minimax_h3_pipeline ¶

def decode_audio(self, audio_rows: np.ndarray, *, num_frames: int) -> np.ndarray:
    """Normalized packed audio rows -> stereo waveform (2, S) fp32 in [-1, 1]."""
    import mlx.core as mx

    from fastvideo.mlx_runtime.minimax_h3_audio_vae import mlx_h3_audio_vae_from_dir

    num_audio_latents = audio_latent_num_frames(align_num_frames(num_frames))
    latents = unpack_audio_tokens(audio_rows, num_audio_latents)
    vae = mlx_h3_audio_vae_from_dir(self.model_root / "audio_vae", include_encoder=False)
    z = vae.denormalize_latents(mx.array(latents))
    waveform = np.asarray(vae.decode(z))[:, 0, :]  # (B, 1, S) -> (B, S)
    del vae, z
    _cleanup_mlx()
    # Keep audio at least as long as the final video packet. Rounding down
    # by a fractional sample makes ffmpeg's ``-shortest`` drop frame 124.
    expected_samples = _audio_sample_count(align_num_frames(num_frames))
    if waveform.shape[-1] < expected_samples:
        waveform = np.pad(waveform, ((0, 0), (0, expected_samples - waveform.shape[-1])))
    return np.clip(waveform[:, :expected_samples], -1.0, 1.0)

def decode_video(self,
                 video_rows: np.ndarray,
                 *,
                 height: int,
                 width: int,
                 num_frames: int,
                 tiled: bool = True) -> np.ndarray:
    """Normalized packed rows -> (T, H, W, 3) uint8 frames."""
    import mlx.core as mx

    from fastvideo.mlx_runtime.minimax_h3_video_vae import mlx_h3_video_vae_from_dir

    geometry = self.resolve_geometry(height, width, num_frames, enforce_duration=False)
    if self.video_decode_backend == "taeh3":
        from fastvideo.mlx_runtime.minimax_h3_taeh3 import decode_latents_taeh3_mlx

        if self._dit_in_channels != 24:
            raise ValueError("TAEH3 requires a 24-channel H3 checkpoint.")
        latents = unpatchify_video_tokens(video_rows, geometry["latent_frame_count"], geometry["latent_height"],
                                          geometry["latent_width"], self._dit_in_channels, self._dit_patch_size)
        pixels = decode_latents_taeh3_mlx(latents,
                                          checkpoint_path=self.taeh3_checkpoint,
                                          dtype=self.vae_dtype,
                                          chunk_size=self.taeh3_chunk_size)
        frames = (pixels[0] * 255.0).astype(np.uint8)
        if frames.shape != (geometry["num_frames"], height, width, 3):
            raise RuntimeError(f"TAEH3 produced unexpected frame shape: {frames.shape}")
        _cleanup_mlx()
        return frames
    vae = mlx_h3_video_vae_from_dir(self.model_root / "vae", include_encoder=False, storage_dtype=self.vae_dtype)
    expected_height = height // vae.spatial_compression_ratio
    expected_width = width // vae.spatial_compression_ratio
    if (geometry["latent_height"], geometry["latent_width"]) != (expected_height, expected_width):
        raise RuntimeError("H3 pipeline/VAE spatial compression mismatch: "
                           f"pipeline={(geometry['latent_height'], geometry['latent_width'])}, "
                           f"VAE={(expected_height, expected_width)}.")
    if vae.latent_channels != self._dit_in_channels:
        raise RuntimeError(
            f"H3 DiT/VAE latent-channel mismatch: DiT={self._dit_in_channels}, VAE={vae.latent_channels}.")
    latents = unpatchify_video_tokens(
        video_rows,
        geometry["latent_frame_count"],
        geometry["latent_height"],
        geometry["latent_width"],
        vae.latent_channels,
        self._dit_patch_size,
    )
    z = mx.array(latents)
    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))
    pixels = np.clip(np.asarray(vae.denormalize_pixels(decoded)), 0.0, 1.0)
    del vae, decoded, z
    _cleanup_mlx()
    frames = (pixels[0].transpose(1, 2, 3, 0) * 255.0).astype(np.uint8)  # (T, H, W, C)
    if frames.shape[0] != geometry["num_frames"]:
        raise RuntimeError(f"decoded {frames.shape[0]} frames, expected {geometry['num_frames']}")
    return frames

