Update app.py

app.py

@@ -12,7 +12,7 @@ import asyncio  # Import asyncio
 os.environ["SAFETENSORS_FAST_GPU"] = "1"
 os.putenv("HF_HUB_ENABLE_HF_TRANSFER", "1")
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+# No longer needed here: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 # Use gr.State to hold the predictor.  Initialize it to None.
 predictor_state = gr.State(None)
@@ -29,12 +29,13 @@ def init_predictor(task_type: str):
             task_type=TaskType.I2V if task_type == "i2v" else TaskType.T2V,
             model_id="Skywork/skyreels-v1-Hunyuan-i2v",  # Adjust model ID as needed
             quant_model=True,
-            is_offload=True,
+            is_offload=True,  # Consider removing if you have enough GPU memory
             offload_config=OffloadConfig(
                 high_cpu_memory=True,
                 parameters_level=True,
             ),
             use_multiprocessing=False,
+            device="cuda:0" if torch.cuda.is_available() else "cpu" # Pass device to the constructor
         )
         return predictor
     except (RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError) as e:
@@ -74,15 +75,18 @@ async def generate_video(prompt, image_file, predictor):
     }
 
     try:
-      kwargs["image"] = load_image(image=image_file.name)
+        # Load the image and move it to the correct device *before* inference
+        image = load_image(image=image_file.name)
+        # No need to manually move to device. SkyReelsVideoInfer should handle it.
+        kwargs["image"] = image
     except Exception as e:
-      return gr.Error(f"image loading error: {e}")
+        return gr.Error(f"Image loading error: {e}")
 
     try:
         output = predictor.inference(kwargs)
         frames = output
     except Exception as e:
-        return gr.Error(f"Inference error: {e}")
+        return gr.Error(f"Inference error: {e}"), None  # Return None for predictor on error
 
     save_dir = "./result/i2v"  # Consistent directory
     os.makedirs(save_dir, exist_ok=True)
@@ -92,7 +96,7 @@ async def generate_video(prompt, image_file, predictor):
     try:
         export_to_video(frames, video_out_file, fps=24)
     except Exception as e:
-        return gr.Error(f"Video export error: {e}")
+        return gr.Error(f"Video export error: {e}"), None # Return None for predictor
 
     return video_out_file, predictor  # Return updated predictor