Update app.py

app.py

@@ -1,24 +1,10 @@
-import spaces
 import gradio as gr
-import sys
-import time
 import os
+import time
 import random
 
-sys.path.append(".")  # Correct path for Hugging Face Space
-from skyreelsinfer import TaskType
-from skyreelsinfer.offload import OffloadConfig
-from skyreelsinfer.skyreels_video_infer import SkyReelsVideoInfer
-from diffusers.utils import export_to_video
-from diffusers.utils import load_image
-import torch
-from huggingface_hub import HfApi
-from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
-# --- Model Loading ---
-predictor = None
-
 def get_transformer_model_id(task_type: str) -> str:
     if task_type == "i2v":
         return "Skywork/skyreels-v1-Hunyuan-i2v"
@@ -26,7 +12,13 @@ def get_transformer_model_id(task_type: str) -> str:
         return "Skywork/skyreels-v1-Hunyuan-t2v"
 
 def init_predictor(task_type: str):
-    global predictor
+    # ALL IMPORTS NOW INSIDE THIS FUNCTION
+    import torch
+    from skyreelsinfer import TaskType
+    from skyreelsinfer.offload import OffloadConfig
+    from skyreelsinfer.skyreels_video_infer import SkyReelsVideoInfer
+    from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError
+
     try:
         predictor = SkyReelsVideoInfer(
             task_type=TaskType.I2V if task_type == "i2v" else TaskType.T2V,
@@ -36,25 +28,27 @@ def init_predictor(task_type: str):
             offload_config=OffloadConfig(
                 high_cpu_memory=True,
                 parameters_level=True,
-                # compiler_transformer=False,  # Keep this consistent
             ),
-            use_multiprocessing=False,  # Still pass this, though it's ignored
+            use_multiprocessing=False,
         )
-        return "Model loaded successfully!"
+        return "Model loaded successfully!", predictor  # Return predictor
 
     except (RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError) as e:
-      return f"Error: Model not found. Details: {e}"
+        return f"Error: Model not found. Details: {e}", None
     except Exception as e:
-        return f"Error loading model: {e}"
+        return f"Error loading model: {e}", None
+
 
-@spaces.GPU(duration=90)
-def generate_video(prompt, seed, image=None, task_type=None):
-    global predictor
+def generate_video(prompt, seed, image, task_type, predictor): # predictor as argument
+    # IMPORTS INSIDE THIS FUNCTION TOO
+    from diffusers.utils import export_to_video
+    from diffusers.utils import load_image
+    import os
 
     if task_type == "i2v" and not isinstance(image, str):
-        return "Error: For i2v, provide a valid image file path.", "{}"
+        return "Error: For i2v, provide image path.", "{}"
     if not isinstance(prompt, str) or not isinstance(seed, (int, float)):
-        return "Error: Invalid input types.", "{}"
+        return "Error: Invalid inputs.", "{}"
 
     if seed == -1:
         random.seed(time.time())
@@ -62,68 +56,70 @@ def generate_video(prompt, seed, image=None, task_type=None):
 
     kwargs = {
         "prompt": prompt,
-        "height": 256,  # Reduced for faster CPU processing
-        "width": 256,   # Reduced for faster CPU processing
-        "num_frames": 64,  # Reduced for faster CPU processing
-        "num_inference_steps": 50,  # Reduced for faster CPU processing
+        "height": 256,
+        "width": 256,
+        "num_frames": 24,
+        "num_inference_steps": 30,
         "seed": int(seed),
-        "guidance_scale": 6.0,
+        "guidance_scale": 7.0,
         "embedded_guidance_scale": 1.0,
-        "negative_prompt": "bad quality", #shorter prompt
+        "negative_prompt": "bad quality, blur",
         "cfg_for": False,
     }
 
     if task_type == "i2v":
         if image is None or not os.path.exists(image):
-            return "Error: Image not provided or not found.", "{}"
+            return "Error: Image not found.", "{}"
         try:
             kwargs["image"] = load_image(image=image)
         except Exception as e:
-              return f"Error loading image: {e}", "{}"
+            return f"Error loading image: {e}", "{}"
 
     try:
         if predictor is None:
-            return "Error: Model not initialized.", "{}"
+            return "Error: Model not init.", "{}"
+
+        output = predictor.inference(kwargs)
+        frames = output
 
-        output = predictor.inference(**kwargs)
         save_dir = f"./result/{task_type}"
         os.makedirs(save_dir, exist_ok=True)
-        video_out_file = f"{save_dir}/{prompt[:100].replace('/','')}_{int(seed)}.mp4"
-        print(f"Generating video, local path: {video_out_file}")
-        export_to_video(output, video_out_file, fps=24)
-        return output, str(kwargs)
+        video_out_file = f"{save_dir}/{prompt[:100]}_{int(seed)}.mp4"
+        print(f"Generating video: {video_out_file}")
+        export_to_video(frames, video_out_file, fps=24)
+        return video_out_file, str(kwargs)
 
     except Exception as e:
-        return f"Error during generation: {e}", "{}"
+        return f"Error: {e}", "{}"
 
 # --- Gradio Interface ---
 with gr.Blocks() as demo:
     with gr.Row():
         task_type_dropdown = gr.Dropdown(
-            choices=["i2v", "t2v"], label="Task Type", value="t2v"
+            choices=["i2v", "t2v"], label="Task", value="t2v"
         )
         load_model_button = gr.Button("Load Model")
-        model_status = gr.Textbox(label="Model Status")
+        model_status = gr.Textbox(label="Status")
     with gr.Row():
         with gr.Column():
-            prompt = gr.Textbox(label="Input Prompt")
-            seed = gr.Number(label="Random Seed", value=-1)
-            image = gr.Image(label="Upload Image (for i2v)", type="filepath")
-            submit_button = gr.Button("Generate Video")
+            prompt = gr.Textbox(label="Prompt")
+            seed = gr.Number(label="Seed", value=-1)
+            image = gr.Image(label="Image (i2v)", type="filepath")
+            submit_button = gr.Button("Generate")
         with gr.Column():
-            output_video = gr.Video(label="Generated Video")
-            output_params = gr.Textbox(label="Output Parameters")
+            output_video = gr.Video(label="Video")
+            output_params = gr.Textbox(label="Params")
 
     load_model_button.click(
         fn=init_predictor,
         inputs=[task_type_dropdown],
-        outputs=[model_status]
+        outputs=[model_status, "state"],  # Output to a hidden state
     )
 
     submit_button.click(
         fn=generate_video,
-        inputs=[prompt, seed, image, task_type_dropdown],
+        inputs=[prompt, seed, image, task_type_dropdown, "state"],  # Input from state
         outputs=[output_video, output_params],
     )
 
-demo.launch() # Don't use demo.launch() inside HuggingFace Spaces.
+demo.launch()