minimax_h3 ¶

def apply_h3_rotary(x, cos, sin):
    """Rotate the RoPE prefix of each head, preserving the rest.

    x: (S, H, D); cos/sin: (S, R) with R <= D. Half-split (GPT-NeoX style)
    rotation inside the rotary prefix, computed in fp32.
    """
    import mlx.core as mx

    rotary_dim = cos.shape[-1]
    x_rotary = x[..., :rotary_dim].astype(mx.float32)
    x_pass = x[..., rotary_dim:]
    cos_b = cos[:, None, :].astype(mx.float32)
    sin_b = sin[:, None, :].astype(mx.float32)
    first_half, second_half = mx.split(x_rotary, 2, axis=-1)
    rotated = mx.concatenate([-second_half, first_half], axis=-1)
    out = x_rotary * cos_b + rotated * sin_b
    return mx.concatenate([out.astype(x.dtype), x_pass], axis=-1)

def build_packed_layout(
    num_text_tokens: int,
    num_latent_frames: int,
    latent_height: int,
    latent_width: int,
    num_audio_latents: int,
    patch_size: tuple[int, int, int] = (1, 2, 2),
    keyframe_anchors: tuple[str, ...] = (),
    text_token_tags: np.ndarray | None = None,
    video_temporal_scale: float = 1.0,
) -> MiniMaxH3PackedLayout:
    """Build the ``[text | condition | audio | video]`` layout (T2VA / FL2VA)."""
    _, patch_h, patch_w = patch_size
    if text_token_tags is None:
        text_token_tags = np.full(num_text_tokens, MINIMAX_H3_TEXT_TAG, dtype=np.int64)
    if text_token_tags.shape != (num_text_tokens, ):
        raise ValueError(f"text_token_tags must have shape ({num_text_tokens},), got {text_token_tags.shape}.")
    if not np.isin(text_token_tags, (MINIMAX_H3_TEXT_TAG, MINIMAX_H3_VIDEO_TAG)).all():
        raise ValueError("text_token_tags may contain only text and vision tags.")
    if not np.isfinite(video_temporal_scale) or video_temporal_scale <= 0:
        raise ValueError(f"video_temporal_scale must be positive and finite, got {video_temporal_scale}.")
    if keyframe_anchors and video_temporal_scale != 1.0:
        raise ValueError("video_temporal_scale is currently supported only for T2VA without keyframe anchors.")

    rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w)
    num_condition_rows = len(keyframe_anchors) * rows_per_frame
    num_audio_rows = num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS
    num_video_rows = num_latent_frames * rows_per_frame
    sequence_length = num_text_tokens + num_condition_rows + num_audio_rows + num_video_rows

    condition_start = num_text_tokens
    audio_start = condition_start + num_condition_rows
    video_start = audio_start + num_audio_rows

    position_ids = np.zeros((sequence_length, 3), dtype=np.float64)
    position_ids[:num_text_tokens, 0] = np.arange(num_text_tokens, dtype=np.float64)

    sqrt_area = float(np.sqrt(latent_height * latent_width))
    height_grid = spatial_position_grid(latent_height, patch_h, sqrt_area)
    width_grid = spatial_position_grid(latent_width, patch_w, sqrt_area)
    mesh_h, mesh_w = np.meshgrid(height_grid, width_grid, indexing="ij")
    frame_grid = np.stack([mesh_h.reshape(-1), mesh_w.reshape(-1)], axis=-1)

    for index, anchor in enumerate(keyframe_anchors):
        if anchor == "first":
            anchor_time = float(num_text_tokens)
        elif anchor == "last":
            anchor_time = (float(num_text_tokens) + _temporal_position_span(num_latent_frames) -
                           MINIMAX_H3_ROPE_FRAME_RESCALE)
        else:
            raise ValueError(f"A keyframe anchor must be 'first' or 'last', got {anchor!r}.")
        rows = slice(condition_start + index * rows_per_frame, condition_start + (index + 1) * rows_per_frame)
        position_ids[rows, 0] = anchor_time
        position_ids[rows, 1:] = frame_grid

    audio_time = float(num_text_tokens) + np.arange(num_audio_latents, dtype=np.float64)
    position_ids[audio_start:video_start, 0] = np.tile(audio_time, MINIMAX_H3_AUDIO_CHANNELS)
    position_ids[audio_start:video_start, 2] = np.concatenate([
        np.full(num_audio_latents, float(width_grid[0]), dtype=np.float64),
        np.full(num_audio_rows - num_audio_latents, float(width_grid[-1]), dtype=np.float64),
    ])

    frame_time = temporal_position_grid(num_latent_frames, 0.0) * video_temporal_scale + float(num_text_tokens)
    position_ids[video_start:, 0] = np.repeat(frame_time, rows_per_frame)
    position_ids[video_start:, 1:] = np.tile(frame_grid, (num_latent_frames, 1))

    video_indices = np.concatenate([
        np.arange(condition_start, audio_start),
        np.arange(video_start, sequence_length),
    ])
    audio_indices = np.arange(audio_start, video_start)
    text_indices = np.arange(num_text_tokens)
    token_tags = np.empty(sequence_length, dtype=np.int64)
    token_tags[text_indices] = text_token_tags
    token_tags[audio_indices] = MINIMAX_H3_AUDIO_TAG
    token_tags[video_indices] = MINIMAX_H3_VIDEO_TAG

    return MiniMaxH3PackedLayout(
        sequence_length=sequence_length,
        position_ids=position_ids,
        token_tags=token_tags,
        video_indices=video_indices,
        audio_indices=audio_indices,
        text_indices=text_indices,
        num_condition_video_rows=num_condition_rows,
        num_condition_audio_rows=0,
        num_video_latent_frames=num_latent_frames,
        latent_height=latent_height,
        latent_width=latent_width,
        num_audio_latents=num_audio_latents,
    )

