# Reduce allocator fragmentation on ZeroGPU's partitioned GPU. Without this the # large contiguous VAE-decode allocation trips the CUDA caching allocator's NVML # query path ("NVML_SUCCESS == r INTERNAL ASSERT FAILED") after denoising finishes. import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces import torch from diffusers import AutoencoderKLWan, WanPipeline, WanImageToVideoPipeline, WanTransformer3DModel, UniPCMultistepScheduler from diffusers.utils import export_to_video from transformers import UMT5EncoderModel from torchao.quantization import quantize_, Int8WeightOnlyConfig import gradio as gr import tempfile import numpy as np from PIL import Image import random MODEL_ID = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers" vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32) # Tiled + sliced VAE decode keeps peak memory small so the full-resolution decode # no longer needs one huge allocation (the step that was crashing on ZeroGPU). vae.enable_tiling() vae.enable_slicing() # TI2V-5B uses the same transformer and text encoder for t2v and i2v; load each # once and share across both pipelines instead of paying for two copies (the # UMT5-XXL encoder alone is ~11GB bf16 — duplicating it doubled the ZeroGPU # pack size and per-task GPU attach time). text_encoder = UMT5EncoderModel.from_pretrained(MODEL_ID, subfolder="text_encoder", torch_dtype=torch.bfloat16) transformer = WanTransformer3DModel.from_pretrained(MODEL_ID, subfolder="transformer", torch_dtype=torch.bfloat16) # Int8 weight-only quant on the text encoder (same recipe as the 14B spaces): # halves its memory with no visible quality impact. quantize_(text_encoder, Int8WeightOnlyConfig()) # Initialize pipelines text_to_video_pipe = WanPipeline.from_pretrained(MODEL_ID, vae=vae, text_encoder=text_encoder, transformer=transformer, torch_dtype=torch.bfloat16) image_to_video_pipe = WanImageToVideoPipeline.from_pretrained(MODEL_ID, vae=vae, text_encoder=text_encoder, transformer=transformer, torch_dtype=torch.bfloat16) for pipe in [text_to_video_pipe, image_to_video_pipe]: pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0) pipe.to("cuda") # Constants MOD_VALUE = 32 DEFAULT_H_SLIDER_VALUE = 896 DEFAULT_W_SLIDER_VALUE = 896 NEW_FORMULA_MAX_AREA = 720 * 1024 SLIDER_MIN_H, SLIDER_MAX_H = 256, 1024 SLIDER_MIN_W, SLIDER_MAX_W = 256, 1024 MAX_SEED = np.iinfo(np.int32).max FIXED_FPS = 24 MIN_FRAMES_MODEL = 25 MAX_FRAMES_MODEL = 193 default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation" default_negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards, watermark, text, signature" def _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area, min_slider_h, max_slider_h, min_slider_w, max_slider_w, default_h, default_w): orig_w, orig_h = pil_image.size if orig_w <= 0 or orig_h <= 0: return default_h, default_w aspect_ratio = orig_h / orig_w calc_h = round(np.sqrt(calculation_max_area * aspect_ratio)) calc_w = round(np.sqrt(calculation_max_area / aspect_ratio)) calc_h = max(mod_val, (calc_h // mod_val) * mod_val) calc_w = max(mod_val, (calc_w // mod_val) * mod_val) new_h = int(np.clip(calc_h, min_slider_h, (max_slider_h // mod_val) * mod_val)) new_w = int(np.clip(calc_w, min_slider_w, (max_slider_w // mod_val) * mod_val)) return new_h, new_w def handle_image_upload_for_dims_wan(uploaded_pil_image, current_h_val, current_w_val): if uploaded_pil_image is None: return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE) try: new_h, new_w = _calculate_new_dimensions_wan( uploaded_pil_image, MOD_VALUE, NEW_FORMULA_MAX_AREA, SLIDER_MIN_H, SLIDER_MAX_H, SLIDER_MIN_W, SLIDER_MAX_W, DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE ) return gr.update(value=new_h), gr.update(value=new_w) except Exception as e: gr.Warning("Error attempting to calculate new dimensions") return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE) def get_duration(input_image, prompt, height, width, negative_prompt, duration_seconds, guidance_scale, steps, seed, randomize_seed, progress): # Scale the GPU request with the actual workload (frames x pixels x steps) # instead of flat 60/75/90 tiers. Calibrated against a measured default run # (4 steps, 49 frames, 896x896: ~35s wall): base overhead ~15s for model # attach + VAE decode + export, ~8s per step at the reference workload. # Over-requesting gets anonymous visitors rejected up front (ZeroGPU checks # the request against their remaining quota) and deprioritizes the queue. frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL) factor = (frames * height * width) / (49 * 896 * 896) if guidance_scale > 1: factor *= 2 # CFG runs two forward passes per step return int(np.clip(15 + int(steps) * 8 * factor, 25, 110)) @spaces.GPU(duration=get_duration) # NO end_image parameter here, deliberately. WanImageToVideoPipeline accepts # `last_image`, but Wan 2.2 5B sets `expand_timesteps=True`, and prepare_latents # takes the `if self.config.expand_timesteps:` branch that builds the video # condition from the first frame alone and discards `last_image`. Verified on # the live Space: two runs at a fixed seed, one with an end image and one # without, came back pixel-identical. Exposing the control would advertise a # feature the model does not have. def generate_video(input_image, prompt, height, width, negative_prompt=default_negative_prompt, duration_seconds=2, guidance_scale=0, steps=4, seed=44, randomize_seed=False, progress=gr.Progress(track_tqdm=True)): target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE) target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE) num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL) current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) if input_image is not None: resized_image = input_image.resize((target_w, target_h)) with torch.inference_mode(): output_frames_list = image_to_video_pipe( image=resized_image, prompt=prompt, negative_prompt=negative_prompt, height=target_h, width=target_w, num_frames=num_frames, guidance_scale=float(guidance_scale), num_inference_steps=int(steps), generator=torch.Generator(device="cuda").manual_seed(current_seed) ).frames[0] else: with torch.inference_mode(): output_frames_list = text_to_video_pipe( prompt=prompt, negative_prompt=negative_prompt, height=target_h, width=target_w, num_frames=num_frames, guidance_scale=float(guidance_scale), num_inference_steps=int(steps), generator=torch.Generator(device="cuda").manual_seed(current_seed) ).frames[0] with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile: video_path = tmpfile.name export_to_video(output_frames_list, video_path, fps=FIXED_FPS) return video_path, current_seed import importlib.metadata as _ilm for _pkg in ["torch","torchvision","transformers","diffusers","peft","accelerate","gradio","gradio_client","spaces","numpy","huggingface_hub","sentencepiece","safetensors","ftfy","imageio","imageio-ffmpeg","opencv-python","einops","hf_transfer"]: try: print(f"PINDUMP {_pkg}=={_ilm.version(_pkg)}", flush=True) except Exception: pass UPSAMPLER_THEME = gr.themes.Soft( primary_hue=gr.themes.colors.indigo, secondary_hue=gr.themes.colors.purple, neutral_hue=gr.themes.colors.slate, ).set( button_primary_background_fill="linear-gradient(135deg, #6366f1, #a855f7)", button_primary_background_fill_hover="linear-gradient(135deg, #5457e5, #9333ea)", button_primary_text_color="#ffffff", button_primary_border_color="*primary_500", ) UPSAMPLER_CSS = """ footer{display:none !important} .gradio-container{max-width:1000px !important; margin:0 auto !important} h1,h2,h3{font-family:system-ui,-apple-system,'Segoe UI',sans-serif} """ _header = """

