Cannot modify batch size and video length

[SingleZombie]

Experience

I wanted to process a video of 11 frames, so I modified the config.

...
input_data:
  ...
  n_sample_frames: 11
...
validation_data:
   ...
  video_length: 11
...

But when I ran the script, I got "CUDA out of memory". I tried to modify the code by spliting the data into batchs manually (in test_vid2vid_zero.py):

bs = 8
for step, batch in enumerate(input_dataloader):
    all_pixel_values = batch["pixel_values"]
    n_batch = all_pixel_values.shape[1]
    samples = []
    for i in range((n_batch - 1) // bs + 1):
        s_id = i * bs
        e_id = (i + 1) * bs
        if e_id > n_batch:
            e_id = n_batch
        pixel_values = all_pixel_values[:, s_id:e_id].to(weight_dtype).to(
            'cuda:0')
     ......

After the modification, I still got errors indicating the shapes of some tensors are not aligned. I found it is because when you pass validation_data to validation_pipeline

                sample = validation_pipeline(
                    prompts,
                    generator=generator,
                    latents=ddim_inv_latent,
                    uncond_embeddings=uncond_embeddings,
                    **validation_data).images

, video_length is also passed to this function. To allow the batch processing, I have to pop video_length from validation_data and pass current batch size to the function. I didn't want to modify the codes anymore.

Summary

Currently, the users cannot set the batch size of testing data. All the frames are processed in a single batch, which may lead to "CUDA out of memory". I think you should refactor the codes in Dataset class and let DataLoader manage the batch prcoessing.