sada9

Running on Zero

sada9 / h3_a2v_blocks.py
Transfer Bot
Moved to Hugging Face automatically
c3da9f3 16 days ago
31.3 kB
"""Audio-conditioned `fl2va` / `t2va` for MiniMax-H3: a **given** soundtrack, with or without a first frame.
MiniMax-H3 denoises video and audio as one packed sequence, under two schedules (`shift = 12` for video, `shift = 3`
for audio). `fl2va` generates both halves: the keyframe anchors the leading video rows and the audio rows are drawn
from noise like every other generated row, so the model invents its own soundtrack. `t2va` is the same thing without
the keyframe.
This module turns those audio rows into *given* rows. The waveform is encoded by the audio VAE and normalized the way
the pipeline normalizes any soundtrack, and the resulting rows are imposed on every step in one of two ways:
* **locked** — the clean rows are written back after every scheduler step, and the layout tells the transformer they
  sit at `t = 1.0` (clean). This is how `ref2va` presents a *reference* soundtrack, except that here the rows are the
  target ones, so they are aligned with the generated frames rather than parked ahead of them. The model then reads
  clean audio beside noisy video, which is a distribution it has seen for references and not for targets.
* **blended** — the rows are re-noised to *that step's* audio sigma, off the audio schedule (shift 3, not the video
  one), and written back after every step. Ordinary diffusion inpainting: at every step the transformer reads audio
  at exactly the noise level its timestep claims, and the last step leaves the clean waveform in place.
`off` runs the plain workflow and is the control leg.
Nothing else about the request changes: the video noise is still the request generator's first draw, so the same seed
gives the same starting video latents under all three modes and the legs are comparable.
None of the three conditioning blocks reads the keyframe — they only ever touch audio rows — so the same three go
into a `t2va`-shaped core as into the `fl2va` one, and `MiniMaxH3A2VAutoDenoiseStep` picks between the two on
`image`, the way main's own `MiniMaxH3AutoDenoiseStep` does.
Written against `diffusers` **main**, whose `fl2va` layout splits the conditioning latents (`PrepareConditionLatents`
/ `FL2VAPrepareLatents`) out of the latent draw and pulls `AfterDenoise` inside the core. Main rather than the
MiniMax-H3 follow-up PR because that PR branched before the LoRA work landed and carries no
`MiniMaxH3LoraLoaderMixin` at all, so no turbo LoRA can be read against it.
"""

from __future__ import annotations

import torch
import torch.nn.functional as F

from diffusers import attention_backend
from diffusers.models import AutoencoderKLMiniMaxH3Audio
from diffusers.modular_pipelines.minimax_h3.before_denoise import (
    MiniMaxH3FL2VAPrepareLatentsStep,
    MiniMaxH3NoKeyframeAnchorsStep,
    MiniMaxH3PrepareConditionLatentsStep,
    MiniMaxH3PrepareLatentsStep,
    MiniMaxH3PrepareLayoutStep,
    MiniMaxH3SetTimestepsStep,
)
from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3ResizeStep
from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
from diffusers.modular_pipelines.minimax_h3.denoise import (
    MiniMaxH3DenoiseLoopWrapper,
    MiniMaxH3LoopDenoiser,
    MiniMaxH3LoopSchedulerStep,
)
from diffusers.modular_pipelines.minimax_h3.encoders import MiniMaxH3KeyframeVaeEncoderStep
from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import MiniMaxH3DecodeStep
from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
from diffusers.modular_pipelines.modular_pipeline import (
    BlockState,
    ConditionalPipelineBlocks,
    ModularPipelineBlocks,
    PipelineState,
    SequentialPipelineBlocks,
)
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam

from h3_long_memory import A2VBlockStop


LOCKED, BLENDED, OFF = "locked", "blended", "off"
AUDIO_CONDITION_MODES = (LOCKED, BLENDED, OFF)


