Update inference/flovd_demo.py

inference/flovd_demo.py

@@ -264,6 +264,27 @@ def save_flow_warped_video(image, flow, filename, fps=16):
         frame_list.append(Image.fromarray(frame))
     
     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, image, num_frames, height, width, batch_size, dtype, generator, do_classifier_free_guidance=False):
+        image_latents = self.vae.encode(image.to(self.device, dtype=dtype)).latent_dist.sample()
+        image_latents = image_latents * self.vae.config.scaling_factor
+
+        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]),
+                device=image_latents.device, dtype=image_latents.dtype
+            )
+            image_latents = torch.cat([image_latents, latent_padding], dim=1)
+
+        noise = torch.randn_like(image_latents, generator=generator)
+        latents = noise.to(self.device, dtype=dtype)
+
+        return latents, image_latents.to(self.device, dtype=dtype)
+
+    CogVideoXImageToVideoPipeline.prepare_latents = new_prepare_latents
 
 def generate_video(
     prompt: str,
@@ -308,19 +329,7 @@ def generate_video(
     - fps (int): The frames per second for the generated video.
     """
 
-    def patch_prepare_latents_with_device():
-        original_prepare_latents = CogVideoXImageToVideoPipeline.prepare_latents
-    
-        def prepare_latents_with_device(self, *args, **kwargs):
-            result = original_prepare_latents(self, *args, **kwargs)
-            # Ensure returned tensors are moved to the correct device
-            if isinstance(result, tuple):
-                result = tuple(t.to(self.device) if isinstance(t, torch.Tensor) else t for t in result)
-            elif isinstance(result, torch.Tensor):
-                result = result.to(self.device)
-            return result
-    
-        CogVideoXImageToVideoPipeline.prepare_latents = types.MethodType(prepare_latents_with_device, CogVideoXImageToVideoPipeline)
+    patch_prepare_latents_safe()
 
     print("at generate video", flush=True)
     local_rank = 'cuda'