rahul7star commited on
Commit
c2716fd
·
verified ·
1 Parent(s): 1d212f5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +148 -364
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
- # Assuming optimize_pipeline_ is a custom function; if not available, define a no-op
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
- transformer_refiner = WanTransformer3DModel.from_pretrained(
95
- "linoyts/Wan2.2-T2V-A14B-Diffusers-BF16",
96
- subfolder="transformer_2",
97
- torch_dtype=torch.bfloat16,
98
- device_map="cuda",
99
- )
100
 
101
- # Assemble pipeline
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
- # Enable offload to CPU when idle
121
- pipe.enable_model_cpu_offload()
122
- print(f"Max VRAM after enable CPU offload reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB")
 
 
 
 
 
 
 
 
 
 
 
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
- import time
138
- import torch
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="cuda",
161
- )
162
-
163
- transformer_refiner = WanTransformer3DModel.from_pretrained(
164
- "linoyts/Wan2.2-T2V-A14B-Diffusers-BF16",
165
- subfolder="transformer_2",
166
  torch_dtype=torch.bfloat16,
167
- device_map="cuda",
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
- def load_i2v_pipeline():
202
- """Load and optimize the I2V pipeline."""
203
- if i2v_pipe_cache[0] is None:
204
- i2v_pipe_cache[0] = WanImageToVideoPipeline.from_pretrained(I2V_MODEL_ID,
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
- def unload_t2v_pipeline():
229
- if t2v_pipe_cache[0] is not None:
230
- t2v_pipe_cache[0].to("cpu")
231
- del t2v_pipe_cache[0]
232
- t2v_pipe_cache[0] = None
233
- clear_memory()
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 int(steps) * 15
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=1,
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
- if mode == "Text-to-Video":
302
- unload_i2v_pipeline() # Unload I2V to free memory
303
- print(t2v_pipe.hf_device_map)
304
- output_frames_list = t2v_pipe_cache[0](
305
- prompt=prompt,
306
- negative_prompt=negative_prompt,
307
- height=LANDSCAPE_HEIGHT,
308
- width=LANDSCAPE_WIDTH,
309
- num_frames=num_frames,
310
- guidance_scale=float(guidance_scale),
311
- guidance_scale_2=float(guidance_scale_2),
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
- clear_memory() # Clean up after generation
339
  return video_path, current_seed
340
 
341
  with gr.Blocks() as demo:
342
- gr.Markdown("# Fast 4 steps Wan 2.2 T2V/I2V (14B) with Lightning LoRA")
343
- gr.Markdown("Run Wan 2.2 in just 4-8 steps, with [Lightning LoRA](https://huggingface.co/Kijai/WanVideo_comfy/tree/main/Wan22-Lightning), fp8 quantization & AoT compilation - compatible with 🧨 diffusers and ZeroGPU⚡️")
344
-
345
- with gr.Tabs() as tabs:
346
- with gr.TabItem("Text-to-Video"):
347
- with gr.Row():
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
- i2v_inputs = [
407
- gr.State(value="Image-to-Video"),
408
- i2v_input_image_component,
409
- i2v_prompt_input,
410
- i2v_negative_prompt_input,
411
- i2v_duration_seconds_input,
412
- i2v_guidance_scale_input,
413
- i2v_guidance_scale_2_input,
414
- i2v_steps_slider,
415
- i2v_seed_input,
416
- i2v_randomize_seed_checkbox
417
- ]
418
- i2v_generate_button.click(fn=generate_video, inputs=i2v_inputs, outputs=[i2v_video_output, i2v_seed_input])
419
-
420
- gr.Examples(
421
- examples=[
422
- ["wan_i2v_input.JPG", "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.", 4],
423
- ["wan22_input_2.jpg", "A sleek lunar vehicle glides into view from left to right, kicking up moon dust as astronauts in white spacesuits hop aboard with characteristic lunar bouncing movements. In the distant background, a VTOL craft descends straight down and lands silently on the surface. Throughout the entire scene, ethereal aurora borealis ribbons dance across the star-filled sky, casting shimmering curtains of green, blue, and purple light that bathe the lunar landscape in an otherworldly, magical glow.", 4],
424
- ["kill_bill.jpeg", "Uma Thurman's character, Beatrix Kiddo, holds her razor-sharp katana blade steady in the cinematic lighting. Suddenly, the polished steel begins to soften and distort, like heated metal starting to lose its structural integrity. The blade's perfect edge slowly warps and droops, molten steel beginning to flow downward in silvery rivulets while maintaining its metallic sheen. The transformation starts subtly at first - a slight bend in the blade - then accelerates as the metal becomes increasingly fluid. The camera holds steady on her face as her piercing eyes gradually narrow, not with lethal focus, but with confusion and growing alarm as she watches her weapon dissolve before her eyes. Her breathing quickens slightly as she witnesses this impossible transformation. The melting intensifies, the katana's perfect form becoming increasingly abstract, dripping like liquid mercury from her grip. Molten droplets fall to the ground with soft metallic impacts. Her expression shifts from calm readiness to bewilderment and concern as her legendary instrument of vengeance literally liquefies in her hands, leaving her defenseless and disoriented.", 6],
425
- ],
426
- inputs=[i2v_input_image_component, i2v_prompt_input, i2v_steps_slider],
427
- outputs=[i2v_video_output, i2v_seed_input],
428
- fn=generate_video,
429
- cache_examples="lazy"
430
- )
 
 
 
 
 
 
 
 
 
 
 
431
 
432
  if __name__ == "__main__":
433
- # Load pipeline once when the app starts
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)