hyperflow-h3-flash / ref2va_int8_convrot /hyperflow_h3_flash_transformer.py

HyperFlow H3 Flash 1.0.0: INT8 convrot and GGUF for fl2va and ref2va

9e9ed14 about 17 hours ago

19.2 kB

# Copyright 2026 The HyperFlow authors. Licensed under the Apache License, Version 2.0.
#
# `HyperFlowH3FlashTransformer3DModel.forward` is derived from `MiniMaxH3Transformer3DModel.forward` in diffusers
# (`models/transformers/transformer_minimax_h3.py`), Copyright 2025 The MiniMax Team and The HuggingFace Team,
# Apache License 2.0. See https://github.com/Video-Rebirth/hyperflow-h3-flash/blob/main/THIRD_PARTY_NOTICES.md.
"""The HyperFlow H3 Flash transformer: MiniMax-H3, compressed and conditioned on two times.
Four changes to `MiniMaxH3Transformer3DModel`, all driven by `config.json`:
* **Two-time conditioning.** An `endpoint_time_embedder` embeds the step's endpoint ``r`` next to the current time
  ``t``; every AdaLN reads ``temb(t, r) = emb_t(t) + hyperflow_gate * (emb_r(r) - emb_t(t))``.
* **Dropped blocks** (`dropped_blocks`): those blocks are identities and have no weights.
* **Low-rank linears** (`lowrank_ranks`): ``W ~= up @ down``, stored as `<name>.down.weight` (rank x in) and
  `<name>.up.weight` (out x rank). Ranks are zero-padded to multiples of 64, in the checkpoint or at load
  (`pad_lowrank_ranks`), so both GEMMs are aligned; `set_lowrank_rank_multiple` re-pads them at run time.
* **INT8 linears** (`quantization`): W8A8 with a Hadamard-rotated input, see `int8.py`.
Embeddings, AdaLN, the token refiner and the output heads are the base model's, in bf16 / fp32. The block
feed-forwards run in token chunks (`ChunkedFeedForward`, `set_ff_chunk`), which bounds their activations;
`set_block_chunk` slices every token-wise op of a block.
"""

from __future__ import annotations

import sys
from typing import Any, ClassVar

import torch
from diffusers import MiniMaxH3Transformer3DModel
from diffusers.configuration_utils import register_to_config
from diffusers.models.embeddings import TimestepEmbedding
from diffusers.models.modeling_utils import get_parameter_dtype
from diffusers.models.transformers.transformer_minimax_h3 import (
    MINIMAX_H3_MODALITY_NUM,
    MiniMaxH3TransformerOutput,
)
from diffusers.utils import apply_lora_scale
from torch import Tensor, nn

from .hyperflow_h3_flash_int8 import Int8RotLinear

#: Tokens per feed-forward chunk; a 5.2 s clip has about 38k, a 14.4 s clip about 105k.
FF_CHUNK_TOKENS = 16384
#: Low-rank ranks are padded to multiples of this at load, so both GEMMs of a low-rank linear are aligned.
RANK_MULTIPLE = 64


