h3_lora [processing/.]

h3_lora.py

@@ -0,0 +1,292 @@
+"""Turbo and style LoRA support for the diffusers MiniMax-H3 transformer: several checkpoint layouts, one fold
+mechanism.
+
+Every LoRA is applied by folding `scale * (lora_B @ lora_A)` into the bf16 weights rather than as runtime
+wrappers, because the AoTI block package (`h3_aoti`) reads each block's live weights and a wrapper module would
+be invisible to it. Deltas are computed in float32 and round once on the way back into bf16. The low-rank factors
+of every loaded LoRA stay resident, so the active one can be switched per request (`set_active`) — unfold the
+old, fold the new, in place, through one bf16 rounding.
+
+The supported checkpoints ship in different layouts:
+
+* `larry` (`larryvrh/MiniMax-H3-Turbo-Lora`) targets the reference (ComfyUI) module tree —
+  `blocks.N.attn.qkv_proj`, `blocks.N.mlp.fc1`, `token_refiner.blocks.N`, `final_layer.adaln_proj.linear` — with
+  `alpha == rank` (scale 1). Each delta gets the same transform the base weights got in the diffusers conversion
+  (scripts/convert_minimax_h3_to_diffusers.py, huggingface/diffusers#14371): fused-QKV row thirds onto
+  `attn.to_q/k/v`, the `SwiGLU` gate/value swap onto `ff.net.0.proj`, `fc2` -> `ff.net.2`, `blocks.` ->
+  `transformer_blocks.`, `token_refiner.blocks.` -> `token_refiner.refiner_blocks.`,
+  `final_layer.adaln_proj.linear` -> `norm_out.linear`. The row transforms are applied to `lora_B` directly
+  (rows of `B @ A` are rows of `B`), so no full delta is ever materialized at load.
+
+* `lightx` / `lightx8` (`lightx2v/Minimax-h3-Turbo`) are PEFT checkpoints against the diffusers tree itself —
+  `transformer_blocks.N.attn.to_q.lora_A.default.weight` and friends — rank 128, `alpha == 8`, so the fold scale
+  is `8 / 128 = 0.0625` (matching `set_adapters(weights=1.0)` in their inference script). Keys map name-for-name.
+  `lightx` is the 4-step file, `lightx8` the 8-step v1.0 file.
+
+* `realism` (`fal/MiniMax-H3-Realism-People-LoRA`) is a style LoRA, not a turbo one — realistic people, trigger
+  word `r34l1sm`. Same reference tree as larry under a `diffusion_model.` prefix, attention only (`qkv_proj` /
+  `out_proj` on the 52 blocks), rank 16, `alpha == rank` (the card's scale 1.0), so it goes through the same
+  `_larry_targets` mapping. It wants the full step count, not 4–6.
+
+* `joyfox` (`joyfox/MiniMax-H3-Turbo`) is another 4-step turbo LoRA, ComfyUI-native like realism but covering
+  more of the tree: attention, MLP, the block and final AdaLN projections, and both output heads
+  (`final_layer.video_out` -> `proj_out`, `final_layer.audio_out` -> `audio_proj_out`). Ranks are mixed
+  (32 attention/MLP, 8 modulation/heads) with a scalar `alpha` per key, so each `lora_B` is prescaled by
+  `alpha / rank` at load (here always 1.0). Its `adaln_proj` entries are skipped: the Comfy-Org checkpoint it was
+  trained against gives the modulation projection an 8-dimensional input (`[96768, 8]`), while the diffusers
+  transformer projects the full 2688-dim time embedding (`[96768, 2688]`), so those deltas have no counterpart
+  to fold into.
+
+* `H3-Facial-Realism-CloseUp` (`prithivMLmods/MiniMax-H3-Facial-Realism-CloseUp`) is a close-up facial realism
+  style LoRA, not a turbo one — it wants the full step count, not 4–6. Its layout is detected at load time:
+  PEFT keys against the diffusers tree map name-for-name, everything else is treated as the reference tree
+  (with or without the `diffusion_model.` prefix) and goes through `_larry_targets`. Per-key scalar `alpha`
+  values are folded into `lora_B` as `alpha / rank` when present, otherwise `alpha == rank` is assumed;
+  reference-tree `adaln_proj` entries whose input dim is not the 2688-dim time embedding of the diffusers tree
+  are skipped, like joyfox's.
+
+`H3_LORA` selects the larry file (`off` skips loading it), `H3_LIGHTX=off` skips lightx, `H3_REALISM=off` skips
+realism, `H3_JOYFOX=off` skips joyfox, `H3_LIGHTX8=off` skips the lightx 8-step file, `H3_FACIAL=off` skips the
+facial realism LoRA, `H3_LORA_DEFAULT` picks which set starts folded, and `H3_LORA_STRENGTH` is the larry card's
+sharpness/artifact dial.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+
+LARRY_REPO = os.environ.get("H3_LORA_REPO", "larryvrh/MiniMax-H3-Turbo-Lora")
+LARRY_FILE = os.environ.get("H3_LORA", "minimax_h3_turbo_v4_step600_ema.safetensors")
+LIGHTX_REPO = os.environ.get("H3_LIGHTX_REPO", "lightx2v/Minimax-h3-Turbo")
+LIGHTX_FILE = os.environ.get("H3_LIGHTX_FILE", "minimax_h3_fl2v_turbo_4step_v0.1.safetensors")
+LIGHTX8_FILE = os.environ.get("H3_LIGHTX8_FILE", "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors")
+LIGHTX_ALPHA = 8
+REALISM_REPO = os.environ.get("H3_REALISM_REPO", "fal/MiniMax-H3-Realism-People-LoRA")
+REALISM_FILE = os.environ.get("H3_REALISM_FILE", "h3-realism-people-t2v-i2v-r2v.safetensors")
+JOYFOX_REPO = os.environ.get("H3_JOYFOX_REPO", "joyfox/MiniMax-H3-Turbo")
+JOYFOX_FILE = os.environ.get("H3_JOYFOX_FILE", "minimax_h3_fl2va_4step_lora.safetensors")
+FACIAL_REPO = os.environ.get("H3_FACIAL_REPO", "prithivMLmods/MiniMax-H3-Facial-Realism-CloseUp")
+FACIAL_FILE = os.environ.get("H3_FACIAL_FILE", "minimax-h3-facial-realism-closeup-cp2000.safetensors")
+FACIAL_NAME = "H3-Facial-Realism-CloseUp"
+LARRY_STRENGTH = float(os.environ.get("H3_LORA_STRENGTH", "1.0"))
+DEFAULT_LORA = os.environ.get("H3_LORA_DEFAULT", "larry")
+
+
+def _larry_targets(name: str, b: torch.Tensor, inner_dim: int) -> list[tuple[str, torch.Tensor]]:
+    """Map one reference-tree base name and its `lora_B` onto diffusers parameter key + row-transformed B."""
+    if name.startswith("token_refiner.blocks."):
+        target = name.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1)
+    elif name.startswith("blocks."):
+        target = name.replace("blocks.", "transformer_blocks.", 1)
+    else:
+        target = name
+    target = target.replace("final_layer.adaln_proj.linear", "norm_out.linear")
+    target = target.replace("final_layer.video_out", "proj_out").replace("final_layer.audio_out", "audio_proj_out")
+
+    if target.endswith(".attn.qkv_proj"):
+        prefix = target.removesuffix("qkv_proj")
+        return [
+            (f"{prefix}to_{kind}.weight", part.contiguous())
+            for kind, part in zip(("q", "k", "v"), b.split(inner_dim, dim=0))
+        ]
+    if target.endswith(".mlp.fc1"):
+        gate, value = b.chunk(2, dim=0)
+        return [(target.replace(".mlp.fc1", ".ff.net.0.proj") + ".weight", torch.cat([value, gate]).contiguous())]
+    if target.endswith(".mlp.fc2"):
+        return [(target.replace(".mlp.fc2", ".ff.net.2") + ".weight", b)]
+    if target.endswith(".attn.out_proj"):
+        return [(target.replace(".attn.out_proj", ".attn.to_out.0") + ".weight", b)]
+    return [(target + ".weight", b)]
+
+
+def _load_larry(inner_dim: int) -> dict:
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+
+    lora = load_file(hf_hub_download(LARRY_REPO, LARRY_FILE))
+    bases = sorted({key.rsplit(".lora_", 1)[0] for key in lora})
+    entries = []
+    for name in bases:
+        a = lora[f"{name}.lora_A.weight"]
+        b = lora[f"{name}.lora_B.weight"]
+        entries.extend((key, a, b_part) for key, b_part in _larry_targets(name, b, inner_dim))
+    return {
+        "label": f"{LARRY_REPO}/{LARRY_FILE}",
+        "scale": LARRY_STRENGTH,
+        "entries": entries,
+    }
+
+
+def _load_lightx(file: str = LIGHTX_FILE) -> dict:
+    """A diffusers-native PEFT checkpoint from `lightx2v/Minimax-h3-Turbo` — keys map name-for-name."""
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+
+    lora = load_file(hf_hub_download(LIGHTX_REPO, file))
+    suffix_a, suffix_b = ".lora_A.default.weight", ".lora_B.default.weight"
+    bases = sorted({key[: -len(suffix_a)] for key in lora if key.endswith(suffix_a)})
+    ranks = {lora[f"{name}{suffix_a}"].shape[0] for name in bases}
+    if len(ranks) != 1:
+        raise ValueError(f"Mixed LoRA ranks in {file}: {sorted(ranks)}")
+    entries = [(f"{name}.weight", lora[f"{name}{suffix_a}"], lora[f"{name}{suffix_b}"]) for name in bases]
+    return {
+        "label": f"{LIGHTX_REPO}/{file}",
+        "scale": LIGHTX_ALPHA / ranks.pop(),
+        "entries": entries,
+    }
+
+
+def _load_realism(inner_dim: int) -> dict:
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+
+    lora = load_file(hf_hub_download(REALISM_REPO, REALISM_FILE))
+    bases = sorted({key.rsplit(".lora_", 1)[0].removeprefix("diffusion_model.") for key in lora})
+    entries = []
+    for name in bases:
+        a = lora[f"diffusion_model.{name}.lora_A.weight"]
+        b = lora[f"diffusion_model.{name}.lora_B.weight"]
+        entries.extend((key, a, b_part) for key, b_part in _larry_targets(name, b, inner_dim))
+    return {
+        "label": f"{REALISM_REPO}/{REALISM_FILE}",
+        "scale": 1.0,
+        "entries": entries,
+    }
+
+
+def _load_joyfox(inner_dim: int) -> dict:
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+
+    lora = load_file(hf_hub_download(JOYFOX_REPO, JOYFOX_FILE))
+    bases = sorted({key.rsplit(".lora_", 1)[0].removeprefix("diffusion_model.") for key in lora if ".lora_" in key})
+    entries = []
+    for name in bases:
+        if "adaln_proj" in name:
+            continue
+        prefixed = f"diffusion_model.{name}"
+        a = lora[f"{prefixed}.lora_A.weight"]
+        b = lora[f"{prefixed}.lora_B.weight"] * (lora[f"{prefixed}.alpha"].item() / a.shape[0])
+        entries.extend((key, a, b_part) for key, b_part in _larry_targets(name, b, inner_dim))
+    return {
+        "label": f"{JOYFOX_REPO}/{JOYFOX_FILE}",
+        "scale": 1.0,
+        "entries": entries,
+    }
+
+
+def _load_facial(inner_dim: int) -> dict:
+    """Close-up facial realism style LoRA. The checkpoint layout is detected at load time: PEFT keys against the
+    diffusers tree map name-for-name, everything else is treated as the reference tree (with or without a
+    `diffusion_model.` prefix) and goes through `_larry_targets`. Per-key scalar `alpha` is folded into `lora_B`
+    as `alpha / rank`; reference-tree `adaln_proj` entries whose input dim is not the 2688-dim time embedding of
+    the diffusers tree are skipped, like joyfox's."""
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+
+    lora = load_file(hf_hub_download(FACIAL_REPO, FACIAL_FILE))
+    label = f"{FACIAL_REPO}/{FACIAL_FILE}"
+
+    suffix_a, suffix_b = ".lora_A.default.weight", ".lora_B.default.weight"
+    if any(key.endswith(suffix_a) for key in lora):
+        entries = []
+        for base in sorted({key[: -len(suffix_a)] for key in lora if key.endswith(suffix_a)}):
+            a = lora[f"{base}{suffix_a}"]
+            b = lora[f"{base}{suffix_b}"]
+            alpha = lora.get(f"{base}.alpha")
+            if alpha is not None:
+                b = b * (alpha.item() / a.shape[0])
+            name = base.removeprefix("transformer.")
+            entries.append((f"{name}.weight", a, b))
+        return {"label": label, "scale": 1.0, "entries": entries}
+
+    prefix = "diffusion_model." if any(key.startswith("diffusion_model.") for key in lora) else ""
+    entries = []
+    for name in sorted({key.rsplit(".lora_", 1)[0].removeprefix(prefix) for key in lora if ".lora_" in key}):
+        full = f"{prefix}{name}"
+        a = lora[f"{full}.lora_A.weight"]
+        if "adaln_proj" in name and a.shape[1] != 2688:
+            continue
+        b = lora[f"{full}.lora_B.weight"]
+        alpha = lora.get(f"{full}.alpha")
+        if alpha is not None:
+            b = b * (alpha.item() / a.shape[0])
+        entries.extend((key, a, b_part) for key, b_part in _larry_targets(name, b, inner_dim))
+    return {"label": label, "scale": 1.0, "entries": entries}
+
+
+def _apply(entries, params, sign: float) -> None:
+    for key, a, b in entries:
+        param = params.get(key)
+        if param is None:
+            raise KeyError(f"LoRA target `{key}` not found in the transformer")
+        delta = sign * (b.to(torch.float32) @ a.to(torch.float32))
+        param.data = (param.data.float() + delta.to(param.device)).to(param.dtype)
+
+
+def available() -> list[str]:
+    """The LoRA sets that were loaded at startup, plus `off`."""
+    state = getattr(_PIPE_TRANSFORMER, "_lora_state", None) if _PIPE_TRANSFORMER is not None else None
+    return sorted(state["sets"]) + ["off"] if state else ["off"]
+
+
+_PIPE_TRANSFORMER = None
+
+
+def apply_lora(transformer) -> str | None:
+    """Load every enabled LoRA set, fold the default one into `transformer`, and stash the factors for per-request
+    switching. Returns a status line, or `None` when everything is disabled."""
+    global _PIPE_TRANSFORMER
+    _PIPE_TRANSFORMER = transformer
+
+    inner_dim = transformer.config.num_attention_heads * transformer.config.attention_head_dim
+    sets = {}
+    if LARRY_FILE.lower() not in ("", "off", "none"):
+        sets["larry"] = _load_larry(inner_dim)
+    if os.environ.get("H3_LIGHTX", "on").lower() not in ("", "off", "none"):
+        sets["lightx"] = _load_lightx()
+    if os.environ.get("H3_LIGHTX8", "on").lower() not in ("", "off", "none"):
+        sets["lightx8"] = _load_lightx(LIGHTX8_FILE)
+    if os.environ.get("H3_REALISM", "on").lower() not in ("", "off", "none"):
+        sets["realism"] = _load_realism(inner_dim)
+    if os.environ.get("H3_JOYFOX", "on").lower() not in ("", "off", "none"):
+        sets["joyfox"] = _load_joyfox(inner_dim)
+    if os.environ.get("H3_FACIAL", "on").lower() not in ("", "off", "none"):
+        sets[FACIAL_NAME] = _load_facial(inner_dim)
+    if not sets:
+        return None
+
+    active = DEFAULT_LORA if DEFAULT_LORA in sets else sorted(sets)[0]
+    params = dict(transformer.named_parameters())
+    _apply(sets[active]["entries"], params, sets[active]["scale"])
+    transformer._lora_state = {"active": active, "sets": sets}
+    return (
+        f"LoRAs loaded: "
+        + ", ".join(f"`{name}` ({spec['label']}, {len(spec['entries'])} weights)" for name, spec in sets.items())
+        + f" · active `{active}`"
+    )
+
+
+def set_active(transformer, name: str) -> str:
+    """Switch the folded LoRA in place. No-op when the state already matches. Returns the active set."""
+    state = getattr(transformer, "_lora_state", None)
+    if state is None:
+        return "off"
+    name = name if name in state["sets"] else "off"
+    if state["active"] == name:
+        return name
+    params = dict(transformer.named_parameters())
+    if state["active"] != "off":
+        old = state["sets"][state["active"]]
+        _apply(old["entries"], params, -old["scale"])
+    if name != "off":
+        _apply(state["sets"][name]["entries"], params, state["sets"][name]["scale"])
+    state["active"] = name
+    return name
+
+
+def set_enabled(transformer, enabled: bool) -> bool:
+    """Backwards-compatible boolean toggle over the default set."""
+    return set_active(transformer, DEFAULT_LORA if enabled else "off") != "off"

