Update app.py
Browse files
app.py
CHANGED
@@ -1,15 +1,11 @@
|
|
|
|
1 |
import os
|
2 |
-
# Set environment variables before any imports to suppress inductor warnings
|
3 |
-
os.environ["TORCHINDUCTOR_CUDA_GRAPHS"] = "0"
|
4 |
-
os.environ["TORCHINDUCTOR_MAX_AUTOTUNE_GEMM"] = "0"
|
5 |
-
|
6 |
-
# Install dependencies as specified
|
7 |
os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
|
8 |
|
|
|
9 |
import spaces
|
10 |
import torch
|
11 |
from diffusers import WanPipeline, AutoencoderKLWan
|
12 |
-
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
|
13 |
from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
|
14 |
from diffusers.utils.export_utils import export_to_video
|
15 |
import gradio as gr
|
@@ -18,256 +14,78 @@ import numpy as np
|
|
18 |
from PIL import Image
|
19 |
import random
|
20 |
import gc
|
|
|
|
|
21 |
|
22 |
-
|
23 |
-
try:
|
24 |
-
from optimization import optimize_pipeline_
|
25 |
-
except ImportError:
|
26 |
-
def optimize_pipeline_(pipe, **kwargs):
|
27 |
-
pass # No-op if optimization is not available
|
28 |
|
29 |
-
# Model configurations
|
30 |
-
T2V_MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
|
31 |
-
I2V_MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
|
32 |
LANDSCAPE_WIDTH = 832
|
33 |
LANDSCAPE_HEIGHT = 480
|
34 |
MAX_SEED = np.iinfo(np.int32).max
|
|
|
35 |
FIXED_FPS = 16
|
36 |
MIN_FRAMES_MODEL = 8
|
37 |
MAX_FRAMES_MODEL = 81
|
38 |
-
MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS, 1)
|
39 |
-
MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS, 1)
|
40 |
-
|
41 |
-
# Cache for pipelines
|
42 |
-
t2v_pipe_cache = [None]
|
43 |
-
i2v_pipe_cache = [None]
|
44 |
-
|
45 |
-
|
46 |
-
def warmup_pipeline(pipe):
|
47 |
-
dummy_prompt = "warmup"
|
48 |
-
with torch.no_grad():
|
49 |
-
pipe(
|
50 |
-
prompt=dummy_prompt,
|
51 |
-
negative_prompt="",
|
52 |
-
height=LANDSCAPE_HEIGHT,
|
53 |
-
width=LANDSCAPE_WIDTH,
|
54 |
-
num_frames=MIN_FRAMES_MODEL,
|
55 |
-
guidance_scale=1.0,
|
56 |
-
guidance_scale_2=1.0,
|
57 |
-
num_inference_steps=1,
|
58 |
-
generator=torch.Generator(device="cuda").manual_seed(0),
|
59 |
-
)
|
60 |
-
def clear_memory():
|
61 |
-
"""Aggressively clear memory and CUDA cache."""
|
62 |
-
for _ in range(3):
|
63 |
-
gc.collect()
|
64 |
-
if torch.cuda.is_available():
|
65 |
-
torch.cuda.empty_cache()
|
66 |
-
torch.cuda.synchronize()
|
67 |
-
|
68 |
-
import time
|
69 |
-
import torch
|
70 |
-
|
71 |
-
def load_t2v_pipeline0():
|
72 |
-
"""Load and optimize the T2V pipeline once, reuse via cache."""
|
73 |
-
if t2v_pipe_cache[0] is None:
|
74 |
-
print("start the T2V pipeline ")
|
75 |
-
start_time = time.time()
|
76 |
-
clear_memory() # clean before loading to avoid fragmentation
|
77 |
-
|
78 |
-
with torch.inference_mode():
|
79 |
-
# Load VAE directly in BF16 on GPU
|
80 |
-
vae = AutoencoderKLWan.from_pretrained(
|
81 |
-
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
82 |
-
subfolder="vae",
|
83 |
-
torch_dtype=torch.bfloat16
|
84 |
-
).to("cuda", non_blocking=True)
|
85 |
-
|
86 |
-
# Load main transformers
|
87 |
-
transformer_main = WanTransformer3DModel.from_pretrained(
|
88 |
-
"linoyts/Wan2.2-T2V-A14B-Diffusers-BF16",
|
89 |
-
subfolder="transformer",
|
90 |
-
torch_dtype=torch.bfloat16,
|
91 |
-
device_map="cuda",
|
92 |
-
)
|
93 |
|
94 |
-
|
95 |
-
|
96 |
-
subfolder="transformer_2",
|
97 |
-
torch_dtype=torch.bfloat16,
|
98 |
-
device_map="cuda",
|
99 |
-
)
|
100 |
|
101 |
-
|
102 |
-
pipe = WanPipeline.from_pretrained(
|
103 |
-
T2V_MODEL_ID,
|
104 |
-
transformer=transformer_main,
|
105 |
-
transformer_2=transformer_refiner,
|
106 |
-
vae=vae,
|
107 |
-
torch_dtype=torch.bfloat16,
|
108 |
-
).to("cuda", non_blocking=True)
|
109 |
|
110 |
-
# Apply optimizations
|
111 |
-
print(f"Max VRAM reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB")
|
112 |
-
optimize_pipeline_(
|
113 |
-
pipe,
|
114 |
-
prompt="prompt",
|
115 |
-
height=LANDSCAPE_HEIGHT,
|
116 |
-
width=LANDSCAPE_WIDTH,
|
117 |
-
num_frames=MAX_FRAMES_MODEL,
|
118 |
-
)
|
119 |
|
120 |
-
|
121 |
-
|
122 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
123 |
|
124 |
-
# Save to cache
|
125 |
-
t2v_pipe_cache[0] = pipe
|
126 |
|
127 |
-
# Log load time and VRAM
|
128 |
-
elapsed = time.time() - start_time
|
129 |
-
print(f"T2V pipeline loaded in {elapsed:.2f}s")
|
130 |
-
print(f"Max VRAM reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB")
|
131 |
|
132 |
-
# Final cleanup
|
133 |
-
clear_memory()
|
134 |
|
135 |
-
return t2v_pipe_cache[0]
|
136 |
|
137 |
-
|
138 |
-
|
139 |
-
|
140 |
-
import time
|
141 |
-
import torch
|
142 |
-
|
143 |
-
def load_t2v_pipeline():
|
144 |
-
"""Always load the T2V pipeline at startup and store in cache."""
|
145 |
-
print("start the T2V pipeline ")
|
146 |
-
start_time = time.time()
|
147 |
-
clear_memory()
|
148 |
-
|
149 |
-
# ✅ Load in normal mode (allows tensor version counters)
|
150 |
-
vae = AutoencoderKLWan.from_pretrained(
|
151 |
-
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
152 |
-
subfolder="vae",
|
153 |
-
torch_dtype=torch.bfloat16
|
154 |
-
).to("cuda", non_blocking=True)
|
155 |
-
|
156 |
-
transformer_main = WanTransformer3DModel.from_pretrained(
|
157 |
-
"linoyts/Wan2.2-T2V-A14B-Diffusers-BF16",
|
158 |
-
subfolder="transformer",
|
159 |
torch_dtype=torch.bfloat16,
|
160 |
-
device_map=
|
161 |
-
)
|
162 |
-
|
163 |
-
|
164 |
-
"linoyts/Wan2.2-T2V-A14B-Diffusers-BF16",
|
165 |
-
subfolder="transformer_2",
|
166 |
torch_dtype=torch.bfloat16,
|
167 |
-
device_map=
|
168 |
-
)
|
|
|
|
|
|
|
169 |
|
170 |
-
pipe = WanPipeline.from_pretrained(
|
171 |
-
T2V_MODEL_ID,
|
172 |
-
transformer=transformer_main,
|
173 |
-
transformer_2=transformer_refiner,
|
174 |
-
vae=vae,
|
175 |
-
torch_dtype=torch.bfloat16,
|
176 |
-
).to("cuda", non_blocking=True)
|
177 |
-
|
178 |
-
print(f"Max VRAM reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB")
|
179 |
-
# ✅ Run optimization before inference mode
|
180 |
-
optimize_pipeline_(
|
181 |
-
pipe,
|
182 |
-
prompt="prompt",
|
183 |
-
height=LANDSCAPE_HEIGHT,
|
184 |
-
width=LANDSCAPE_WIDTH,
|
185 |
-
num_frames=MAX_FRAMES_MODEL,
|
186 |
-
)
|
187 |
-
|
188 |
-
pipe.enable_model_cpu_offload()
|
189 |
-
warmup_pipeline(pipe)
|
190 |
-
print("T2V pipeline warmed up with dummy inference.")
|
191 |
-
print(f"Max VRAM after offload reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB")
|
192 |
-
t2v_pipe_cache[0] = pipe
|
193 |
-
|
194 |
-
elapsed = time.time() - start_time
|
195 |
-
print(f"T2V pipeline preloaded in {elapsed:.2f}s")
|
196 |
-
print(f"Max VRAM reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB")
|
197 |
-
|
198 |
-
clear_memory()
|
199 |
-
return t2v_pipe_cache[0]
|
200 |
|
201 |
-
|
202 |
-
|
203 |
-
|
204 |
-
|
205 |
-
transformer=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
|
206 |
-
subfolder='transformer',
|
207 |
-
torch_dtype=torch.bfloat16,
|
208 |
-
device_map='cuda',
|
209 |
-
),
|
210 |
-
transformer_2=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
|
211 |
-
subfolder='transformer_2',
|
212 |
-
torch_dtype=torch.bfloat16,
|
213 |
-
device_map='cuda',
|
214 |
-
),
|
215 |
-
torch_dtype=torch.bfloat16,
|
216 |
-
).to('cuda')
|
217 |
-
optimize_pipeline_(i2v_pipe_cache[0],
|
218 |
-
image=Image.new('RGB', (LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT)),
|
219 |
-
prompt='prompt',
|
220 |
-
height=LANDSCAPE_HEIGHT,
|
221 |
-
width=LANDSCAPE_WIDTH,
|
222 |
-
num_frames=MAX_FRAMES_MODEL,
|
223 |
-
)
|
224 |
-
i2v_pipe_cache[0].enable_model_cpu_offload() # Enable CPU offload for memory optimization
|
225 |
-
clear_memory()
|
226 |
-
return i2v_pipe_cache[0]
|
227 |
|
228 |
-
|
229 |
-
|
230 |
-
|
231 |
-
|
232 |
-
|
233 |
-
|
234 |
|
235 |
-
def unload_i2v_pipeline():
|
236 |
-
if i2v_pipe_cache[0] is not None:
|
237 |
-
i2v_pipe_cache[0].to("cpu")
|
238 |
-
del i2v_pipe_cache[0]
|
239 |
-
i2v_pipe_cache[0] = None
|
240 |
-
clear_memory()
|
241 |
|
242 |
-
# Default prompts
|
243 |
default_prompt_t2v = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
|
244 |
-
default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
|
245 |
default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作��, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"
|
246 |
|
247 |
-
def resize_image(image: Image.Image) -> Image.Image:
|
248 |
-
if image.height > image.width:
|
249 |
-
transposed = image.transpose(Image.Transpose.ROTATE_90)
|
250 |
-
resized = resize_image_landscape(transposed)
|
251 |
-
return resized.transpose(Image.Transpose.ROTATE_270)
|
252 |
-
return resize_image_landscape(image)
|
253 |
-
|
254 |
-
def resize_image_landscape(image: Image.Image) -> Image.Image:
|
255 |
-
target_aspect = LANDSCAPE_WIDTH / LANDSCAPE_HEIGHT
|
256 |
-
width, height = image.size
|
257 |
-
in_aspect = width / height
|
258 |
-
if in_aspect > target_aspect:
|
259 |
-
new_width = round(height * target_aspect)
|
260 |
-
left = (width - new_width) // 2
|
261 |
-
image = image.crop((left, 0, left + new_width, height))
|
262 |
-
else:
|
263 |
-
new_height = round(width / target_aspect)
|
264 |
-
top = (height - new_height) // 2
|
265 |
-
image = image.crop((0, top, width, top + new_height))
|
266 |
-
return image.resize((LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT), Image.LANCZOS)
|
267 |
|
268 |
def get_duration(
|
269 |
-
mode,
|
270 |
-
input_image,
|
271 |
prompt,
|
272 |
negative_prompt,
|
273 |
duration_seconds,
|
@@ -278,161 +96,127 @@ def get_duration(
|
|
278 |
randomize_seed,
|
279 |
progress,
|
280 |
):
|
281 |
-
return
|
282 |
|
283 |
@spaces.GPU(duration=get_duration)
|
284 |
-
@torch.no_grad()
|
285 |
def generate_video(
|
286 |
-
mode,
|
287 |
-
input_image,
|
288 |
prompt,
|
289 |
negative_prompt=default_negative_prompt,
|
290 |
-
duration_seconds=MAX_DURATION,
|
291 |
-
guidance_scale=1,
|
292 |
-
guidance_scale_2=
|
293 |
-
steps=4,
|
294 |
-
seed=42,
|
295 |
-
randomize_seed=False,
|
296 |
progress=gr.Progress(track_tqdm=True),
|
297 |
):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
298 |
num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
|
299 |
current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
|
300 |
|
301 |
-
|
302 |
-
|
303 |
-
|
304 |
-
|
305 |
-
|
306 |
-
|
307 |
-
|
308 |
-
|
309 |
-
|
310 |
-
|
311 |
-
|
312 |
-
num_inference_steps=int(steps),
|
313 |
-
generator=torch.Generator(device="cuda").manual_seed(current_seed),
|
314 |
-
).frames[0]
|
315 |
-
else: # Image-to-Video
|
316 |
-
unload_t2v_pipeline() # Unload T2V to free memory
|
317 |
-
pipe = load_i2v_pipeline()
|
318 |
-
if input_image is None:
|
319 |
-
raise gr.Error("Please upload an input image.")
|
320 |
-
resized_image = resize_image(input_image)
|
321 |
-
output_frames_list = pipe(
|
322 |
-
image=resized_image,
|
323 |
-
prompt=prompt,
|
324 |
-
negative_prompt=negative_prompt,
|
325 |
-
height=resized_image.height,
|
326 |
-
width=resized_image.width,
|
327 |
-
num_frames=num_frames,
|
328 |
-
guidance_scale=float(guidance_scale),
|
329 |
-
guidance_scale_2=float(guidance_scale_2),
|
330 |
-
num_inference_steps=int(steps),
|
331 |
-
generator=torch.Generator(device="cuda").manual_seed(current_seed),
|
332 |
-
).frames[0]
|
333 |
|
334 |
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
|
335 |
video_path = tmpfile.name
|
336 |
|
337 |
export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
|
338 |
-
|
339 |
return video_path, current_seed
|
340 |
|
341 |
with gr.Blocks() as demo:
|
342 |
-
gr.Markdown("# Fast 4 steps Wan 2.2 T2V
|
343 |
-
gr.Markdown("
|
344 |
-
|
345 |
-
|
346 |
-
|
347 |
-
|
348 |
-
with gr.Column():
|
349 |
-
t2v_prompt_input = gr.Textbox(label="Prompt", value=default_prompt_t2v)
|
350 |
-
t2v_duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=MAX_DURATION, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
|
351 |
-
with gr.Accordion("Advanced Settings", open=False):
|
352 |
-
t2v_negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
|
353 |
-
t2v_seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
|
354 |
-
t2v_randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
|
355 |
-
t2v_steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=4, label="Inference Steps")
|
356 |
-
t2v_guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage")
|
357 |
-
t2v_guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=3, label="Guidance Scale 2 - low noise stage")
|
358 |
-
t2v_generate_button = gr.Button("Generate Video", variant="primary")
|
359 |
-
with gr.Column():
|
360 |
-
t2v_video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
|
361 |
-
|
362 |
-
t2v_inputs = [
|
363 |
-
gr.State(value="Text-to-Video"),
|
364 |
-
gr.State(value=None),
|
365 |
-
t2v_prompt_input,
|
366 |
-
t2v_negative_prompt_input,
|
367 |
-
t2v_duration_seconds_input,
|
368 |
-
t2v_guidance_scale_input,
|
369 |
-
t2v_guidance_scale_2_input,
|
370 |
-
t2v_steps_slider,
|
371 |
-
t2v_seed_input,
|
372 |
-
t2v_randomize_seed_checkbox
|
373 |
-
]
|
374 |
-
t2v_generate_button.click(fn=generate_video, inputs=t2v_inputs, outputs=[t2v_video_output, t2v_seed_input])
|
375 |
-
|
376 |
-
gr.Examples(
|
377 |
-
examples=[
|
378 |
-
["POV selfie video, white cat with sunglasses standing on surfboard, relaxed smile, tropical beach behind (clear water, green hills, blue sky with clouds). Surfboard tips, cat falls into ocean, camera plunges underwater with bubbles and sunlight beams. Brief underwater view of cat’s face, then cat resurfaces, still filming selfie, playful summer vacation mood."],
|
379 |
-
["Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."],
|
380 |
-
["A cinematic shot of a boat sailing on a calm sea at sunset."],
|
381 |
-
["Drone footage flying over a futuristic city with flying cars."],
|
382 |
-
],
|
383 |
-
inputs=[t2v_prompt_input],
|
384 |
-
outputs=[t2v_video_output, t2v_seed_input],
|
385 |
-
fn=generate_video,
|
386 |
-
cache_examples="lazy"
|
387 |
-
)
|
388 |
-
|
389 |
-
with gr.TabItem("Image-to-Video"):
|
390 |
-
with gr.Row():
|
391 |
-
with gr.Column():
|
392 |
-
i2v_input_image_component = gr.Image(type="pil", label="Input Image (auto-resized to target H/W)")
|
393 |
-
i2v_prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v)
|
394 |
-
i2v_duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=MAX_DURATION, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
|
395 |
-
with gr.Accordion("Advanced Settings", open=False):
|
396 |
-
i2v_negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
|
397 |
-
i2v_seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
|
398 |
-
i2v_randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
|
399 |
-
i2v_steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=6, label="Inference Steps")
|
400 |
-
i2v_guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage")
|
401 |
-
i2v_guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale 2 - low noise stage")
|
402 |
-
i2v_generate_button = gr.Button("Generate Video", variant="primary")
|
403 |
-
with gr.Column():
|
404 |
-
i2v_video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
|
405 |
|
406 |
-
|
407 |
-
gr.
|
408 |
-
|
409 |
-
|
410 |
-
|
411 |
-
|
412 |
-
|
413 |
-
|
414 |
-
|
415 |
-
|
416 |
-
|
417 |
-
|
418 |
-
|
419 |
-
|
420 |
-
|
421 |
-
|
422 |
-
|
423 |
-
|
424 |
-
|
425 |
-
|
426 |
-
|
427 |
-
|
428 |
-
|
429 |
-
|
430 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
431 |
|
432 |
if __name__ == "__main__":
|
433 |
-
|
434 |
-
# At the top of app.py
|
435 |
-
print("Loading T2V pipeline into GPU memory...")
|
436 |
-
t2v_pipe = load_t2v_pipeline()
|
437 |
-
print("T2V pipeline loaded and ready!")
|
438 |
-
demo.queue().launch(mcp_server=True)
|
|
|
1 |
+
# PyTorch 2.8 (temporary hack)
|
2 |
import os
|
|
|
|
|
|
|
|
|
|
|
3 |
os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
|
4 |
|
5 |
+
# Actual demo code
|
6 |
import spaces
|
7 |
import torch
|
8 |
from diffusers import WanPipeline, AutoencoderKLWan
|
|
|
9 |
from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
|
10 |
from diffusers.utils.export_utils import export_to_video
|
11 |
import gradio as gr
|
|
|
14 |
from PIL import Image
|
15 |
import random
|
16 |
import gc
|
17 |
+
from optimization import optimize_pipeline_
|
18 |
+
|
19 |
|
20 |
+
MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
|
|
|
|
|
|
|
|
|
|
|
21 |
|
|
|
|
|
|
|
22 |
LANDSCAPE_WIDTH = 832
|
23 |
LANDSCAPE_HEIGHT = 480
|
24 |
MAX_SEED = np.iinfo(np.int32).max
|
25 |
+
|
26 |
FIXED_FPS = 16
|
27 |
MIN_FRAMES_MODEL = 8
|
28 |
MAX_FRAMES_MODEL = 81
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
29 |
|
30 |
+
MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS,1)
|
31 |
+
MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS,1)
|
|
|
|
|
|
|
|
|
32 |
|
33 |
+
vae = AutoencoderKLWan.from_pretrained("Wan-AI/Wan2.2-T2V-A14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
34 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
35 |
|
36 |
+
# pipe = WanPipeline.from_pretrained(MODEL_ID,
|
37 |
+
# transformer=WanTransformer3DModel.from_pretrained('rahul7star/wan2.2',
|
38 |
+
# subfolder='Wan2.2-T2V-A14B-Diffusers-BF16/transformer',
|
39 |
+
# torch_dtype=torch.bfloat16,
|
40 |
+
# device_map='cuda',
|
41 |
+
# ),
|
42 |
+
# transformer_2=WanTransformer3DModel.from_pretrained('rahul7star/wan2.2',
|
43 |
+
# subfolder='Wan2.2-T2V-A14B-Diffusers-BF16/transformer_2',
|
44 |
+
# torch_dtype=torch.bfloat16,
|
45 |
+
# device_map='cuda',
|
46 |
+
# ),
|
47 |
+
# vae=vae,
|
48 |
+
# torch_dtype=torch.bfloat16,
|
49 |
+
# ).to('cuda')
|
50 |
|
|
|
|
|
51 |
|
|
|
|
|
|
|
|
|
52 |
|
|
|
|
|
53 |
|
|
|
54 |
|
55 |
+
pipe = WanPipeline.from_pretrained(MODEL_ID,
|
56 |
+
transformer=WanTransformer3DModel.from_pretrained('linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
|
57 |
+
subfolder='transformer',
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
58 |
torch_dtype=torch.bfloat16,
|
59 |
+
device_map='cuda',
|
60 |
+
),
|
61 |
+
transformer_2=WanTransformer3DModel.from_pretrained('linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
|
62 |
+
subfolder='transformer_2',
|
|
|
|
|
63 |
torch_dtype=torch.bfloat16,
|
64 |
+
device_map='cuda',
|
65 |
+
),
|
66 |
+
vae=vae,
|
67 |
+
torch_dtype=torch.bfloat16,
|
68 |
+
).to('cuda')
|
69 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
70 |
|
71 |
+
for i in range(3):
|
72 |
+
gc.collect()
|
73 |
+
torch.cuda.synchronize()
|
74 |
+
torch.cuda.empty_cache()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
75 |
|
76 |
+
optimize_pipeline_(pipe,
|
77 |
+
prompt='prompt',
|
78 |
+
height=LANDSCAPE_HEIGHT,
|
79 |
+
width=LANDSCAPE_WIDTH,
|
80 |
+
num_frames=MAX_FRAMES_MODEL,
|
81 |
+
)
|
82 |
|
|
|
|
|
|
|
|
|
|
|
|
|
83 |
|
|
|
84 |
default_prompt_t2v = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
|
|
|
85 |
default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作��, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"
|
86 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
87 |
|
88 |
def get_duration(
|
|
|
|
|
89 |
prompt,
|
90 |
negative_prompt,
|
91 |
duration_seconds,
|
|
|
96 |
randomize_seed,
|
97 |
progress,
|
98 |
):
|
99 |
+
return steps * 15
|
100 |
|
101 |
@spaces.GPU(duration=get_duration)
|
|
|
102 |
def generate_video(
|
|
|
|
|
103 |
prompt,
|
104 |
negative_prompt=default_negative_prompt,
|
105 |
+
duration_seconds = MAX_DURATION,
|
106 |
+
guidance_scale = 1,
|
107 |
+
guidance_scale_2 = 3,
|
108 |
+
steps = 4,
|
109 |
+
seed = 42,
|
110 |
+
randomize_seed = False,
|
111 |
progress=gr.Progress(track_tqdm=True),
|
112 |
):
|
113 |
+
"""
|
114 |
+
Generate a video from a text prompt using the Wan 2.2 14B T2V model with Lightning LoRA.
|
115 |
+
|
116 |
+
This function takes an input prompt and generates a video animation based on the provided
|
117 |
+
prompt and parameters. It uses an FP8 qunatized Wan 2.2 14B Text-to-Video model with Lightning LoRA
|
118 |
+
for fast generation in 4-8 steps.
|
119 |
+
|
120 |
+
Args:
|
121 |
+
prompt (str): Text prompt describing the desired animation or motion.
|
122 |
+
negative_prompt (str, optional): Negative prompt to avoid unwanted elements.
|
123 |
+
Defaults to default_negative_prompt (contains unwanted visual artifacts).
|
124 |
+
duration_seconds (float, optional): Duration of the generated video in seconds.
|
125 |
+
Defaults to 2. Clamped between MIN_FRAMES_MODEL/FIXED_FPS and MAX_FRAMES_MODEL/FIXED_FPS.
|
126 |
+
guidance_scale (float, optional): Controls adherence to the prompt. Higher values = more adherence.
|
127 |
+
Defaults to 1.0. Range: 0.0-20.0.
|
128 |
+
guidance_scale_2 (float, optional): Controls adherence to the prompt. Higher values = more adherence.
|
129 |
+
Defaults to 1.0. Range: 0.0-20.0.
|
130 |
+
steps (int, optional): Number of inference steps. More steps = higher quality but slower.
|
131 |
+
Defaults to 4. Range: 1-30.
|
132 |
+
seed (int, optional): Random seed for reproducible results. Defaults to 42.
|
133 |
+
Range: 0 to MAX_SEED (2147483647).
|
134 |
+
randomize_seed (bool, optional): Whether to use a random seed instead of the provided seed.
|
135 |
+
Defaults to False.
|
136 |
+
progress (gr.Progress, optional): Gradio progress tracker. Defaults to gr.Progress(track_tqdm=True).
|
137 |
+
|
138 |
+
Returns:
|
139 |
+
tuple: A tuple containing:
|
140 |
+
- video_path (str): Path to the generated video file (.mp4)
|
141 |
+
- current_seed (int): The seed used for generation (useful when randomize_seed=True)
|
142 |
+
|
143 |
+
Raises:
|
144 |
+
gr.Error: If input_image is None (no image uploaded).
|
145 |
+
|
146 |
+
Note:
|
147 |
+
- The function automatically resizes the input image to the target dimensions
|
148 |
+
- Frame count is calculated as duration_seconds * FIXED_FPS (24)
|
149 |
+
- Output dimensions are adjusted to be multiples of MOD_VALUE (32)
|
150 |
+
- The function uses GPU acceleration via the @spaces.GPU decorator
|
151 |
+
- Generation time varies based on steps and duration (see get_duration function)
|
152 |
+
"""
|
153 |
+
|
154 |
num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
|
155 |
current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
|
156 |
|
157 |
+
output_frames_list = pipe(
|
158 |
+
prompt=prompt,
|
159 |
+
negative_prompt=negative_prompt,
|
160 |
+
height=480,
|
161 |
+
width=832,
|
162 |
+
num_frames=num_frames,
|
163 |
+
guidance_scale=float(guidance_scale),
|
164 |
+
guidance_scale_2=float(guidance_scale_2),
|
165 |
+
num_inference_steps=int(steps),
|
166 |
+
generator=torch.Generator(device="cuda").manual_seed(current_seed),
|
167 |
+
).frames[0]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
168 |
|
169 |
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
|
170 |
video_path = tmpfile.name
|
171 |
|
172 |
export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
|
173 |
+
|
174 |
return video_path, current_seed
|
175 |
|
176 |
with gr.Blocks() as demo:
|
177 |
+
gr.Markdown("# Fast 4 steps Wan 2.2 T2V (14B) with Lightning LoRA")
|
178 |
+
gr.Markdown("run Wan 2.2 in just 4-8 steps, with [Wan 2.2 Lightning LoRA](https://huggingface.co/Kijai/WanVideo_comfy/tree/main/Wan22-Lightning), fp8 quantization & AoT compilation - compatible with 🧨 diffusers and ZeroGPU⚡️")
|
179 |
+
with gr.Row():
|
180 |
+
with gr.Column():
|
181 |
+
prompt_input = gr.Textbox(label="Prompt", value=default_prompt_t2v)
|
182 |
+
duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=MAX_DURATION, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
183 |
|
184 |
+
with gr.Accordion("Advanced Settings", open=False):
|
185 |
+
negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
|
186 |
+
seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
|
187 |
+
randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
|
188 |
+
steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=4, label="Inference Steps")
|
189 |
+
guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage")
|
190 |
+
guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=3, label="Guidance Scale 2 - low noise stage")
|
191 |
+
|
192 |
+
generate_button = gr.Button("Generate Video", variant="primary")
|
193 |
+
with gr.Column():
|
194 |
+
video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
|
195 |
+
|
196 |
+
ui_inputs = [
|
197 |
+
prompt_input,
|
198 |
+
negative_prompt_input, duration_seconds_input,
|
199 |
+
guidance_scale_input, guidance_scale_2_input, steps_slider, seed_input, randomize_seed_checkbox
|
200 |
+
]
|
201 |
+
generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
|
202 |
+
|
203 |
+
gr.Examples(
|
204 |
+
examples=[
|
205 |
+
[
|
206 |
+
"POV selfie video, white cat with sunglasses standing on surfboard, relaxed smile, tropical beach behind (clear water, green hills, blue sky with clouds). Surfboard tips, cat falls into ocean, camera plunges underwater with bubbles and sunlight beams. Brief underwater view of cat’s face, then cat resurfaces, still filming selfie, playful summer vacation mood.",
|
207 |
+
],
|
208 |
+
[
|
209 |
+
"Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
|
210 |
+
],
|
211 |
+
[
|
212 |
+
"A cinematic shot of a boat sailing on a calm sea at sunset.",
|
213 |
+
],
|
214 |
+
[
|
215 |
+
"Drone footage flying over a futuristic city with flying cars.",
|
216 |
+
],
|
217 |
+
],
|
218 |
+
inputs=[prompt_input], outputs=[video_output, seed_input], fn=generate_video, cache_examples="lazy"
|
219 |
+
)
|
220 |
|
221 |
if __name__ == "__main__":
|
222 |
+
demo.queue().launch(mcp_server=True)
|
|
|
|
|
|
|
|
|
|