Update inference/flovd_demo.py

inference/flovd_demo.py

@@ -265,8 +265,6 @@ def save_flow_warped_video(image, flow, filename, fps=16):
     
     export_to_video(frame_list, filename, fps=fps)
 
-    
-from diffusers.pipelines.cogvideo.pipeline_cogvideox_image2video import CogVideoXImageToVideoPipeline
 def patch_prepare_latents_safe():
     def new_prepare_latents(
         self,
@@ -281,16 +279,20 @@ def patch_prepare_latents_safe():
         generator,
         latents=None,
     ):
-        # Ensure 5D input: [B, C, F=1, H, W]
         image_5d = image.unsqueeze(2) if image.ndim == 4 else image
-
         image_latents = self.vae.encode(image_5d.to(device, dtype=dtype)).latent_dist.sample()
         image_latents = image_latents * self.vae.config.scaling_factor
 
-        # Pad temporal dimension if needed
+        # Pad frame dim if needed
         if image_latents.shape[2] != num_frames:
             latent_padding = torch.zeros(
-                (image_latents.shape[0], latent_channels, num_frames - image_latents.shape[2], height, width),
+                (
+                    image_latents.shape[0],
+                    image_latents.shape[1],
+                    num_frames - image_latents.shape[2],
+                    image_latents.shape[3],
+                    image_latents.shape[4],
+                ),
                 device=image_latents.device,
                 dtype=image_latents.dtype
             )