Skip to content

Commit 9a337a0

Browse files
kevin314kevin314
committed
Add scripts
1 parent 2f91701 commit 9a337a0

5 files changed

Lines changed: 337 additions & 26 deletions

File tree

‎examples/inference/optimizations/FastWan_QAD_TAEHV.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@
5252

5353
# TAEHV checkpoint for Wan2.1. Clone https://github.com/madebyollin/taehv to get
5454
# ``taew2_1.pth`` (Wan 2.1 / Wan 2.2-14B / Qwen-Image all use this VAE).
55-
DEFAULT_TAEHV_CHECKPOINT = "/root/taehv/taew2_1.pth"
55+
DEFAULT_TAEHV_CHECKPOINT = "/workspace/taehv/taew2_1.pth"
5656

5757
PROMPT = (
5858
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "

‎examples/inference/optimizations/fp4_qad_wan2_1_1_3b.py‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
"""NVFP4 QAD inference example.
2+
3+
Runs Wan2.1-T2V-1.3B with the FastWan-QAD-1.3B distilled checkpoint and
4+
NVFP4QATConfig quantization. Uses ATTN_QAT_INFER attention backend.
5+
6+
Requirements:
7+
- GPU: Blackwell (B200/B300, sm100a+) for the FP4 linear path
8+
- TAEHV (optional): Follow install instructions at https://github.com/madebyollin/taehv
9+
10+
Usage:
11+
python fp4_qad_wan2_1_1_3b.py # NVFP4 QAD (default)
12+
python fp4_qad_wan2_1_1_3b.py --bf16 # BF16 baseline
13+
python fp4_qad_wan2_1_1_3b.py --taehv-checkpoint /path/to/taew2_1.pth
14+
"""
15+
16+
import argparse
17+
import os
18+
import sys
19+
import time
20+
21+
import torch
22+
23+
OUTPUT_PATH = "video_samples"
24+
25+
26+
def load_taehv(checkpoint_path, device="cuda", dtype=torch.float16):
27+
repo_dir = os.path.dirname(checkpoint_path)
28+
if repo_dir not in sys.path:
29+
sys.path.insert(0, repo_dir)
30+
from taehv import TAEHV
31+
print(f"Loading TAEHV from {checkpoint_path}...")
32+
model = TAEHV(checkpoint_path=checkpoint_path).to(device, dtype)
33+
print("TAEHV loaded.")
34+
return model
35+
36+
37+
@torch.no_grad() # type: ignore[misc]
38+
def decode_with_taehv(taehv_model, latents):
39+
latents = latents.permute(0, 2, 1, 3, 4)
40+
latents = latents.to(device=next(taehv_model.parameters()).device,
41+
dtype=next(taehv_model.parameters()).dtype)
42+
decoded = taehv_model.decode_video(latents, parallel=False, show_progress_bar=False)
43+
frames = []
44+
for frame in decoded[0]:
45+
frame_np = (frame.clamp(0, 1) * 255).byte().cpu().permute(1, 2, 0).numpy()
46+
frames.append(frame_np)
47+
return frames
48+
49+
50+
def main():
51+
parser = argparse.ArgumentParser(description="NVFP4 QAD video generation benchmark")
52+
parser.add_argument("--bf16", action="store_true",
53+
help="BF16 baseline (no NVFP4 quantization)")
54+
parser.add_argument("--taehv-checkpoint", default=None, metavar="PATH",
55+
help="Path to taew2_1.pth; enables TAEHV tiny autoencoder decoding")
56+
parser.add_argument("--model", default="FastVideo/FastWan-QAD-1.3B",
57+
help="Model path or HuggingFace ID")
58+
parser.add_argument("--no-compile", action="store_true", help="Disable torch.compile for the DiT")
59+
parser.add_argument("--num_gpus", type=int, default=1)
60+
parser.add_argument("--infer_steps", type=int, default=3)
61+
args = parser.parse_args()
62+
63+
os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_INFER")
64+
os.environ["FASTVIDEO_DISABLE_ATTENTION_COMPILE"] = "0"
65+
os.environ["FLASHINFER_CUDA_ARCH_LIST"] = "12.0a"
66+
os.environ["FLASHINFER_EXTRA_CFLAGS"] = "-DCCCL_DISABLE_CTK_COMPATIBILITY_CHECK"
67+
os.environ["FLASHINFER_EXTRA_CUDAFLAGS"] = "-DCCCL_DISABLE_CTK_COMPATIBILITY_CHECK"
68+
# os.environ["CUDA_HOME"] = "/root/miniconda3/envs/fastvideo/lib/python3.12/site-packages/nvidia/cu13"
69+
70+
from fastvideo import VideoGenerator
71+
from fastvideo.configs.pipelines.base import PipelineConfig
72+
73+
mode = "bf16" if args.bf16 else "nvfp4_qad"
74+
if not args.no_compile:
75+
mode += "_compile"
76+
use_taehv = args.taehv_checkpoint is not None
77+
print(f"Mode: {mode.upper()}" + (" decoder=TAEHV" if use_taehv else " decoder=VAE"))
78+
79+
taehv_model = load_taehv(args.taehv_checkpoint) if use_taehv else None
80+
81+
pipeline_config = PipelineConfig.from_pretrained(args.model)
82+
pipeline_config.text_encoder_precisions = ("bf16",)
83+
if not args.bf16:
84+
from fastvideo.layers.quantization.nvfp4_qat_config import NVFP4QATConfig
85+
pipeline_config.dit_config.quant_config = NVFP4QATConfig()
86+
87+
generator = VideoGenerator.from_pretrained(
88+
args.model,
89+
pipeline_config=pipeline_config,
90+
num_gpus=args.num_gpus,
91+
use_fsdp_inference=False,
92+
dit_cpu_offload=False,
93+
dit_layerwise_offload=False,
94+
vae_cpu_offload=use_taehv,
95+
text_encoder_cpu_offload=False,
96+
pin_cpu_memory=False,
97+
enable_torch_compile=not args.no_compile,
98+
enable_torch_compile_text_encoder=not args.no_compile,
99+
enable_torch_compile_vae=not args.no_compile and not use_taehv,
100+
output_type="latent" if use_taehv else "pil",
101+
)
102+
103+
prompt = (
104+
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
105+
"wide with interest. The playful yet serene atmosphere is complemented by soft "
106+
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
107+
)
108+
109+
n_warmup = 2 if not args.no_compile else 0
110+
for _ in range(n_warmup):
111+
warmup_result = generator.generate(request={"prompt": prompt, "sampling": {"num_inference_steps": 3, "guidance_scale": 1.0},
112+
"output": {"save_video": False}})
113+
if use_taehv:
114+
decode_with_taehv(taehv_model, warmup_result.samples)
115+
116+
os.makedirs(OUTPUT_PATH, exist_ok=True)
117+
video_path = os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")
118+
if use_taehv:
119+
import imageio
120+
result = generator.generate(request={
121+
"prompt": prompt,
122+
"sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0},
123+
"output": {"save_video": False},
124+
})
125+
denoise_elapsed = result.generation_time
126+
torch.cuda.synchronize()
127+
t_decode = time.perf_counter()
128+
frames = decode_with_taehv(taehv_model, result.samples)
129+
torch.cuda.synchronize()
130+
decode_elapsed = time.perf_counter() - t_decode
131+
total = denoise_elapsed + decode_elapsed
132+
imageio.mimsave(video_path, frames, fps=16, format="mp4")
133+
print(f"Saved TAEHV-decoded video to: {video_path}")
134+
print(f"[{mode.upper()}] {args.infer_steps} steps in {total:.3f}s "
135+
f"(denoise {denoise_elapsed:.3f}s + decode {decode_elapsed:.3f}s)")
136+
else:
137+
result = generator.generate(request={
138+
"prompt": prompt,
139+
"sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0},
140+
"output": {"save_video": True, "output_path": video_path},
141+
})
142+
elapsed = result.generation_time
143+
print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.3f}s "
144+
f"({args.infer_steps / elapsed:.2f} it/s)")
145+
146+
generator.shutdown()
147+
148+
149+
if __name__ == "__main__":
150+
main()

