"""The denoising half of a split MiniMax-H3 deployment, for the `ref2va` partition.

`MiniMaxH3Blocks` is cut at its `text_encoder` step: the 62 GiB Qwen3-VL conditioner runs in
`multimodalart/qwen3vl-conditioner` (a public embedding service this Space calls over the gradio API), everything
else — the `transformer_ref` partition, both autoencoders, the two schedulers — runs here. `prompt_embeds` +
`text_token_tags` is the whole wire format between the two halves.

Same layout as the character-swap H3 Space (`hugging-apps/minimax-h3-character-swap-lora`), which ships the same
subclass against the same diffusers classes.
"""

import torch

from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
from diffusers.modular_pipelines.minimax_h3.encoders import MiniMaxH3Ref2VAReferenceEncoderStep
from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks
from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
    MiniMaxH3DecodeStep,
    MiniMaxH3Ref2VACoreDenoiseStep,
)
from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam


class MiniMaxH3Ref2VAGeneratorBlocks(SequentialPipelineBlocks):
    """The denoising half of a split `ref2va`: the `ref2va` branch with its `text_encoder` step removed.

    `reference_encoder` stays here, next to the two autoencoders it runs: its output shapes are where every
    reference block's geometry in the packed layout comes from.
    """

    model_name = "minimax-h3"
    # `after_denoise` is deliberately absent: in diffusers 0.41.0 `MiniMaxH3Ref2VACoreDenoiseStep` already ends with
    # its own `MiniMaxH3AfterDenoiseStep`, so a second copy here would run the unpack twice and crash on the
    # audio reshape (`shape '[2, 207, 207]' is invalid for input of size 13248`). The pre-merge PR commit the
    # sibling deployment pinned did not nest it, which is where the duplication came from.
    block_classes = [
        MiniMaxH3Ref2VASetupStep,
        MiniMaxH3Ref2VAReferenceEncoderStep,
        MiniMaxH3Ref2VACoreDenoiseStep,
        MiniMaxH3DecodeStep,
    ]
    block_names = ["setup", "reference_encoder", "denoise", "decode"]

    @property
    def description(self):
        return (
            "The denoising half of a split MiniMax-H3 `ref2va` deployment: the `ref2va` branch of `MiniMaxH3Blocks` "
            "without its text-encoder step, so `prompt_embeds` and `text_token_tags` come in as inputs and the "
            "62 GiB Qwen3-VL conditioner is never loaded here. The transformer is the `transformer_ref` partition."
        )

    @property
    def outputs(self):
        # Inlined from `MiniMaxH3Blocks.outputs` (diffusers 0.41.0): the pre-merge PR exposed this list through a
        # `_generation_outputs()` helper that the release dropped, so the same three params are spelled out here.
        return [
            OutputParam.template("videos", description="The generated video."),
            OutputParam(
                "audio",
                type_hint=torch.Tensor,
                description="The generated soundtrack, of shape `(1, 2, num_samples)`.",
            ),
            OutputParam("sampling_rate", type_hint=int, description="Sample rate of the generated soundtrack in Hz."),
        ]