82 lines
2.6 KiB
Python
82 lines
2.6 KiB
Python
"""torch.compile A/B example for FastVideo.
|
|
|
|
`enable_torch_compile=True` compiles the DiT submodules that declare
|
|
`_compile_conditions` for a substantial end-to-end speedup (e.g.
|
|
Wan2.1-T2V-1.3B on A100: ~-24% e2e). It is off by default.
|
|
|
|
The first compiled generation pays a one-time graph-build cost; it
|
|
amortizes over later generations with the same input shapes. This
|
|
script does one un-measured warmup then a measured run so the reported
|
|
number is steady-state, not graph-build — measuring the warmup is the
|
|
most common way to wrongly conclude compile is slower.
|
|
|
|
Usage:
|
|
# baseline (eager)
|
|
python torch_compile_example.py
|
|
# compiled
|
|
python torch_compile_example.py --compile
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import time
|
|
|
|
from fastvideo import VideoGenerator
|
|
|
|
PROMPT = ("A high-definition video of a robotic arm welding a metal structure, "
|
|
"bright sparks and smoke, industrial setting.")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="torch.compile A/B")
|
|
parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
|
|
parser.add_argument("--compile", action="store_true", help="Enable torch.compile for the DiT")
|
|
parser.add_argument("--num_gpus", type=int, default=1)
|
|
args = parser.parse_args()
|
|
|
|
mode = "COMPILE" if args.compile else "BASELINE"
|
|
print(f"Mode: {mode} (enable_torch_compile={args.compile})")
|
|
|
|
os.makedirs("video_samples", exist_ok=True)
|
|
|
|
generator = VideoGenerator.from_pretrained(
|
|
args.model,
|
|
num_gpus=args.num_gpus,
|
|
enable_torch_compile=args.compile,
|
|
)
|
|
|
|
def _run(tag: str) -> float:
|
|
save = tag == "measured"
|
|
# Modern typed-request API (generate_video is deprecated). Same
|
|
# prompt/seed/shapes both runs so the compiled graph is reused.
|
|
request: dict = {
|
|
"prompt": PROMPT,
|
|
"sampling": {
|
|
"seed": 1024
|
|
},
|
|
"output": {
|
|
"save_video": save
|
|
},
|
|
}
|
|
if save:
|
|
request["output"]["output_path"] = (f"video_samples/torch_compile_{tag}.mp4")
|
|
t0 = time.perf_counter()
|
|
generator.generate(request)
|
|
return time.perf_counter() - t0
|
|
|
|
try:
|
|
# Warmup: pays the one-time graph build when --compile. Discarded.
|
|
w = _run("warmup")
|
|
print(f"warmup: {w:.2f}s "
|
|
f"({'incl. graph build' if args.compile else 'cold start'})")
|
|
|
|
# Measured: steady state, compiled graph reused (same shapes/seed).
|
|
m = _run("measured")
|
|
print(f"=== {mode} steady-state e2e: {m:.2f}s ===")
|
|
finally:
|
|
generator.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|