‎examples/inference/optimizations/fp4_sa2_wan2_1_1_3b.py‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
"""NVFP4 QAD inference example with SageAttention 2 backend.
2+
3+
Runs Wan2.1-T2V-1.3B with the FastWan-QAD-1.3B-SA2 distilled checkpoint and
4+
NVFP4QATConfig quantization. Uses the SAGE_ATTN attention backend.
5+
6+
Requirements:
7+
- GPU: sm89+ (H100, L40S, RTX 4090, Ada Lovelace, or newer)
8+
- sageattention: pip install sageattention
9+
- TAEHV (optional): Follow install instructions at https://github.com/madebyollin/taehv
10+
11+
Usage:
12+
python fp4_sa2_wan2_1_1_3b.py # NVFP4 + SageAttn2 (default)
13+
python fp4_sa2_wan2_1_1_3b.py --bf16 # BF16 baseline
14+
python fp4_sa2_wan2_1_1_3b.py --taehv-checkpoint /path/to/taew2_1.pth
15+
"""
16+
17+
import argparse
18+
import os
19+
import sys
20+
import time
21+
22+
import torch
23+
24+
OUTPUT_PATH = "video_samples"
25+
26+
27+
def load_taehv(checkpoint_path, device="cuda", dtype=torch.float16):
28+
repo_dir = os.path.dirname(checkpoint_path)
29+
if repo_dir not in sys.path:
30+
sys.path.insert(0, repo_dir)
31+
from taehv import TAEHV
32+
print(f"Loading TAEHV from {checkpoint_path}...")
33+
model = TAEHV(checkpoint_path=checkpoint_path).to(device, dtype)
34+
print("TAEHV loaded.")
35+
return model
36+
37+
38+
@torch.no_grad() # type: ignore[misc]
39+
def decode_with_taehv(taehv_model, latents):
40+
latents = latents.permute(0, 2, 1, 3, 4)
41+
latents = latents.to(device=next(taehv_model.parameters()).device,
42+
dtype=next(taehv_model.parameters()).dtype)
43+
decoded = taehv_model.decode_video(latents, parallel=False, show_progress_bar=False)
44+
frames = []
45+
for frame in decoded[0]:
46+
frame_np = (frame.clamp(0, 1) * 255).byte().cpu().permute(1, 2, 0).numpy()
47+
frames.append(frame_np)
48+
return frames
49+
50+
51+
def main():
52+
parser = argparse.ArgumentParser(description="NVFP4 QAD + SageAttention2 video generation benchmark")
53+
parser.add_argument("--bf16", action="store_true",
54+
help="BF16 baseline (no NVFP4 quantization)")
55+
parser.add_argument("--taehv-checkpoint", default=None, metavar="PATH",
56+
help="Path to taew2_1.pth; enables TAEHV tiny autoencoder decoding")
57+
parser.add_argument("--model", default="FastVideo/FastWan-QAD-1.3B-SA2",
58+
help="Model path or HuggingFace ID")
59+
parser.add_argument("--no-compile", action="store_true", help="Disable torch.compile for the DiT")
60+
parser.add_argument("--num_gpus", type=int, default=1)
61+
parser.add_argument("--infer_steps", type=int, default=3)
62+
args = parser.parse_args()
63+
64+
os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")
65+
os.environ["FASTVIDEO_DISABLE_ATTENTION_COMPILE"] = "0"
66+
os.environ["FLASHINFER_CUDA_ARCH_LIST"] = "12.0a"
67+
os.environ["FLASHINFER_EXTRA_CFLAGS"] = "-DCCCL_DISABLE_CTK_COMPATIBILITY_CHECK"
68+
os.environ["FLASHINFER_EXTRA_CUDAFLAGS"] = "-DCCCL_DISABLE_CTK_COMPATIBILITY_CHECK"
69+
70+
from fastvideo import VideoGenerator
71+
from fastvideo.configs.pipelines.base import PipelineConfig
72+
73+
mode = "bf16" if args.bf16 else "nvfp4_sa2"
74+
if not args.no_compile:
75+
mode += "_compile"
76+
use_taehv = args.taehv_checkpoint is not None
77+
print(f"Mode: {mode.upper()}" + (" decoder=TAEHV" if use_taehv else " decoder=VAE"))
78+
79+
taehv_model = load_taehv(args.taehv_checkpoint) if use_taehv else None
80+
81+
pipeline_config = PipelineConfig.from_pretrained(args.model)
82+
pipeline_config.text_encoder_precisions = ("bf16",)
83+
if not args.bf16:
84+
from fastvideo.layers.quantization.nvfp4_qat_config import NVFP4QATConfig
85+
pipeline_config.dit_config.quant_config = NVFP4QATConfig()
86+
87+
generator = VideoGenerator.from_pretrained(
88+
args.model,
89+
pipeline_config=pipeline_config,
90+
num_gpus=args.num_gpus,
91+
use_fsdp_inference=False,
92+
dit_cpu_offload=False,
93+
dit_layerwise_offload=False,
94+
vae_cpu_offload=use_taehv,
95+
text_encoder_cpu_offload=False,
96+
pin_cpu_memory=False,
97+
enable_torch_compile=not args.no_compile,
98+
enable_torch_compile_text_encoder=not args.no_compile,
99+
enable_torch_compile_vae=not args.no_compile and not use_taehv,
100+
output_type="latent" if use_taehv else "pil",
101+
)
102+
103+
prompt = (
104+
"A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
105+
"wide with interest. The playful yet serene atmosphere is complemented by soft "
106+
"natural light filtering through the petals. Mid-shot, warm and cheerful tones."
107+
)
108+
109+
n_warmup = 2 if not args.no_compile else 0
110+
for _ in range(n_warmup):
111+
warmup_result = generator.generate(request={"prompt": prompt, "sampling": {"num_inference_steps": 3, "guidance_scale": 1.0},
112+
"output": {"save_video": False}})
113+
if use_taehv:
114+
decode_with_taehv(taehv_model, warmup_result.samples)
115+
116+
os.makedirs(OUTPUT_PATH, exist_ok=True)
117+
video_path = os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")
118+
if use_taehv:
119+
import imageio
120+
result = generator.generate(request={
121+
"prompt": prompt,
122+
"sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0},
123+
"output": {"save_video": False},
124+
})
125+
denoise_elapsed = result.generation_time
126+
torch.cuda.synchronize()
127+
t_decode = time.perf_counter()
128+
frames = decode_with_taehv(taehv_model, result.samples)
129+
torch.cuda.synchronize()
130+
decode_elapsed = time.perf_counter() - t_decode
131+
total = denoise_elapsed + decode_elapsed
132+
imageio.mimsave(video_path, frames, fps=16, format="mp4")
133+
print(f"Saved TAEHV-decoded video to: {video_path}")
134+
print(f"[{mode.upper()}] {args.infer_steps} steps in {total:.3f}s "
135+
f"(denoise {denoise_elapsed:.3f}s + decode {decode_elapsed:.3f}s)")
136+
else:
137+
result = generator.generate(request={
138+
"prompt": prompt,
139+
"sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0},
140+
"output": {"save_video": True, "output_path": video_path},
141+
})
142+
elapsed = result.generation_time
143+
print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.3f}s "
144+
f"({args.infer_steps / elapsed:.2f} it/s)")
145+
146+
generator.shutdown()
147+
148+
149+
if __name__ == "__main__":
150+
main()

0 commit comments

Comments
 (0)