import gradio as gr
import torch
import gc
import os
import time
from diffusers.utils import load_image
from skyreels_v2_infer.modules import download_model
from skyreels_v2_infer.pipelines import Text2VideoPipeline

# Descargar y preparar modelo
MODEL_ID = "Skywork/SkyReels-V2-T2V-14B-540P"
model_path = download_model(MODEL_ID)
pipe = Text2VideoPipeline(model_path=model_path, dit_path=model_path)

def generate_video(prompt):
    print(f"Generando video para: {prompt}")
    height, width = 544, 960
    seed = int(time.time()) % 4294967294
    generator = torch.Generator(device="cuda").manual_seed(seed)

    kwargs = {
        "prompt": prompt,
        "negative_prompt": "ugly, blurry, bad quality, watermark",
        "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]

    output_path = f"/tmp/video_{seed}.mp4"
    import imageio
    imageio.mimwrite(output_path, video_frames, fps=24, quality=8)

    gc.collect()
    torch.cuda.empty_cache()

    return output_path

with gr.Blocks() as demo:
    gr.Markdown("# SkyReels V2")
    with gr.Row():
        prompt_input = gr.Textbox(label="Prompt para generar video")
    with gr.Row():
        generate_btn = gr.Button("Generar video")
    with gr.Row():
        output_video = gr.Video(label="Video generado")

    generate_btn.click(fn=generate_video, inputs=prompt_input, outputs=output_video)

demo.launch()