Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 import spaces
 import gradio as gr
 import argparse
@@ -6,17 +7,16 @@ import os
 import random
 import subprocess
 from PIL import Image
-import numpy as np  # Import NumPy
+import numpy as np
 
-
-subprocess.run(['sh', './sky.sh'])  # Keep if needed
+subprocess.run(['sh', './sky.sh'])
 sys.path.append("./SkyReels-V1")
 
-# Corrected Relative Imports
 from skyreelsinfer import TaskType
 from skyreelsinfer.offload import OffloadConfig
-from skyreelsinfer.skyreels_video_infer import SkyReelsVideoSingleGpuInfer  # Import the class
+from skyreelsinfer.skyreels_video_infer import SkyReelsVideoSingleGpuInfer
 from diffusers.utils import export_to_video
+
 import torch
 import logging
 
@@ -29,13 +29,12 @@ torch.backends.cudnn.benchmark = False
 torch.set_float32_matmul_precision("highest")
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-# --- Dummy Classes (Moved to skyreelsinfer/__init__.py) ---
 logger = logging.getLogger(__name__)
-# --- Global Variables and Argument Parsing ---
+
 
 _predictor = None
-task_type = TaskType.I2V  # Default task type.
+task_type = TaskType.I2V
+
 @spaces.GPU(duration=90)
 def init_predictor():
     global _predictor
@@ -46,7 +45,7 @@ def init_predictor():
         if task_type == TaskType.I2V:
             model_id = "Skywork/SkyReels-V1-Hunyuan-I2V"
         elif task_type == TaskType.T2V:
-            model_id = "your_t2v_model_id"  # Replace
+            model_id = "your_t2v_model_id"
         else:
             raise ValueError(f"Invalid task_type: {task_type}")
 
@@ -94,16 +93,15 @@ def generate_video(prompt, seed, image=None):
     elif task_type == TaskType.T2V:
         pass
     else:
-      raise ValueError(f"Invalid task_type: {task_type}")
+      raise ValueError(f"Invalid Tasktype")
 
     if _predictor is None:
         init_predictor()
 
     output = _predictor.infer(**kwargs)
 
-    # --- Convert to NumPy, move to CPU, scale, and change dtype ---
     output = (output.cpu().numpy() * 255).astype(np.uint8)
-    output = output.transpose(0, 2, 3, 4, 1) #Correct transpose.
+    output = output.transpose(0, 2, 3, 4, 1)
 
     save_dir = f"./result/{task_type.name}"
     os.makedirs(save_dir, exist_ok=True)