Class WanMotionTensorDatasetWiRefMotion

datasets/video_datasets.py:8–204

class WanMotionTensorDatasetWiRefMotion(torch.utils.data.Dataset):
    def __init__(self, train_json_file_list, test_json_file_list, motion_mean_path, motion_std_path, min_motion_length, max_motion_length, null_context_path, uncond_prob=0.5, is_test=False, use_global_orient=True, text_key='video_text_annot', duplicate_meta=None, t2m_prob_low_quality=0.4, t2m_prob_default=0.8, **kwargs):
        self.data_list = self._load_data_list(
            train_json_file_list=train_json_file_list,
            test_json_file_list=test_json_file_list,
            duplicate_meta=duplicate_meta,
            is_test=is_test,
        )
        print(len(self.data_list), "tensors cached in metadata.")

        motion_mean = np.load(motion_mean_path)
        motion_std = np.load(motion_std_path)

        self.motion_mean = torch.from_numpy(motion_mean).float()
        self.motion_std = torch.from_numpy(motion_std).float()
        self.motion_dim = motion_mean.shape[-1]
        self.min_motion_length = min_motion_length
        self.max_motion_length = max_motion_length
        self.uncond_prob = uncond_prob
        self.null_context = torch.load(null_context_path, weights_only=True, map_location="cpu")   # [226, 4096]
        self.is_test = is_test
        self.use_global_orient = use_global_orient
        self.text_key = text_key
        self.prompt_emb_key = f'{text_key}_wanvideot5_embed_path'
        # motion31/mogen_db captions are noisier, so use lower T2M probability by default.
        self.t2m_prob_low_quality = t2m_prob_low_quality
        # Other datasets default to a higher T2M probability.
        self.t2m_prob_default = t2m_prob_default
        total_samples = len(self.data_list)
        missing_prompt = 0
        for sample in self.data_list:
            prompt_path = sample.get(self.prompt_emb_key)
            if prompt_path is None or not os.path.exists(prompt_path):
                missing_prompt += 1
        print(
            f'[{self.__class__.__name__}] prompt embeddings with key "{self.prompt_emb_key}": '
            f'{total_samples - missing_prompt}/{total_samples} available'
        )

    def __getitem__(self, index):
        sample_index, data, motion, motion_path = self._fetch_motion_sample(index)
        motion_length = motion.shape[0]

        normalized_motion = self._normalize_motion_tensor(motion)

        data_dict = {
            "motion": normalized_motion,
            "motion_length": motion_length,
        }
        motion_dim_mask = torch.ones(self.motion_dim)

        ref_motion_dict, motion_dim_mask, attend_to_text_mask = self._prepare_ref_motion(
            data=data,
            motion=normalized_motion,
            motion_dim_mask=motion_dim_mask,
            motion_path=motion_path,
        )
        data_dict.update(ref_motion_dict)