Function collate_fn_motion_wanvideo

datasets/dataloader.py:109–167

(batch)

def collate_fn_motion_wanvideo(batch):
    # filter out None
    batch = [x for x in batch if x is not None]

    # collate motion
    motions = [x.pop('motion') for x in batch]
    motions = collate_tensors(motions)
    motion_lengths = [x.get('motion_length') for x in batch]
    motion_mask = lengths_to_mask(motion_lengths, motions.device, motions.shape[1])

    with_ref_motion = False
    if 'ref_motion' in batch[0]:
        with_ref_motion = True
        ref_motions = [x.pop('ref_motion') for x in batch]
        ref_motions_original = [x.pop('ref_motion_original') for x in batch]
        ref_motions = collate_tensors(ref_motions)
        ref_motions_original = collate_tensors(ref_motions_original)
        ref_motion_lengths = [x.pop('ref_motion_length') for x in batch]
        ref_motion_mask = lengths_to_mask(ref_motion_lengths, ref_motions.device, ref_motions.shape[1])

    # collate prompt
    with_prompt_emb = False
    if 'prompt_emb' in batch[0]:
        with_prompt_emb = True
        prompt_emb = [x.pop('prompt_emb') for x in batch]
        prompt_emb = collate_tensors(prompt_emb)
        prompt_lengths = [x.pop('prompt_length') for x in batch]
        prompt_emb_mask = lengths_to_mask(prompt_lengths, prompt_emb.device, prompt_emb.shape[1])

    texts = None
    if 'text' in batch[0]:
        texts = [x.pop('text') for x in batch]

    sample_ids = None
    if 'test_sample_id' in batch[0]:
        sample_ids = [x.pop('test_sample_id') for x in batch]

    try:
        ret = torch.utils.data.default_collate(batch)
    except Exception as e:
        print(f'Failed to collate batch in collate_fn_motion_wanvideo: {e}')
        for idx, sample in enumerate(batch[:2]):
            print(f'sample {idx} keys: {list(sample.keys())}')
        raise
    
    ret['motion'] = motions
    ret['motion_mask'] = motion_mask
    if with_ref_motion:
        ret['ref_motion'] = ref_motions
        ret['ref_motion_mask'] = ref_motion_mask
        ret['ref_motion_original'] = ref_motions_original
    if with_prompt_emb:
        ret['prompt_emb'] = prompt_emb
        ret['prompt_emb_mask'] = prompt_emb_mask
    if texts is not None:
        ret['text'] = texts
    if sample_ids is not None:
        ret['test_sample_id'] = sample_ids