import os # Video DiTs have transient allocation spikes; expandable segments avoid NVML allocator asserts. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # noqa: E402 (must be imported before torch / diffusers) import math # noqa: E402 import random # noqa: E402 import tempfile # noqa: E402 import time # noqa: E402 import gradio as gr # noqa: E402 import torch # noqa: E402 from diffusers import AutoencoderKLWan, WanImageToVideoPipeline, WanPipeline # noqa: E402 from diffusers.utils import export_to_video # noqa: E402 from PIL import Image, ImageOps # noqa: E402 # -------------------------------------------------------------------------------------- # Model # -------------------------------------------------------------------------------------- MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" FPS = 24 MOD = 32 # VAE spatial stride (16) x transformer patch size (2): sizes must be multiples of 32 MIN_SIDE = 256 # The VAE runs in fp32 and the transformer / text encoder in bf16, as in the model card. print("Loading Wan2.2 VAE on CPU...", flush=True) vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32) print("Loading Wan2.2 text-to-video pipeline on CPU...", flush=True) t2v_pipe = WanPipeline.from_pretrained(MODEL_ID, vae=vae, torch_dtype=torch.bfloat16) # TI2V-5B needs per-token timesteps ("expand_timesteps") for both modes. if not t2v_pipe.config.get("expand_timesteps", False): t2v_pipe.register_to_config(expand_timesteps=True) # Image-to-video uses the same weights through the I2V pipeline (first-frame conditioning). print("Preparing image-to-video pipeline...", flush=True) i2v_pipe = WanImageToVideoPipeline(**t2v_pipe.components, expand_timesteps=True) print("Model components loaded on CPU; waiting for a GPU request.", flush=True) # -------------------------------------------------------------------------------------- # Settings # -------------------------------------------------------------------------------------- # Official negative prompt from the model card. DEFAULT_NEGATIVE_PROMPT = ( "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，" "JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，" "手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" ) PRESETS = { "832×480 · fast (default)": (832, 480), "480×832 · portrait, fast": (480, 832), "1280×704 · 720p": (1280, 704), "704×1280 · 720p portrait": (704, 1280), } DEFAULT_RES = next(iter(PRESETS)) # --- GPU-time budgeting ----------------------------------------------------------------- # ZeroGPU quotas are small (free accounts get a few minutes per day), so the request is # sized to a realistic runtime instead of a fixed worst case. These constants are a # first-order FLOPs model, NOT measurements: calibrate them from the "GPU time" line shown # under each video (and `hf spaces logs`), then adjust EFFECTIVE_FLOPS / OVERHEAD_S. MAX_GPU_SECONDS = 180 # hard ceiling requested from ZeroGPU for one call SAFETY = 1.2 # reserved duration = estimate * SAFETY MIN_STEPS = 20 LIN_FLOPS_PER_TOKEN = 1.0e10 # ~2 * 5B params per token per forward ATTN_FLOPS_COEF = 4 * 30 * 3072 # 4 * layers * dim, multiplied by tokens**2 EFFECTIVE_FLOPS = 80e12 # conservative sustained throughput on half a GPU OVERHEAD_S = 25.0 # text encoding, weight streaming on cold start, export DECODE_PX_PER_S = 6e6 # VAE decode throughput (pixels of video per second) def resolve_size(image, resolution): """Output (width, height). With an input image, keep its aspect ratio at the preset's area.""" w0, h0 = PRESETS.get(resolution, PRESETS[DEFAULT_RES]) if image is None: return w0, h0 aspect = image.height / image.width area = w0 * h0 h = max(MIN_SIDE, round(math.sqrt(area * aspect)) // MOD * MOD) w = max(MIN_SIDE, round(math.sqrt(area / aspect)) // MOD * MOD) return w, h def estimate_seconds(w, h, frames, steps, guidance): tokens = (w // MOD) * (h // MOD) * ((frames - 1) // 4 + 1) per_forward = LIN_FLOPS_PER_TOKEN * tokens + ATTN_FLOPS_COEF * tokens**2 cfg = 2 if guidance > 1.0 else 1 denoise = steps * cfg * per_forward / EFFECTIVE_FLOPS decode = frames * w * h / DECODE_PX_PER_S return OVERHEAD_S + denoise + decode def plan_job(w, h, seconds, steps, guidance): """Shrink steps, then duration, until the estimated GPU time fits the ZeroGPU budget.""" frames = int(seconds) * FPS + 1 # always 4k + 1, as the VAE requires steps = int(steps) budget = MAX_GPU_SECONDS / SAFETY while estimate_seconds(w, h, frames, steps, guidance) > budget and steps > MIN_STEPS: steps = max(MIN_STEPS, steps - 2) while estimate_seconds(w, h, frames, steps, guidance) > budget and frames > FPS + 1: frames -= FPS return frames, steps, estimate_seconds(w, h, frames, steps, guidance) def _gpu_duration(prompt, image=None, resolution=DEFAULT_RES, seconds=2, steps=30, guidance_scale=5.0, *args, **kwargs): # Gradio also passes the remaining inputs and `progress`, hence the *args / **kwargs. try: w, h = resolve_size(image, resolution) _, _, est = plan_job(w, h, seconds, steps, guidance_scale) return int(min(MAX_GPU_SECONDS, math.ceil(est * SAFETY))) except Exception: return 120 # -------------------------------------------------------------------------------------- # Inference # -------------------------------------------------------------------------------------- @spaces.GPU(duration=_gpu_duration) def generate( prompt: str, image: Image.Image | None = None, resolution: str = DEFAULT_RES, seconds: int = 2, steps: int = 30, guidance_scale: float = 5.0, negative_prompt: str = DEFAULT_NEGATIVE_PROMPT, seed: int = 0, randomize_seed: bool = True, enable_tiling: bool = True, progress=gr.Progress(track_tqdm=True), ): """Generate a 24 fps video with Wan2.2 TI2V-5B. Text-to-video when no image is given; image-to-video when an image is provided (the image becomes the first frame and its aspect ratio is kept). Args: prompt: Description of the video. Detailed prompts (subject, motion, camera, lighting, style) work best. image: Optional first-frame image for image-to-video. resolution: One of the preset sizes. With an image, only the preset's pixel area is used. seconds: Video length in seconds (24 fps). Long or large videos may be shortened to fit the GPU time limit. steps: Number of denoising steps. May be lowered to fit the GPU time limit. guidance_scale: Classifier-free guidance scale. negative_prompt: What to avoid. The default is the official Wan negative prompt. seed: Random seed. randomize_seed: Pick a new random seed for each run (overrides `seed`). Returns: The generated MP4 file, the seed used, and a short status line. """ start = time.perf_counter() t2v_pipe.to("cuda") if enable_tiling: vae.enable_tiling( tile_sample_min_height=192, tile_sample_min_width=192, tile_sample_stride_height=160, tile_sample_stride_width=160, ) else: vae.disable_tiling() if randomize_seed: seed = random.randint(0, 2**31 - 1) seed = int(seed) w, h = resolve_size(image, resolution) frames, n_steps, _ = plan_job(w, h, seconds, steps, guidance_scale) generator = torch.Generator(device="cuda").manual_seed(seed) kwargs = dict( prompt=prompt, negative_prompt=negative_prompt or None, height=h, width=w, num_frames=frames, num_inference_steps=n_steps, guidance_scale=float(guidance_scale), generator=generator, ) if image is not None: first_frame = ImageOps.exif_transpose(image).convert("RGB").resize((w, h), Image.LANCZOS) output = i2v_pipe(image=first_frame, **kwargs) else: output = t2v_pipe(**kwargs) with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f: path = f.name export_to_video(output.frames[0], path, fps=FPS) elapsed = time.perf_counter() - start mode = "image-to-video" if image is not None else "text-to-video" note = "" if n_steps != int(steps) or frames != int(seconds) * FPS + 1: note = " · settings reduced to fit the GPU time limit" info = ( f"{mode} · {w}×{h} · {frames} frames ({frames / FPS:.1f}s) · {n_steps} steps · " f"seed {seed} · GPU time {elapsed:.0f}s{note}" ) print(info, flush=True) return path, seed, info def check_inputs(prompt, image): if not (prompt and prompt.strip()): raise gr.Error("Please enter a prompt describing the video.") # -------------------------------------------------------------------------------------- # UI # -------------------------------------------------------------------------------------- HERE = os.path.dirname(os.path.abspath(__file__)) CAT_BOXING_PROMPT = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." CAT_SURF_PROMPT = ( "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline " "gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring " "crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally " "relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's " "intricate details and the refreshing atmosphere of the seaside." ) CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks(title="Wan2.2 TI2V-5B") as demo: with gr.Column(elem_id="col-container"): gr.Markdown( "# 🎬 Wan2.2 TI2V-5B\n" "Text-to-video **and** image-to-video at up to 720p / 24 fps with " "[Wan2.2-TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers). " "Leave the image empty for text-to-video, or upload one to use as the first frame. " "[Paper](https://arxiv.org/abs/2503.20314) · [GitHub](https://github.com/Wan-Video/Wan2.2)\n\n" "_Runs on ZeroGPU, which has small daily quotas, so the defaults are kept short " "(832×480, 2 s). Long or 720p runs are automatically trimmed to fit the GPU time limit._" ) with gr.Row(): with gr.Column(): prompt = gr.Textbox( label="Prompt", lines=4, placeholder="Describe the subject, the motion, the camera and the lighting…", ) image = gr.Image(label="First frame (optional, for image-to-video)", type="pil", sources=["upload", "clipboard"]) run = gr.Button("Generate", variant="primary") with gr.Accordion("Advanced settings", open=False): resolution = gr.Dropdown(list(PRESETS), value=DEFAULT_RES, label="Resolution") seconds = gr.Slider(1, 5, value=2, step=1, label="Duration (seconds, 24 fps)") steps = gr.Slider(10, 50, value=30, step=1, label="Inference steps") guidance_scale = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Guidance scale") negative_prompt = gr.Textbox(label="Negative prompt", value=DEFAULT_NEGATIVE_PROMPT, lines=3) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) seed = gr.Number(label="Seed", value=0, precision=0) enable_tiling = gr.Checkbox(label="Enable VAE Spatial Tiling (16 tiles: 192px / 160px stride)", value=True) with gr.Column(): video = gr.Video(label="Result", autoplay=True, loop=True) info = gr.Markdown() gr.Examples( examples=[ [CAT_BOXING_PROMPT, None], [CAT_SURF_PROMPT, os.path.join(HERE, "examples", "i2v_input.JPG")], ], inputs=[prompt, image], outputs=[video, seed, info], fn=generate, cache_examples=True, cache_mode="lazy", ) all_inputs = [prompt, image, resolution, seconds, steps, guidance_scale, negative_prompt, seed, randomize_seed, enable_tiling] gr.on( triggers=[run.click, prompt.submit], fn=check_inputs, inputs=[prompt, image], outputs=None, queue=False, api_name=False, ).success( fn=generate, inputs=all_inputs, outputs=[video, seed, info], api_name="generate", ) demo.queue(max_size=20) if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)