Spaces:
Running
Running
File size: 2,912 Bytes
b45529f c8eea54 5562446 b45529f 5562446 c8eea54 b45529f 5562446 b45529f ad7a082 b45529f ad7a082 b45529f c8eea54 5562446 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import gradio as gr
import torch
from diffusers.utils import export_to_video
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
import os
from uuid import uuid4
# Check if CUDA is available and set device accordingly
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}") # Print device information
# Load model on startup
try:
print("Loading model...")
model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
scheduler = UniPCMultistepScheduler(
prediction_type='flow_prediction',
use_flow_sigmas=True,
num_train_timesteps=1000,
flow_shift=5.0
)
pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
pipe.scheduler = scheduler
pipe.to(device) # Move model to CUDA if available, otherwise CPU
print("Model loaded successfully.")
except Exception as e:
print(f"Error loading model: {e}")
raise e
# Define the generation function
def generate_video(prompt, negative_prompt="", height=720, width=1280, num_frames=81, guidance_scale=5.0):
try:
print(f"Generating video for prompt: {prompt}")
output = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
height=height,
width=width,
num_frames=num_frames,
guidance_scale=guidance_scale,
).frames[0]
output_filename = f"{uuid4()}.mp4"
output_path = os.path.join("outputs", output_filename)
os.makedirs("outputs", exist_ok=True)
export_to_video(output, output_path, fps=16)
print(f"Video generated successfully: {output_path}")
return output_path # Gradio returns this as a downloadable file/video
except Exception as e:
print(f"Error generating video: {e}")
return None # Return None in case of error
# Gradio Interface with API support
iface = gr.Interface(
fn=generate_video,
inputs=[
gr.Textbox(label="Prompt"),
gr.Textbox(label="Negative Prompt", value=""),
gr.Number(label="Height", value=720),
gr.Number(label="Width", value=1280),
gr.Number(label="Number of Frames", value=81),
gr.Number(label="Guidance Scale", value=5.0)
],
outputs=gr.File(label="Generated Video"),
title="Wan2.1 Video Generator",
description="Generate realistic videos from text prompts using the Wan2.1 T2V model.",
api=True # This enables the API
)
try:
print("Launching Gradio interface...")
iface.launch(share=True) # `share=True` will allow others to access your app via a public link
print("Gradio interface launched successfully.")
except Exception as e:
print(f"Error launching Gradio interface: {e}")
|