Update inference/flovd_demo.py

inference/flovd_demo.py

@@ -37,6 +37,8 @@ import numpy as np
 from PIL import Image
 
 import torch
+import types
+from diffusers.pipelines.cogvideo.pipeline_cogvideox_image2video import CogVideoXImageToVideoPipeline
 
 from diffusers import (
     CogVideoXDPMScheduler,
@@ -305,6 +307,21 @@ def generate_video(
     - seed (int): The seed for reproducibility.
     - 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)
+
     print("at generate video", flush=True)
     local_rank = 'cuda'