import gc import inspect import gradio as gr import torch from diffusers import DiffusionPipeline from diffusers.utils import export_to_video from PIL import Image # --------------------------------------------------------------------------- # Preset models # --------------------------------------------------------------------------- PRESET_MODELS = [ ("Wan 2.2 T2V 14B", "Wan-AI/Wan2.2-T2V-A14B-Diffusers"), ("Wan 2.2 I2V 14B", "Wan-AI/Wan2.2-I2V-A14B-Diffusers"), ("Wan 2.2 TI2V 5B", "Wan-AI/Wan2.2-TI2V-5B-Diffusers"), ("Wan 2.2 Animate 14B", "Wan-AI/Wan2.2-Animate-14B-Diffusers"), ("LTX Video 0.9.8 13B Distilled", "Lightricks/LTX-Video-0.9.8-13B-distilled"), ("LTX Video 0.9.7 Distilled", "Lightricks/LTX-Video-0.9.7-distilled"), ("HunyuanVideo 1.5 T2V 720p", "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v"), ("HunyuanVideo 1.5 I2V 720p", "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v"), ("HunyuanVideo T2V", "hunyuanvideo-community/HunyuanVideo"), ("Wan 2.1 T2V 1.3B (small/fast)", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"), ] PRESET_CHOICES = [(label, model_id) for label, model_id in PRESET_MODELS] # --------------------------------------------------------------------------- # Global pipeline state # --------------------------------------------------------------------------- current_pipeline = None current_model_id = None # --------------------------------------------------------------------------- # Model loading # --------------------------------------------------------------------------- def load_model(model_id: str, dtype_str: str, cpu_offload: bool) -> str: global current_pipeline, current_model_id if not model_id: return "Please select or enter a model ID." if model_id == current_model_id and current_pipeline is not None: return f"Already loaded: {model_id}" # Unload previous model if current_pipeline is not None: del current_pipeline current_pipeline = None current_model_id = None gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() dtype_map = { "bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32, } torch_dtype = dtype_map[dtype_str] try: current_pipeline = DiffusionPipeline.from_pretrained( model_id, torch_dtype=torch_dtype ) if cpu_offload: current_pipeline.enable_model_cpu_offload() else: current_pipeline.to("cuda") current_model_id = model_id pipe_class = current_pipeline.__class__.__name__ return ( f"Loaded: {model_id}\n" f"Pipeline: {pipe_class}\n" f"Dtype: {dtype_str}\n" f"CPU offload: {cpu_offload}" ) except Exception as e: current_pipeline = None current_model_id = None return f"Failed to load: {e}" # --------------------------------------------------------------------------- # Video generation # --------------------------------------------------------------------------- def generate( prompt: str, image, negative_prompt: str, num_steps: int, guidance_scale: float, num_frames: int, height: int, width: int, seed: int, fps: int, ) -> str: if current_pipeline is None: raise gr.Error("Load a model first!") generator = ( torch.Generator(device="cpu").manual_seed(seed) if seed >= 0 else None ) # Inspect the pipeline's __call__ signature so we only pass supported args sig = inspect.signature(current_pipeline.__call__) params = { name for name, p in sig.parameters.items() if p.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) } kwargs: dict = {} if "prompt" in params: kwargs["prompt"] = prompt if negative_prompt and "negative_prompt" in params: kwargs["negative_prompt"] = negative_prompt if "num_inference_steps" in params: kwargs["num_inference_steps"] = num_steps if "guidance_scale" in params: kwargs["guidance_scale"] = guidance_scale if "num_frames" in params: kwargs["num_frames"] = num_frames if "height" in params: kwargs["height"] = height if "width" in params: kwargs["width"] = width if "generator" in params: kwargs["generator"] = generator # Image input for I2V models if image is not None: pil_image = Image.open(image) if isinstance(image, str) else image if "image" in params: kwargs["image"] = pil_image output = current_pipeline(**kwargs) frames = output.frames[0] video_path = export_to_video(frames, fps=fps) return video_path # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- with gr.Blocks(title="Video Gen Playground") as demo: gr.Markdown( "# Video Gen Playground\n" "Load any diffusers video model from the Hub, configure settings, generate.\n\n" "*GPU hardware is configured in Space Settings.*" ) # --- Model loading section --- with gr.Row(): with gr.Column(scale=2): model_dropdown = gr.Dropdown( choices=PRESET_CHOICES, label="Model", allow_custom_value=True, info="Pick a preset or type any HF model ID", ) with gr.Column(scale=1): dtype_dropdown = gr.Dropdown( choices=["bfloat16", "float16", "float32"], value="bfloat16", label="Dtype", ) with gr.Column(scale=1): cpu_offload_checkbox = gr.Checkbox( value=True, label="Enable model CPU offload", ) with gr.Column(scale=1): load_btn = gr.Button("Load Model") status_box = gr.Textbox(label="Status", interactive=False, lines=4) load_btn.click( fn=load_model, inputs=[model_dropdown, dtype_dropdown, cpu_offload_checkbox], outputs=[status_box], ) gr.Markdown("---") # --- Inputs + settings + output --- with gr.Row(): # Left column: inputs with gr.Column(scale=1): prompt_box = gr.Textbox( label="Prompt", lines=3, placeholder="A cinematic shot of a cat riding a skateboard...", ) image_input = gr.Image( label="Image (optional, for I2V models)", type="pil", ) negative_prompt_box = gr.Textbox( label="Negative prompt", lines=2, placeholder="blurry, low quality, distorted", ) # Middle column: generation settings with gr.Column(scale=1): num_steps_slider = gr.Slider( minimum=1, maximum=100, value=30, step=1, label="Inference steps", ) guidance_slider = gr.Slider( minimum=0.0, maximum=20.0, value=6.0, step=0.5, label="Guidance scale", ) num_frames_slider = gr.Slider( minimum=1, maximum=257, value=49, step=4, label="Number of frames (4k+1 recommended)", ) height_slider = gr.Slider( minimum=256, maximum=1080, value=480, step=8, label="Height", ) width_slider = gr.Slider( minimum=256, maximum=1920, value=832, step=8, label="Width", ) seed_input = gr.Number( value=-1, label="Seed (-1 = random)", precision=0, ) fps_slider = gr.Slider( minimum=1, maximum=60, value=24, step=1, label="Output FPS", ) generate_btn = gr.Button("Generate", variant="primary") # Right column: output with gr.Column(scale=1): video_output = gr.Video(label="Generated video") generate_btn.click( fn=generate, inputs=[ prompt_box, image_input, negative_prompt_box, num_steps_slider, guidance_slider, num_frames_slider, height_slider, width_slider, seed_input, fps_slider, ], outputs=[video_output], ) # --------------------------------------------------------------------------- demo.launch()