36df317293
Comment-only pass over the Python this PR touches: drop what the code already says, collapse multi-line explanations that still read on one line, and keep the reasoning that is not recoverable from the code. No code, docstring semantics or behaviour changes; verified with an AST comparison against the previous revision, and the backend suite is unchanged (same 37 environment failures as before: the API integration tests that need a live keyed server, the flash-attn install hooks, and the GPU memory fields).
305 lines
11 KiB
Python
305 lines
11 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Numerical + lifecycle tests for the per-arch eager fusions (``diffusion_arch_patches``).
|
|
|
|
Each per-arch patch only fuses ``a + b*c`` -> ``torch.addcmul`` (1-ULP, more accurate), so
|
|
the patched method/forward must match the stock diffusers one within fp tolerance. We also
|
|
check install/uninstall reversibility + idempotency, the kill-switch, and the body-drift
|
|
guard (a diffusers whose block body changed is left unpatched). Runs on CPU.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import types
|
|
|
|
import pytest
|
|
|
|
torch = pytest.importorskip("torch")
|
|
pytest.importorskip("diffusers")
|
|
|
|
from core.inference import diffusion_arch_patches as ap # noqa: E402
|
|
|
|
|
|
@pytest.fixture(autouse = True)
|
|
def _clean():
|
|
ap.uninstall_arch_patches()
|
|
yield
|
|
ap.uninstall_arch_patches()
|
|
|
|
|
|
# ── qwen-image _modulate (modulation addcmul) ───────────────────────────────────
|
|
|
|
|
|
def test_qwen_modulate_matches_stock_global_and_indexed():
|
|
from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock as Q
|
|
|
|
B, L, D = 2, 16, 64
|
|
x = torch.randn(B, L, D)
|
|
mod = torch.randn(B, 3 * D)
|
|
# _modulate uses no real `self` state, so call it unbound with self=None.
|
|
ref_x, ref_g = Q._modulate(None, x, mod)
|
|
got_x, got_g = ap._qwen_modulate(None, x, mod)
|
|
torch.testing.assert_close(got_x, ref_x, atol = 1e-5, rtol = 1e-4)
|
|
assert torch.equal(got_g, ref_g)
|
|
|
|
# per-token `index` branch (mod batch is 2*B).
|
|
idx = torch.randint(0, 2, (B, L))
|
|
mod2 = torch.randn(2 * B, 3 * D)
|
|
ref2_x, ref2_g = Q._modulate(None, x, mod2, idx)
|
|
got2_x, got2_g = ap._qwen_modulate(None, x, mod2, idx)
|
|
torch.testing.assert_close(got2_x, ref2_x, atol = 1e-5, rtol = 1e-4)
|
|
assert torch.equal(got2_g, ref2_g)
|
|
|
|
|
|
# ── z-image block forward (gated-residual addcmul) ──────────────────────────────
|
|
|
|
|
|
class _AttnStub(torch.nn.Module):
|
|
"""Deterministic stand-in for ZImageAttention so the block forward runs without RoPE /
|
|
freqs_cis (the patch only changes the residual adds, which is what we validate)."""
|
|
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.proj = torch.nn.Linear(dim, dim, bias = False)
|
|
|
|
def forward(self, h, **kwargs):
|
|
return self.proj(h)
|
|
|
|
|
|
def _zimage_block(dim = 64, heads = 4):
|
|
from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock
|
|
|
|
blk = ZImageTransformerBlock(
|
|
layer_id = 0,
|
|
dim = dim,
|
|
n_heads = heads,
|
|
n_kv_heads = heads,
|
|
norm_eps = 1e-5,
|
|
qk_norm = True,
|
|
modulation = True,
|
|
).eval()
|
|
blk.attention = _AttnStub(dim).eval()
|
|
return blk
|
|
|
|
|
|
def _adaln_dim(dim):
|
|
from diffusers.models.transformers.transformer_z_image import ADALN_EMBED_DIM
|
|
return min(dim, ADALN_EMBED_DIM)
|
|
|
|
|
|
def test_zimage_forward_matches_stock_global_modulation():
|
|
from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock
|
|
|
|
torch.manual_seed(0)
|
|
blk = _zimage_block()
|
|
B, L, D = 2, 16, 64
|
|
x = torch.randn(B, L, D)
|
|
adaln = torch.randn(B, _adaln_dim(D))
|
|
|
|
with torch.inference_mode():
|
|
ref = ZImageTransformerBlock.forward(blk, x, None, None, adaln_input = adaln).clone()
|
|
got = ap._zimage_forward(blk, x, None, None, adaln_input = adaln)
|
|
torch.testing.assert_close(got, ref, atol = 1e-5, rtol = 1e-4)
|
|
|
|
|
|
def test_zimage_forward_matches_stock_per_token_modulation():
|
|
from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock
|
|
|
|
torch.manual_seed(1)
|
|
blk = _zimage_block()
|
|
B, L, D = 2, 16, 64
|
|
x = torch.randn(B, L, D)
|
|
ad = _adaln_dim(D)
|
|
adaln_noisy = torch.randn(B, ad)
|
|
adaln_clean = torch.randn(B, ad)
|
|
noise_mask = torch.randint(0, 2, (B, L))
|
|
|
|
with torch.inference_mode():
|
|
ref = ZImageTransformerBlock.forward(
|
|
blk,
|
|
x,
|
|
None,
|
|
None,
|
|
noise_mask = noise_mask,
|
|
adaln_noisy = adaln_noisy,
|
|
adaln_clean = adaln_clean,
|
|
).clone()
|
|
got = ap._zimage_forward(
|
|
blk,
|
|
x,
|
|
None,
|
|
None,
|
|
noise_mask = noise_mask,
|
|
adaln_noisy = adaln_noisy,
|
|
adaln_clean = adaln_clean,
|
|
)
|
|
torch.testing.assert_close(got, ref, atol = 1e-5, rtol = 1e-4)
|
|
|
|
|
|
# ── flux.1 / flux.2 block forwards (modulation + gated-residual addcmul) ─────────
|
|
|
|
|
|
class _Tuple2AttnStub(torch.nn.Module):
|
|
"""Double-stream attention stub -> (img_out, ctx_out)."""
|
|
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.pi = torch.nn.Linear(dim, dim, bias = False)
|
|
self.pc = torch.nn.Linear(dim, dim, bias = False)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states,
|
|
encoder_hidden_states = None,
|
|
**kwargs,
|
|
):
|
|
return self.pi(hidden_states), self.pc(encoder_hidden_states)
|
|
|
|
|
|
class _SingleAttnStub(torch.nn.Module):
|
|
"""Single-stream attention stub -> tensor."""
|
|
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.p = torch.nn.Linear(dim, dim, bias = False)
|
|
|
|
def forward(self, hidden_states, **kwargs):
|
|
return self.p(hidden_states)
|
|
|
|
|
|
def _close_any(got, ref):
|
|
if isinstance(ref, tuple):
|
|
assert len(got) == len(ref)
|
|
for g, r in zip(got, ref):
|
|
torch.testing.assert_close(g, r, atol = 1e-5, rtol = 1e-4)
|
|
else:
|
|
torch.testing.assert_close(got, ref, atol = 1e-5, rtol = 1e-4)
|
|
|
|
|
|
D, H = 64, 4
|
|
B, L, LC = 2, 16, 8
|
|
|
|
|
|
def test_flux_double_forward_matches_stock():
|
|
from diffusers.models.transformers.transformer_flux import FluxTransformerBlock
|
|
|
|
torch.manual_seed(0)
|
|
blk = FluxTransformerBlock(dim = D, num_attention_heads = H, attention_head_dim = D // H).eval()
|
|
blk.attn = _Tuple2AttnStub(D).eval()
|
|
hs, ehs, temb = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, D)
|
|
with torch.inference_mode():
|
|
ref = FluxTransformerBlock.forward(blk, hs, ehs, temb)
|
|
got = ap._flux_double_forward(blk, hs, ehs, temb)
|
|
_close_any(got, ref)
|
|
|
|
|
|
def test_flux_single_forward_matches_stock():
|
|
from diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock
|
|
|
|
torch.manual_seed(1)
|
|
blk = FluxSingleTransformerBlock(dim = D, num_attention_heads = H, attention_head_dim = D // H).eval()
|
|
blk.attn = _SingleAttnStub(D).eval()
|
|
hs, ehs, temb = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, D)
|
|
with torch.inference_mode():
|
|
ref = FluxSingleTransformerBlock.forward(blk, hs, ehs, temb)
|
|
got = ap._flux_single_forward(blk, hs, ehs, temb)
|
|
_close_any(got, ref)
|
|
|
|
|
|
def test_flux2_double_forward_matches_stock():
|
|
from diffusers.models.transformers.transformer_flux2 import Flux2TransformerBlock
|
|
|
|
torch.manual_seed(2)
|
|
blk = Flux2TransformerBlock(dim = D, num_attention_heads = H, attention_head_dim = D // H).eval()
|
|
blk.attn = _Tuple2AttnStub(D).eval()
|
|
hs, ehs = torch.randn(B, L, D), torch.randn(B, LC, D)
|
|
tmi, tmt = torch.randn(B, 6 * D), torch.randn(B, 6 * D)
|
|
with torch.inference_mode():
|
|
ref = Flux2TransformerBlock.forward(blk, hs, ehs, tmi, tmt)
|
|
got = ap._flux2_double_forward(blk, hs, ehs, tmi, tmt)
|
|
_close_any(got, ref)
|
|
|
|
|
|
def test_flux2_single_forward_matches_stock():
|
|
from diffusers.models.transformers.transformer_flux2 import Flux2SingleTransformerBlock
|
|
|
|
torch.manual_seed(3)
|
|
blk = Flux2SingleTransformerBlock(
|
|
dim = D, num_attention_heads = H, attention_head_dim = D // H
|
|
).eval()
|
|
blk.attn = _SingleAttnStub(D).eval()
|
|
hs, ehs, tm = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, 3 * D)
|
|
with torch.inference_mode():
|
|
ref = Flux2SingleTransformerBlock.forward(blk, hs, ehs, tm)
|
|
got = ap._flux2_single_forward(blk, hs, ehs, tm)
|
|
_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 ───────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_install_idempotent_and_reversible():
|
|
from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock as Q
|
|
from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock as Z
|
|
|
|
q_orig, z_orig = Q._modulate, Z.forward
|
|
n1 = ap.install_arch_patches()
|
|
n2 = ap.install_arch_patches() # idempotent
|
|
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()
|
|
|
|
ap.uninstall_arch_patches()
|
|
assert not ap.is_installed()
|
|
assert Q._modulate is q_orig and Z.forward is z_orig # exact restore
|
|
ap.uninstall_arch_patches() # idempotent
|
|
|
|
|
|
def test_kill_switch(monkeypatch):
|
|
from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock as Q
|
|
|
|
monkeypatch.setenv("UNSLOTH_DIFFUSION_ARCH_PATCHES", "0")
|
|
orig = Q._modulate
|
|
assert ap.install_arch_patches() == 0
|
|
assert not ap.is_installed()
|
|
assert Q._modulate is orig
|
|
|
|
|
|
def test_body_drift_guard_skips_changed_block(monkeypatch):
|
|
# A resolver whose body-check fails (diffusers changed the lines we rewrite) is skipped. Force
|
|
# the qwen resolver to see a drifted body.
|
|
monkeypatch.setattr(ap, "_body_has", lambda fn, *needles: False)
|
|
assert ap.install_arch_patches() == 0
|
|
assert not ap.is_installed()
|