def build_row_timesteps(
    layout: MiniMaxH3PackedLayout,
    video_timestep: float,
    audio_timestep: float,
    condition_video_timestep: float = 1.0,
    condition_audio_timestep: float = 1.0,
) -> tuple[np.ndarray, np.ndarray]:
    """Per-row timesteps -> (unique timesteps sorted ascending, per-row indices).

    Matches ``torch.unique(..., sorted=True, return_inverse=True)``.
    """
    row_timesteps = np.full(layout.sequence_length, video_timestep, dtype=np.float64)
    if layout.num_condition_video_rows:
        row_timesteps[layout.video_indices[:layout.num_condition_video_rows]] = condition_video_timestep
    row_timesteps[layout.audio_indices[layout.num_condition_audio_rows:]] = audio_timestep
    if layout.num_condition_audio_rows:
        row_timesteps[layout.audio_indices[:layout.num_condition_audio_rows]] = condition_audio_timestep
    unique, inverse = np.unique(row_timesteps, return_inverse=True)
    return unique.astype(np.float32), inverse.astype(np.int64)

def linear(x, weight, bias=None):
    """Run an H3 linear with its measured wide-row affine dispatch enabled."""
    return _shared_linear(x, weight, bias, use_affine_dq_gemm=True)

def load_mlx_h3_checkpoint(checkpoint_dir: str | Path) -> MLXMiniMaxH3DiT:
    """Rebuild an H3 DiT saved by :func:`save_mlx_h3_checkpoint`.

    Refuses a quantization grid the installed MLX cannot execute, loudly —
    never silently dequantizes onto a different grid.
    """
    import mlx.core as mx

    checkpoint_dir = Path(checkpoint_dir)
    manifest_path = checkpoint_dir / H3_MANIFEST_FILENAME
    weights_path = checkpoint_dir / H3_WEIGHTS_FILENAME
    if not manifest_path.exists() or not weights_path.exists():
        raise FileNotFoundError(f"Not an MLX H3 checkpoint directory: {checkpoint_dir} "
                                f"(expected {H3_MANIFEST_FILENAME} and {H3_WEIGHTS_FILENAME}).")

    manifest = json.loads(manifest_path.read_text())
    version = manifest.get("format_version")
    if version != H3_FORMAT_VERSION:
        raise ValueError(f"MLX H3 checkpoint {checkpoint_dir} has format_version={version}; "
                         f"this build reads version {H3_FORMAT_VERSION}. Re-export the checkpoint.")

    spec = None
    if manifest["quantization"] is not None:
        spec = MLXQuantizationSpec(**manifest["quantization"])
        ensure_quantization_supported(spec)

    arrays = mx.load(str(weights_path))
    quantized_keys: dict[str, dict[str, Any]] = manifest["quantized_keys"]

    def rebuild(key: str):
        if key not in quantized_keys:
            return arrays[key]
        info = quantized_keys[key]
        assert spec is not None, f"Quantized key '{key}' in a checkpoint without a quantization spec"
        return QuantizedMatrix(
            weight=arrays[key],
            scales=arrays[f"{key}.scales"],
            biases=arrays[f"{key}.biases"] if info["has_biases"] else None,
            spec=spec,
            dequantized_dtype=_name_to_dtype(info["dequantized_dtype"]),
        )

    weights: dict[str, Any] = {}
    blocks: list[dict[str, Any] | None] = [None] * int(manifest["num_blocks"])
    refiner: list[dict[str, Any] | None] = [None] * int(manifest["num_refiner_blocks"])
    for key in arrays:
        if key.startswith("__adaln_cache."):
            continue
        if key.endswith(".scales") or key.endswith(".biases"):
            continue
        value = rebuild(key)
        if key.startswith("blocks."):
            _, index_str, sub = key.split(".", 2)
            index = int(index_str)
            block = blocks[index]
            if block is None:
                block = {}
                blocks[index] = block
            block[sub] = value
        elif key.startswith("refiner."):
            _, index_str, sub = key.split(".", 2)
            index = int(index_str)
            refiner_block = refiner[index]
            if refiner_block is None:
                refiner_block = {}
                refiner[index] = refiner_block
            refiner_block[sub] = value
        else:
            weights[key] = value

    if any(block is None for block in blocks) or any(block is None for block in refiner):
        raise ValueError(f"MLX H3 checkpoint {checkpoint_dir} is missing block weights.")
    dit = MLXMiniMaxH3DiT(
        weights,
        [block for block in blocks if block is not None],
        [block for block in refiner if block is not None],
        manifest["config"],
    )
    cache_info = manifest.get("adaln_cache")
    if cache_info is not None:
        block_tables = []
        for block_index in range(int(cache_info["num_blocks"])):
            block_tables.append(
                tuple(arrays[f"__adaln_cache.block.{block_index}.table.{table_index}"]
                      for table_index in range(int(cache_info["tables_per_block"]))))
        dit._adaln_cache = MiniMaxH3StepCache(
            timesteps=np.asarray(cache_info["timesteps"], dtype=np.float32),
            block_tables=block_tables,
            norm_out_shift=arrays["__adaln_cache.norm_out_shift"],
            norm_out_scale=arrays["__adaln_cache.norm_out_scale"],
        )
    return dit

