app.py

import gradio as gr
from diffusers import CogVideoXTransformer3DModel, DiffusionPipeline
from diffusers.utils import export_to_video
import torch
import tempfile
import os
import spaces
# Lista de modelos disponibles
TRANSFORMER_MODELS = [
    "sayakpaul/pika-dissolve-v0",
    "finetrainers/crush-smol-v0", 
    "finetrainers/3dgs-v0",
    "finetrainers/cakeify-v0"
]

@spaces.GPU
def generate_video(transformer_model, prompt, negative_prompt):
    # Cargar el modelo del transformer seleccionado
    transformer = CogVideoXTransformer3DModel.from_pretrained(
        transformer_model, 
        torch_dtype=torch.bfloat16
    )
    
    # Inicializar el pipeline
    pipeline = DiffusionPipeline.from_pretrained(
        "THUDM/CogVideoX-5b", 
        transformer=transformer, 
        torch_dtype=torch.bfloat16
    ).to("cuda")
    
    # Generar el video
    video_frames = pipeline(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_frames=81,
        height=512, 
        width=768,
        num_inference_steps=50
    ).frames[0]
    
    # Guardar el video en un archivo temporal
    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_file:
        export_to_video(video_frames, tmp_file.name, fps=25)
        return tmp_file.name

# Crear la interfaz de Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Generador de Videos con CogVideoX")
    
    with gr.Row():
        with gr.Column():
            # Entradas
            model_dropdown = gr.Dropdown(
                choices=TRANSFORMER_MODELS,
                value=TRANSFORMER_MODELS[0],
                label="Modelo Transformer"
            )
            prompt_input = gr.Textbox(
                lines=5,
                label="Prompt",
                placeholder="Describe el video que quieres generar..."
            )
            negative_prompt_input = gr.Textbox(
                lines=2,
                label="Prompt Negativo",
                value="inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs"
            )
            generate_btn = gr.Button("Generar Video")
        
        with gr.Column():
            # Salida
            video_output = gr.Video(label="Video Generado")
    
    # Conectar la función
    generate_btn.click(
        fn=generate_video,
        inputs=[model_dropdown, prompt_input, negative_prompt_input],
        outputs=video_output
    )

# Lanzar la aplicación
demo.launch()