import spaces
import torch
import gradio as gr
from diffusers import CogVideoXPipeline
from diffusers.utils import export_to_video

BASE_MODEL = "zai-org/CogVideoX-2b"
LORA_REPO = "BinaryLight1011/Cinemotion"

# Carrega o pipeline base uma vez (fica em CPU até a inferência)
pipe = CogVideoXPipeline.from_pretrained(BASE_MODEL, torch_dtype=torch.float16)
pipe.load_lora_weights(LORA_REPO)
pipe.to("cuda")

pipe.vae.enable_slicing()
pipe.vae.enable_tiling()


@spaces.GPU(duration=120)
def generate(prompt, negative_prompt, num_frames, guidance_scale, steps, seed):
    generator = torch.Generator(device="cuda").manual_seed(int(seed))
    video = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt or None,
        num_videos_per_prompt=1,
        num_inference_steps=int(steps),
        num_frames=int(num_frames),
        guidance_scale=float(guidance_scale),
        generator=generator,
    ).frames[0]

    output_path = "output.mp4"
    export_to_video(video, output_path, fps=8)
    return output_path


with gr.Blocks(title="Cinemotion") as demo:
    gr.Markdown("# 🎬 Cinemotion — Text-to-Video (CogVideoX-2b + LoRA)")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=4, placeholder="Descreva a cena...")
            negative_prompt = gr.Textbox(label="Negative prompt (opcional)", lines=2)
            num_frames = gr.Slider(9, 49, value=49, step=8, label="Número de frames")
            guidance_scale = gr.Slider(1.0, 15.0, value=6.0, step=0.5, label="Guidance scale")
            steps = gr.Slider(10, 100, value=50, step=5, label="Inference steps")
            seed = gr.Number(value=42, label="Seed")
            btn = gr.Button("Gerar vídeo", variant="primary")

        with gr.Column():
            output_video = gr.Video(label="Resultado")

    btn.click(
        fn=generate,
        inputs=[prompt, negative_prompt, num_frames, guidance_scale, steps, seed],
        outputs=output_video,
    )

demo.launch()