Fix CUDA device mismatch - resize image and add autocast

handler.py

@@ -64,6 +64,10 @@ async def predict(request: dict):
         else:
             image_bytes = base64.b64decode(image_data)
             image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        
+        # Resize to expected dimensions for Cosmos Video2World
+        image = image.resize((1280, 704))
+        
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Failed to load image: {str(e)}")
 
@@ -73,20 +77,22 @@ async def predict(request: dict):
     guidance_scale = inputs.get("guidance_scale", 7.0)
     seed = inputs.get("seed")
 
+    # Create generator on correct device
     generator = None
     if seed is not None:
         generator = torch.Generator(device="cuda").manual_seed(int(seed))
 
     try:
-        output = pipe(
-            image=image,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_frames=num_frames,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            generator=generator,
-        )
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            output = pipe(
+                image=image,
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_frames=num_frames,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                generator=generator,
+            )
 
         video_path = "/tmp/output.mp4"
         export_to_video(output.frames[0], video_path, fps=16)
@@ -97,6 +103,8 @@ async def predict(request: dict):
         return {"video": video_b64, "content_type": "video/mp4"}
 
     except Exception as e:
+        import traceback
+        traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"Inference failed: {str(e)}")
 
 @app.get("/health")