Update app.py

app.py

@@ -1,123 +1,84 @@
 import spaces
 import gradio as gr
+import argparse
 import sys
 import time
 import os
 import random
-from PIL import Image
-import torch
-import asyncio  # Import asyncio
-from skyreelsinfer import TaskType
 from skyreelsinfer.offload import OffloadConfig
-from skyreelsinfer.skyreels_video_infer import SkyReelsVideoInfer
-from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError
+from skyreelsinfer import TaskType
+from skyreelsinfer.skyreels_video_infer import SkyReelsVideoSingleGpuInfer
 from diffusers.utils import export_to_video
 from diffusers.utils import load_image
 
-# os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Uncomment if needed
-os.environ["SAFETENSORS_FAST_GPU"] = "1"
-os.putenv("HF_HUB_ENABLE_HF_TRANSFER", "1")
-
-# No longer needed here: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+#predictor = None
+#task_type = None
 
-# Use gr.State to hold the predictor.  Initialize it to None.
-predictor_state = gr.State(None)
-device="cuda:0" if torch.cuda.is_available() else "cpu" # Pass device to the constructor
-
-def init_predictor(task_type: str):
-    try:
-        predictor = SkyReelsVideoInfer(
-            task_type=TaskType.I2V if task_type == "i2v" else TaskType.T2V,
-            model_id="Skywork/skyreels-v1-Hunyuan-i2v",  # Adjust model ID as needed
-            quant_model=True,
-            is_offload=False,  # Consider removing if you have enough GPU memory
-            offload_config=None, #OffloadConfig(
-            #    high_cpu_memory=True,
-            #    parameters_level=True,
-            #),
-            use_multiprocessing=False,
+#@spaces.GPU(duration=120)
+def init_predictor():
+    global predictor
+    predictor = SkyReelsVideoSingleGpuInfer(
+        task_type= TaskType.I2V,
+        model_id="Skywork/SkyReels-V1-Hunyuan-I2V",
+        quant_model=False,
+        is_offload=False,
+        offload_config=OffloadConfig(
+            high_cpu_memory=True,
+            parameters_level=True,
+            compiler_transformer=False,
         )
-        return predictor
-    except (RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError) as e:
-        print(f"Error: Model not found. Details: {e}")
-        return None
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        return None
-
-# Make generate_video async
-async def generate_video(prompt, image_file, predictor):
-    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))
+    )
+    
+@spaces.GPU(duration=80)
+def generate_video(prompt, seed, image=None):
+    print(f"image:{type(image)}")
+    if seed == -1:
+        random.seed(time.time())
+        seed = int(random.randrange(4294967294))
     kwargs = {
         "prompt": prompt,
-        "height": 256,
-        "width": 256,
-        "num_frames": 24,
+        "height": 512,
+        "width": 512,
+        "num_frames": 97,
         "num_inference_steps": 30,
-        "seed": int(seed),
-        "guidance_scale": 7.0,
+        "seed": seed,
+        "guidance_scale": 6.0,
         "embedded_guidance_scale": 1.0,
-        "negative_prompt": "bad quality, blur",
+        "negative_prompt": "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion",
         "cfg_for": False,
     }
-    try:
-        # Load the image and move it to the correct device *before* inference
-        image = load_image(image=image_file.name)
-        # No need to manually move to device. SkyReelsVideoInfer should handle it.
-        kwargs["image"] = image
-    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}"), None  # Return None for predictor on error
-    save_dir = "./result/i2v"  # Consistent directory
+    assert image is not None, "please input image"
+    kwargs["image"] = load_image(image=image)
+    #global predictor
+    output = predictor.inference(kwargs)
+    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")
-    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 None for predictor
-    return video_out_file, predictor  # Return updated predictor
+    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
 
-def display_image(file):
-    if file is not None:
-        return Image.open(file.name)
-    else:
-        return None
-
-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()
+def create_gradio_interface():
+        with gr.Blocks() as demo:
+            with gr.Row():
+                image = gr.Image(label="Upload Image", type="filepath")
+                prompt = gr.Textbox(label="Input Prompt")
+                seed = gr.Number(label="Random Seed", value=-1)
+            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],
+                outputs=[output_video, output_params],
+            )
+        return demo
+    
+#init_predictor()
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    #import multiprocessing
+    #multiprocessing.freeze_support()
+    init_predictor()
+    demo = create_gradio_interface()
+    demo.launch()