127 lines
3.9 KiB
Python
127 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Profile one FastH3 Preview generation with NVTX after warmup."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
from fastvideo import VideoGenerator
|
|
from fastvideo.api import (
|
|
CompileConfig,
|
|
ComponentConfig,
|
|
EngineConfig,
|
|
GenerationRequest,
|
|
GeneratorConfig,
|
|
OffloadConfig,
|
|
OutputConfig,
|
|
ParallelismConfig,
|
|
PipelineSelection,
|
|
QuantizationConfig,
|
|
SamplingConfig,
|
|
)
|
|
from fastvideo.profiler import nvtx_range
|
|
|
|
VARIANT_BACKENDS = {
|
|
"dense-datafree": "FLASH_ATTN",
|
|
"vsa-datafree": "VIDEO_SPARSE_ATTN_H3",
|
|
"vsa-synthetic-step1300": "VIDEO_SPARSE_ATTN_H3",
|
|
"vsa-synthetic-step1900": "VIDEO_SPARSE_ATTN_H3",
|
|
}
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("variant", choices=VARIANT_BACKENDS)
|
|
parser.add_argument("--prompt", required=True)
|
|
parser.add_argument("--output", required=True)
|
|
parser.add_argument("--num-frames", type=int, default=345)
|
|
parser.add_argument("--warmup-runs", type=int, default=3)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
"""Warm the selected FastH3 recipe, then profile one identical generation."""
|
|
args = parse_args()
|
|
output_dir = Path(args.output)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
attention_backend = VARIANT_BACKENDS[args.variant]
|
|
experimental = {
|
|
"attention_backend": attention_backend,
|
|
"inference_torch_compile": attention_backend == "FLASH_ATTN",
|
|
"vae_parallel_decode": True,
|
|
"vae_parallel_decode_strategy": "gather",
|
|
}
|
|
if attention_backend == "VIDEO_SPARSE_ATTN_H3":
|
|
experimental.update({
|
|
"VSA_sparsity": 0.9,
|
|
"VSA_tile_size": 64,
|
|
})
|
|
|
|
adapter_path = hf_hub_download(
|
|
repo_id="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA",
|
|
filename=f"{args.variant}/adapter_model.safetensors",
|
|
)
|
|
generator = VideoGenerator.from_config(
|
|
GeneratorConfig(
|
|
model_path="MiniMaxAI/MiniMax-H3",
|
|
pipeline=PipelineSelection(
|
|
components=ComponentConfig(lora_path=adapter_path, lora_strength=1.0),
|
|
experimental=experimental,
|
|
),
|
|
engine=EngineConfig(
|
|
num_gpus=4,
|
|
parallelism=ParallelismConfig(tp_size=1, sp_size=4),
|
|
offload=OffloadConfig(dit=False, dit_layerwise=False),
|
|
compile=CompileConfig(vae_enabled=True),
|
|
quantization=QuantizationConfig(transformer_quant="MXFP8"),
|
|
),
|
|
)
|
|
)
|
|
|
|
request = GenerationRequest(
|
|
prompt=args.prompt,
|
|
negative_prompt="",
|
|
sampling=SamplingConfig(
|
|
height=768,
|
|
width=1344,
|
|
num_frames=args.num_frames,
|
|
fps=24,
|
|
num_inference_steps=5,
|
|
guidance_scale=1.0,
|
|
batch_cfg=False,
|
|
seed=1000,
|
|
),
|
|
output=OutputConfig(
|
|
output_path=str(output_dir / "fasth3_profile.mp4"),
|
|
save_video=True,
|
|
return_frames=False,
|
|
),
|
|
)
|
|
|
|
try:
|
|
for warmup_index in range(args.warmup_runs):
|
|
warmup_result = generator.generate(request)
|
|
print(f"Warmup {warmup_index + 1}/{args.warmup_runs}: {warmup_result.video_path}")
|
|
|
|
torch.cuda.profiler.start()
|
|
try:
|
|
with nvtx_range("fasth3.profiled_generation"):
|
|
measured_result = generator.generate(request)
|
|
finally:
|
|
torch.cuda.profiler.stop()
|
|
|
|
print(f"Profiled output: {measured_result.video_path}")
|
|
if measured_result.generation_time is not None:
|
|
print(f"Generation time: {measured_result.generation_time:.2f}s")
|
|
finally:
|
|
generator.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|