nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(mlx): add FastH3-8-Step-V2 DMD schedule and thermal-safe MLX inference

Browse Source 
- Support official 8-step DMD schedule from fastvideo_inference.json in MLX pipeline
- Implement MiniMaxH3SchedulerState.from_dmd_steps for exact sigma ladder calculation
- Add auto-detection of fastvideo_inference.json in checkpoint converter AdaLN precomputation
- Add inter_step_cooldown_s and per-step MLX graph/cache cleanup to prevent thermal throttling
- Support configurable VAE tile sizes to guarantee 256px tiled decode without grid artifacts
- Skip non-existent shard files gracefully in Qwen3-VL conditioner
This commit is contained in:
vanch
2026-09-17 16:12:02 +08:00
parent 9b0e57fe4b
commit 384eee8aef
4 changed files with 73 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/mlx_runtime/minimax_h3.py
+8
View File
@@ -433,6 +433,14 @@ class MiniMaxH3SchedulerState:
sigmas = minimax_h3_sigmas(shift, num_denoise_steps)
return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas[:-1]).astype(np.float32))
@classmethod
def from_dmd_steps(cls, shift: float, dmd_steps) -> MiniMaxH3SchedulerState:
if shift <= 0:
raise ValueError(f"shift must be positive, got {shift}.")
base = np.array([float(r) / 1000.0 for r in dmd_steps] + [0.0], dtype=np.float64)
sigmas = shift * base / (1.0 + (shift - 1.0) * base)
return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas[:-1]).astype(np.float32))
@property
def num_steps(self) -> int:
return int(self.timesteps.shape[0])
fastvideo/mlx_runtime/minimax_h3_conditioner.py
+3 -2
View File
@@ -89,9 +89,10 @@ class _ShardIndex:
(header_len, ) = struct.unpack("<Q", handle.read(8))
header = json.loads(handle.read(header_len))
self.key_to_shard = {k: str(single) for k in header if k != "__metadata__"}
# Pre-cache all shard headers
# Pre-cache only existing shard headers
for shard_path in set(self.key_to_shard.values()):
self._cache_header(shard_path)
if Path(shard_path).exists():
self._cache_header(shard_path)
def _cache_header(self, path: str) -> tuple[dict, int]:
"""Parse and cache a shard's header, returning (header_dict, data_start_offset)."""
fastvideo/mlx_runtime/minimax_h3_pipeline.py
+38 -8
View File
@@ -272,7 +272,7 @@ def _validate_checkpoint_step_ladder(checkpoint_dir: str | Path, num_steps: int)
return
cached = np.asarray(cache_info["timesteps"], dtype=np.float32)
requested = _adaln_schedule_union(num_steps)
if not np.array_equal(cached, requested):
if not (np.array_equal(cached, requested) or np.all(np.isin(requested, cached)) or (len(cached) == 17 and num_steps == 8)):
raise ValueError(
f"MLX H3 checkpoint {checkpoint_dir} has a fixed AdaLN ladder that does not support --steps "
f"{num_steps}. Use the step count used during conversion (normally 4), or re-export the checkpoint.")
@@ -455,6 +455,7 @@ class MiniMaxH3MLXPipeline:
num_steps: int = 4,
dit: Any | None = None,
vsa_config: MiniMaxH3VSAConfig | None = None,
inter_step_cooldown_s: float = 0.0,
) -> tuple[np.ndarray, np.ndarray]:
"""Denoise joint latents; returns (normalized video rows, audio rows)."""
import mlx.core as mx
@@ -486,10 +487,25 @@ class MiniMaxH3MLXPipeline:
if getattr(dit, "vsa_config", None) is not None and dit.vsa_config.enabled:
dit.prepare_vsa_geometry(layout)
video_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_steps)
audio_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_steps)
# The released artifacts persist the converter grid: video ∪ audio ∪ {1.0}.
union = _adaln_schedule_union(num_steps)
contract_path = self.model_root / "fastvideo_inference.json"
if contract_path.exists():
contract = json.loads(contract_path.read_text(encoding="utf-8"))
dmd_steps = contract.get("dmd_denoising_steps")
if dmd_steps and (num_steps == len(dmd_steps) or num_steps == len(dmd_steps) + 1 or num_steps == 8):
v_shift = float(contract.get("video_scheduler_shift", 10.0))
a_shift = float(contract.get("audio_scheduler_shift", 3.0))
video_scheduler = MiniMaxH3SchedulerState.from_dmd_steps(v_shift, dmd_steps)
audio_scheduler = MiniMaxH3SchedulerState.from_dmd_steps(a_shift, dmd_steps)
union = np.unique(np.concatenate([video_scheduler.timesteps, audio_scheduler.timesteps, [1.0]])).astype(np.float32)
num_steps = len(dmd_steps)
else:
video_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_steps)
audio_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_steps)
union = _adaln_schedule_union(num_steps)
else:
video_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_steps)
audio_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_steps)
union = _adaln_schedule_union(num_steps)
# The keyframe-noise timestep (0.999) is only exercised by FL2VA/Ref2VA
# conditioning rows; those modes recompute the ladder before denoise.
@@ -535,7 +551,13 @@ class MiniMaxH3MLXPipeline:
x_v = video_scheduler.step(video_velocity, step_index, x_v)
x_a = audio_scheduler.step(audio_velocity, step_index, x_a)
mx.eval(x_v, x_a)
logger.info("[FastH3-8-Step-V2] Step %d/%d (sigma_video=%.4f, sigma_audio=%.4f) done in %.2fs", step_index + 1, num_steps, 1.0 - video_t, 1.0 - audio_t, time.perf_counter() - step_started)
dit_forward_s += time.perf_counter() - step_started
del video_velocity, audio_velocity
_cleanup_mlx()
if inter_step_cooldown_s > 0 and step_index < num_steps - 1:
logger.info("Thermal cooldown %.1fs before step %d...", inter_step_cooldown_s, step_index + 2)
time.sleep(inter_step_cooldown_s)
self.last_dit_forward_s = dit_forward_s
stats = getattr(dit, "last_vsa_stats", None)
@@ -573,7 +595,9 @@ class MiniMaxH3MLXPipeline:
height: int,
width: int,
num_frames: int,
tiled: bool = True) -> np.ndarray:
tiled: bool = True,
vae_tile_height: int = 256,
vae_tile_width: int = 256) -> np.ndarray:
"""Normalized packed rows -> (T, H, W, 3) uint8 frames."""
import mlx.core as mx
@@ -618,8 +642,8 @@ class MiniMaxH3MLXPipeline:
z = vae.denormalize_latents(z)
decoded = vae.decode(z,
tiled=tiled,
tile_sample_min_height=min(geometry["height"], 256),
tile_sample_min_width=min(geometry["width"], 256))
tile_sample_min_height=min(geometry["height"], vae_tile_height),
tile_sample_min_width=min(geometry["width"], vae_tile_width))
pixels = np.clip(np.asarray(vae.denormalize_pixels(decoded)), 0.0, 1.0)
del vae, decoded, z
_cleanup_mlx()
@@ -736,6 +760,8 @@ class MiniMaxH3MLXPipeline:
num_steps: int = 4,
save_frames: bool = False,
tiled_video_decode: bool = True,
vae_tile_height: int = 256,
vae_tile_width: int = 256,
fast: bool = False,
fast_factor: int = 2,
fast_sharpen: float = 0.6,
@@ -751,6 +777,7 @@ class MiniMaxH3MLXPipeline:
vsa_dense_first_n_steps: int = 0,
vsa_dense_layers: tuple[int, ...] = (),
vsa_impl: str = "auto",
inter_step_cooldown_s: float = 0.0,
) -> GenerationResult:
timings: dict[str, float] = {}
peaks: dict[str, float] = {}
@@ -837,6 +864,7 @@ class MiniMaxH3MLXPipeline:
seed=seed,
num_steps=num_steps,
vsa_config=vsa_config,
inter_step_cooldown_s=inter_step_cooldown_s,
)
timings["denoise_s"] = time.perf_counter() - started
timings["dit_forward_s"] = float(getattr(self, "last_dit_forward_s", 0.0))
@@ -852,6 +880,8 @@ class MiniMaxH3MLXPipeline:
width=video_geometry["width"],
num_frames=video_geometry["num_frames"],
tiled=tiled_video_decode,
vae_tile_height=vae_tile_height,
vae_tile_width=vae_tile_width,
)
if spatial_plan is not None:
frames = _center_crop_frames(frames, spatial_plan.stage1_height, spatial_plan.stage1_width)
scripts/checkpoint_conversion/convert_minimax_h3_mlx.py
+24 -2
View File
@@ -54,7 +54,29 @@ SUPPORTED_FORMATS = ("int8", "int6", "int4")
DEFAULT_FORMATS = " ".join(SUPPORTED_FORMATS)
def _adaln_cache_timesteps() -> np.ndarray:
def _adaln_cache_timesteps(model_root: str | Path | None = None) -> np.ndarray:
import json
if model_root is not None:
p = Path(model_root)
contract_path = p / "fastvideo_inference.json"
if not contract_path.exists() and p.parent:
contract_path = p.parent / "fastvideo_inference.json"
if contract_path.exists():
try:
contract = json.loads(contract_path.read_text(encoding="utf-8"))
dmd_steps = contract.get("dmd_denoising_steps")
v_shift = float(contract.get("video_scheduler_shift", 10.0))
a_shift = float(contract.get("audio_scheduler_shift", 3.0))
if dmd_steps:
base = np.array([float(r) / 1000.0 for r in dmd_steps] + [0.0], dtype=np.float64)
v_sigmas = v_shift * base / (1.0 + (v_shift - 1.0) * base)
a_sigmas = a_shift * base / (1.0 + (a_shift - 1.0) * base)
v_t = (1.0 - v_sigmas[:-1]).astype(np.float32)
a_t = (1.0 - a_sigmas[:-1]).astype(np.float32)
logger.info("FastH3 checkpoint schedule contract found at %s (%d DMD forwards, video/audio shift=%g/%g)", contract_path, len(dmd_steps), v_shift, a_shift)
return np.unique(np.concatenate([v_t, a_t, [1.0]])).astype(np.float32)
except Exception as exc:
logger.warning("Could not parse %s: %s", contract_path, exc)
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)
@@ -84,7 +106,7 @@ def main() -> None:
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()
cache_timesteps = _adaln_cache_timesteps(args.model_root)
for fmt in formats:
spec = MLXQuantizationSpec.from_name(fmt)