Update app.py

app.py

@@ -8,12 +8,15 @@ pipe = CogVideoXImageToVideoPipeline.from_pretrained(
         "THUDM/CogVideoX-5b-I2V",
         torch_dtype=torch.bfloat16
     )
-    
-pipe.vae.enable_tiling()
-pipe.vae.enable_slicing()
 
 @spaces.GPU(duration=120)
 def generate_video(prompt, image):
+    # Ensure the generator is on the same device as the model
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    pipe.vae.enable_tiling()
+    pipe.vae.enable_slicing()
+    generator = torch.Generator(device=device).manual_seed(42)
+    
     video = pipe(
         prompt=prompt,
         image=image,
@@ -21,7 +24,7 @@ def generate_video(prompt, image):
         num_inference_steps=50,
         num_frames=49,
         guidance_scale=6,
-        generator=torch.Generator(device="cuda").manual_seed(42),
+        generator=generator,
     ).frames[0]
     
     video_path = "output.mp4"