Blame for h3_long_memory.py

"""Keep a long (up to 60 s), full-canvas MiniMax-H3 pass inside one card — without lowering the quality.

A 28 s 768x1344 request is ~200k rows in the packed sequence and 60 s at the same canvas is ~430k. Stock
`MiniMaxH3TransformerBlock.forward` materializes, for *every* row at once, the modulated norm, six per-row AdaLN
tensors (`index_select` over the table), the feed-forward's hidden layer (SwiGLU, several times the model width) and
the LoRA branch next to it. That is what ran the allocator dry (`NVML_SUCCESS == r INTERNAL ASSERT FAILED` in
`CUDACachingAllocator` is how an exhausted card surfaces under ZeroGPU), inside `ff`.

Everything in a block except the attention core is **per row**: RMSNorm, the AdaLN modulation, the q/k/v projections,
q/k norms, the rotary rotation, the output projection, the gates and the whole feed-forward. So they are run here in
chunks of rows written into preallocated buffers. Only the attention core itself sees the full sequence, and the
fused cuDNN kernel needs no `S x S` matrix. The result is the same computation row for row — same weights, same
LoRA, same dtypes, same order of operations — just not all rows at the same time.

The same goes for the two ends of the model and for the decode:

* `norm_out` is computed chunk by chunk straight into the float32 buffer the output heads read.
* The video decode streams: the VAE already decodes in temporal clips; each finished clip is denormalized,
  converted to uint8 and moved to the CPU immediately, instead of the whole float32 video being built on the card
  (18 GB for 60 s at 768x1344 before the clamp copies it again).

`patched(pipe)` installs all of it for the duration of one window and restores the stock methods afterwards; the
one-shot `/generate` path never sees it.

**The decode can be split too.** A ZeroGPU visitor who is not signed in has 2 minutes a day — about 60 s of real
`xlarge` time — so even the video decode may not fit next to the setup a window pays. The decode already runs clip by
clip; `DecodeControl` lets a window stop on a clip boundary and a later window pick up at the next clip. The loop's
whole state at a clip boundary is the frames emitted so far (uint8, on the CPU) and the cross-fade tail `overlap`, so
that is all that is handed over.
"""

from __future__ import annotations

import contextlib
import os
import types

CHUNK_ROWS = int(os.environ.get("H3_CHUNK_ROWS", "16384"))


def _chunks(total: int):
    for start in range(0, total, CHUNK_ROWS):
        yield start, min(total, start + CHUNK_ROWS)


class A2VBlockStop(BaseException):
    """Raised from inside a transformer forward, on a *block* boundary, when the window is out of time.

    A `BaseException` on purpose: diffusers' `loop_step` logs every `Exception` passing through it as a block error
    with a full traceback, and this is a planned stop, not an error.

    `block` is the index of the first block NOT yet run; `hidden` is the packed sequence entering it (on the CPU),
    or None when this window had not run a single block of the step yet (nothing worth keeping).
    """

    def __init__(self, block: int, hidden):
        super().__init__(f"window stopped before transformer block {block}")
        self.block = int(block)
        self.hidden = hidden


class A2VDecodeStop(BaseException):
    """Raised from inside the streaming video decode, on a *clip* boundary, when the window is out of time.

    A `BaseException` for the same reason as `A2VBlockStop`: it is a planned stop, not an error.

    `clip` is the first clip NOT yet decoded; `frames` the uint8 frames THIS window decoded and has not handed over,
    `(F, H, W, 3)` on the CPU, or None when it decoded none; `overlap` the cross-fade tail the next clip blends with
    (on the CPU), or None.
    """

    def __init__(self, clip: int, frames, overlap):
        super().__init__(f"window stopped before video decode clip {clip}")
        self.clip = int(clip)
        self.frames = frames
        self.overlap = overlap


class DecodeControl:
    """Lets the VIDEO DECODE be split across ZeroGPU windows, clip by clip.

    * `resume_clip` / `resume_overlap` — start at that clip with that cross-fade tail instead of at clip 0.
    * `should_stop(clip, pending_frames)` — asked before every clip that would actually run; `pending_frames` is how
      many frames this window has decoded so far (they still have to be written out if it stops).

    After the decode, `num_clips` is how many clips the whole video has and `pad_frames` how many trailing frames the
    last clip's latent padding produced. A resumable decode never trims them itself — the caller joins the parts of
    all windows and trims once, because the tail may live in a different window than the padding's origin.
    """

    def __init__(self, resume_clip: int = 0, resume_overlap=None, should_stop=None):
        self.resume_clip = int(resume_clip or 0)
        self.resume_overlap = resume_overlap if self.resume_clip > 0 else None
        self.should_stop = should_stop
        self.num_clips = 0
        self.pad_frames = 0
        self.decoded = 0  # clips this window actually decoded