def normalize_waveform(waveform: torch.Tensor, sample_rate: int, target_sample_rate: int) -> torch.Tensor:
    """A `(channels, num_samples)` waveform as MiniMax-H3's stereo, float32, `target_sample_rate` soundtrack."""
    waveform = torch.as_tensor(waveform).to(torch.float32)
    if waveform.ndim == 1:
        waveform = waveform[None]
    if waveform.ndim != 2:
        raise ValueError(f"The soundtrack must be a `(channels, num_samples)` waveform, got {tuple(waveform.shape)}.")
    if waveform.shape[0] > 2:
        # Anything past stereo is downmixed rather than refused: an uploaded file is not a curated reference.
        waveform = waveform.mean(dim=0, keepdim=True)
    if sample_rate != target_sample_rate:
        import torchaudio

        waveform = torchaudio.functional.resample(waveform, sample_rate, target_sample_rate)
    if waveform.shape[0] == 1:
        waveform = waveform.expand(2, -1).contiguous()
    return waveform


class MiniMaxH3AudioConditionStep(ModularPipelineBlocks):
    model_name = "minimax-h3"

    @property
    def description(self) -> str:
        return (
            "Encodes the given soundtrack into the audio rows of the packed sequence, and keeps the noise those rows "
            "were drawn from. It runs after the noise draw so the request's generator is untouched: the video noise "
            "is still its first draw, and the audio draw still happens, which is what makes the three conditioning "
            "modes comparable at one seed."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("audio_vae", AutoencoderKLMiniMaxH3Audio)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                name="given_audio",
                type_hint=torch.Tensor,
                description="The given soundtrack, a `(channels, num_samples)` waveform. None runs plain `fl2va`.",
            ),
            InputParam(
                name="given_audio_sample_rate",
                type_hint=int,
                description="The rate `given_audio` carries its samples at. Left out, the audio VAE's own.",
            ),
            InputParam(
                name="audio_condition_mode",
                type_hint=str,
                default=BLENDED,
                description="`locked`, `blended` or `off`.",
            ),
            InputParam(
                name="num_audio_latents",
                type_hint=int,
                required=True,
                description="Number of audio latents per channel.",
            ),
            InputParam(
                name="audio_latents",
                type_hint=torch.Tensor,
                required=True,
                description="The audio rows of the packed sequence, as drawn from the request's generator.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "audio_latents",
                type_hint=torch.Tensor,
                description="The audio rows the loop starts from: the given ones when locked, the drawn noise else.",
            ),
            OutputParam(
                "given_audio_rows",
                type_hint=torch.Tensor,
                description="The encoded, normalized soundtrack in channel-major row layout, or None.",
            ),
            OutputParam(
                "audio_noise_rows",
                type_hint=torch.Tensor,
                description="The noise the audio rows were drawn from, which `blended` re-noises against.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        if block_state.audio_condition_mode not in AUDIO_CONDITION_MODES:
            raise ValueError(
                f"`audio_condition_mode` must be one of {AUDIO_CONDITION_MODES}, got "
                f"{block_state.audio_condition_mode!r}."
            )

        block_state.audio_noise_rows = block_state.audio_latents.clone()
        block_state.given_audio_rows = None

        if block_state.given_audio is not None and block_state.audio_condition_mode != OFF:
            sample_rate = block_state.given_audio_sample_rate or components.audio_sampling_rate
            waveform = normalize_waveform(block_state.given_audio, sample_rate, components.audio_sampling_rate)

            # The audio VAE right-pads to a whole hop, so a waveform of exactly `num_audio_latents` hops encodes to
            # exactly the rows the layout reserved. Silence pads a soundtrack shorter than the video.
            num_samples = block_state.num_audio_latents * components.audio_vae.hop_length
            if waveform.shape[-1] < num_samples:
                waveform = F.pad(waveform, (0, num_samples - waveform.shape[-1]))
            waveform = waveform[:, :num_samples]

            # Channel-major rows: the two stereo channels are two batch items of the mono audio VAE, and soundtracks
            # take the posterior *mean* rather than a sample, the way the reference encoder does.
            #
            # Under `native` attention explicitly: the audio VAE's encoder carries a causal-attention `pre_block`, it
            # is float32 by `_keep_in_fp32_modules`, and the backend a pipeline selects for its transformer is global
            # — cuDNN has no float32 SDPA kernel and aborts here. Nothing else in MiniMax-H3 runs this path, since
            # only *encoding* audio reaches it.
            with attention_backend("native"):
                posterior = components.audio_vae.encode(waveform.to(device)[:, None], return_dict=False)[0]
            latents = posterior.mode().float().transpose(1, 2)
            mean = torch.tensor(components.audio_vae.config.latents_mean, device=latents.device).view(1, 1, -1)
            std = torch.tensor(components.audio_vae.config.latents_std, device=latents.device).view(1, 1, -1)
            rows = ((latents - mean) / std).reshape(-1, components.audio_latent_channels)
            if rows.shape != block_state.audio_latents.shape:
                raise ValueError(
                    f"The soundtrack encodes to {tuple(rows.shape)} audio rows but the layout holds "
                    f"{tuple(block_state.audio_latents.shape)}."
                )
            block_state.given_audio_rows = rows.to(block_state.audio_latents.dtype)

            if block_state.audio_condition_mode == LOCKED:
                # Clean from the first forward on, matching the `t = 1.0` the timestep plan will claim for these rows.
                block_state.audio_latents = block_state.given_audio_rows.clone()

        self.set_block_state(state, block_state)
        return components, state


class MiniMaxH3A2VSetTimestepsStep(MiniMaxH3SetTimestepsStep):
    model_name = "minimax-h3"

    @property
    def description(self) -> str:
        return (
            "The `fl2va` timestep plan, with the audio rows presented as clean (`t = 1.0`) when the soundtrack is "
            "locked. That is the presentation `ref2va` gives a reference soundtrack; `blended` and `off` leave the "
            "audio rows on their own schedule, where their content really is at the level their timestep claims."
        )

    @property
    def inputs(self) -> list[InputParam]:
        return super().inputs + [
            InputParam(name="audio_condition_mode", type_hint=str, default=BLENDED, description="See the block above."),
            InputParam(name="given_audio_rows", type_hint=torch.Tensor, description="The encoded soundtrack, or None."),
        ]

    @torch.no_grad()
    def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
        components.audio_scheduler.set_timesteps(block_state.num_inference_steps, device=device)
        block_state.timesteps = components.scheduler.timesteps
        block_state.audio_timesteps = components.audio_scheduler.timesteps

        locked = block_state.audio_condition_mode == LOCKED and block_state.given_audio_rows is not None
        # The rows are *given*, not *prepended*: `num_condition_audio_rows` stays 0 everywhere else, so the loop still
        # steps them (the write-back overrules it) and the unpack step still keeps them.
        num_clean_audio_rows = int(block_state.audio_indices.numel()) if locked else 0

        block_state.row_timestep_plan = [
            tuple(
                tensor.to(device)
                for tensor in self.build_row_timesteps(
                    block_state.video_indices,
                    block_state.audio_indices,
                    block_state.num_condition_video_rows,
                    num_clean_audio_rows,
                    block_state.text_indices.numel(),
                    float(timestep),
                    float(audio_timestep),
                    max(float(timestep), components.keyframe_noise_aug),
                    1.0,
                )
            )
            for timestep, audio_timestep in zip(block_state.timesteps, block_state.audio_timesteps)
        ]

        self.set_block_state(state, block_state)
        return components, state


class MiniMaxH3A2VLoopSchedulerStep(MiniMaxH3LoopSchedulerStep):
    model_name = "minimax-h3"

    @property
    def description(self) -> str:
        return (
            "Steps both modalities, then imposes the given soundtrack on the audio rows: clean when locked, re-noised "
            "to the audio schedule's next sigma when blended. The scheduler is still stepped over those rows so its "
            "step index keeps up with the loop; the write-back is what the next forward reads."
        )

    @property
    def inputs(self) -> list[InputParam]:
        return super().inputs + [
            InputParam(name="audio_condition_mode", type_hint=str, default=BLENDED, description="See the block above."),
            InputParam(name="given_audio_rows", type_hint=torch.Tensor, description="The encoded soundtrack, or None."),
            InputParam(
                name="audio_noise_rows",
                type_hint=torch.Tensor,
                description="The noise the audio rows were drawn from.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components: MiniMaxH3ModularPipeline, state: BlockState, i: int, t: torch.Tensor):
        components, block_state = super().__call__(components, state, i, t)

        given = block_state.given_audio_rows
        if given is None or block_state.audio_condition_mode == OFF:
            return components, block_state

        if block_state.audio_condition_mode == LOCKED:
            # `.clone()` is load-bearing: `Tensor.to` is a no-op when the dtype and device already match, so without
            # it the rows *are* `given_audio_rows`, and the next step's in-place `audio_latents[0:] = ...` overwrites
            # the soundtrack itself. The locked leg then drifts into the model's own audio from step 1 on.
            block_state.audio_latents = given.to(block_state.audio_latents).clone()
        else:
            # The rows the *next* forward reads are at the next entry of the audio schedule; past the last one the
            # schedule has reached sigma = 0, i.e. `t = 1`, which returns the soundtrack itself.
            audio_timesteps = block_state.audio_timesteps
            next_timestep = float(audio_timesteps[i + 1]) if i + 1 < audio_timesteps.numel() else 1.0
            block_state.audio_latents = components.audio_scheduler.scale_noise(
                given.to(block_state.audio_latents),
                next_timestep,
                block_state.audio_noise_rows.to(block_state.audio_latents),
            )
        return components, block_state


class A2VWindowStop(BaseException):
    """Raised inside the denoise loop when a resumable window is out of time or steps.
    A `BaseException` so diffusers' per-block `except Exception` handlers don't log this planned stop as an
    "Error in block" with a traceback at every level; `run_window` catches it by name."""


class A2VWindow:
    """One ZeroGPU window of a resumable render, handed to the loop through the `a2v_window` pipeline input.
    * `start` — the absolute schedule index this window resumes at. Iterations before it are skipped outright:
      `MiniMaxH3Scheduler.step` finds its own index from the timestep it is handed when `_step_index` is unset, and
      the audio scheduler is stepped with `audio_timesteps[i]`, so both schedules land on the right entry.
    * `resume_rows` — the target video rows (float32, conditioning rows excluded) as of `start`. Written over the
      freshly drawn noise right before the first executed step. The audio rows need no carrying: in `locked` mode
      they are the given soundtrack at every step.
    * `stop(i, total)` — asked before every step; `True` ends the window *before* doing that step's work.
    * `tick(i)` — told after every completed step (per-step timing lives on the caller's side).
    * `after_last()` — asked once the final step has landed; `True` ends the window *before* the decode, so a long
      render's decode can get a window of its own (the next window resumes at `start == total`: no steps, decode).
    After a stop, `rows` holds the target rows to checkpoint and `stopped_at` the index the next window starts at.
    * `keep_final_rows` — when every step has landed (here or in an earlier window) and the loop ends *normally*, keep
      a CPU copy of the final target rows in `final_rows`. The decode that follows may itself stop between two clips,
      and the window that picks it up needs those rows to rebuild the latents it decodes.
    """

    def __init__(self, start=0, resume_rows=None, stop=None, tick=None, after_last=None, keep_final_rows=False):
        self.start = int(start or 0)
        self.resume_rows = resume_rows
        self.stop = stop
        self.tick = tick
        self.after_last = after_last
        self.keep_final_rows = bool(keep_final_rows)
        self.final_rows = None
        self.rows = None
        self.stopped_at = None
        self.executed = 0
        # Set when the stop landed between two transformer blocks of step `stopped_at`.
        self.partial_block = None
        self.partial_hidden = None


class MiniMaxH3A2VDenoiseStep(MiniMaxH3DenoiseLoopWrapper):
    model_name = "minimax-h3"
    block_classes = [MiniMaxH3LoopDenoiser, MiniMaxH3A2VLoopSchedulerStep]
    block_names = ["denoiser", "update"]

    @property
    def description(self) -> str:
        return (
            "Runs the `fl2va` denoising loop with the given soundtrack imposed on the audio rows every step. With an "
            "`a2v_window` it runs only a slice of the schedule — resuming from saved rows and stopping on a step "
            "boundary — so one render can be spread over several ZeroGPU windows."
        )

    @property
    def loop_inputs(self) -> list[InputParam]:
        return super().loop_inputs + [
            InputParam(
                name="a2v_window",
                type_hint=object,
                description="An `A2VWindow`, or None for an ordinary uninterrupted run.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        window = getattr(block_state, "a2v_window", None)
        if window is None:
            # Exactly the upstream loop.
            with self.progress_bar(total=len(block_state.timesteps)) as progress_bar:
                for i, t in enumerate(block_state.timesteps):
                    components, block_state = self.loop_step(components, block_state, i=i, t=t)
                    progress_bar.update()
            self.set_block_state(state, block_state)
            return components, state

        if block_state.audio_condition_mode != LOCKED:
            raise ValueError("Resumable windows only carry the video rows, so they need the `locked` audio mode.")

        num_condition_video_rows = int(block_state.num_condition_video_rows or 0)
        total = len(block_state.timesteps)
        if window.start > total:
            raise ValueError(f"The window starts at step {window.start}, past the schedule's {total}.")

        if window.start > 0:
            if window.resume_rows is None:
                raise ValueError("A window that resumes past step 0 needs the saved video rows.")
            target = block_state.latents[num_condition_video_rows:]
            rows = window.resume_rows.to(device=target.device, dtype=target.dtype)
            if rows.shape != target.shape:
                raise ValueError(
                    f"The saved video rows {tuple(rows.shape)} do not match this request's {tuple(target.shape)} — "
                    "the state file belongs to a different canvas or segment length."
                )
            block_state.latents[num_condition_video_rows:] = rows

        # Fresh schedulers for this call: `set_timesteps` already cleared their step index, and the first `step`
        # re-derives it from the timestep it is handed.
        with self.progress_bar(total=max(0, total - window.start)) as progress_bar:
            for i, t in enumerate(block_state.timesteps):
                if i < window.start:
                    continue
                if window.stop is not None and window.stop(i, total):
                    window.rows = block_state.latents[num_condition_video_rows:].detach().float().cpu().clone()
                    window.stopped_at = i
                    raise A2VWindowStop()
                try:
                    components, block_state = self.loop_step(components, block_state, i=i, t=t)
                except A2VBlockStop as stop:
                    # Out of time *inside* this step, between two transformer blocks. The latents have not been
                    # touched yet (the scheduler update runs after the forward), so the rows are still this step's
                    # input; the sequence entering `stop.block` goes with them so the next window finishes the step.
                    window.rows = block_state.latents[num_condition_video_rows:].detach().float().cpu().clone()
                    window.stopped_at = i
                    window.partial_block = stop.block if stop.hidden is not None else None
                    window.partial_hidden = stop.hidden
                    raise A2VWindowStop() from None
                window.executed += 1
                if window.tick is not None:
                    window.tick(i)
                progress_bar.update()
                if i == total - 1 and window.after_last is not None and window.after_last():
                    window.rows = block_state.latents[num_condition_video_rows:].detach().float().cpu().clone()
                    window.stopped_at = total
                    raise A2VWindowStop()
        if window.keep_final_rows:
            # The loop ended normally: these are the final target rows. A decode-only window never ran a step, so the
            # rows it was handed ARE the final ones.
            window.final_rows = (
                window.resume_rows
                if window.start >= total
                else block_state.latents[num_condition_video_rows:].detach().float().cpu().clone()
            )
        if window.start >= total:
            # A decode-only window (every step landed in earlier windows): no forward ran here, so the loop's
            # outputs were never written. `set_block_state` insists on them — give them explicit empties — and put
            # the soundtrack rows back exactly as the last step's update left them (`locked`: the soundtrack itself).
            for output in self.intermediate_outputs:
                if not hasattr(block_state, output.name):
                    setattr(block_state, output.name, None)
            given = block_state.given_audio_rows
            if given is not None and getattr(block_state, "audio_latents", None) is not None:
                block_state.audio_latents = given.to(block_state.audio_latents).clone()
        self.set_block_state(state, block_state)
        return components, state


class MiniMaxH3A2VCoreDenoiseStep(SequentialPipelineBlocks):
    model_name = "minimax-h3"
    block_classes = [
        MiniMaxH3PrepareLayoutStep,
        MiniMaxH3PrepareConditionLatentsStep,
        MiniMaxH3PrepareLatentsStep,
        MiniMaxH3FL2VAPrepareLatentsStep,
        MiniMaxH3AudioConditionStep,
        MiniMaxH3A2VSetTimestepsStep,
        MiniMaxH3A2VDenoiseStep,
        MiniMaxH3AfterDenoiseStep,
    ]
    block_names = [
        "prepare_layout",
        "prepare_condition_latents",
        "prepare_latents",
        "prepare_latents_fl2va",
        "audio_condition",
        "set_timesteps",
        "denoise",
        "after_denoise",
    ]

    @property
    def description(self) -> str:
        return (
            "Core denoising workflow for audio-conditioned `fl2va`: `MiniMaxH3FL2VACoreDenoiseStep` with the given "
            "soundtrack encoded into the audio rows and imposed on every step."
        )


class MiniMaxH3A2VT2VACoreDenoiseStep(SequentialPipelineBlocks):
    """`MiniMaxH3CoreDenoiseStep` with the same three audio-conditioning blocks inserted at the same points.
    None of the three reads the keyframe — they only ever touch audio rows — so the difference between this and the
    `fl2va` core is entirely the core they are inserted into: no `prepare_condition_latents`, no
    `MiniMaxH3FL2VAPrepareLatentsStep`, and `MiniMaxH3NoKeyframeAnchorsStep` in front to say the request anchors
    none. The packed sequence is `[text | target audio | target video]` rather than
    `[text | keyframe conditions | target audio | target video]`.
    """

    model_name = "minimax-h3"
    block_classes = [
        MiniMaxH3NoKeyframeAnchorsStep,
        MiniMaxH3PrepareLayoutStep,
        MiniMaxH3PrepareLatentsStep,
        MiniMaxH3AudioConditionStep,
        MiniMaxH3A2VSetTimestepsStep,
        MiniMaxH3A2VDenoiseStep,
        MiniMaxH3AfterDenoiseStep,
    ]
    block_names = [
        "no_keyframe_anchors",
        "prepare_layout",
        "prepare_latents",
        "audio_condition",
        "set_timesteps",
        "denoise",
        "after_denoise",
    ]

    @property
    def description(self) -> str:
        return (
            "Core denoising workflow for audio-conditioned `t2va`: `MiniMaxH3CoreDenoiseStep` with the given "
            "soundtrack encoded into the audio rows and imposed on every step."
        )


class MiniMaxH3A2VAutoBeforeEncodeStep(ConditionalPipelineBlocks):
    model_name = "minimax-h3"
    block_classes = [MiniMaxH3ResizeStep]
    block_names = ["keyframes"]
    block_trigger_inputs = ["image", "last_image"]
    default_block_name = None

    def select_block(self, **kwargs) -> str | None:
        if kwargs.get("image") is not None or kwargs.get("last_image") is not None:
            return "keyframes"
        return None

    @property
    def description(self) -> str:
        return (
            "Media preparation block.\n"
            " - `MiniMaxH3ResizeStep` runs when a keyframe is provided (`fl2va`), putting it onto the target canvas.\n"
            " - an audio-only request (`t2va`) skips this block, and the layout step resolves the canvas from the "
            "requested height and width alone."
        )


class MiniMaxH3A2VAutoVaeEncoderStep(ConditionalPipelineBlocks):
    model_name = "minimax-h3"
    block_classes = [MiniMaxH3KeyframeVaeEncoderStep]
    block_names = ["keyframes"]
    block_trigger_inputs = ["image", "last_image"]
    default_block_name = None

    def select_block(self, **kwargs) -> str | None:
        if kwargs.get("image") is not None or kwargs.get("last_image") is not None:
            return "keyframes"
        return None

    @property
    def description(self) -> str:
        return (
            "VAE encoder block.\n"
            " - `MiniMaxH3KeyframeVaeEncoderStep` runs when a keyframe is provided (`fl2va`).\n"
            " - an audio-only request (`t2va`) skips this block."
        )


class MiniMaxH3A2VAutoDenoiseStep(ConditionalPipelineBlocks):
    model_name = "minimax-h3"
    block_classes = [MiniMaxH3A2VCoreDenoiseStep, MiniMaxH3A2VT2VACoreDenoiseStep]
    block_names = ["fl2va", "t2va"]
    block_trigger_inputs = ["image", "last_image"]
    default_block_name = "t2va"

    def select_block(self, **kwargs) -> str | None:
        if kwargs.get("image") is not None or kwargs.get("last_image") is not None:
            return "fl2va"
        return None

    @property
    def description(self) -> str:
        return (
            "Denoise block, audio-conditioned either way.\n"
            " - the `fl2va` core runs when a keyframe is provided, against the `transformer` partition.\n"
            " - the `t2va` core runs otherwise, against the same partition.\n"
            "There is no `ref2va` branch: it would declare `transformer_ref`, the second checkpoint partition, which "
            "this Space must never load."
        )


class MiniMaxH3A2VGeneratorBlocks(SequentialPipelineBlocks):
    """The denoising half of the split deployment, audio-conditioned.
    `MiniMaxH3GeneratorBlocks` with its core denoise swapped for the audio-conditioned one, so the 62.14 GiB Qwen3-VL
    still lives in the conditioner Space and `prompt_embeds` / `text_token_tags` are still the whole wire format.
    """

    model_name = "minimax-h3"
    # Conditional wrappers of the same shape as main's own, branching on `image` the way `MiniMaxH3AutoDenoiseStep`
    # does — but two-way rather than three-way. The upstream auto wrappers carry a `ref2va` branch, which declares
    # `transformer_ref`, the second checkpoint partition; this Space must never load it, so the branch is not here.
    block_classes = [
        MiniMaxH3A2VAutoBeforeEncodeStep,
        MiniMaxH3A2VAutoVaeEncoderStep,
        MiniMaxH3A2VAutoDenoiseStep,
        MiniMaxH3DecodeStep,
    ]
    block_names = ["resize", "vae_encoder", "denoise", "decode"]

    @property
    def description(self) -> str:
        return (
            "The denoising half of a split MiniMax-H3 deployment, conditioned on a given soundtrack: the `fl2va` "
            "branch when a keyframe is given and the `t2va` branch when one is not, without their text-encoder step, "
            "with the audio rows imposed rather than generated in either."
        )

    @property
    def outputs(self):
        return [
            OutputParam.template("videos", description="The generated video."),
            OutputParam(
                "audio",
                type_hint=torch.Tensor,
                description="The soundtrack of the packed sequence, of shape `(1, 2, num_samples)`.",
            ),
            OutputParam("sampling_rate", type_hint=int, description="Sample rate of the soundtrack in Hz."),
        ]