nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
f9b2d4957bd60830bc35181df2fde10b59cd418b
hao-ai-lab-FastVideo/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py
T

194 lines
8.7 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
"""Convert the FastH3 DiT into pre-quantized (or dense BF16) MLX checkpoints.
Each output dir is a self-contained, shippable artifact in the
`mlx_h3_dit.safetensors` format: weights already cast/quantized, optional
AdaLN tables dropped at load time by the runtime. Conversion streams source
weights so the full BF16 DiT is not resident at once.
Usage (on the target Mac):
python scripts/checkpoint_conversion/convert_minimax_h3_mlx.py \\
--model-root ~/models/FastH3-Preview-v0.2/transformer \\
--out ~/models/FastH3-MLX \\
--formats "int8 int6 int4"
Dense conversion drops the 50 ``attn.to_gate_compress`` matrices (~3.6 GiB BF16).
Add ``--include-vsa`` to keep them, quantize them with the selected affine
INT8/INT6/INT4 grid, and record ``vsa.capable`` in the manifest. Write VSA
checkpoints to a new directory — do not overwrite an existing dense export.
python scripts/checkpoint_conversion/convert_minimax_h3_mlx.py \\
--model-root ~/models/FastH3-Preview-v0.2/transformer \\
--out ~/models/FastH3-MLX-vsa \\
--formats int6 --include-vsa
``bf16`` writes a dense, unquantized checkpoint (~38 GiB for the 33B body; AdaLN
table baked and its projection weights dropped, exactly like the quantized
formats). It targets large-memory Macs (96 GB+) where MLX runs dense BF16
matmul fastest. Conversion holds the dense body resident while saving, so it
needs roughly the output size in free memory.
python scripts/checkpoint_conversion/convert_minimax_h3_mlx.py \\
--model-root ~/models/FastH3/transformer_ref \\
--out ~/models/FastH3-MLX-vsa \\
--formats bf16 --include-vsa
Output layout: `~/models/FastH3-MLX/<format>/mlx_h3_dit.safetensors` + manifest.
"""
from __future__ import annotations
import argparse
import gc
import time
from pathlib import Path
import numpy as np
from fastvideo.logger import init_logger
from fastvideo.mlx_runtime.fastwan import MLXQuantizationSpec, ensure_quantization_supported
from fastvideo.mlx_runtime.minimax_h3 import (
H3_MANIFEST_FILENAME,
H3_WEIGHTS_FILENAME,
MINIMAX_H3_AUDIO_SHIFT,
MINIMAX_H3_VIDEO_SHIFT,
MiniMaxH3SchedulerState,
adaln_timestep_union,
load_fasth3_inference_contract,
mlx_h3_checkpoint_vsa_capable,
mlx_h3_dit_from_diffusers_safetensors,
minimax_h3_sigmas,
pdd_adaln_timesteps,
pdd_plan_from_contract,
save_mlx_h3_checkpoint,
)
logger = init_logger(__name__)
DENSE_FORMAT = "bf16"
QUANTIZED_FORMATS = ("int8", "int6", "int4")
SUPPORTED_FORMATS = (*QUANTIZED_FORMATS, DENSE_FORMAT)
DEFAULT_FORMATS = " ".join(QUANTIZED_FORMATS)
def _adaln_cache_timesteps(model_root: str | Path | None = None) -> np.ndarray:
"""Four-step uniform grid, unless the snapshot declares its own shifts.
FastH3 8-Step V2 carries ``video_scheduler_shift`` / ``audio_scheduler_shift``.
Four-step sidecars do not, so those conversions keep the existing cache.
A shifted contract that fails to parse aborts conversion.
"""
contract = load_fasth3_inference_contract(model_root)
if contract is None:
video = 1.0 - minimax_h3_sigmas(MINIMAX_H3_VIDEO_SHIFT, 4)[:-1]
audio = 1.0 - minimax_h3_sigmas(MINIMAX_H3_AUDIO_SHIFT, 4)[:-1]
return np.unique(np.concatenate([video, audio, [1.0]])).astype(np.float32)
if contract["kind"] == "pdd":
plan = pdd_plan_from_contract(contract)
video_state = MiniMaxH3SchedulerState.from_pdd_node_sigmas(contract["video_scheduler_shift"],
plan.node_sigmas["video"])
audio_state = MiniMaxH3SchedulerState.from_pdd_node_sigmas(contract["audio_scheduler_shift"],
plan.node_sigmas["audio"])
logger.info("FastH3 PDD schedule contract: %d fused blocks of a %d-interval grid, video/audio shift=%g/%g",
plan.num_steps, plan.pdd_steps, contract["video_scheduler_shift"],
contract["audio_scheduler_shift"])
return pdd_adaln_timesteps(contract, video_state, audio_state)
video_state = MiniMaxH3SchedulerState.from_dmd_steps(contract["video_scheduler_shift"],
contract["dmd_denoising_steps"])
audio_state = MiniMaxH3SchedulerState.from_dmd_steps(contract["audio_scheduler_shift"],
contract["dmd_denoising_steps"])
logger.info("FastH3 shifted schedule contract: %d DMD forwards, video/audio shift=%g/%g",
len(contract["dmd_denoising_steps"]), contract["video_scheduler_shift"],
contract["audio_scheduler_shift"])
return adaln_timestep_union(video_state, audio_state)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model-root", required=True, help="transformer/ dir of the diffusers snapshot")
parser.add_argument("--out", required=True, help="base dir for the per-format MLX checkpoints")
parser.add_argument("--formats", default=DEFAULT_FORMATS)
parser.add_argument(
"--include-vsa",
action="store_true",
help=("retain and quantize transformer_blocks.*.attn.to_gate_compress.weight "
"(required for MLX VSA inference; omitted by dense conversion)"),
)
parser.add_argument("--nvfp4-conditioner-root", type=Path,
help="also cache the released packed encoder in native MLX layout, without requantization")
parser.add_argument("--nvfp4-conditioner-out", type=Path,
help="empty encoder-cache output directory; defaults to OUT/nvfp4-encoder")
return parser.parse_args()
def main() -> None:
import mlx.core as mx
args = parse_args()
formats = args.formats.split()
unsupported = sorted(set(formats) - set(SUPPORTED_FORMATS))
if unsupported:
raise ValueError(f"Unsupported H3 MLX formats: {unsupported}. Choose from {SUPPORTED_FORMATS}.")
out_base = Path(args.out)
out_base.mkdir(parents=True, exist_ok=True)
cache_timesteps = _adaln_cache_timesteps(args.model_root)
for fmt in formats:
spec = None if fmt == DENSE_FORMAT else MLXQuantizationSpec.from_name(fmt)
if spec is None and fmt != DENSE_FORMAT:
raise ValueError(f"Expected a quantized H3 format, got {fmt}.")
ensure_quantization_supported(spec)
dtype = "bf16"
out_dir = out_base / fmt
already = (out_dir / H3_MANIFEST_FILENAME).exists() and (out_dir / H3_WEIGHTS_FILENAME).exists()
if already:
capable = mlx_h3_checkpoint_vsa_capable(out_dir)
if bool(capable) == bool(args.include_vsa):
print(f"[skip] {fmt} already converted at {out_dir} (vsa.capable={capable})", flush=True)
continue
logger.warning(
"Skipping %s: existing checkpoint has vsa.capable=%s, requested include_vsa=%s. "
"Use a new output directory to convert this format in the requested mode.",
out_dir, capable, args.include_vsa,
)
continue
print(f"[convert] {fmt} (dtype={dtype}, spec={spec}, include_vsa={args.include_vsa})", flush=True)
if hasattr(mx, "reset_peak_memory"):
mx.reset_peak_memory()
t0 = time.perf_counter()
dit = mlx_h3_dit_from_diffusers_safetensors(
args.model_root,
quantization=spec,
dtype=dtype,
adaln_cache_timesteps=cache_timesteps,
include_vsa=args.include_vsa,
)
mx.eval()
save_mlx_h3_checkpoint(dit, out_dir)
peak_gib = float(getattr(mx, "get_peak_memory", lambda: 0)()) / 2**30
print(f"[done] {fmt} in {time.perf_counter() - t0:.1f}s | peak {peak_gib:.1f} GiB -> {out_dir}", flush=True)
del dit
gc.collect()
if hasattr(mx, "clear_cache"):
mx.clear_cache()
if args.nvfp4_conditioner_root is not None:
from fastvideo.mlx_runtime.minimax_h3_conditioner import (
MLX_NVFP4_ENCODER_MANIFEST,
export_mlx_h3_nvfp4_encoder,
)
encoder_out = args.nvfp4_conditioner_out or out_base / "nvfp4-encoder"
if (encoder_out / MLX_NVFP4_ENCODER_MANIFEST).exists() and (encoder_out / "model.safetensors").exists():
print(f"[skip] NVFP4 encoder already cached at {encoder_out}", flush=True)
else:
started = time.perf_counter()
export_mlx_h3_nvfp4_encoder(args.nvfp4_conditioner_root, encoder_out)
print(f"[encoder] cached packed NVFP4 encoder in {time.perf_counter() - started:.1f}s at {encoder_out}",
flush=True)
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink