Update app.py

app.py

@@ -4,30 +4,27 @@ import sys
 import time
 import os
 import random
-from PIL import Image
-import torch
-import asyncio
-
-# os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Uncomment if needed
+from PIL import Image 
+ # os.environ["CUDA_VISIBLE_DEVICES"] = ""
 os.environ["SAFETENSORS_FAST_GPU"] = "1"
-os.putenv("HF_HUB_ENABLE_HF_TRANSFER", "1")
-
+os.putenv("HF_HUB_ENABLE_HF_TRANSFER","1")
+import torch
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# Create the gr.State component *outside* the gr.Blocks context.  This is CORRECT.
-predictor_state = gr.State(None)  # Use a consistent name
-task_type_state = gr.State("i2v")  # Store the task type
+# Create the gr.State component *outside* the gr.Blocks context
+
+predictor = gr.State(None)
 
 def init_predictor(task_type: str):
     from skyreelsinfer import TaskType
     from skyreelsinfer.offload import OffloadConfig
     from skyreelsinfer.skyreels_video_infer import SkyReelsVideoInfer
     from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError
-
+    global predictor
     try:
         predictor = SkyReelsVideoInfer(
             task_type=TaskType.I2V if task_type == "i2v" else TaskType.T2V,
-            model_id="Skywork/skyreels-v1-Hunyuan-i2v",  # Or a different model for T2V
+            model_id="Skywork/skyreels-v1-Hunyuan-i2v",
             quant_model=True,
             is_offload=True,
             offload_config=OffloadConfig(
@@ -38,27 +35,24 @@ def init_predictor(task_type: str):
         )
         return predictor
     except (RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError) as e:
-        print(f"Error: Model not found. Details: {e}")  # Log the error
-        return None  # Return None if model loading fails
+        return f"Error: Model not found. Details: {e}", None
     except Exception as e:
-        print(f"Error loading model: {e}")  # Log the error
-        return None # Return None if model loading fails
-
-
+        return f"Error loading model: {e}", None
+        
+predictor = init_predictor('i2v')
+global predictor
 
 @spaces.GPU(duration=80)
-async def generate_video(prompt, image_file, predictor):
+def generate_video(prompt, image, predictor):
     from diffusers.utils import export_to_video
     from diffusers.utils import load_image
-    if image_file is None:
-        return gr.Error("Error: For i2v, provide image path.") , None # Use gr.Error
-    if not isinstance(prompt, str) or not prompt.strip():  # Check for empty prompt
-        return gr.Error("Error: No prompt."), None # Use gr.Error
-
-    #if seed == -1: #removed seed parameter
+    if image == None:
+        return "Error: For i2v, provide image path.", "{}"
+    if not isinstance(prompt, str):
+        return "Error: No prompt.", "{}"
+    #if seed == -1:
     random.seed(time.time())
     seed = int(random.randrange(4294967294))
-
     kwargs = {
         "prompt": prompt,
         "height": 256,
@@ -71,63 +65,42 @@ async def generate_video(prompt, image_file, predictor):
         "negative_prompt": "bad quality, blur",
         "cfg_for": False,
     }
-    # Check if predictor is initialized and the task type is i2v.
-
-    kwargs["image"] = load_image(image=image_file)
-
-    if predictor is None:
-        return gr.Error("Predictor not initialized."), None
-
-    try:
-        output = predictor.inference(kwargs)
-    except Exception as e:
-        return gr.Error(f"Inference error: {e}"), None
 
+    kwargs["image"] = load_image(image=image)
+    output = predictor.inference(kwargs)
     frames = output
-    save_dir = f"./result/i2v"  # Use a consistent directory
+    save_dir = f"./result/{task_type}"
     os.makedirs(save_dir, exist_ok=True)
-    video_out_file = os.path.join(save_dir, f"{prompt[:100]}_{int(seed)}.mp4")
+    video_out_file = f"{save_dir}/{prompt[:100]}_{int(seed)}.mp4"
     print(f"Generating video: {video_out_file}")
-    try:
-        export_to_video(frames, video_out_file, fps=24)
-    except Exception as e:
-        return gr.Error(f"Video export error: {e}"), None
-    return video_out_file, predictor  # Return the updated predictor
-
-
+    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 was removed
+        return Image.open(file.name)
     else:
         return None
-
-async def load_model_async(task_type):
-  predictor = init_predictor(task_type)
-  return predictor
-
-async def main():
-    with gr.Blocks() as demo:
-        # Remove the local predictor variable.  We'll use the global predictor_state.
-        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.Text(label="Prompt")
-        generate_button = gr.Button("Generate")
-        output_video = gr.Video(label="Output Video")
-
-        image_file.change(
-            display_image,
-            inputs=[image_file],
-            outputs=[image_file_preview]
-        )
-        # Use predictor_state in inputs and outputs, and return it from the function.
-        generate_button.click(
-            fn=generate_video,
-            inputs=[prompt_textbox, image_file, predictor_state],
-            outputs=[output_video, predictor_state],  # Output the predictor_state
-        )
-        predictor_state.value = await load_model_async("i2v")
-
-    await demo.launch()
-
-if __name__ == "__main__":
-    asyncio.run(main())
+        
+with gr.Blocks() as demo:
+    #predictor = gr.State({})  # Initialize as an empty dictionary
+
+    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.Text(label="Prompt")
+    generate_button = gr.Button("Generate")
+    output_video = gr.Video(label="Output Video")
+
+    image_file.change(
+        display_image,
+        inputs=[image_file],
+        outputs=[image_file_preview]
+    )
+
+    generate_button.click(
+        fn=generate_video,
+        inputs=[prompt_textbox, image_file, predictor],
+        outputs=[output_video],
+    )
+
+demo.launch()