diff --git a/studio/backend/core/inference/diffusion_arch_patches.py b/studio/backend/core/inference/diffusion_arch_patches.py index e079e09603..f73085197e 100644 --- a/studio/backend/core/inference/diffusion_arch_patches.py +++ b/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. diff --git a/studio/backend/tests/test_diffusion_arch_patches.py b/studio/backend/tests/test_diffusion_arch_patches.py index 77276a0fb1..02e975c5a8 100644 --- a/studio/backend/tests/test_diffusion_arch_patches.py +++ b/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()