class BlockControl:
    """Lets ONE denoising step be split across ZeroGPU windows, block by block.

    A 60 s full-canvas step costs more GPU time than a single ZeroGPU booking can hold, so a step can no longer be
    the smallest unit of work. Nothing changes numerically: the packed sequence entering block `k` is exactly what the
    uninterrupted forward would have handed it, so resuming there with it gives the same step, bit for bit.

    * `resume_block` / `resume_hidden` — skip blocks `< resume_block` of the first forward and start from the saved
      sequence instead. The embeddings in front of block 0 are recomputed (cheap, deterministic) and discarded.
    * `should_stop(index)` — asked before every block that would actually run.
    """

    def __init__(self, resume_block: int = 0, resume_hidden=None, should_stop=None, decode: "DecodeControl | None" = None):
        self.resume_block = int(resume_block or 0)
        self.resume_hidden = resume_hidden if self.resume_block > 0 else None
        self.should_stop = should_stop
        self.decode = decode  # the same window's control over the video decode, if it may stop there too
        self.first_block = self.resume_block  # the first block this window will actually compute
        self.computed = 0


CONTROL: BlockControl | None = None


def _control_gate(block, hidden_states):
    """`(skip, hidden_states)` for one block under `CONTROL`: skip a block an earlier window already ran (handing
    over the saved sequence at block 0), or raise `A2VBlockStop` when the window is out of time."""
    control = CONTROL
    if control is None:
        return False, hidden_states
    index = int(getattr(block, "_a2v_index", -1))
    if control.resume_block > 0 and 0 <= index < control.resume_block:
        if index == 0:
            saved = control.resume_hidden
            if tuple(saved.shape) != tuple(hidden_states.shape):
                raise ValueError(
                    f"The saved mid-step sequence {tuple(saved.shape)} does not match {tuple(hidden_states.shape)}."
                )
            hidden_states = saved.to(device=hidden_states.device, dtype=hidden_states.dtype)
        if index == control.resume_block - 1:
            control.resume_block, control.resume_hidden = 0, None
        return True, hidden_states
    if control.should_stop is not None and control.should_stop(index):
        # Block 0's input is only the embeddings (recomputed for free), and stopping at the very block this window
        # resumed at, before computing anything, has nothing new to keep either.
        keep = index > 0 and not (index == control.first_block and control.computed == 0)
        raise A2VBlockStop(index, hidden_states.detach().to("cpu") if keep else None)
    control.computed += 1
    return False, hidden_states


def _head_major_backend(backend) -> bool:
    """Whether this attention backend re-lays q/k/v out as `(batch, heads, rows, dim)` before the kernel.

    diffusers' native backends (`_native_cudnn` here, SDPA in general) do `x.permute(0, 2, 1, 3).contiguous()` on
    each of q, k and v. On a `(batch, rows, heads, dim)` buffer that is a full COPY — at 60 s full canvas three extra
    ~6.3 GB tensors, alive next to the originals, right at the memory peak. That copy is what ran the card out of
    memory (`NVML_SUCCESS == r INTERNAL ASSERT FAILED` in `CUDACachingAllocator`, raised inside that `contiguous`).
    """
    name = str(getattr(backend, "value", backend) or "").lower()
    return name == "" or "native" in name


def _attention_buffer(like, batch: int, total: int, head_major: bool):
    """A full-sequence q/k/v buffer, returned as a `(batch, rows, heads, dim)` view.

    With `head_major` the storage itself is `(batch, heads, rows, dim)`, so the backend's permute lands on a tensor
    that is already contiguous and `.contiguous()` returns it without copying. Same values, same kernel, same
    result — only the copy is gone.
    """
    heads, dim = like.shape[2], like.shape[3]
    if head_major:
        return like.new_empty((batch, heads, total, dim)).transpose(1, 2)
    return like.new_empty((batch, total, heads, dim))


