SolitaryThinker
bbbb7ab021
Move all examples/inference and scripts off the legacy
VideoGenerator.from_pretrained(model, **kwargs) +
generate_video(prompt, **kwargs) surface onto the typed
from_config(GeneratorConfig(...)) + generate(GenerationRequest(...))
convention, including the Kandinsky-5 and DreamX-World examples added
on main after the original migration. git-mv api/compat.py to
api/translation.py ('compat' implied a temporary shim; the forward
translation is its honest permanent role) and repoint all importers,
tests, and doc links.
93 lines
2.8 KiB
Python
93 lines
2.8 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
|
|
from fastvideo.api import (
|
|
CompileConfig,
|
|
EngineConfig,
|
|
GenerationRequest,
|
|
GeneratorConfig,
|
|
OutputConfig,
|
|
SamplingConfig,
|
|
)
|
|
|
|
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_config(
|
|
GeneratorConfig(
|
|
model_path=args.model,
|
|
engine=EngineConfig(
|
|
num_gpus=args.num_gpus,
|
|
compile=CompileConfig(enabled=args.compile),
|
|
),
|
|
)
|
|
)
|
|
|
|
def _run(tag: str) -> float:
|
|
save = tag == "measured"
|
|
# Same prompt/seed/shapes both runs so the compiled graph is reused.
|
|
output = OutputConfig(save_video=save)
|
|
if save:
|
|
output.output_path = f"video_samples/torch_compile_{tag}.mp4"
|
|
request = GenerationRequest(
|
|
prompt=PROMPT,
|
|
sampling=SamplingConfig(seed=1024),
|
|
output=output,
|
|
)
|
|
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()
|