def minimax_h3_sigmas(shift: float, num_denoise_steps: int) -> np.ndarray:
    """Sigma grid of length ``num_denoise_steps + 1`` (descending, ending at 0).

    The reference ``set_timesteps(num_inference_steps=N)`` builds N sigmas and
    N-1 timesteps; this wrapper takes the *denoise step count* directly, which
    is what a sampler actually schedules.
    """
    if num_denoise_steps < 1:
        raise ValueError(f"num_denoise_steps must be >= 1, got {num_denoise_steps}.")
    base = np.linspace(1.0, 0.0, num_denoise_steps + 1, dtype=np.float64)
    sigmas = shift * base / (1.0 + (shift - 1.0) * base)
    return _unique_consecutive(sigmas).astype(np.float32)

def mlx_h3_checkpoint_vsa_capable(checkpoint_dir: str | Path) -> bool:
    """True when a saved MLX H3 checkpoint retained the VSA gate projections."""
    manifest_path = Path(checkpoint_dir) / H3_MANIFEST_FILENAME
    if not manifest_path.is_file():
        return False
    manifest = json.loads(manifest_path.read_text())
    if bool((manifest.get("vsa") or {}).get("capable")):
        return True
    quantized_keys = manifest.get("quantized_keys") or {}
    return any("attn.to_gate_compress.weight" in key for key in quantized_keys)