def _block_forward(self, hidden_states, temb, adaln_indices, rotary_emb, attention_mask=None):
    import torch
    from diffusers.models.attention_dispatch import dispatch_attention_fn
    from diffusers.models.transformers.transformer_minimax_h3 import _apply_rotary_emb

    skip, hidden_states = _control_gate(self, hidden_states)
    if skip:
        return hidden_states

    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(temb)
    attn = self.attn
    batch, total, _ = hidden_states.shape

    # --- attention inputs, row chunk by row chunk -------------------------------------------------------------
    processor = attn.processor
    head_major = _head_major_backend(getattr(processor, "_attention_backend", None))
    query = key = value = None
    for start, end in _chunks(total):
        idx = adaln_indices[start:end]
        norm = self.norm1(hidden_states[:, start:end])
        norm = norm * (1.0 + scale_msa.index_select(0, idx)) + shift_msa.index_select(0, idx)
        if getattr(attn, "fused_projections", False):
            q, k, v = attn.to_qkv(norm).chunk(3, dim=-1)
        else:
            q, k, v = attn.to_q(norm), attn.to_k(norm), attn.to_v(norm)
        del norm
        q = attn.norm_q(q.unflatten(-1, (attn.heads, -1)))
        k = attn.norm_k(k.unflatten(-1, (attn.heads, -1)))
        v = v.unflatten(-1, (attn.heads, -1))
        if rotary_emb is not None:
            cos, sin = (t[start:end] for t in rotary_emb)
            q = _apply_rotary_emb(q, cos, sin)
            k = _apply_rotary_emb(k, cos, sin)
        if query is None:
            query, key, value = (_attention_buffer(t, batch, total, head_major) for t in (q, k, v))
        query[:, start:end], key[:, start:end], value[:, start:end] = q, k, v
        del q, k, v

    # --- the one full-sequence operation ----------------------------------------------------------------------
    out = dispatch_attention_fn(
        query,
        key,
        value,
        attn_mask=attention_mask,
        dropout_p=0.0,
        is_causal=False,
        backend=getattr(processor, "_attention_backend", None),
        parallel_config=getattr(processor, "_parallel_config", None),
    )
    dtype = query.dtype
    del query, key, value

    # --- attention output + residual, then the feed-forward + residual, row chunk by row chunk ----------------
    # In place: under no_grad nothing else holds the block's input, and the model rebinds it to the return value.
    for start, end in _chunks(total):
        idx = adaln_indices[start:end]
        projected = attn.to_out[1](attn.to_out[0](out[:, start:end].flatten(2, 3).type(dtype)))
        hidden_states[:, start:end] += gate_msa.index_select(0, idx) * projected
        del projected
    del out
    for start, end in _chunks(total):
        idx = adaln_indices[start:end]
        norm = self.norm2(hidden_states[:, start:end])
        norm = norm * (1.0 + scale_mlp.index_select(0, idx)) + shift_mlp.index_select(0, idx)
        hidden_states[:, start:end] += gate_mlp.index_select(0, idx) * self.ff(norm)
        del norm
    return hidden_states


def _norm_out_forward(self, hidden_states, temb, timestep_indices):
    """`norm_out.forward` (released or pruned), chunked, returning float32 — what the model casts it to next anyway."""
    import torch
    from diffusers.models.modeling_utils import get_parameter_dtype

    dtype = get_parameter_dtype(self.linear)
    if hasattr(self, "lora_adaln_bias"):
        # The pruned checkpoint (`MiniMaxH3PrunedAdaLayerNormOut`, multimodalart/MiniMax-H3-Pruned): `temb` is already
        # the table's timestep coordinates, so there is NO `silu`, and the modulation is `linear(temb)` plus the float32
        # `folded_bias` — which carries almost the whole modulation — plus the turbo LoRA's AdaLN constant term.
        # Leaving those out (the stock formula below) breaks the last layer on every step: that is what burned the
        # colours/contrast of every `animate_a2v` render while the one-shot `/generate` path stayed clean.
        projected = self.linear(temb.to(dtype))
        shift, scale = (projected.float() + self.lora_adaln_bias()).to(dtype).chunk(2, dim=-1)
    else:
        # The released checkpoint's `MiniMaxH3AdaLayerNormOut`.
        shift, scale = self.linear(torch.nn.functional.silu(temb).to(dtype)).chunk(2, dim=-1)
    out = torch.empty(hidden_states.shape, device=hidden_states.device, dtype=torch.float32)
    for start, end in _chunks(hidden_states.shape[1]):
        idx = timestep_indices[start:end]
        chunk = self.norm(hidden_states[:, start:end])
        out[:, start:end] = (chunk * (1.0 + scale.index_select(0, idx)) + shift.index_select(0, idx)).float()
    return out


