app.py

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)