AiCoderv2's picture
Update app.py
971a591 verified
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()