Spaces:
Running on Zero
Running on Zero
| 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 | |
| # -------------------------------------------------------------------------------------- | |
| 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) | |