def mlx_h3_dit_from_diffusers_safetensors(
    transformer_path: str | Path,
    config: dict[str, Any] | None = None,
    *,
    dtype: str = "fp16",
    num_blocks: int | None = None,
    quantization: str | MLXQuantizationSpec | None = None,
    adaln_cache_timesteps: np.ndarray | None = None,
    include_vsa: bool = False,
) -> MLXMiniMaxH3DiT:
    """Load the released H3 transformer (diffusers layout) into MLX.

    ``transformer_path`` is the ``transformer/`` directory of the HF repo (or
    a single safetensors file, e.g. a student checkpoint). ``config`` defaults
    to ``config.json`` next to the weights. fp32-release modules stay fp32;
    attention/FFN matrices are quantized when ``quantization`` is set.

    When ``adaln_cache_timesteps`` is provided, AdaLN tables are computed one
    block at a time while the checkpoint is read. The 13B projection weights
    are never retained together and are omitted from the returned model. This
    is the memory-bounded build path for the released 33B student.

    ``include_vsa`` keeps ``attn.to_gate_compress`` matrices and quantizes them
    with the same affine grid as the other linear weights. Dense conversion
    continues to drop them.
    """
    import mlx.core as mx
    transformer_path = Path(transformer_path)
    if config is None:
        config_path = transformer_path / "config.json" if transformer_path.is_dir() else transformer_path.with_name(
            "config.json")
        config = json.loads(Path(config_path).read_text())
    total_blocks = int(config["num_layers"])
    if num_blocks is None:
        num_blocks = total_blocks
    num_refiner = int(config["num_refiner_layers"])
    cast_dtype = {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[dtype]
    spec = MLXQuantizationSpec.from_name(quantization) if (quantization is None
                                                           or isinstance(quantization, str)) else quantization
    ensure_quantization_supported(spec)

    weights: dict[str, Any] = {}
    blocks: list[dict[str, Any] | None] = [None] * num_blocks
    refiner: list[dict[str, Any] | None] = [None] * num_refiner
    cache_timesteps = None
    temb = None
    cached_block_tables: list[tuple[Any, ...] | None] | None = None
    pending_adaln: dict[int, dict[str, Any]] = {}

    def assign(key: str, value) -> None:
        if key.startswith("transformer_blocks."):
            _, index_str, sub = key.split(".", 2)
            index = int(index_str)
            if index < num_blocks:
                block = blocks[index]
                if block is None:
                    block = {}
                    blocks[index] = block
                block[sub] = value
        elif key.startswith("token_refiner.refiner_blocks."):
            _, _, index_str, sub = key.split(".", 3)
            index = int(index_str)
            refiner_block = refiner[index]
            if refiner_block is None:
                refiner_block = {}
                refiner[index] = refiner_block
            refiner_block[sub] = value
        else:
            weights[key] = value

    if adaln_cache_timesteps is not None:
        cache_timesteps = np.unique(np.asarray(adaln_cache_timesteps, dtype=np.float32))
        for shard in _safetensors_shards(transformer_path):
            shard_arrays = mx.load(str(shard))
            for key, source in shard_arrays.items():
                if not key.startswith("time_embedder."):
                    continue
                keep_fp32 = key.split(".", 1)[0] in FP32_MODULE_PREFIXES
                target_dtype = mx.float32 if keep_fp32 else cast_dtype
                assign(key, _load_array(source, target_dtype))
            del shard_arrays
        required_time_keys = {
            "time_embedder.linear_1.weight",
            "time_embedder.linear_1.bias",
            "time_embedder.linear_2.weight",
            "time_embedder.linear_2.bias",
        }
        missing_time_keys = sorted(required_time_keys - weights.keys())
        if missing_time_keys:
            raise KeyError(f"Missing time-embedder weights needed for the AdaLN cache: {missing_time_keys}")
        timestep_rows = mx.array(cache_timesteps)
        t_freq = timestep_embedding(timestep_rows, int(config["freq_dim"])).astype(
            weight_dtype(weights["time_embedder.linear_1.weight"]))
        temb = linear(t_freq, weights["time_embedder.linear_1.weight"], weights["time_embedder.linear_1.bias"])
        temb = linear(silu(temb), weights["time_embedder.linear_2.weight"], weights["time_embedder.linear_2.bias"])
        mx.eval(temb)
        cached_block_tables = [None] * num_blocks

    for shard in _safetensors_shards(transformer_path):
        shard_arrays = mx.load(str(shard))
        for key, source in shard_arrays.items():
            if _is_ignored_dense_key(key, include_vsa=include_vsa):
                continue
            if temb is not None and key.startswith("time_embedder."):
                continue
            if key.startswith("transformer_blocks."):
                index = int(key.split(".")[1])
                if index >= num_blocks:
                    continue
            if key.startswith("rope."):
                continue  # non-persistent analytic buffer, rebuilt on the fly
            keep_fp32 = key.split(".", 1)[0] in FP32_MODULE_PREFIXES
            target_dtype = mx.float32 if keep_fp32 else cast_dtype
            array = _load_array(source, target_dtype)
            if temb is not None and ".adaln_proj.linear." in key:
                _, index_str, sub = key.split(".", 2)
                index = int(index_str)
                block_pending = pending_adaln.setdefault(index, {})
                block_pending[sub] = array
                if {"adaln_proj.linear.weight", "adaln_proj.linear.bias"} <= block_pending.keys():
                    tables = _adaln_tables(block_pending, temb)
                    mx.eval(tables)
                    assert cached_block_tables is not None
                    cached_block_tables[index] = tables
                    block = blocks[index]
                    if block is None:
                        block = {}
                        blocks[index] = block
                    block["adaln_proj.linear.weight"] = None
                    block["adaln_proj.linear.bias"] = None
                    del pending_adaln[index]
                continue
            if spec is not None and _is_quantizable(key, include_vsa=include_vsa) and not keep_fp32:
                value = quantize_matrix(array, spec)
                del array
            else:
                value = array
            _eval_value(value)
            assign(key, value)
        del shard_arrays

    cache = None
    if temb is not None:
        if pending_adaln:
            raise KeyError(f"Incomplete AdaLN projection pairs for blocks {sorted(pending_adaln)}")
        assert cache_timesteps is not None and cached_block_tables is not None
        missing_cache_blocks = [index for index, tables in enumerate(cached_block_tables) if tables is None]
        if missing_cache_blocks:
            raise KeyError(f"Missing AdaLN cache tables for blocks {missing_cache_blocks}")
        shift_scale = linear(
            silu(temb).astype(weight_dtype(weights["norm_out.linear.weight"])),
            weights["norm_out.linear.weight"],
            weights["norm_out.linear.bias"],
        )
        norm_out_shift, norm_out_scale = mx.split(shift_scale, 2, axis=-1)
        mx.eval(norm_out_shift, norm_out_scale)
        cache = MiniMaxH3StepCache(
            timesteps=cache_timesteps,
            block_tables=[tables for tables in cached_block_tables if tables is not None],
            norm_out_shift=norm_out_shift,
            norm_out_scale=norm_out_scale,
        )

    if any(block is None for block in blocks):
        missing = [i for i, block in enumerate(blocks) if block is None]
        raise KeyError(f"Missing transformer block weights for indices {missing}.")
    if any(block is None for block in refiner):
        raise KeyError("Missing token refiner weights.")
    loaded_blocks = [block for block in blocks if block is not None]
    if include_vsa:
        missing_gates = [
            index for index, block in enumerate(loaded_blocks) if "attn.to_gate_compress.weight" not in block
        ]
        if missing_gates:
            raise KeyError("VSA conversion requested but gate projections are missing for transformer blocks "
                           f"{missing_gates}. The source checkpoint must contain "
                           f"`transformer_blocks.*.{VSA_GATE_KEY_SUFFIX}`.")
        logger.info("Retained %d VSA gate projections (%s).", len(loaded_blocks), VSA_GATE_KEY_SUFFIX)
    dit = MLXMiniMaxH3DiT(
        weights,
        loaded_blocks,
        [block for block in refiner if block is not None],
        config,
    )
    dit._adaln_cache = cache
    return dit