class LowRankLinear(nn.Module):
    """``y = up(down(x))``, the factorized replacement of a bias-free `nn.Linear`."""

    def __init__(self, in_features: int, out_features: int, rank: int) -> None:
        super().__init__()
        self.in_features, self.out_features, self.rank = in_features, out_features, rank
        self.down = nn.Linear(in_features, rank, bias=False)
        self.up = nn.Linear(rank, out_features, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        return self.up(self.down(x))

    @torch.no_grad()
    def pad_rank(self, multiple: int = RANK_MULTIPLE) -> bool:
        """Zero-pad the rank up to a multiple of ``multiple``: new rows of ``down`` and columns of ``up`` are zero, so
        ``up @ down`` is unchanged. Returns whether the rank changed."""
        rank = -(-self.rank // multiple) * multiple
        if rank == self.rank:
            return False
        extra = rank - self.rank
        down, up = self.down.weight, self.up.weight
        self.down.weight = nn.Parameter(torch.cat([down, down.new_zeros(extra, down.shape[1])]), requires_grad=False)
        self.up.weight = nn.Parameter(torch.cat([up, up.new_zeros(up.shape[0], extra)], dim=1), requires_grad=False)
        self.down.out_features = self.up.in_features = self.rank = rank
        return True


class ChunkedFeedForward(nn.Module):
    """A block's feed-forward run over at most ``chunk_tokens`` tokens at a time; it keeps the block's ``net``.
    The feed-forward is token-wise, so the chunks compute the same function; only the rounding of bf16 GEMMs can
    change with their row count. It bounds the SwiGLU input, ``2 * ffn_dim`` wide per token (6 GiB for a 14.4 s
    clip), to one chunk.
    """

    def __init__(self, net: nn.ModuleList, out_features: int, chunk_tokens: int) -> None:
        super().__init__()
        self.net, self.out_features, self.chunk_tokens = net, out_features, chunk_tokens

    def _forward(self, hidden_states: Tensor) -> Tensor:
        for module in self.net:
            hidden_states = module(hidden_states)
        return hidden_states

    def forward(self, hidden_states: Tensor) -> Tensor:
        if hidden_states.shape[-2] <= self.chunk_tokens:
            return self._forward(hidden_states)
        out = hidden_states.new_empty(*hidden_states.shape[:-1], self.out_features)
        for start in range(0, hidden_states.shape[-2], self.chunk_tokens):
            stop = start + self.chunk_tokens
            out[..., start:stop, :] = self._forward(hidden_states[..., start:stop, :])
        return out


class DroppedBlock(nn.Module):
    """A removed transformer block: the hidden states go straight to the next block.
    The signature is the block's own: context parallel hooks address `adaln_indices` by name.
    """

    def forward(self, hidden_states: Tensor, temb=None, adaln_indices=None, rotary_emb=None, attention_mask=None):
        return hidden_states


def _replace_linear(model: nn.Module, name: str, make: type[nn.Module], *args: Any) -> None:
    owner_name, _, attr = name.rpartition(".")
    owner = model.get_submodule(owner_name)
    dense = owner[int(attr)] if isinstance(owner, nn.ModuleList) else getattr(owner, attr)
    if type(dense) is not nn.Linear or dense.bias is not None:
        raise TypeError(f"{name} is not a bias-free nn.Linear")
    replacement = make(dense.in_features, dense.out_features, *args)
    if isinstance(owner, nn.ModuleList):
        owner[int(attr)] = replacement
    else:
        setattr(owner, attr, replacement)


class HyperFlowH3FlashTransformer3DModel(MiniMaxH3Transformer3DModel):
    """`MiniMaxH3Transformer3DModel` conditioned on ``(t, r)``, with dropped blocks, low-rank and INT8 linears."""

    # Matched against whole dotted name parts at load time, so `time_embedder` does not cover the endpoint embedder.
    _keep_in_fp32_modules: ClassVar[list[str]] = [
        *MiniMaxH3Transformer3DModel._keep_in_fp32_modules,
        "endpoint_time_embedder",
        "wscale",
    ]

    @register_to_config
    def __init__(
        self,
        num_attention_heads: int = 56,
        attention_head_dim: int = 128,
        hidden_size: int = 5376,
        num_layers: int = 50,
        num_refiner_layers: int = 2,
        ffn_dim: int = 14336,
        in_channels: int = 24,
        audio_in_channels: int = 32,
        patch_size: tuple[int, int, int] = (1, 2, 2),
        text_dim: int = 5120,
        freq_dim: int = 256,
        time_embed_hidden_dim: int = 5376,
        time_embed_dim: int = 16,
        rope_freq_dim: int = 16,
        rope_theta: float = 10000.0,
        norm_eps: float = 1e-5,
        qk_norm_eps: float = 1e-5,
        final_norm_eps: float = 1e-5,
        hyperflow_gate: float = 0.25,
        dropped_blocks: list[int] | None = None,
        lowrank_ranks: dict[str, int] | None = None,
        quantization: dict[str, Any] | None = None,
    ) -> None:
        super().__init__(
            num_attention_heads=num_attention_heads,
            attention_head_dim=attention_head_dim,
            hidden_size=hidden_size,
            num_layers=num_layers,
            num_refiner_layers=num_refiner_layers,
            ffn_dim=ffn_dim,
            in_channels=in_channels,
            audio_in_channels=audio_in_channels,
            patch_size=patch_size,
            text_dim=text_dim,
            freq_dim=freq_dim,
            time_embed_hidden_dim=time_embed_hidden_dim,
            time_embed_dim=time_embed_dim,
            rope_freq_dim=rope_freq_dim,
            rope_theta=rope_theta,
            norm_eps=norm_eps,
            qk_norm_eps=qk_norm_eps,
            final_norm_eps=final_norm_eps,
        )
        self.endpoint_time_embedder = TimestepEmbedding(
            in_channels=freq_dim, time_embed_dim=time_embed_hidden_dim, out_dim=time_embed_dim
        )

        dropped = set(dropped_blocks or [])
        for index, block in enumerate(self.transformer_blocks):
            if index in dropped:
                self.transformer_blocks[index] = DroppedBlock()
            else:
                block.ff = ChunkedFeedForward(block.ff.net, block.ff.net[-1].out_features, FF_CHUNK_TOKENS)
        for name, rank in (lowrank_ranks or {}).items():
            if int(name.split(".")[1]) not in dropped:
                _replace_linear(self, name, LowRankLinear, int(rank))
        if quantization:
            if quantization.get("method") != "int8_convrot":
                raise ValueError(f"unsupported quantization method {quantization.get('method')!r}")
            for name in quantization["linears"]:
                _replace_linear(self, name, Int8RotLinear, int(quantization["group"]))
        for block in self.transformer_blocks:
            attn = getattr(block, "attn", None)
            for name in ("to_q", "to_k"):
                if isinstance(layer := getattr(attn, name, None), Int8RotLinear):
                    layer.keep_input = True

    def pad_lowrank_ranks(self, multiple: int = RANK_MULTIPLE) -> int:
        """Zero-pad every low-rank linear to a rank multiple of ``multiple`` (see `LowRankLinear.pad_rank`) for a
        checkpoint stored at its unpadded ranks. Returns how many linears changed."""
        ranks = dict(self.config.lowrank_ranks or {})
        padded = 0
        for name, module in self.named_modules():
            if isinstance(module, LowRankLinear) and module.pad_rank(multiple):
                ranks[name] = module.rank
                padded += 1
        if padded:
            self.register_to_config(lowrank_ranks=ranks)
        return padded

    def time_embedding(self, timestep: Tensor, endpoint_timestep: Tensor) -> Tensor:
        """``temb`` of every distinct ``(t, r)`` pair, at the time embedders' precision (fp32)."""
        if endpoint_timestep.shape != timestep.shape:
            raise ValueError(f"endpoint_timestep {list(endpoint_timestep.shape)} != timestep {list(timestep.shape)}")
        dtype = get_parameter_dtype(self.time_embedder)
        current = self.time_embedder(self.time_proj(timestep).to(dtype))
        endpoint = self.endpoint_time_embedder(self.time_proj(endpoint_timestep).to(dtype))
        return current + self.config.hyperflow_gate * (endpoint - current)

    @apply_lora_scale("attention_kwargs")
    def forward(
        self,
        hidden_states: Tensor,
        audio_hidden_states: Tensor,
        encoder_hidden_states: Tensor,
        timestep: Tensor,
        timestep_indices: Tensor,
        token_tags: Tensor,
        position_ids: Tensor,
        video_indices: Tensor,
        audio_indices: Tensor,
        text_indices: Tensor,
        endpoint_timestep: Tensor | None = None,
        attention_kwargs: dict[str, Any] | None = None,
        return_dict: bool = True,
    ) -> MiniMaxH3TransformerOutput | tuple[Tensor, Tensor]:
        """The base forward plus ``endpoint_timestep``: ``(timestep[i], endpoint_timestep[i])`` is the ``i``-th pair.
        ``timestep_indices`` maps every row of the packed sequence to its pair. Build the pairs with
        `hyperflow_h3_flash.schedule.build_row_time_pairs`, or run the model through `flash_blocks(...)`.
        """
        if endpoint_timestep is None:
            raise ValueError(
                "HyperFlow H3 Flash needs `endpoint_timestep`, the endpoint `r` of every `(t, r)` pair: it is a "
                "few-step model and has no single-time mode. Run it through `hyperflow_h3_flash.flash_blocks(...)`."
            )
        if position_ids.ndim != 2 or position_ids.shape[-1] != 3:
            raise ValueError(f"`position_ids` must be a `(seq_len, 3)` tensor, got {list(position_ids.shape)}.")
        sequence_length = position_ids.shape[0]
        if token_tags.shape != (sequence_length,) or timestep_indices.shape != (sequence_length,):
            raise ValueError("`token_tags` and `timestep_indices` must both be `(seq_len,)` tensors.")

        rotary_emb = self.rope(position_ids)

        video_embeds = self.proj_in(hidden_states.to(get_parameter_dtype(self.proj_in)))
        audio_embeds = self.audio_proj_in(audio_hidden_states.to(get_parameter_dtype(self.audio_proj_in)))
        text_embeds = self.context_embedder(encoder_hidden_states.to(get_parameter_dtype(self.context_embedder)))
        text_embeds = self.token_refiner(text_embeds)

        hidden_states = text_embeds.new_zeros((text_embeds.shape[0], sequence_length, text_embeds.shape[-1]))
        hidden_states = hidden_states.index_copy(1, text_indices, text_embeds)
        hidden_states = hidden_states.index_copy(1, video_indices, video_embeds.to(text_embeds.dtype))
        hidden_states = hidden_states.index_copy(1, audio_indices, audio_embeds.to(text_embeds.dtype))

        temb = self.time_embedding(timestep, endpoint_timestep)
        adaln_indices = timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags
        for block in self.transformer_blocks:
            hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb)

        hidden_states = self.norm_out(hidden_states, temb, timestep_indices).to(get_parameter_dtype(self.proj_out))
        video_output = self.proj_out(hidden_states).index_select(1, video_indices)
        audio_output = self.audio_proj_out(hidden_states).index_select(1, audio_indices)
        if not return_dict:
            return (video_output, audio_output)
        return MiniMaxH3TransformerOutput(sample=video_output, audio_sample=audio_output)


def _used_rank(layer: LowRankLinear) -> int:
    """The rank without the zero padding (trailing zero rows of ``down`` and columns of ``up``)."""
    used = (layer.down.weight.detach().abs().amax(1) != 0) | (layer.up.weight.detach().abs().amax(0) != 0)
    nonzero = used.nonzero()
    return int(nonzero[-1]) + 1 if len(nonzero) else 1


def set_lowrank_rank_multiple(model: nn.Module, multiple: int) -> int:
    """Re-pad every low-rank linear to a rank multiple of ``multiple`` (0: no padding); returns the layers changed.
    Zero rows and columns add exact zeros, so any multiple computes the same function; the GEMM shapes change, and with
    them the bf16 rounding (8 is 2 to 7 % faster than 64 on H200). Run time only: the config still lists the stored
    ranks, so do not `save_pretrained` the re-padded model.
    """
    changed = 0
    for layer in model.modules():
        if not isinstance(layer, LowRankLinear):
            continue
        rank = _used_rank(layer)
        target = -(-rank // multiple) * multiple if multiple > 0 else rank
        if target == layer.rank:
            continue
        down, up = layer.down.weight.detach(), layer.up.weight.detach()
        if target < layer.rank:
            down, up = down[:target], up[:, :target]
        else:
            down = torch.cat([down, down.new_zeros(target - layer.rank, down.shape[1])])
            up = torch.cat([up, up.new_zeros(up.shape[0], target - layer.rank)], dim=1)
        layer.down.weight = nn.Parameter(down.contiguous(), requires_grad=False)
        layer.up.weight = nn.Parameter(up.contiguous(), requires_grad=False)
        layer.rank = layer.down.out_features = layer.up.in_features = target
        changed += 1
    return changed


def set_ff_chunk(model: nn.Module, tokens: int) -> None:
    """Run every feed-forward over at most ``tokens`` tokens at a time (0: whole sequence)."""
    for module in model.modules():
        if isinstance(module, ChunkedFeedForward):
            module.chunk_tokens = tokens if tokens > 0 else sys.maxsize


def chunk_block(block: nn.Module, chunk: int) -> None:
    """Run every token-wise op of a transformer block over ``chunk``-token slices; only the attention sees all of q /
    k / v.
    The unchunked block holds about a dozen ``(seq, hidden)`` tensors at once (each modulation picked per token, the
    modulated norm, the residual sums), 1.2 GB apiece at 108k tokens. Per-row results are the same up to the bf16 GEMM
    rounding for the smaller row count. Apply before context parallelism and offloading, which hook the block.
    """
    import diffusers.models.transformers.transformer_minimax_h3 as h3

    attn = block.attn

    def forward(hidden_states, temb, adaln_indices, rotary_emb, attention_mask=None):
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = block.adaln_proj(temb)
        batch, seq, _ = hidden_states.shape
        cos, sin = rotary_emb
        qkv = None
        for start in range(0, seq, chunk):
            rows = slice(start, min(start + chunk, seq))
            idx = adaln_indices[rows]
            h = block.norm1(hidden_states[:, rows])
            h = h * (1.0 + scale_msa.index_select(0, idx)) + shift_msa.index_select(0, idx)
            parts = []
            for proj, norm in ((attn.to_q, attn.norm_q), (attn.to_k, attn.norm_k), (attn.to_v, None)):
                x = proj(h).unflatten(-1, (attn.heads, -1))
                if norm is not None:
                    x = h3._apply_rotary_emb(norm(x), cos[rows], sin[rows])
                parts.append(x)
            if qkv is None:
                qkv = [x.new_empty((batch, seq, *x.shape[2:])) for x in parts]
            for full, x in zip(qkv, parts, strict=True):
                full[:, rows] = x
            del h, parts, x
        out = h3.dispatch_attention_fn(
            *qkv,
            attn_mask=attention_mask,
            dropout_p=0.0,
            is_causal=False,
            backend=attn.processor._attention_backend,
            parallel_config=attn.processor._parallel_config,
        )
        dtype = qkv[0].dtype
        del qkv
        result = torch.empty_like(hidden_states)
        for start in range(0, seq, chunk):
            rows = slice(start, min(start + chunk, seq))
            idx = adaln_indices[rows]
            a = attn.to_out[1](attn.to_out[0](out[:, rows].flatten(2, 3).to(dtype)))
            h = hidden_states[:, rows] + gate_msa.index_select(0, idx) * a
            n = block.norm2(h)
            n = n * (1.0 + scale_mlp.index_select(0, idx)) + shift_mlp.index_select(0, idx)
            result[:, rows] = h + gate_mlp.index_select(0, idx) * block.ff(n)
            del a, h, n
        return result

    block.forward = forward


def set_block_chunk(model: nn.Module, tokens: int) -> None:
    """`chunk_block` on every kept block (0: off). Call once, before context parallelism and offloading."""
    if tokens > 0:
        for block in model.transformer_blocks:
            if not isinstance(block, DroppedBlock):
                chunk_block(block, tokens)


__all__ = [
    "RANK_MULTIPLE",
    "ChunkedFeedForward",
    "DroppedBlock",
    "HyperFlowH3FlashTransformer3DModel",
    "LowRankLinear",
    "chunk_block",
    "set_block_chunk",
    "set_ff_chunk",
    "set_lowrank_rank_multiple",
]

	# Copyright 2026 The HyperFlow authors. Licensed under the Apache License, Version 2.0.
	#
	# `HyperFlowH3FlashTransformer3DModel.forward` is derived from `MiniMaxH3Transformer3DModel.forward` in diffusers
	# (`models/transformers/transformer_minimax_h3.py`), Copyright 2025 The MiniMax Team and The HuggingFace Team,
	# Apache License 2.0. See https://github.com/Video-Rebirth/hyperflow-h3-flash/blob/main/THIRD_PARTY_NOTICES.md.
	"""The HyperFlow H3 Flash transformer: MiniMax-H3, compressed and conditioned on two times.

	Four changes to `MiniMaxH3Transformer3DModel`, all driven by `config.json`:

	* Two-time conditioning. An `endpoint_time_embedder` embeds the step's endpoint ``r`` next to the current time
	``t``; every AdaLN reads ``temb(t, r) = emb_t(t) + hyperflow_gate * (emb_r(r) - emb_t(t))``.
	* Dropped blocks (`dropped_blocks`): those blocks are identities and have no weights.
	* Low-rank linears (`lowrank_ranks`): ``W ~= up @ down``, stored as `<name>.down.weight` (rank x in) and
	`<name>.up.weight` (out x rank). Ranks are zero-padded to multiples of 64, in the checkpoint or at load
	(`pad_lowrank_ranks`), so both GEMMs are aligned; `set_lowrank_rank_multiple` re-pads them at run time.
	* INT8 linears (`quantization`): W8A8 with a Hadamard-rotated input, see `int8.py`.

	Embeddings, AdaLN, the token refiner and the output heads are the base model's, in bf16 / fp32. The block
	feed-forwards run in token chunks (`ChunkedFeedForward`, `set_ff_chunk`), which bounds their activations;
	`set_block_chunk` slices every token-wise op of a block.
	"""

	from __future__ import annotations

	import sys
	from typing import Any, ClassVar

	import torch
	from diffusers import MiniMaxH3Transformer3DModel
	from diffusers.configuration_utils import register_to_config
	from diffusers.models.embeddings import TimestepEmbedding
	from diffusers.models.modeling_utils import get_parameter_dtype
	from diffusers.models.transformers.transformer_minimax_h3 import (
	MINIMAX_H3_MODALITY_NUM,
	MiniMaxH3TransformerOutput,
	)
	from diffusers.utils import apply_lora_scale
	from torch import Tensor, nn

	from .hyperflow_h3_flash_int8 import Int8RotLinear

	#: Tokens per feed-forward chunk; a 5.2 s clip has about 38k, a 14.4 s clip about 105k.
	FF_CHUNK_TOKENS = 16384
	#: Low-rank ranks are padded to multiples of this at load, so both GEMMs of a low-rank linear are aligned.
	RANK_MULTIPLE = 64


	class LowRankLinear(nn.Module):
	"""``y = up(down(x))``, the factorized replacement of a bias-free `nn.Linear`."""

	def __init__(self, in_features: int, out_features: int, rank: int) -> None:
	super().__init__()
	self.in_features, self.out_features, self.rank = in_features, out_features, rank
	self.down = nn.Linear(in_features, rank, bias=False)
	self.up = nn.Linear(rank, out_features, bias=False)

	def forward(self, x: Tensor) -> Tensor:
	return self.up(self.down(x))

	@torch.no_grad()
	def pad_rank(self, multiple: int = RANK_MULTIPLE) -> bool:
	"""Zero-pad the rank up to a multiple of ``multiple``: new rows of ``down`` and columns of ``up`` are zero, so
	``up @ down`` is unchanged. Returns whether the rank changed."""
	rank = -(-self.rank // multiple) * multiple
	if rank == self.rank:
	return False
	extra = rank - self.rank
	down, up = self.down.weight, self.up.weight
	self.down.weight = nn.Parameter(torch.cat([down, down.new_zeros(extra, down.shape[1])]), requires_grad=False)
	self.up.weight = nn.Parameter(torch.cat([up, up.new_zeros(up.shape[0], extra)], dim=1), requires_grad=False)
	self.down.out_features = self.up.in_features = self.rank = rank
	return True


	class ChunkedFeedForward(nn.Module):
	"""A block's feed-forward run over at most ``chunk_tokens`` tokens at a time; it keeps the block's ``net``.

	The feed-forward is token-wise, so the chunks compute the same function; only the rounding of bf16 GEMMs can
	change with their row count. It bounds the SwiGLU input, ``2 * ffn_dim`` wide per token (6 GiB for a 14.4 s
	clip), to one chunk.
	"""

	def __init__(self, net: nn.ModuleList, out_features: int, chunk_tokens: int) -> None:
	super().__init__()
	self.net, self.out_features, self.chunk_tokens = net, out_features, chunk_tokens

	def _forward(self, hidden_states: Tensor) -> Tensor:
	for module in self.net:
	hidden_states = module(hidden_states)
	return hidden_states

	def forward(self, hidden_states: Tensor) -> Tensor:
	if hidden_states.shape[-2] <= self.chunk_tokens:
	return self._forward(hidden_states)
	out = hidden_states.new_empty(*hidden_states.shape[:-1], self.out_features)
	for start in range(0, hidden_states.shape[-2], self.chunk_tokens):
	stop = start + self.chunk_tokens
	out[..., start:stop, :] = self._forward(hidden_states[..., start:stop, :])
	return out


	class DroppedBlock(nn.Module):
	"""A removed transformer block: the hidden states go straight to the next block.

	The signature is the block's own: context parallel hooks address `adaln_indices` by name.
	"""

	def forward(self, hidden_states: Tensor, temb=None, adaln_indices=None, rotary_emb=None, attention_mask=None):
	return hidden_states


	def _replace_linear(model: nn.Module, name: str, make: type[nn.Module], *args: Any) -> None:
	owner_name, _, attr = name.rpartition(".")
	owner = model.get_submodule(owner_name)
	dense = owner[int(attr)] if isinstance(owner, nn.ModuleList) else getattr(owner, attr)
	if type(dense) is not nn.Linear or dense.bias is not None:
	raise TypeError(f"{name} is not a bias-free nn.Linear")
	replacement = make(dense.in_features, dense.out_features, *args)
	if isinstance(owner, nn.ModuleList):
	owner[int(attr)] = replacement
	else:
	setattr(owner, attr, replacement)


	class HyperFlowH3FlashTransformer3DModel(MiniMaxH3Transformer3DModel):
	"""`MiniMaxH3Transformer3DModel` conditioned on ``(t, r)``, with dropped blocks, low-rank and INT8 linears."""

	# Matched against whole dotted name parts at load time, so `time_embedder` does not cover the endpoint embedder.
	_keep_in_fp32_modules: ClassVar[list[str]] = [
	*MiniMaxH3Transformer3DModel._keep_in_fp32_modules,
	"endpoint_time_embedder",
	"wscale",
	]

	@register_to_config
	def __init__(
	self,
	num_attention_heads: int = 56,
	attention_head_dim: int = 128,
	hidden_size: int = 5376,
	num_layers: int = 50,
	num_refiner_layers: int = 2,
	ffn_dim: int = 14336,
	in_channels: int = 24,
	audio_in_channels: int = 32,
	patch_size: tuple[int, int, int] = (1, 2, 2),
	text_dim: int = 5120,
	freq_dim: int = 256,
	time_embed_hidden_dim: int = 5376,
	time_embed_dim: int = 16,
	rope_freq_dim: int = 16,
	rope_theta: float = 10000.0,
	norm_eps: float = 1e-5,
	qk_norm_eps: float = 1e-5,
	final_norm_eps: float = 1e-5,
	hyperflow_gate: float = 0.25,
	dropped_blocks: list[int] \| None = None,
	lowrank_ranks: dict[str, int] \| None = None,
	quantization: dict[str, Any] \| None = None,
	) -> None:
	super().__init__(
	num_attention_heads=num_attention_heads,
	attention_head_dim=attention_head_dim,
	hidden_size=hidden_size,
	num_layers=num_layers,
	num_refiner_layers=num_refiner_layers,
	ffn_dim=ffn_dim,
	in_channels=in_channels,
	audio_in_channels=audio_in_channels,
	patch_size=patch_size,
	text_dim=text_dim,
	freq_dim=freq_dim,
	time_embed_hidden_dim=time_embed_hidden_dim,
	time_embed_dim=time_embed_dim,
	rope_freq_dim=rope_freq_dim,
	rope_theta=rope_theta,
	norm_eps=norm_eps,
	qk_norm_eps=qk_norm_eps,
	final_norm_eps=final_norm_eps,
	)
	self.endpoint_time_embedder = TimestepEmbedding(
	in_channels=freq_dim, time_embed_dim=time_embed_hidden_dim, out_dim=time_embed_dim
	)

	dropped = set(dropped_blocks or [])
	for index, block in enumerate(self.transformer_blocks):
	if index in dropped:
	self.transformer_blocks[index] = DroppedBlock()
	else:
	block.ff = ChunkedFeedForward(block.ff.net, block.ff.net[-1].out_features, FF_CHUNK_TOKENS)
	for name, rank in (lowrank_ranks or {}).items():
	if int(name.split(".")[1]) not in dropped:
	_replace_linear(self, name, LowRankLinear, int(rank))
	if quantization:
	if quantization.get("method") != "int8_convrot":
	raise ValueError(f"unsupported quantization method {quantization.get('method')!r}")
	for name in quantization["linears"]:
	_replace_linear(self, name, Int8RotLinear, int(quantization["group"]))
	for block in self.transformer_blocks:
	attn = getattr(block, "attn", None)
	for name in ("to_q", "to_k"):
	if isinstance(layer := getattr(attn, name, None), Int8RotLinear):
	layer.keep_input = True

	def pad_lowrank_ranks(self, multiple: int = RANK_MULTIPLE) -> int:
	"""Zero-pad every low-rank linear to a rank multiple of ``multiple`` (see `LowRankLinear.pad_rank`) for a
	checkpoint stored at its unpadded ranks. Returns how many linears changed."""
	ranks = dict(self.config.lowrank_ranks or {})
	padded = 0
	for name, module in self.named_modules():
	if isinstance(module, LowRankLinear) and module.pad_rank(multiple):
	ranks[name] = module.rank
	padded += 1
	if padded:
	self.register_to_config(lowrank_ranks=ranks)
	return padded

	def time_embedding(self, timestep: Tensor, endpoint_timestep: Tensor) -> Tensor:
	"""``temb`` of every distinct ``(t, r)`` pair, at the time embedders' precision (fp32)."""
	if endpoint_timestep.shape != timestep.shape:
	raise ValueError(f"endpoint_timestep {list(endpoint_timestep.shape)} != timestep {list(timestep.shape)}")
	dtype = get_parameter_dtype(self.time_embedder)
	current = self.time_embedder(self.time_proj(timestep).to(dtype))
	endpoint = self.endpoint_time_embedder(self.time_proj(endpoint_timestep).to(dtype))
	return current + self.config.hyperflow_gate * (endpoint - current)

	@apply_lora_scale("attention_kwargs")
	def forward(
	self,
	hidden_states: Tensor,
	audio_hidden_states: Tensor,
	encoder_hidden_states: Tensor,
	timestep: Tensor,
	timestep_indices: Tensor,
	token_tags: Tensor,
	position_ids: Tensor,
	video_indices: Tensor,
	audio_indices: Tensor,
	text_indices: Tensor,
	endpoint_timestep: Tensor \| None = None,
	attention_kwargs: dict[str, Any] \| None = None,
	return_dict: bool = True,
	) -> MiniMaxH3TransformerOutput \| tuple[Tensor, Tensor]:
	"""The base forward plus ``endpoint_timestep``: ``(timestep[i], endpoint_timestep[i])`` is the ``i``-th pair.

	``timestep_indices`` maps every row of the packed sequence to its pair. Build the pairs with
	`hyperflow_h3_flash.schedule.build_row_time_pairs`, or run the model through `flash_blocks(...)`.
	"""
	if endpoint_timestep is None:
	raise ValueError(
	"HyperFlow H3 Flash needs `endpoint_timestep`, the endpoint `r` of every `(t, r)` pair: it is a "
	"few-step model and has no single-time mode. Run it through `hyperflow_h3_flash.flash_blocks(...)`."
	)
	if position_ids.ndim != 2 or position_ids.shape[-1] != 3:
	raise ValueError(f"`position_ids` must be a `(seq_len, 3)` tensor, got {list(position_ids.shape)}.")
	sequence_length = position_ids.shape[0]
	if token_tags.shape != (sequence_length,) or timestep_indices.shape != (sequence_length,):
	raise ValueError("`token_tags` and `timestep_indices` must both be `(seq_len,)` tensors.")

	rotary_emb = self.rope(position_ids)

	video_embeds = self.proj_in(hidden_states.to(get_parameter_dtype(self.proj_in)))
	audio_embeds = self.audio_proj_in(audio_hidden_states.to(get_parameter_dtype(self.audio_proj_in)))
	text_embeds = self.context_embedder(encoder_hidden_states.to(get_parameter_dtype(self.context_embedder)))
	text_embeds = self.token_refiner(text_embeds)

	hidden_states = text_embeds.new_zeros((text_embeds.shape[0], sequence_length, text_embeds.shape[-1]))
	hidden_states = hidden_states.index_copy(1, text_indices, text_embeds)
	hidden_states = hidden_states.index_copy(1, video_indices, video_embeds.to(text_embeds.dtype))
	hidden_states = hidden_states.index_copy(1, audio_indices, audio_embeds.to(text_embeds.dtype))

	temb = self.time_embedding(timestep, endpoint_timestep)
	adaln_indices = timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags
	for block in self.transformer_blocks:
	hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb)

	hidden_states = self.norm_out(hidden_states, temb, timestep_indices).to(get_parameter_dtype(self.proj_out))
	video_output = self.proj_out(hidden_states).index_select(1, video_indices)
	audio_output = self.audio_proj_out(hidden_states).index_select(1, audio_indices)
	if not return_dict:
	return (video_output, audio_output)
	return MiniMaxH3TransformerOutput(sample=video_output, audio_sample=audio_output)


	def _used_rank(layer: LowRankLinear) -> int:
	"""The rank without the zero padding (trailing zero rows of ``down`` and columns of ``up``)."""
	used = (layer.down.weight.detach().abs().amax(1) != 0) \| (layer.up.weight.detach().abs().amax(0) != 0)
	nonzero = used.nonzero()
	return int(nonzero[-1]) + 1 if len(nonzero) else 1


	def set_lowrank_rank_multiple(model: nn.Module, multiple: int) -> int:
	"""Re-pad every low-rank linear to a rank multiple of ``multiple`` (0: no padding); returns the layers changed.

	Zero rows and columns add exact zeros, so any multiple computes the same function; the GEMM shapes change, and with
	them the bf16 rounding (8 is 2 to 7 % faster than 64 on H200). Run time only: the config still lists the stored
	ranks, so do not `save_pretrained` the re-padded model.
	"""
	changed = 0
	for layer in model.modules():
	if not isinstance(layer, LowRankLinear):
	continue
	rank = _used_rank(layer)
	target = -(-rank // multiple) * multiple if multiple > 0 else rank
	if target == layer.rank:
	continue
	down, up = layer.down.weight.detach(), layer.up.weight.detach()
	if target < layer.rank:
	down, up = down[:target], up[:, :target]
	else:
	down = torch.cat([down, down.new_zeros(target - layer.rank, down.shape[1])])
	up = torch.cat([up, up.new_zeros(up.shape[0], target - layer.rank)], dim=1)
	layer.down.weight = nn.Parameter(down.contiguous(), requires_grad=False)
	layer.up.weight = nn.Parameter(up.contiguous(), requires_grad=False)
	layer.rank = layer.down.out_features = layer.up.in_features = target
	changed += 1
	return changed


	def set_ff_chunk(model: nn.Module, tokens: int) -> None:
	"""Run every feed-forward over at most ``tokens`` tokens at a time (0: whole sequence)."""
	for module in model.modules():
	if isinstance(module, ChunkedFeedForward):
	module.chunk_tokens = tokens if tokens > 0 else sys.maxsize


	def chunk_block(block: nn.Module, chunk: int) -> None:
	"""Run every token-wise op of a transformer block over ``chunk``-token slices; only the attention sees all of q /
	k / v.

	The unchunked block holds about a dozen ``(seq, hidden)`` tensors at once (each modulation picked per token, the
	modulated norm, the residual sums), 1.2 GB apiece at 108k tokens. Per-row results are the same up to the bf16 GEMM
	rounding for the smaller row count. Apply before context parallelism and offloading, which hook the block.
	"""
	import diffusers.models.transformers.transformer_minimax_h3 as h3

	attn = block.attn

	def forward(hidden_states, temb, adaln_indices, rotary_emb, attention_mask=None):
	shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = block.adaln_proj(temb)
	batch, seq, _ = hidden_states.shape
	cos, sin = rotary_emb
	qkv = None
	for start in range(0, seq, chunk):
	rows = slice(start, min(start + chunk, seq))
	idx = adaln_indices[rows]
	h = block.norm1(hidden_states[:, rows])
	h = h * (1.0 + scale_msa.index_select(0, idx)) + shift_msa.index_select(0, idx)
	parts = []
	for proj, norm in ((attn.to_q, attn.norm_q), (attn.to_k, attn.norm_k), (attn.to_v, None)):
	x = proj(h).unflatten(-1, (attn.heads, -1))
	if norm is not None:
	x = h3._apply_rotary_emb(norm(x), cos[rows], sin[rows])
	parts.append(x)
	if qkv is None:
	qkv = [x.new_empty((batch, seq, *x.shape[2:])) for x in parts]
	for full, x in zip(qkv, parts, strict=True):
	full[:, rows] = x
	del h, parts, x
	out = h3.dispatch_attention_fn(
	*qkv,
	attn_mask=attention_mask,
	dropout_p=0.0,
	is_causal=False,
	backend=attn.processor._attention_backend,
	parallel_config=attn.processor._parallel_config,
	)
	dtype = qkv[0].dtype
	del qkv
	result = torch.empty_like(hidden_states)
	for start in range(0, seq, chunk):
	rows = slice(start, min(start + chunk, seq))
	idx = adaln_indices[rows]
	a = attn.to_out[1](attn.to_out[0](out[:, rows].flatten(2, 3).to(dtype)))
	h = hidden_states[:, rows] + gate_msa.index_select(0, idx) * a
	n = block.norm2(h)
	n = n * (1.0 + scale_mlp.index_select(0, idx)) + shift_mlp.index_select(0, idx)
	result[:, rows] = h + gate_mlp.index_select(0, idx) * block.ff(n)
	del a, h, n
	return result

	block.forward = forward


	def set_block_chunk(model: nn.Module, tokens: int) -> None:
	"""`chunk_block` on every kept block (0: off). Call once, before context parallelism and offloading."""
	if tokens > 0:
	for block in model.transformer_blocks:
	if not isinstance(block, DroppedBlock):
	chunk_block(block, tokens)


	__all__ = [
	"RANK_MULTIPLE",
	"ChunkedFeedForward",
	"DroppedBlock",
	"HyperFlowH3FlashTransformer3DModel",
	"LowRankLinear",
	"chunk_block",
	"set_block_chunk",
	"set_ff_chunk",
	"set_lowrank_rank_multiple",
	]