import gradio as gr
from diffusers import DiffusionPipeline
import torch
import imageio
import numpy as np

# Detectar dispositivo
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

print(f"Usando dispositivo: {device} - dtype: {dtype}")

# Cargar modelo
pipe = DiffusionPipeline.from_pretrained(
    "cerspense/zeroscope_v2_576w",
    torch_dtype=dtype
)

# Mover al dispositivo correcto
if device == "cuda":
    pipe.to("cuda")
else:
    pipe.to("cpu")

def generate_video(prompt: str):

    if not prompt.strip():
        return None, "El prompt no puede estar vacío."

    num_frames = 16
    height = 320
    width = 576
    steps = 15

    try:
        result = pipe(
            prompt,
            num_inference_steps=steps,
            num_frames=num_frames,
            height=height,
            width=width
        )

        frames = result.frames[0] if isinstance(result.frames, list) else result.frames

        # Convertir PIL a numpy
        if hasattr(frames[0], 'size'):
            frames = [np.array(frame) for frame in frames]

        output_path = "/tmp/video.mp4"

        imageio.mimsave(output_path, frames, fps=8)

        return output_path, "✅ Video generado correctamente."

    except Exception as e:
        return None, f"❌ Error: {str(e)}"

# Interfaz
with gr.Blocks(title="Zeroscope AI Video Generator") as demo:

    gr.Markdown("# 🎬 Generador de Videos IA")
    gr.Markdown("Genera videos cortos desde texto usando Zeroscope.")

    with gr.Row():

        with gr.Column():
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Ej: un gato cyberpunk caminando bajo lluvia neon",
                lines=3
            )

            generate_btn = gr.Button(
                "🎥 Generar Video",
                variant="primary"
            )

        with gr.Column():
            video_output = gr.Video(label="Resultado")
            status_text = gr.Textbox(label="Estado")

    generate_btn.click(
        fn=generate_video,
        inputs=prompt_input,
        outputs=[video_output, status_text]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)