Upload h3_eros.py

h3_eros.py

@@ -0,0 +1,139 @@
+"""10Eros-Max (TenStrip) beta_5 support: replace the diffusers transformer's weights with the Eros checkpoint.
+
+`TenStrip/10Eros-Max` ships full bf16 checkpoints in the *reference* (ComfyUI) module tree, not LoRAs, so the update is
+a straight weight replacement rather than a fold. Each tensor is renamed and row-transformed exactly like the
+`larryvrh` LoRA deltas in `h3_lora` (fused-QKV thirds onto `attn.to_q/k/v`, the `SwiGLU` gate/value swap onto
+`ff.net.0.proj`, `blocks.` -> `transformer_blocks.`, ...). The mapping was checked against the base weights: block 0
+q/k/v/out_proj/fc1 all sit within ~1% cosine of MiniMax-H3 after that transform, and fc1 does not match without the swap.
+
+AdaLN: the reference file stores the projections rank-8 compressed over the 1025-step timestep grid — `adaln_t_table[i] =
+adaln_basis @ (silu(temb(i / 1024)) - adaln_mean)` and `*.adaln_proj.linear` of shape `[out, 8]`. Diffusers computes the
+full `linear(silu(temb))`, so each projection is rebuilt as `W = W_ref @ basis`, `b = b_ref - W @ mean`. Checked against
+the base model (time embedder from MiniMax-H3, which the file does not carry): the table matches at cos 0.9998 and the
+rebuilt output matches the reference form at cos 1.0000 (0.9993-0.9999 vs. base weights). Exact on the sampler's
+timestep grid, a rank-8 approximation off it; the time embedder itself stays base. `rope.inv_freq` is skipped.
+
+Tensors are streamed one at a time with `safe_open`, so peak CPU memory is one tensor rather than the 40 GB file.
+
+`H3_EROS=off` skips it, `H3_EROS_REPO` / `H3_EROS_FILE` pick the checkpoint. A `TURBO` file has the 4-step delta baked
+in, which is why `app.py` then defaults the LoRA selector to `off` — stacking larry/lightx on top double-applies it.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+
+EROS_REPO = os.environ.get("H3_EROS_REPO", "TenStrip/10Eros-Max")
+EROS_FILE = os.environ.get("H3_EROS_FILE", "10Eros_Max_h3_TURBO-hybrid_beta5.safetensors")
+
+_SKIP_EXACT = {"rope.inv_freq", "adaln_basis", "adaln_mean", "adaln_t_table"}
+_DIRECT = {
+    "video_patch_proj": "proj_in",
+    "audio_patch_proj": "audio_proj_in",
+    "condition_proj": "context_embedder",
+    "final_layer.video_out": "proj_out",
+    "final_layer.audio_out": "audio_proj_out",
+    "final_layer.norm": "norm_out.norm",
+}
+
+
+def enabled() -> bool:
+    return EROS_FILE.lower() not in ("", "off", "none") and os.environ.get("H3_EROS", "on").lower() not in ("", "off", "none")
+
+
+def turbo_baked() -> bool:
+    return enabled() and "turbo" in EROS_FILE.lower()
+
+
+def adaln_target(name: str) -> str | None:
+    """Diffusers module prefix for a compressed reference AdaLN projection tensor, else `None`."""
+    if name.startswith("final_layer.adaln_proj.linear."):
+        return "norm_out.linear"
+    if name.startswith("blocks.") and ".adaln_proj.linear." in name:
+        return name.replace("blocks.", "transformer_blocks.", 1).rsplit(".", 1)[0]
+    return None
+
+
+def rebuild_adaln(w_ref: torch.Tensor, b_ref: torch.Tensor, basis: torch.Tensor, mean: torch.Tensor):
+    """Rank-8 reference projection -> full `[out, time_embed_dim]` weight and bias (float32)."""
+    w = w_ref.float() @ basis.float()
+    return w, b_ref.float() - w @ mean.float()
+
+
+def map_tensor(name: str, tensor: torch.Tensor, inner_dim: int) -> list[tuple[str, torch.Tensor]]:
+    """Reference-tree tensor -> [(diffusers key, tensor)]. Returns [] for tensors with no diffusers counterpart."""
+    if name in _SKIP_EXACT or adaln_target(name) is not None:
+        return []
+    stem, _, leaf = name.rpartition(".")  # leaf is `weight` / `bias`
+    if stem in _DIRECT:
+        return [(f"{_DIRECT[stem]}.{leaf}", tensor)]
+
+    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  # token_refiner.final_norm.weight
+
+    if target.endswith(".attn.qkv_proj.weight"):
+        prefix = target.removesuffix("qkv_proj.weight")
+        return [
+            (f"{prefix}to_{kind}.weight", part.contiguous())
+            for kind, part in zip(("q", "k", "v"), tensor.split(inner_dim, dim=0))
+        ]
+    if target.endswith(".mlp.fc1.weight"):
+        gate, value = tensor.chunk(2, dim=0)
+        return [(target.replace(".mlp.fc1.weight", ".ff.net.0.proj.weight"), torch.cat([value, gate]).contiguous())]
+    if target.endswith(".mlp.fc2.weight"):
+        return [(target.replace(".mlp.fc2.weight", ".ff.net.2.weight"), tensor)]
+    if target.endswith(".attn.out_proj.weight"):
+        return [(target.replace(".attn.out_proj.weight", ".attn.to_out.0.weight"), tensor)]
+    if target.endswith(".attn.q_norm.weight"):
+        return [(target.replace(".attn.q_norm.weight", ".attn.norm_q.weight"), tensor)]
+    if target.endswith(".attn.k_norm.weight"):
+        return [(target.replace(".attn.k_norm.weight", ".attn.norm_k.weight"), tensor)]
+    return [(target, tensor)]  # norm1 / norm2 / token_refiner.final_norm
+
+
+def apply_eros(transformer) -> str | None:
+    """Replace the transformer's weights in place with the Eros checkpoint. Returns a status line, or `None` if off."""
+    if not enabled():
+        return None
+    from huggingface_hub import hf_hub_download
+    from safetensors import safe_open
+
+    inner_dim = transformer.config.num_attention_heads * transformer.config.attention_head_dim
+    params = dict(transformer.named_parameters())
+    path = hf_hub_download(EROS_REPO, EROS_FILE)
+
+    replaced, skipped = 0, []  # skipped: reference-only tensors (basis / mean / table / rope)
+    with safe_open(path, framework="pt", device="cpu") as f:
+        basis, mean = f.get_tensor("adaln_basis"), f.get_tensor("adaln_mean")
+        for name in f.keys():
+            if name.endswith("adaln_proj.linear.weight") and (prefix := adaln_target(name)):
+                w, b = rebuild_adaln(f.get_tensor(name), f.get_tensor(name.replace(".weight", ".bias")), basis, mean)
+                for key, tensor in ((f"{prefix}.weight", w), (f"{prefix}.bias", b)):
+                    if key not in params or tuple(params[key].shape) != tuple(tensor.shape):
+                        raise ValueError(f"AdaLN `{key}`: transformer vs Eros {tuple(tensor.shape)} mismatch")
+                    params[key].data.copy_(tensor.to(params[key].dtype))
+                    replaced += 1
+                continue
+            targets = map_tensor(name, f.get_tensor(name), inner_dim)
+            if not targets:
+                skipped.append(name)
+                continue
+            for key, tensor in targets:
+                param = params.get(key)
+                if param is None:
+                    raise KeyError(f"Eros tensor `{name}` -> `{key}` not found in the transformer")
+                if tuple(param.shape) != tuple(tensor.shape):
+                    raise ValueError(f"`{key}`: transformer {tuple(param.shape)} vs Eros {tuple(tensor.shape)}")
+                param.data.copy_(tensor.to(param.dtype))
+                replaced += 1
+
+    return (
+        f"Eros `{EROS_REPO}/{EROS_FILE}`: {replaced} weights replaced, {len(skipped)} reference-only tensors skipped "
+        f"(AdaLN rebuilt from rank-8){' · turbo baked in' if turbo_baked() else ''}"
+    )