File size: 6,539 Bytes
cf8a54e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """CPU tests for the Krea 2 / Qwen-Image 2.1 UltraTex port (tiny random weights)."""
import math
import sys
import os
import pytest
import torch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from ultratex.backbones import common # noqa: E402
from ultratex.backbones.common import block_causal_attention, dense_attention, foreground_indices, gather_sequence # noqa: E402
from ultratex.backbones.krea2.model import Krea2, Krea2Params # noqa: E402
from ultratex.backbones.qwen21.model import Qwen21Params, QwenImage21 # noqa: E402
from ultratex.backbones.qwen21.util import pack_image as qwen_pack # noqa: E402
from ultratex.backbones.krea2.util import pack_image as krea_pack # noqa: E402
from ultratex import mv # noqa: E402
torch.manual_seed(0)
def tiny_krea2():
p = Krea2Params(features=64, tdim=16, txtdim=32, heads=4, kvheads=2, multiplier=1, layers=2, patch=2, channels=4,
txtlayers=3, txtheads=2, txtkvheads=2)
return Krea2(p).float()
def tiny_qwen21():
p = Qwen21Params(in_channels=8, out_channels=8, num_layers=2, attention_head_dim=16, num_attention_heads=4,
context_in_dim=24, mlp_ratio=2, axes_dims_rope=(4, 6, 6))
return QwenImage21(p).float()
def test_block_causal_matches_masked_sdpa():
b, h, n, d = 1, 2, 23, 8
q, k, v = (torch.randn(b, h, n, d) for _ in range(3))
segments = [(0, 5, True), (5, 12, False), (12, 17, False), (17, 23, False)]
out = block_causal_attention(q, k, v, segments)
mask = torch.zeros(n, n, dtype=torch.bool)
for start, end, causal in segments:
for i in range(start, end):
mask[i, :end] = True
if causal:
mask[i, i + 1:end] = False
ref = dense_attention(q, k, v, mask)
assert torch.allclose(out, ref, atol=1e-5)
def test_krea2_forward_and_timestep_zero():
m = tiny_krea2()
latent = torch.randn(1, 4, 8, 12)
tokens, ids = krea_pack(latent, 0, patch=2)
ref_tokens, ref_ids = krea_pack(torch.randn(1, 4, 8, 12), 1, patch=2)
assert tokens.shape == (1, 24, 16) and ids.shape == (1, 24, 3)
assert ids[0, -1].tolist() == [0.0, 3.0, 5.0]
ctx = torch.randn(1, 5, 3 * 32)
img = torch.cat([tokens, ref_tokens], 1)
img_ids = torch.cat([ids, ref_ids], 1)
t = torch.tensor([0.4])
out = m(img, img_ids, t, ctx)
assert out.shape == (1, 48, 16)
out0 = m(img, img_ids, t, ctx, timestep_zero_index=24)
assert out0.shape == out.shape and not torch.allclose(out, out0)
# dropping background tokens keeps the survivors' ids
keep = torch.tensor([[0, 3, 7, 20]])
kept_ids = gather_sequence(ids, keep)
assert kept_ids[0, 2].tolist() == ids[0, 7].tolist()
def test_qwen21_pack_centres_on_target():
latent = torch.randn(1, 8, 4, 6)
tokens, ids = qwen_pack(latent, frame=7, target_hw=(4, 6))
assert tokens.shape == (1, 24, 8)
assert ids[0, 0].tolist() == [7.0, -2.0, -3.0]
assert ids[0, -1].tolist() == [7.0, 1.0, 2.0]
# odd reference on an even target gets the half-token parity shift
_, ids_odd = qwen_pack(torch.randn(1, 8, 3, 5), frame=1, target_hw=(4, 6))
assert ids_odd[0, 0, 1].item() == pytest.approx(-2 + 0.5)
def test_qwen21_prefix_cache_matches_full_pass():
m = tiny_qwen21().eval()
target_latent = torch.randn(1, 8, 4, 6)
ref_latent = torch.randn(1, 8, 4, 6)
ctx = torch.randn(1, 3, 24)
txt_len = ctx.shape[1]
ref_tokens, ref_ids = qwen_pack(ref_latent, frame=txt_len, target_hw=(4, 6))
tgt_tokens, tgt_ids = qwen_pack(target_latent, frame=txt_len + 6, target_hw=(4, 6))
# background-prune the reference and the target
ref_keep = torch.tensor([[1, 2, 5, 9, 13]])
tgt_keep = torch.tensor([[0, 4, 8, 9, 10, 20]])
ref_tokens, ref_ids = gather_sequence(ref_tokens, ref_keep), gather_sequence(ref_ids, ref_keep)
tgt_tokens, tgt_ids = gather_sequence(tgt_tokens, tgt_keep), gather_sequence(tgt_ids, tgt_keep)
img = torch.cat([ref_tokens, tgt_tokens], 1)
ids = torch.cat([ref_ids, tgt_ids], 1)
t = torch.tensor([0.7])
with torch.no_grad():
full = m(img, ids, t, ctx, ref_lengths=[5], attention="block_causal")
cache = {}
filled = m(img, ids, t, ctx, ref_lengths=[5], attention="block_causal", prefix_cache=cache)
assert "blocks" in cache and cache["prefix_len"] == txt_len + 5
cached = m(tgt_tokens, tgt_ids, torch.tensor([0.2]), ctx, attention="block_causal", prefix_cache=cache)
again = m(img, ids, torch.tensor([0.2]), ctx, ref_lengths=[5], attention="block_causal")
assert full.shape == (1, 6, 8)
assert torch.allclose(full, filled, atol=1e-5)
# the prefix is timestep-independent, so cached target rows equal a fresh full pass at the new t
assert torch.allclose(cached, again, atol=1e-4)
# and differ from bidirectional attention
bidir = m(img, ids, torch.tensor([0.2]), ctx, ref_lengths=[5], attention="full")
assert not torch.allclose(bidir, again, atol=1e-4)
def test_pipeline_layout_orders_and_ids():
spec = mv.get_spec("qwen21")
pipe = mv.MVPipeline(spec, vae=None, text_cache=None, device="cpu", weight_dtype=torch.float32)
target = torch.randn(1, 64, 4, 6)
refs = [torch.randn(1, 64, 4, 6), torch.randn(1, 64, 2, 2)]
tgt_tokens, tgt_ids, packed_refs, token_hw = pipe._layout(target, refs, txt_len=10)
assert token_hw == (4, 6)
assert packed_refs[0][1][0, 0, 0].item() == 10.0
assert packed_refs[1][1][0, 0, 0].item() == 16.0
assert tgt_ids[0, 0, 0].item() == 18.0
spec = mv.get_spec("krea2")
pipe = mv.MVPipeline(spec, vae=None, text_cache=None, device="cpu", weight_dtype=torch.float32)
target = torch.randn(1, 16, 8, 12)
tgt_tokens, tgt_ids, packed_refs, token_hw = pipe._layout(target, [torch.randn(1, 16, 8, 12)], txt_len=10)
assert token_hw == (4, 6) and tgt_tokens.shape == (1, 24, 64)
assert tgt_ids[0, 0, 0].item() == 0.0 and packed_refs[0][1][0, 0, 0].item() == 1.0
def test_foreground_indices_and_schedule():
mask = torch.zeros(1, 4, 6)
mask[0, 1:3, 2:5] = 1
idx = foreground_indices(mask, 24, "target")
assert idx.tolist() == [[8, 9, 10, 14, 15, 16]]
with pytest.raises(ValueError):
foreground_indices(torch.zeros(1, 4, 6), 24, "target")
sched = common.shifted_schedule(25, common.qwen21_mu(16384))
assert sched[0] == 1.0 and sched[-1] == 0.0 and all(a > b for a, b in zip(sched, sched[1:]))
assert math.isclose(common.qwen21_mu(4096), 0.6935, abs_tol=1e-3)
|