| """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 |
| from ultratex.backbones.common import block_causal_attention, dense_attention, foreground_indices, gather_sequence |
| from ultratex.backbones.krea2.model import Krea2, Krea2Params |
| from ultratex.backbones.qwen21.model import Qwen21Params, QwenImage21 |
| from ultratex.backbones.qwen21.util import pack_image as qwen_pack |
| from ultratex.backbones.krea2.util import pack_image as krea_pack |
| from ultratex import mv |
|
|
| 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) |
| |
| 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] |
| |
| _, 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)) |
| |
| 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) |
| |
| assert torch.allclose(cached, again, atol=1e-4) |
| |
| 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) |
|
|
