Update app.py

app.py

@@ -10,15 +10,7 @@ import random
 # Create the gr.State component *outside* the gr.Blocks context
 predictor_state = gr.State(None)
 
-def get_transformer_model_id(task_type: str) -> str:
-    if task_type == "i2v":
-        return "Skywork/skyreels-v1-Hunyuan-i2v"
-    else:
-        return "Skywork/skyreels-v1-Hunyuan-t2v"
-        
- # @spaces.GPU(duration=120)
-def init_predictor(task_type: str):
-    # ALL IMPORTS NOW INSIDE THIS FUNCTION
+ def init_predictor(task_type: str):
     import torch
     from skyreelsinfer import TaskType
     from skyreelsinfer.offload import OffloadConfig
@@ -28,7 +20,7 @@ def init_predictor(task_type: str):
     try:
         predictor = SkyReelsVideoInfer(
             task_type=TaskType.I2V if task_type == "i2v" else TaskType.T2V,
-            model_id=get_transformer_model_id(task_type),
+            model_id="Skywork/skyreels-v1-Hunyuan-i2v",
             quant_model=True,
             is_offload=True,
             offload_config=OffloadConfig(
@@ -47,20 +39,19 @@ def init_predictor(task_type: str):
 init_predictor('i2v')
 
 @spaces.GPU(duration=80)
-def generate_video(prompt, seed, image, task_type, predictor): # predictor as argument
-    # IMPORTS INSIDE THIS FUNCTION TOO
+def generate_video(prompt, image, predictor):
     from diffusers.utils import export_to_video
     from diffusers.utils import load_image
     import os
 
-    if task_type == "i2v" and not isinstance(image, str):
+    if image == None:
         return "Error: For i2v, provide image path.", "{}"
-    if not isinstance(prompt, str) or not isinstance(seed, (int, float)):
-        return "Error: Invalid inputs.", "{}"
+    if not isinstance(prompt, str):
+        return "Error: No prompt.", "{}"
 
-    if seed == -1:
-        random.seed(time.time())
-        seed = int(random.randrange(4294967294))
+    #if seed == -1:
+    random.seed(time.time())
+    seed = int(random.randrange(4294967294))
 
     kwargs = {
         "prompt": prompt,
@@ -75,51 +66,41 @@ def generate_video(prompt, seed, image, task_type, predictor): # predictor as ar
         "cfg_for": False,
     }
 
-    if task_type == "i2v":
-        if image is None or not os.path.exists(image):
-            return "Error: Image not found.", "{}"
-        try:
-            kwargs["image"] = load_image(image=image)
-        except Exception as e:
-            return f"Error loading image: {e}", "{}"
-
-    try:
-        if predictor is None:
-            return "Error: Model not init.", "{}"
-
-        output = predictor.inference(kwargs)
-        frames = output
-
-        save_dir = f"./result/{task_type}"
-        os.makedirs(save_dir, exist_ok=True)
-        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
-
-    except Exception as e:
-        return f"Error: {e}", "{}"
-
+    kwargs["image"] = load_image(image=image)
+    output = predictor.inference(kwargs)
+    frames = output
+    save_dir = f"./result/{task_type}"
+    os.makedirs(save_dir, exist_ok=True)
+    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
+
+ 
+def display_image(file):
+    if file is not None:
+        return Image.open(file.name)
+    else:
+        return None
+        
 # --- Minimal Gradio Interface ---
 with gr.Blocks() as demo:
-    task_type_dropdown = gr.Dropdown(
-        choices=["i2v", "t2v"], label="Task", value="t2v", elem_id="task_type"
-    )
-    load_model_button = gr.Button("Load Model")
+
+    image_file = gr.File(label="Image Prompt (Required)", file_types=["image"])
+    image_file_preview = gr.Image(label="Image Prompt Preview", interactive=False)
     prompt_textbox = gr.Textbox(label="Prompt")
     generate_button = gr.Button("Generate")
-    output_textbox = gr.Textbox(label="Output") # Just a textbox
     output_video = gr.Video(label="Output Video") # Just a textbox
 
-    load_model_button.click(
-        fn=init_predictor,
-        inputs=[task_type_dropdown],
-        outputs=[output_textbox, predictor_state], # Correct order of outputs
+    image_file.change(
+        display_image,
+        inputs=[image_file],
+        outputs=[image_file_preview]
     )
 
     generate_button.click(
         fn=generate_video,
-        inputs=[prompt_textbox, task_type_dropdown, predictor_state],
+        inputs=[prompt_textbox, image, predictor_state],
         outputs=[output_video],
     )