def denoise(
    self,
    text_rows: np.ndarray,
    token_tags: np.ndarray,
    *,
    height: int,
    width: int,
    num_frames: int,
    audio_num_frames: int | None = None,
    video_temporal_scale: float = 1.0,
    seed: int,
    num_steps: int = 4,
    dit: Any | None = None,
    vsa_config: MiniMaxH3VSAConfig | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Denoise joint latents; returns (normalized video rows, audio rows)."""
    import mlx.core as mx

    geometry = self.resolve_geometry(height, width, num_frames, enforce_duration=audio_num_frames is None)
    audio_frames = geometry["num_frames"] if audio_num_frames is None else align_num_frames(audio_num_frames)

    owned_dit = dit is None
    if owned_dit:
        _validate_checkpoint_step_ladder(self.dit_checkpoint, num_steps)
        t0 = time.perf_counter()
        dit = load_mlx_h3_checkpoint(self.dit_checkpoint)
        logger.info("Loaded MLX H3 DiT from %s in %.1fs", self.dit_checkpoint, time.perf_counter() - t0)
    if vsa_config is not None:
        dit.configure_vsa(vsa_config)
    if hasattr(dit, "reset_vsa_stats"):
        dit.reset_vsa_stats()

    layout = build_packed_layout(
        len(token_tags),
        geometry["latent_frame_count"],
        geometry["latent_height"],
        geometry["latent_width"],
        audio_latent_num_frames(audio_frames),
        patch_size=dit.patch_size,
        text_token_tags=np.asarray(token_tags, dtype=np.int64),
        video_temporal_scale=video_temporal_scale,
    )
    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)
    # The keyframe-noise timestep (0.999) is only exercised by FL2VA/Ref2VA
    # conditioning rows; those modes recompute the ladder before denoise.

    cache = getattr(dit, "_adaln_cache", None)
    if cache is None:
        dit.precompute_adaln(union, drop_weights=True)
    elif not np.array_equal(cache.timesteps.astype(np.float32), union):
        extra = np.setdiff1d(union, cache.timesteps)
        logger.info("Recomputing AdaLN cache for %d-step ladder (extra timesteps %s).", num_steps, extra)
        dit.precompute_adaln(union, drop_weights=True)

    video_key, audio_key = mx.random.split(mx.random.key(seed))
    target_video_rows = int(layout.video_indices.shape[0] - layout.num_condition_video_rows)
    target_audio_rows = int(layout.audio_indices.shape[0] - layout.num_condition_audio_rows)
    x_v = mx.random.normal((target_video_rows, dit.patch_dim), key=video_key)
    x_a = mx.random.normal((target_audio_rows, dit.audio_in_channels), key=audio_key)
    text = mx.array(text_rows.astype(np.float32))
    dit_forward_s = 0.0

    for step_index in range(num_steps):
        video_t = float(video_scheduler.timesteps[step_index])
        audio_t = float(audio_scheduler.timesteps[step_index])
        unique, inverse = build_row_timesteps(
            layout,
            video_timestep=video_t,
            audio_timestep=audio_t,
            condition_video_timestep=max(video_t, MINIMAX_H3_KEYFRAME_NOISE_AUG),
            condition_audio_timestep=1.0,
        )
        step_started = time.perf_counter()
        video_velocity, audio_velocity = dit.forward_with_cache(
            x_v,
            x_a,
            text,
            layout=layout,
            step_timesteps=unique,
            row_timestep_inverse=inverse,
            step_index=step_index,
        )
        # Only target rows are being denoised (no conditions in T2VA).
        video_velocity = video_velocity[layout.num_condition_video_rows:]
        audio_velocity = audio_velocity[layout.num_condition_audio_rows:]
        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)
        dit_forward_s += time.perf_counter() - step_started

    self.last_dit_forward_s = dit_forward_s
    stats = getattr(dit, "last_vsa_stats", None)
    self.last_vsa_stats = None if stats is None else {
        "enabled": bool(getattr(dit.vsa_config, "enabled", False)),
        "configured_sparsity": stats.configured_sparsity,
        "layer_sparsity": stats.layer_sparsity,
        "achieved_sparsity": stats.achieved_sparsity,
        "tile_size": stats.tile_size,
        "prefix_mode": stats.prefix_mode,
        "impl": stats.impl,
        "num_prefix_tiles": stats.num_prefix_tiles,
        "num_video_tiles": stats.num_video_tiles,
        "video_keep": stats.video_keep,
        "dense_fallback_reason": stats.dense_fallback_reason,
        "attention_calls": stats.attention_calls,
        "sparse_calls": stats.sparse_calls,
        "impl_counts": stats.impl_counts,
        "fallback_reasons": stats.fallback_reasons,
        "checkpoint_vsa_capable": bool(getattr(dit, "vsa_capable", False)),
    }

    video_rows = np.asarray(x_v, dtype=np.float32)
    audio_rows = np.asarray(x_a, dtype=np.float32)
    if owned_dit:
        del dit
        _cleanup_mlx()
    return video_rows, audio_rows

def encode_prompt(self, prompt: str) -> tuple[np.ndarray, np.ndarray]:
    """Returns (hidden states (S, hidden), token tags). Uses the cache or
    the streamed conditioner."""
    cache_key = None
    if self.prompt_cache_dir is not None:
        cache_key = prompt_cache_path(self.prompt_cache_dir, self.model_root, prompt)
        if cache_key.exists():
            data = np.load(cache_key)
            logger.info("Loaded prompt embeddings from cache %s", cache_key)
            return data["hidden_states"], data["token_tags"]

    conditioner = self._load_conditioner()
    hidden, tags = conditioner.encode_prompt(prompt)
    conditioner.close()
    _cleanup_mlx()
    if cache_key is not None:
        cache_key.parent.mkdir(parents=True, exist_ok=True)
        tmp_cache = cache_key.with_name(f".{cache_key.name}.tmp")
        try:
            with tmp_cache.open("wb") as handle:
                np.savez(handle, hidden_states=hidden, token_tags=tags)
            tmp_cache.replace(cache_key)
        finally:
            tmp_cache.unlink(missing_ok=True)
    return hidden, tags

def mux(self,
        frames: np.ndarray,
        waveform: np.ndarray,
        output_path: str | Path,
        fps: int = MINIMAX_H3_FPS,
        sample_rate: int = 32000) -> Path:
    """H.264 video + AAC stereo audio, A/V durations within one frame."""
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    tmp_video = output_path.with_suffix(".tmp.mp4")
    tmp_audio = output_path.with_suffix(".tmp.wav")

    pcm = (np.clip(waveform.T, -1.0, 1.0) * 32767.0).astype("<i2")  # (S, 2)
    import wave

    try:
        with wave.open(str(tmp_audio), "wb") as handle:
            handle.setnchannels(2)
            handle.setsampwidth(2)
            handle.setframerate(sample_rate)
            handle.writeframes(pcm.tobytes())

        height, width = frames.shape[1:3]
        ffmpeg = shutil.which("ffmpeg")
        if ffmpeg is None:
            raise RuntimeError("ffmpeg is required for MP4 muxing.")
        subprocess.run(
            [
                ffmpeg,
                "-y",
                "-loglevel",
                "error",
                "-f",
                "rawvideo",
                "-pix_fmt",
                "rgb24",
                "-s",
                f"{width}x{height}",
                "-r",
                str(fps),
                "-i",
                "-",
                "-i",
                str(tmp_audio),
                "-c:v",
                "libx264",
                "-preset",
                "medium",
                "-crf",
                "18",
                "-pix_fmt",
                "yuv420p",
                "-c:a",
                "aac",
                "-b:a",
                "192k",
                "-shortest",
                "-movflags",
                "+faststart",
                str(tmp_video),
            ],
            input=frames.tobytes(),
            check=True,
        )
        tmp_video.replace(output_path)
    finally:
        tmp_audio.unlink(missing_ok=True)
        tmp_video.unlink(missing_ok=True)
    return output_path

@staticmethod
def resolve_geometry(
    height: int,
    width: int,
    num_frames: int,
    *,
    enforce_duration: bool = True,
) -> dict[str, int]:
    """Explicit canvases pass through (positive multiples of 32); the
    aspect-ratio resolver only applies when dimensions are omitted."""
    if height <= 0 or width <= 0 or height % 32 or width % 32:
        raise ValueError(f"H3 canvas must be positive multiples of 32, got {height}x{width}.")
    aligned_frames = align_num_frames(num_frames)
    latent_frames = video_latent_num_frames(aligned_frames)
    duration = aligned_frames / MINIMAX_H3_FPS
    if enforce_duration and not 5.0 <= duration <= 15.0:
        raise ValueError(f"H3 generates 5-15 s at {MINIMAX_H3_FPS} fps; {aligned_frames} frames "
                         f"is {duration:.2f} s.")
    return {
        "height": height,
        "width": width,
        "num_frames": aligned_frames,
        "latent_frame_count": latent_frames,
        "latent_height": height // 16,
        "latent_width": width // 16,
    }