def _streaming_video_decode(self, components, state):
    """`MiniMaxH3VideoDecodeStep.__call__` with the VAE's own clip loop, emitting uint8 frames to the CPU per clip.

    With a `DecodeControl` installed (resumable windows) the loop can start at any clip and stop before any clip; the
    returned video then holds only the frames THIS window decoded and the trailing padding frames are NOT trimmed —
    `control.pad_frames` says how many to cut once the parts are joined. Without one it is the stock decode, trimmed.
    """
    import torch
    from diffusers.models.modeling_utils import get_parameter_dtype

    block_state = self.get_block_state(state)
    device = components._execution_device
    vae = components.vae
    control = getattr(CONTROL, "decode", None) if CONTROL is not None else None

    latents_mean = torch.tensor(vae.config.latents_mean, device=device).view(1, -1, 1, 1, 1)
    latents_std = torch.tensor(vae.config.latents_std, device=device).view(1, -1, 1, 1, 1)
    z = (block_state.latents * latents_std + latents_mean).to(get_parameter_dtype(vae.decoder))
    pixel_mean = torch.tensor(components.pixel_mean, device=device).view(1, -1, 1, 1, 1)
    pixel_std = torch.tensor(components.pixel_std, device=device).view(1, -1, 1, 1, 1)

    frames = []

    def emit(chunk):
        video = (chunk.float() * pixel_std + pixel_mean).clamp(0, 1)
        frames.append((video[0] * 255).round().to(torch.uint8).permute(1, 2, 3, 0).contiguous().cpu())

    # The loop below is `AutoencoderKLMiniMaxH3._decode`, line for line, with `decoded_chunks.append` -> `emit`.
    tokens_chunk_size = vae.tokens_chunk_size
    token_drop = vae.config.token_drop
    temporal_ratio = vae.temporal_compression_ratio
    chunk_num_frames = tokens_chunk_size * temporal_ratio
    num_tokens = z.shape[2] + token_drop
    pad_tokens = (-num_tokens) % tokens_chunk_size
    num_chunks = (num_tokens + pad_tokens) // tokens_chunk_size - int(token_drop > 0)
    if pad_tokens > 0:
        z = torch.cat([z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2)

    # The frames the latent padding produced past the last requested one (`_decode` cuts them off at the very end).
    pad_frames = 0
    if pad_tokens > 0:
        intra_tail = vae.config.clip_length % temporal_ratio
        before_pad = z.shape[2] - pad_tokens
        pad_frames = sum(
            intra_tail if intra_tail and (before_pad + k) % tokens_chunk_size == 0 else temporal_ratio
            for k in range(pad_tokens)
        )

    first_clip, overlap = 0, None
    if control is not None:
        control.num_clips, control.pad_frames, control.decoded = int(num_chunks), int(pad_frames), 0
        first_clip = min(int(control.resume_clip), int(num_chunks))
        if control.resume_overlap is not None:
            overlap = control.resume_overlap.to(device)

    with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"):
        for i in range(first_clip, num_chunks):
            if control is not None and control.should_stop is not None:
                if control.should_stop(i, sum(int(f.shape[0]) for f in frames)):
                    raise A2VDecodeStop(
                        i,
                        torch.cat(frames) if frames else None,
                        overlap.detach().to("cpu") if overlap is not None else None,
                    )
            start = i * tokens_chunk_size
            clip = vae._decode_clip(z[:, :, start:start + tokens_chunk_size + vae.token_overlap])
            for j in range(int(token_drop > 0) + 1):
                frame_start = j * chunk_num_frames
                chunk = clip[:, :, frame_start:frame_start + chunk_num_frames]
                chunk = chunk[:, :, vae.frame_pre_padding:]
                if j == 0:
                    if overlap is not None:
                        chunk = vae._blend(overlap, chunk, vae.frame_overlap, dim=-3)
                    emit(chunk)
                else:
                    overlap = chunk
            del clip
            if control is not None:
                control.decoded += 1
        if overlap is not None:
            emit(overlap)

    video = torch.cat(frames)
    if control is None and pad_frames > 0:
        video = video[:-pad_frames]
    block_state.videos = [video]  # (F, H, W, 3) uint8 on the CPU
    self.set_block_state(state, block_state)
    return components, state


@contextlib.contextmanager
def patched(pipe, control: BlockControl | None = None):
    """Install the chunked block/norm forwards on this transformer and the streaming decode, then undo them.

    `control`, when given, lets the window stop between two transformer blocks and resume there later."""
    global CONTROL
    from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3VideoDecodeStep

    transformer = pipe.transformer
    undo = []
    for index, block in enumerate(getattr(transformer, "transformer_blocks", [])):
        if all(hasattr(block, name) for name in ("adaln_proj", "attn", "ff", "norm1", "norm2")):
            block._a2v_index = index
            block.forward = types.MethodType(_block_forward, block)
            undo.append(lambda block=block: block.__dict__.pop("forward", None))
    norm_out = getattr(transformer, "norm_out", None)
    if norm_out is not None and hasattr(norm_out, "linear") and hasattr(norm_out, "norm"):
        norm_out.forward = types.MethodType(_norm_out_forward, norm_out)
        undo.append(lambda: norm_out.__dict__.pop("forward", None))
    original_decode = MiniMaxH3VideoDecodeStep.__call__
    MiniMaxH3VideoDecodeStep.__call__ = _streaming_video_decode
    undo.append(lambda: setattr(MiniMaxH3VideoDecodeStep, "__call__", original_decode))
    print(f"[a2v-memory] chunked {len(undo) - 2} blocks at {CHUNK_ROWS} rows, streaming decode on", flush=True)
    CONTROL = control
    try:
        yield
    finally:
        CONTROL = None
        for step in reversed(undo):
            step()