Update app.py

app.py

@@ -4,27 +4,30 @@ import sys
 import time
 import os
 import random
-from PIL import Image 
- # os.environ["CUDA_VISIBLE_DEVICES"] = ""
-os.environ["SAFETENSORS_FAST_GPU"] = "1"
-os.putenv("HF_HUB_ENABLE_HF_TRANSFER","1")
+from PIL import Image
 import torch
+import asyncio  # Import asyncio
+
+# os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Uncomment if needed
+os.environ["SAFETENSORS_FAST_GPU"] = "1"
+os.putenv("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# Create the gr.State component *outside* the gr.Blocks context
+# Use gr.State to hold the predictor.  Initialize it to None.
+predictor_state = gr.State(None)
 
-global predictor
 
 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",
+            model_id="Skywork/skyreels-v1-Hunyuan-i2v",  # Adjust model ID as needed
             quant_model=True,
             is_offload=True,
             offload_config=OffloadConfig(
@@ -35,23 +38,28 @@ def init_predictor(task_type: str):
         )
         return predictor
     except (RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError) as e:
-        return f"Error: Model not found. Details: {e}", None
+        print(f"Error: Model not found. Details: {e}")
+        return None
     except Exception as e:
-        return f"Error loading model: {e}", None
-        
-predictor = init_predictor('i2v')
+        print(f"Error loading model: {e}")
+        return None
+
 
-@spaces.GPU(duration=80)
-def generate_video(prompt, image, predictor):
+# Make generate_video async
+async def generate_video(prompt, image_file, predictor):
     from diffusers.utils import export_to_video
     from diffusers.utils import load_image
-    if image == None:
-        return "Error: For i2v, provide image path.", "{}"
-    if not isinstance(prompt, str):
-        return "Error: No prompt.", "{}"
-    #if seed == -1:
+
+    if image_file is None:
+        return gr.Error("Error: For i2v, provide an image.")
+    if not isinstance(prompt, str) or not prompt.strip():
+        return gr.Error("Error: Please provide a prompt.")
+    if predictor is None:
+        return gr.Error("Error: Model not loaded.")
+
     random.seed(time.time())
     seed = int(random.randrange(4294967294))
+
     kwargs = {
         "prompt": prompt,
         "height": 256,
@@ -65,41 +73,62 @@ def generate_video(prompt, image, predictor):
         "cfg_for": False,
     }
 
-    kwargs["image"] = load_image(image=image)
-    output = predictor.inference(kwargs)
-    frames = output
-    save_dir = f"./result/{task_type}"
+    try:
+      kwargs["image"] = load_image(image=image_file.name)
+    except Exception as e:
+      return gr.Error(f"image loading error: {e}")
+
+    try:
+        output = predictor.inference(kwargs)
+        frames = output
+    except Exception as e:
+        return gr.Error(f"Inference error: {e}")
+
+    save_dir = "./result/i2v"  # Consistent directory
     os.makedirs(save_dir, exist_ok=True)
-    video_out_file = f"{save_dir}/{prompt[:100]}_{int(seed)}.mp4"
+    video_out_file = os.path.join(save_dir, f"{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
- 
+
+    try:
+        export_to_video(frames, video_out_file, fps=24)
+    except Exception as e:
+        return gr.Error(f"Video export error: {e}")
+
+    return video_out_file, predictor  # Return updated predictor
+
+
 def display_image(file):
     if file is not None:
         return Image.open(file.name)
     else:
         return None
-        
-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()
+
+async def load_model():
+    predictor = init_predictor('i2v')
+    return predictor
+
+async def main():
+    with gr.Blocks() as demo:
+        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_state],
+            outputs=[output_video, predictor_state],  # Output predictor_state
+        )
+        predictor_state.value = await load_model() # load and set predictor
+
+    await demo.launch()
+
+if __name__ == "__main__":
+    asyncio.run(main())