import os; os.system('pip install --upgrade --no-deps spaces') # Reduce allocator fragmentation on ZeroGPU's partitioned GPU so the VAE-decode # allocation does not trip the CUDA caching allocator's NVML query path. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces import torch from diffusers import WanPipeline from diffusers.models.transformers.transformer_wan import WanTransformer3DModel from diffusers.utils.export_utils import export_to_video import gradio as gr import tempfile import numpy as np import random import gc from torchao.quantization import quantize_ from torchao.quantization import Float8DynamicActivationFloat8WeightConfig from torchao.quantization import Int8WeightOnlyConfig import aoti MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Wan 2.2 14B native resolution band (480p). Kept as a few clean presets. MULTIPLE_OF = 16 ASPECT_RATIOS = { "21:9 (976x416)": (976, 416), "16:9 (848x480)": (848, 480), "4:3 (768x576)": (768, 576), "1:1 (640x640)": (640, 640), "9:21 (624x1456)": (624, 1456), "9:21 (416x976)": (416, 976), "9:21 (288x656)": (288, 656), "9:16 (720x1280)": (720, 1280), "9:16 (480x848)": (480, 848), "9:16 (320x576)": (320, 576), "3:4 (576x768)": (576, 768), } DEFAULT_RATIO = "9:16 (480x848)" MAX_SEED = np.iinfo(np.int32).max FIXED_FPS = 16 MIN_FRAMES_MODEL = 8 MAX_FRAMES_MODEL = 240 MIN_DURATION = round(MIN_FRAMES_MODEL / FIXED_FPS, 1) MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1) LIGHTNING_LORA_REPO = "Kijai/WanVideo_comfy" LIGHTNING_LORA_FILE = "Lightx2v/lightx2v_T2V_14B_cfg_step_distill_v2_lora_rank128_bf16.safetensors" # Stage the two 14B MoE experts through host RAM one at a time: load -> # fuse the Lightning LoRA -> fp8-quantize (halves it) -> collect, and only # THEN load the second expert. With both experts resident in bf16 plus the # fuse copies, peak RAM sits at ZeroGPU's 104G startup cap and boots become # a coin flip ("Memory limit exceeded (104.0G)", no traceback). pipe = WanPipeline.from_pretrained(MODEL_ID, transformer=WanTransformer3DModel.from_pretrained(MODEL_ID, subfolder='transformer', torch_dtype=torch.bfloat16, device_map='cuda', low_cpu_mem_usage=True, ), transformer_2=None, torch_dtype=torch.bfloat16, ).to('cuda') quantize_(pipe.text_encoder, Int8WeightOnlyConfig()) pipe.load_lora_weights( LIGHTNING_LORA_REPO, weight_name=LIGHTNING_LORA_FILE, adapter_name="lightx2v" ) pipe.fuse_lora(adapter_names=["lightx2v"], lora_scale=3., components=["transformer"]) pipe.unload_lora_weights() quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig()) gc.collect() torch.cuda.empty_cache() pipe.register_modules( transformer_2=WanTransformer3DModel.from_pretrained(MODEL_ID, subfolder='transformer_2', torch_dtype=torch.bfloat16, device_map='cuda', low_cpu_mem_usage=True, ), ) pipe.load_lora_weights( LIGHTNING_LORA_REPO, weight_name=LIGHTNING_LORA_FILE, adapter_name="lightx2v_2", load_into_transformer_2=True ) pipe.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1., components=["transformer_2"]) pipe.unload_lora_weights() quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig()) gc.collect() torch.cuda.empty_cache() # The repeated transformer blocks are architecturally identical between the Wan 2.2 # A14B T2V and I2V experts (same hidden dim), so the I2V-compiled AOTI package works # here too and keeps generation fast enough for cold anonymous callers. spaces.aoti_load( module=pipe.transformer, repo_id='cbensimon/WanTransformer3DModel-sm120-cu130-raa', ) spaces.aoti_load( module=pipe.transformer_2, repo_id='cbensimon/WanTransformer3DModel-sm120-cu130-raa', ) # Tiled + sliced VAE decode keeps peak memory low on ZeroGPU. pipe.vae.enable_tiling() pipe.vae.enable_slicing() default_prompt_t2v = "" default_negative_prompt = "" def get_num_frames(duration_seconds: float): return 1 + int(np.clip( int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL, )) def get_duration(prompt, aspect_ratio, steps, negative_prompt, duration_seconds, GPU_time, guidance_scale, guidance_scale_2, seed, randomize_seed, progress=None): if GPU_time == 0: width, height = ASPECT_RATIOS.get(aspect_ratio, ASPECT_RATIOS[DEFAULT_RATIO]) BASE_FRAMES_HEIGHT_WIDTH = 81 * 832 * 624 BASE_STEP_DURATION = 11 frames = get_num_frames(duration_seconds) factor = frames * width * height / BASE_FRAMES_HEIGHT_WIDTH step_duration = BASE_STEP_DURATION * factor ** 1.5 estimate = steps * step_duration estimate = min(max(estimate, 10), 120) if guidance_scale > 1 or guidance_scale_2 > 1: estimate *= 2 # CFG runs two forward passes per step else: estimate = GPU_time/1.5 gr.Info(f"GPU time = {estimate * 1.5}s") return estimate @spaces.GPU(duration=get_duration) #@spaces.GPU(duration=120) def generate_video( prompt, aspect_ratio=DEFAULT_RATIO, steps=6, negative_prompt=default_negative_prompt, duration_seconds=MAX_DURATION, GPU_time=0, # Added this parameter guidance_scale=1.5, guidance_scale_2=1.5, seed=42, randomize_seed=True, progress=gr.Progress(track_tqdm=True), ): """ Generate a video from a text prompt using the Wan 2.2 14B T2V model with a 4-step Lightning LoRA, fp8 quantization and AoT-compiled transformer blocks. """ if not prompt or not prompt.strip(): raise gr.Error("Please enter a prompt.") width, height = ASPECT_RATIOS.get(aspect_ratio, ASPECT_RATIOS[DEFAULT_RATIO]) num_frames = get_num_frames(duration_seconds) current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) output_frames_list = pipe( prompt=prompt, negative_prompt=negative_prompt, height=height, width=width, num_frames=num_frames, guidance_scale=float(guidance_scale), guidance_scale_2=float(guidance_scale_2), 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 with gr.Blocks(theme=gr.Theme.from_hub("26A1/_")) as demo: with gr.Row(): with gr.Column(): prompt_input = gr.Textbox(label="Prompt", value=default_prompt_t2v, lines=3) aspect_ratio_input = gr.Dropdown(choices=list(ASPECT_RATIOS.keys()), value=DEFAULT_RATIO, label="Aspect ratio") duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=6, label="Duration (s)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.") GPU_time_input = gr.Slider(value=180,minimum=0,maximum=180,step=1,label="GPU time (s)",info="0:Auto") 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) steps_slider = gr.Slider(minimum=1, maximum=12, step=1, value=6, label="Inference Steps", info="Lightning-distilled: 4-8 steps is the sweet spot.") guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.25, value=1.25, label="Guidance Scale - high noise stage") guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.25, value=2.5, label="Guidance Scale 2 - low noise stage") generate_button = gr.Button("Generate Video", variant="primary") with gr.Column(): video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False) ui_inputs = [ prompt_input, aspect_ratio_input, steps_slider, negative_prompt_input, duration_seconds_input, GPU_time_input, # Fixed: added GPU_time_input, removed duplicate guidance_scale_input, guidance_scale_2_input, seed_input, randomize_seed_checkbox ] generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input], api_name="generate_video") if __name__ == "__main__": demo.queue().launch(ssr_mode=False, show_error=True)