app.py

import os; os.system('pip install --upgrade --no-deps spaces')
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

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:21 (416x976)"

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)

MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
LIGHTNING_LORA_REPO = "Kijai/WanVideo_comfy"
LORA_FILE = "LoRAs/Wan22-Lightning/Wan22_A14B_T2V_LOW_Lightning_4steps_lora_250928_rank64_fp16.safetensors"
LORA_FILE_2 = "LoRAs/Wan22-Lightning/Wan22_A14B_T2V_LOW_Lightning_4steps_lora_250928_rank64_fp16.safetensors"
lora_scale = 2.0
lora_scale_2 = 1.0

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=LORA_FILE, adapter_name="lora_adapter"
)
pipe.fuse_lora(adapter_names=["lora_adapter"], lora_scale=lora_scale, 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=LORA_FILE_2,
    adapter_name="lora_adapter_2", load_into_transformer_2=True
)
pipe.fuse_lora(adapter_names=["lora_adapter_2"], lora_scale=lora_scale_2, components=["transformer_2"])
pipe.unload_lora_weights()
quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
gc.collect()
torch.cuda.empty_cache()

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',
)

pipe.vae.enable_tiling()
pipe.vae.enable_slicing()

def get_num_frames(duration_seconds: float):
    raw_frames = int(round(duration_seconds * FIXED_FPS))
    raw_frames = np.clip(raw_frames, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
    
    raw_frames_adjusted = raw_frames - 1
    remainder = raw_frames_adjusted % 4
    
    if remainder == 0:
        adjusted_frames = raw_frames
    elif remainder <= 2:
        adjusted_frames = raw_frames - remainder
    else:
        adjusted_frames = raw_frames + (4 - remainder)
    
    adjusted_frames = max(MIN_FRAMES_MODEL, min(adjusted_frames, MAX_FRAMES_MODEL))
    
    if (adjusted_frames - 1) % 4 != 0:
        adjusted_frames = ((adjusted_frames - 1) // 4) * 4 + 1
    
    return adjusted_frames

def get_duration(prompt, aspect_ratio, steps, negative_prompt, duration_seconds, GPU_time,
                 guidance_scale, guidance_scale_2, seed, randomize_seed, progress=None):

    GPU_time = float(GPU_time)
    
    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(float(duration_seconds))
        factor = frames * width * height / BASE_FRAMES_HEIGHT_WIDTH
        step_duration = BASE_STEP_DURATION * factor ** 1.5
        estimate = int(steps) * step_duration
        estimate = min(max(estimate, 10), 120)
    
        if float(guidance_scale) > 1 or float(guidance_scale_2) > 1:
            estimate *= 2          
    else:      
        estimate = GPU_time / 1.5

    gr.Info(f"GPU time = {estimate * 1.5}s")  
    return estimate

@spaces.GPU(duration=get_duration)

def generate_video(
    prompt,
    aspect_ratio,
    steps,
    negative_prompt,
    duration_seconds,
    GPU_time,
    guidance_scale,
    guidance_scale_2,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    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(float(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]

    video_filename = f"{current_seed}_{guidance_scale}_{guidance_scale_2}.mp4"
    video_path = os.path.join(tempfile.gettempdir(), video_filename)
    export_to_video(output_frames_list, video_path, fps=FIXED_FPS, quality=7)

    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="", lines=2)
            duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=4.4, label="Duration (s)")
            GPU_time_input = gr.Slider(value=90.0,minimum=0.0,maximum=300.0,step=1.0,label="GPU time (s)")
            randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
            seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, interactive=True)           
            generate_button = gr.Button("Generate Video", variant="primary")

            with gr.Accordion("Advanced Settings", open=True):
                negative_prompt_input = gr.Textbox(label="Negative Prompt", value="", lines=2)                
                aspect_ratio_input = gr.Dropdown(choices=list(ASPECT_RATIOS.keys()), value=DEFAULT_RATIO, label="Aspect ratio")
             
                steps_slider = gr.Slider(minimum=1, maximum=12, step=1, value=6, label="Inference Steps")
                guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=1.5, label="Guidance Scale - high noise stage")
                guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=1.5, label="Guidance Scale 2 - low noise stage")
                
        with gr.Column():
            video_output = gr.Video(label="Generated Video", autoplay=False, interactive=False)

    ui_inputs = [
        prompt_input, aspect_ratio_input, steps_slider, negative_prompt_input, 
        duration_seconds_input, GPU_time_input, 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)