h3_split_blocks.py

@@ -0,0 +1,147 @@
+"""The halves of a **split** MiniMax-H3 deployment, for both of its checkpoint partitions.
+
+MiniMax-H3 is 195.9 GiB in bfloat16 and a ZeroGPU Space is evicted at 150 GB of storage, so `MiniMaxH3Blocks` is cut
+at its `text_encoder` step: the 62.14 GiB Qwen3-VL runs in the conditioner Space, everything else in a generator
+Space, and `prompt_embeds` + `text_token_tags` is the whole wire format between them.
+
+`resize` / `setup` run on **both** sides: they own no pretrained component, and each half needs the canvas and the
+prepared keyframes or normalized references. Both conditioner halves also return the resolved `height` / `width` /
+`num_frames`, which the generating half pins rather than re-deriving.
+
+Two things the blocks leave to the caller: a keyframe reaches them EXIF-transposed and in RGB, and the `t2va` / `fl2va`
+frame count is aligned to `17 * n + 5` before the call, since that arithmetic lives on the denoising side of the cut.
+"""
+
+from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
+from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
+from diffusers.modular_pipelines.minimax_h3.encoders import (
+    MiniMaxH3Ref2VAReferenceEncoderStep,
+    MiniMaxH3Ref2VATextEncoderStep,
+    MiniMaxH3TextEncoderStep,
+)
+from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
+    MiniMaxH3AutoKeyframeVaeEncoderStep,
+    MiniMaxH3AutoResizeStep,
+    MiniMaxH3CoreDenoiseStep,
+    MiniMaxH3DecodeStep,
+    MiniMaxH3Ref2VACoreDenoiseStep,
+    _generation_outputs,
+)
+from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam
+
+
+def _wire_outputs(num_frames: bool = True) -> list[OutputParam]:
+    """The wire format of the split. `num_frames` is declared by the `ref2va` half alone, whose setup resolves one."""
+    return [
+        OutputParam.template("prompt_embeds"),
+        OutputParam("text_token_tags", description="The per-row modality tag of every row of `prompt_embeds`."),
+        OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."),
+        OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."),
+        *(
+            [OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5.")]
+            if num_frames
+            else []
+        ),
+    ]
+
+
+class MiniMaxH3ConditionerBlocks(SequentialPipelineBlocks):
+    """The conditioner half of a split MiniMax-H3: the keyframes on the canvas plus the Qwen3-VL read at layer 50."""
+
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3AutoResizeStep, MiniMaxH3TextEncoderStep]
+    block_names = ["resize", "text_encoder"]
+
+    @property
+    def description(self):
+        return (
+            "The conditioner half of a split MiniMax-H3 deployment: puts the keyframes onto the target canvas and "
+            "encodes MiniMax-H3's presentation of the request into the `prompt_embeds` / `text_token_tags` pair the "
+            "denoising half consumes. The frame count is the caller's to align."
+        )
+
+    @property
+    def outputs(self):
+        return _wire_outputs(num_frames=False)
+
+
+class MiniMaxH3GeneratorBlocks(SequentialPipelineBlocks):
+    """The denoising half of a split MiniMax-H3: `MiniMaxH3Blocks` with its `text_encoder` step removed."""
+
+    model_name = "minimax-h3"
+    block_classes = [
+        MiniMaxH3AutoResizeStep,
+        MiniMaxH3AutoKeyframeVaeEncoderStep,
+        MiniMaxH3CoreDenoiseStep,
+        MiniMaxH3AfterDenoiseStep,
+        MiniMaxH3DecodeStep,
+    ]
+    block_names = ["resize", "vae_encoder", "denoise", "after_denoise", "decode"]
+
+    @property
+    def description(self):
+        return (
+            "The denoising half of a split MiniMax-H3 deployment: the `t2va` / `fl2va` branch of `MiniMaxH3Blocks` "
+            "without its text-encoder step, so `prompt_embeds` and `text_token_tags` come in as inputs and the "
+            "62.14 GiB Qwen3-VL conditioner is never loaded here."
+        )
+
+    @property
+    def outputs(self):
+        return _generation_outputs()
+
+
+class MiniMaxH3Ref2VAConditionerBlocks(SequentialPipelineBlocks):
+    """The conditioner half of a split `ref2va`: the resolved plan plus the Qwen3-VL read at its 50th layer.
+
+    Component for component this is `MiniMaxH3ConditionerBlocks`, so one conditioner Space serves both partitions.
+    What differs is the presentation: `ref2va` prepends a label per reference and a vision block per image and per
+    merged video frame pair, so the references themselves have to reach this half.
+    """
+
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3Ref2VASetupStep, MiniMaxH3Ref2VATextEncoderStep]
+    block_names = ["setup", "text_encoder"]
+
+    @property
+    def description(self):
+        return (
+            "The conditioner half of a split MiniMax-H3 `ref2va` deployment: resolves the request plan (canvas, frame "
+            "count, references normalized onto MiniMax-H3's own rates and resolutions) and encodes MiniMax-H3's "
+            "presentation of it into the `prompt_embeds` / `text_token_tags` pair the denoising half consumes."
+        )
+
+    @property
+    def outputs(self):
+        return _wire_outputs()
+
+
+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"
+    block_classes = [
+        MiniMaxH3Ref2VASetupStep,
+        MiniMaxH3Ref2VAReferenceEncoderStep,
+        MiniMaxH3Ref2VACoreDenoiseStep,
+        MiniMaxH3AfterDenoiseStep,
+        MiniMaxH3DecodeStep,
+    ]
+    block_names = ["setup", "reference_encoder", "denoise", "after_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.14 GiB Qwen3-VL conditioner is never loaded here. The transformer is the `transformer_ref` partition."
+        )
+
+    @property
+    def outputs(self):
+        return _generation_outputs()