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