Update inference/flovd_demo.py

inference/flovd_demo.py

@@ -281,17 +281,20 @@ def patch_prepare_latents_safe():
         generator,
         latents=None,
     ):
-        image_latents = self.vae.encode(image.to(device, dtype=dtype)).latent_dist.sample()
+        # 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
         if image_latents.shape[2] != num_frames:
             latent_padding = torch.zeros(
-                (image_latents.shape[0], num_frames - image_latents.shape[2], image_latents.shape[3], image_latents.shape[4]),
+                (image_latents.shape[0], latent_channels, num_frames - image_latents.shape[2], height, width),
                 device=image_latents.device,
                 dtype=image_latents.dtype
             )
-            image_latents = torch.cat([image_latents, latent_padding], dim=1)
+            image_latents = torch.cat([image_latents, latent_padding], dim=2)
 
         if latents is None:
             noise = torch.randn_like(image_latents, generator=generator)
@@ -303,6 +306,7 @@ def patch_prepare_latents_safe():
     CogVideoXImageToVideoPipeline.prepare_latents = new_prepare_latents
 
 
+
 def generate_video(
     prompt: str,
     fvsm_path: str,