import gradio as gr
import torch
from diffusers import DiffusionPipeline

# Define a dictionary of available models
models = {
    "QuantStack/Wan2.2-Fun-A14B-InP-GGUF"
    # Add more models as needed
}

# Load models into a dict for quick access
loaded_pipelines = {}

def load_model(model_id):
    if model_id not in loaded_pipelines:
        pipe = DiffusionPipeline.from_pretrained(model_id)
        pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
        loaded_pipelines[model_id] = pipe
    return loaded_pipelines[model_id]

def generate_video(model_name, prompt):
    pipe = load_model(models[model_name])
    # Call your model's video generation method
    result = pipe(prompt)
    # Adjust based on actual output (assuming it returns a video)
    video = result.videos[0]  # or result['videos'][0]
    video_path = "output.mp4"
    video.save(video_path)
    return video_path

with gr.Blocks() as demo:
    gr.Markdown("# Video Generation from Text")
    with gr.Row():
        model_choice = gr.Dropdown(choices=list(models.keys()), label="Select Model")
        prompt_input = gr.Textbox(label="Enter your prompt", lines=2, placeholder="A spaceship in space, neon colors")
    generate_button = gr.Button("Generate Video")
    video_output = gr.Video(label="Generated Video")
    
    generate_button.click(
        fn=generate_video,
        inputs=[model_choice, prompt_input],
        outputs=video_output
    )

demo.launch()