4f80c5a3da
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com> Co-authored-by: Aryan <aryan@aryan5v.dev>
109 lines
4.9 KiB
Python
109 lines
4.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Regenerate a 768P MiniMax-H3 draft video at 2560x1440 with Open-H3-Regenerate-2K.
|
|
|
|
The draft is VAE-encoded (the first 121 frames at 24 fps, 1344x768), its latent is upscaled, and three refine steps
|
|
of the H3 transformer with the released LoRA run at 2560x1440. The output video is silent. See
|
|
docs/inference/open_h3_regenerate_2k.md.
|
|
|
|
python examples/inference/basic/basic_open_h3_regenerate_2k.py --video draft.mp4 --prompt "..." --output out/
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from fastvideo import VideoGenerator
|
|
from fastvideo.api import (
|
|
ComponentConfig,
|
|
EngineConfig,
|
|
GenerationRequest,
|
|
GeneratorConfig,
|
|
InputConfig,
|
|
OffloadConfig,
|
|
OutputConfig,
|
|
PipelineSelection,
|
|
QuantizationConfig,
|
|
SamplingConfig,
|
|
)
|
|
|
|
DEFAULT_MODEL = "Efficient-Large-Model/Open-H3-Regenerate-2K"
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
parser.add_argument("--model-path", default=DEFAULT_MODEL, help="release repo id or a local copy of it")
|
|
parser.add_argument("--video", default=None, help="768P draft video (URL-free local path)")
|
|
parser.add_argument("--prompt", required=True, help="text the refine steps are conditioned on")
|
|
parser.add_argument("--output", type=Path, default=Path("outputs/open_h3_regenerate_2k"))
|
|
parser.add_argument("--steps", type=int, choices=(2, 3), default=3, help="refine steps (3 is the released setting)")
|
|
parser.add_argument("--seed", type=int, default=1766580932)
|
|
parser.add_argument("--attention-backend", default="TORCH_SDPA")
|
|
parser.add_argument("--draft-size", choices=("768p", "576p"), default="768p",
|
|
help="size the draft is encoded at; 576p for drafts from other video models")
|
|
parser.add_argument("--vae-decoder", default=None,
|
|
help="NVFP4 decoder-nvfp4.safetensors that replaces the VAE decoder at load "
|
|
"(Blackwell only, see docs/inference/open_h3_regenerate_2k.md)")
|
|
parser.add_argument("--refine-transformer", default=None,
|
|
help="refine on another MiniMax-H3 DiT (e.g. a FastH3 export) instead of the release's own")
|
|
parser.add_argument("--refine-transformer-quant", default=None,
|
|
help="NVFP4 layer profile for a packed refine DiT, e.g. h3_dit_vsa")
|
|
parser.add_argument("--draft-latent", default=None,
|
|
help="native draft latent .pt to refine instead of encoding --video")
|
|
parser.add_argument("--runs", type=int, default=1, help="repeat the request to measure warm timings")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if (args.video is None) == (args.draft_latent is None):
|
|
raise SystemExit("pass exactly one of --video or --draft-latent")
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
experimental = {
|
|
"attention_backend": args.attention_backend,
|
|
"draft_size": args.draft_size,
|
|
}
|
|
if args.vae_decoder:
|
|
experimental["h3_vae_nvfp4_decoder"] = args.vae_decoder
|
|
experimental["refine_transformer_dir"] = args.refine_transformer
|
|
experimental["draft_latent_path"] = args.draft_latent
|
|
engine = EngineConfig(num_gpus=1,
|
|
offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=False))
|
|
if args.refine_transformer_quant:
|
|
engine.quantization = QuantizationConfig(transformer_quant="NVFP4", layer_profile=args.refine_transformer_quant)
|
|
generator = VideoGenerator.from_config(
|
|
GeneratorConfig(
|
|
model_path=args.model_path,
|
|
engine=engine,
|
|
pipeline=PipelineSelection(
|
|
components=ComponentConfig(),
|
|
experimental=experimental,
|
|
),
|
|
))
|
|
try:
|
|
for run in range(args.runs):
|
|
started = time.perf_counter()
|
|
result = generator.generate(
|
|
GenerationRequest(
|
|
prompt=args.prompt,
|
|
inputs=InputConfig(video_path=args.video),
|
|
sampling=SamplingConfig(height=1440,
|
|
width=2560,
|
|
num_frames=124,
|
|
fps=24,
|
|
num_inference_steps=args.steps,
|
|
guidance_scale=1.0,
|
|
seed=args.seed),
|
|
output=OutputConfig(output_path=str(args.output / f"regen2k_{run}.mp4"),
|
|
save_video=True,
|
|
return_frames=False),
|
|
))
|
|
print(f"run {run}: {time.perf_counter() - started:.1f}s -> {result.video_path}")
|
|
finally:
|
|
generator.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|