Add the Krea 2 per-arch eager fusion (Krea2TransformerBlock addcmul patch)

Fuses the block's two inline modulations (1 + scale) * norm(x) + shift and two
gated residuals x + gate * out to torch.addcmul, matching the existing qwen /
z-image / flux fusions (compile-safe, 1-ULP more accurate, body-drift guarded).
Stock-vs-patched equivalence test included; install count is now 7.

studio/backend/core/inference/diffusion_arch_patches.py

@@ -27,13 +27,14 @@ source-body check (``_body_has``) confirms the exact lines we rewrite are still
 future diffusers that changed the block body simply leaves the block UNPATCHED (correctness
 first) rather than running a stale copy. Kill-switch: ``UNSLOTH_DIFFUSION_ARCH_PATCHES=0``.
 
-Implemented for all four families (extend by adding entries to ``_SPECS``):
+Implemented for all five families (extend by adding entries to ``_SPECS``):
   * qwen-image  ``QwenImageTransformerBlock._modulate`` -- modulation addcmul (all 4 sites).
   * z-image     ``ZImageTransformerBlock.forward``      -- the 2 gated-residual addcmuls.
   * flux.1      ``FluxTransformerBlock.forward`` (inline norm2 modulation + 4 gated residuals)
                 + ``FluxSingleTransformerBlock.forward`` (residual + gate*proj_out).
   * flux.2-klein ``Flux2TransformerBlock.forward`` (4 inline modulations + 4 gated residuals)
                 + ``Flux2SingleTransformerBlock.forward`` (inline modulation + gated residual).
+  * krea-2      ``Krea2TransformerBlock.forward`` (2 inline modulations + 2 gated residuals).
 """
 
 from __future__ import annotations
@@ -468,6 +469,53 @@ def _spec_flux2_single():
     return (cls, "forward", _flux2_single_forward)
 
 
+# =====================================================================================
+# krea-2: Krea2TransformerBlock.forward  (2 inline modulations + 2 gated residuals)
+# =====================================================================================
+def _krea2_block_forward(
+    self,
+    hidden_states,
+    temb,
+    image_rotary_emb,
+    attention_mask = None,
+):
+    """diffusers 0.39 ``Krea2TransformerBlock.forward`` with the two inline modulations
+    ``(1 + scale) * norm(x) + shift`` and the two gated residuals ``x + gate * out`` each
+    fused to one ``torch.addcmul``."""
+    # temb: (B, 1, 6 * hidden_size), shared across all blocks; each block only learns an
+    # additive table.
+    modulation = temb.unflatten(-1, (6, -1)) + self.scale_shift_table
+    prescale, preshift, pregate, postscale, postshift, postgate = modulation.unbind(-2)
+
+    norm1 = self.norm1(hidden_states)
+    attn_out = self.attn(
+        torch.addcmul(preshift, norm1, 1 + prescale),
+        attention_mask = attention_mask,
+        image_rotary_emb = image_rotary_emb,
+    )
+    hidden_states = torch.addcmul(hidden_states, pregate, attn_out)
+    norm2 = self.norm2(hidden_states)
+    ff_out = self.ff(torch.addcmul(postshift, norm2, 1 + postscale))
+    return torch.addcmul(hidden_states, postgate, ff_out)
+
+
+def _spec_krea2_forward():
+    try:
+        from diffusers.models.transformers.transformer_krea2 import Krea2TransformerBlock as cls
+    except Exception:  # noqa: BLE001
+        return None
+    orig = getattr(cls, "forward", None)
+    if orig is None or not _body_has(
+        orig,
+        "(1.0 + prescale) * self.norm1(hidden_states) + preshift",
+        "hidden_states = hidden_states + pregate * attn_out",
+        "(1.0 + postscale) * self.norm2(hidden_states) + postshift",
+        "hidden_states = hidden_states + postgate * ff_out",
+    ):
+        return None
+    return (cls, "forward", _krea2_block_forward)
+
+
 # =====================================================================================
 # registry + lifecycle
 # =====================================================================================
@@ -480,6 +528,7 @@ _SPECS: tuple[Callable[[], Optional[tuple]], ...] = (
     _spec_flux_single,
     _spec_flux2_double,
     _spec_flux2_single,
+    _spec_krea2_forward,
 )
 
 # (cls, attr) pairs we successfully patched, for an exact reverse.

studio/backend/tests/test_diffusion_arch_patches.py

@@ -236,6 +236,34 @@ def test_flux2_single_forward_matches_stock():
     _close_any(got, ref)
 
 
+def test_krea2_forward_matches_stock():
+    from diffusers.models.transformers.transformer_krea2 import (
+        Krea2RotaryPosEmbed,
+        Krea2TransformerBlock,
+    )
+    from diffusers.pipelines.krea2.pipeline_krea2 import Krea2Pipeline
+
+    torch.manual_seed(4)
+    blk = Krea2TransformerBlock(
+        hidden_size = D, intermediate_size = 2 * D, num_heads = H, num_kv_heads = H // 2,
+        norm_eps = 1e-6,
+    ).eval()
+    # Give the zero-init modulation table real values so all six scale/shift/gate
+    # branches contribute to the output.
+    with torch.no_grad():
+        blk.scale_shift_table.normal_()
+    # A [text + 2x2 image grid] sequence with the real rotary embed (axes sum to head_dim).
+    position_ids = Krea2Pipeline.prepare_position_ids(4, 2, 2, torch.device("cpu"))
+    rope = Krea2RotaryPosEmbed(theta = 10000, axes_dim = [D // H // 2, D // H // 4, D // H // 4])
+    image_rotary_emb = rope(position_ids)
+    hs = torch.randn(B, position_ids.shape[0], D)
+    tm = torch.randn(B, 1, 6 * D)
+    with torch.inference_mode():
+        ref = Krea2TransformerBlock.forward(blk, hs, tm, image_rotary_emb).clone()
+        got = ap._krea2_block_forward(blk, hs, tm, image_rotary_emb)
+    torch.testing.assert_close(got, ref, atol = 1e-5, rtol = 1e-4)
+
+
 # ── lifecycle ───────────────────────────────────────────────────────────────────
 
 
@@ -246,7 +274,7 @@ def test_install_idempotent_and_reversible():
     q_orig, z_orig = Q._modulate, Z.forward
     n1 = ap.install_arch_patches()
     n2 = ap.install_arch_patches()  # idempotent
-    assert n1 == 6 and n2 == n1  # qwen + z-image + flux.1 x2 + flux.2 x2
+    assert n1 == 7 and n2 == n1  # qwen + z-image + flux.1 x2 + flux.2 x2 + krea-2
     assert Q._modulate is not q_orig and Z.forward is not z_orig
     assert ap.is_installed()