Spaces:
Running
Running
| 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}") | |