Update app.py

app.py

@@ -1,22 +1,22 @@
 import spaces
 import gradio as gr
-import argparse  # Import argparse
+import argparse
 import sys
 import os
 import random
 import subprocess
-from PIL import Image  # Keep PIL import
+from PIL import Image
+import numpy as np  # Import NumPy
 
-subprocess.run(['sh', './sky.sh'])  # Keep if needed
 
+# subprocess.run(['sh', './sky.sh'])  # Keep if needed
 sys.path.append("./SkyReels-V1")
 
 # Corrected Relative Imports
-from skyreelsinfer import TaskType  # Now imported correctly
-from skyreelsinfer.offload import OffloadConfig
-from skyreelsinfer.skyreels_video_infer import SkyReelsVideoSingleGpuInfer  # Import the class
+from SkyReels-V1.skyreelsinfer import TaskType
+from SkyReels-V1.skyreelsinfer.offload import OffloadConfig
+from SkyReels-V1.skyreelsinfer.skyreels_video_infer import SkyReelsVideoSingleGpuInfer  # Import the class
 from diffusers.utils import export_to_video
-
 import torch
 import logging
 
@@ -30,26 +30,24 @@ torch.set_float32_matmul_precision("highest")
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-logger = logging.getLogger(__name__)
-
 # --- 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.  IMPORTANT: Set a default.
+task_type = TaskType.I2V  # Default task type.
 
 @spaces.GPU(duration=90)
 def init_predictor():
     global _predictor
-    global task_type  # Access the global task_type
+    global task_type
     logger = logging.getLogger(__name__)
 
     if _predictor is None:
         if task_type == TaskType.I2V:
             model_id = "Skywork/SkyReels-V1-Hunyuan-I2V"
         elif task_type == TaskType.T2V:
-            model_id = "your_t2v_model_id"  # Replace with your T2V model ID
+            model_id = "your_t2v_model_id"  # Replace
         else:
             raise ValueError(f"Invalid task_type: {task_type}")
 
@@ -69,13 +67,13 @@ def init_predictor():
     else:
         logger.warning("Predictor already initialized (should be rare).")
 
-@spaces.GPU(duration=90) # Needed, because we are saving a file
+@spaces.GPU(duration=90)
 def generate_video(prompt, seed, image=None):
     global _predictor
     global task_type
 
     if seed == -1:
-        random.seed()  # Use system time for randomness if seed is -1
+        random.seed()
         seed = int(random.randrange(4294967294))
 
     kwargs = {
@@ -93,37 +91,40 @@ def generate_video(prompt, seed, image=None):
 
     if task_type == TaskType.I2V:
         assert image is not None, "Please input an image for I2V task."
-        kwargs["image"] = Image.open(image)  # Use PIL.Image.open
+        kwargs["image"] = Image.open(image)
     elif task_type == TaskType.T2V:
-      pass # No image needed.
+        pass # No image
     else:
-        raise ValueError("Invalid Tasktype")
+      raise ValueError(f"Invalid task_type: {task_type}")
 
     if _predictor is None:
         init_predictor()
 
     output = _predictor.infer(**kwargs)
 
-    save_dir = f"./result/{task_type.name}"  # Use task_type.name for directory
+    # --- Convert to NumPy, move to CPU, scale, and change dtype ---
+    output = (output.cpu().numpy() * 255).astype(np.uint8)
+    # --- Convert from B, C, T, H, W to B, T, H, W, C
+    output = output.transpose(0, 2, 3, 4, 1)
+    save_dir = f"./result/{task_type.name}"
     os.makedirs(save_dir, exist_ok=True)
-    video_out_file = f"{save_dir}/{prompt[:100].replace('/', '')}_{seed}.mp4"
+    video_out_file = f"{save_dir}/{prompt[:100].replace('/','')}_{seed}.mp4"
     print(f"generate video, local path: {video_out_file}")
-    export_to_video(output, video_out_file, fps=24)
-    return video_out_file, kwargs # Return the file path
+    export_to_video(output, video_out_file, fps=24) # Pass fps
+    return video_out_file, kwargs
 
 
 def create_gradio_interface():
     with gr.Blocks() as demo:
         with gr.Row():
-            with gr.Column():
-                image = gr.Image(label="Upload Image", type="filepath")
-                prompt = gr.Textbox(label="Input Prompt")
-                seed = gr.Number(label="Random Seed", value=-1)  # Default to -1
-            with gr.Column():
-                submit_button = gr.Button("Generate Video")
-                output_video = gr.Video(label="Generated Video")
-                output_params = gr.Textbox(label="Output Parameters")
-
+          with gr.Column():
+            image = gr.Image(label="Upload Image", type="filepath")
+            prompt = gr.Textbox(label="Input Prompt")
+            seed = gr.Number(label="Random Seed", value=-1)
+          with gr.Column():
+            submit_button = gr.Button("Generate Video")
+            output_video = gr.Video(label="Generated Video")
+            output_params = gr.Textbox(label="Output Parameters")
         submit_button.click(
             fn=generate_video,
             inputs=[prompt, seed, image],
@@ -138,12 +139,11 @@ if __name__ == "__main__":
                         help="Task type, 't2v' for text-to-video, 'i2v' for image-to-video.")
     args = parser.parse_args()
 
-    # Set the global task_type based on command-line arguments
     if args.task_type == "t2v":
         task_type = TaskType.T2V
     elif args.task_type == "i2v":
         task_type = TaskType.I2V
-    # No else needed, default is already set
+    # No else, default is already set.
 
     demo = create_gradio_interface()
-    demo.queue().launch() # Add queue
+    demo.queue().launch()