Wan 2.2 5B

Fast text and image to video with Wan 2.2 5B.

""" _footer = """

Wan 2.2 5B is Alibaba's compact open video model, handling both text-to-video and image-to-video in a single checkpoint. A distilled few-step sampler generates smooth clips in seconds rather than minutes, with custom aspect ratios, which makes it a great free way to prototype shots before committing to a bigger model.

Maintained by Upsampler. Check out the free AI video generator, no sign-up required.

""" with gr.Blocks(theme=UPSAMPLER_THEME, css=UPSAMPLER_CSS) as demo: gr.HTML(_header) with gr.Row(): with gr.Column(): input_image_component = gr.Image(type="pil", label="Input Image (optional, auto-resized to target H/W)") prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v) duration_seconds_input = gr.Slider(minimum=round(MIN_FRAMES_MODEL/FIXED_FPS,1), maximum=round(MAX_FRAMES_MODEL/FIXED_FPS,1), step=0.1, value=2, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.") with gr.Accordion("Advanced Settings", open=False): negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3) seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True) randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True) with gr.Row(): height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label=f"Output Height (multiple of {MOD_VALUE})") width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label=f"Output Width (multiple of {MOD_VALUE})") steps_slider = gr.Slider(minimum=1, maximum=8, step=1, value=4, label="Inference Steps") guidance_scale_input = gr.Slider(minimum=0.0, maximum=5.0, step=0.01, value=0.0, label="Guidance Scale") generate_button = gr.Button("Generate Video", variant="primary") with gr.Column(): video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False) input_image_component.upload( fn=handle_image_upload_for_dims_wan, inputs=[input_image_component, height_input, width_input], outputs=[height_input, width_input] ) input_image_component.clear( fn=handle_image_upload_for_dims_wan, inputs=[input_image_component, height_input, width_input], outputs=[height_input, width_input] ) ui_inputs = [ input_image_component, prompt_input, height_input, width_input, negative_prompt_input, duration_seconds_input, guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox ] generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input], api_name="generate_video") gr.HTML(_footer) if __name__ == "__main__": demo.queue().launch(ssr_mode=False)