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 = 1.7 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)