Update app.py

app.py

@@ -1,26 +1,59 @@
 import gradio as gr
-import subprocess
+import torch
 import os
-import random
 import time
+import imageio
+from diffusers.utils import load_image
+from skyreels_v2_infer.modules import download_model
+from skyreels_v2_infer.pipelines import Text2VideoPipeline
 
-def generate(prompt):
-    seed = random.randint(0, 999999)
-    command = f"python generate_video.py --prompt \"{prompt}\" --resolution 540P --fps 24 --outdir=outputs --seed {seed}"
-    os.makedirs("outputs", exist_ok=True)
-    subprocess.run(command, shell=True)
-    
-    # Buscar el último vídeo generado
-    files = sorted(os.listdir("result/outputs"), key=lambda x: os.path.getmtime(os.path.join("result/outputs", x)))
-    if files:
-        video_path = os.path.join("result/outputs", files[-1])
-        return video_path
-    else:
-        return "No video generated"
-
-iface = gr.Interface(fn=generate, 
-                     inputs=gr.Textbox(label="Prompt para generar video"), 
-                     outputs=gr.Video(label="Video generado"))
-
-iface.launch()
+def generate_video(prompt):
+    model_id = "Skywork/SkyReels-V2-T2V-14B-540P"
+    model_path = download_model(model_id)
 
+    pipe = Text2VideoPipeline(
+        model_path=model_path,
+        dit_path=model_path,
+        use_usp=False,
+        offload=True,
+    )
+
+    height = 544
+    width = 960
+    seed = int(time.time())
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    kwargs = {
+        "prompt": prompt,
+        "negative_prompt": "static, blurred, low quality",
+        "num_frames": 97,
+        "num_inference_steps": 30,
+        "guidance_scale": 6.0,
+        "shift": 8.0,
+        "generator": generator,
+        "height": height,
+        "width": width,
+    }
+
+    with torch.cuda.amp.autocast(dtype=pipe.transformer.dtype), torch.no_grad():
+        video_frames = pipe(**kwargs)[0]
+
+    timestamp = time.strftime("%Y%m%d_%H%M%S")
+    filename = f"output_{timestamp}.mp4"
+    os.makedirs("video_out", exist_ok=True)
+    path = os.path.join("video_out", filename)
+    imageio.mimwrite(path, video_frames, fps=24, quality=8)
+    return path
+
+with gr.Blocks() as demo:
+    gr.Markdown("# SkyReels V2")
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Prompt para generar video", placeholder="Describe la escena que quieres ver")
+    with gr.Row():
+        run_button = gr.Button("Generar video")
+    with gr.Row():
+        output_video = gr.Video(label="Video generado")
+
+    run_button.click(fn=generate_video, inputs=prompt_input, outputs=output_video)
+
+demo.launch()