import math import sys from types import SimpleNamespace import unittest from unittest.mock import patch import torch from torch import nn from torch.nn import functional as F from h3_blackwell_runtime.packing import FRAME_RESCALE, H3PromptPacker, _video_t_spans from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, run_attention from h3_blackwell_runtime.qwen3vl_vision import ( TEXT_HEAD_DIM, TEXT_ROPE_DIMS, TEXT_ROPE_THETA, VISION_HIDDEN, Qwen3VL32BVision, _VisionAttention, _VisionPatchMerger, _apply_rope_vision, _text_run_ids, mrope_freqs_cis, ) from h3_blackwell_runtime.runtime import H3HotRuntime from h3_blackwell_runtime.vae_encoder import MiniMaxH3VideoVAEEncoder, _downsample class Fl2vaVAEContracts(unittest.TestCase): def test_quant_conv_is_a_required_checkpoint_weight(self): names = MiniMaxH3VideoVAEEncoder()._required_encoder_names() self.assertIn("quant_conv.weight", names) self.assertIn("quant_conv.bias", names) def test_single_frame_is_encoded_without_temporal_prepad(self): encoder = MiniMaxH3VideoVAEEncoder(tiling=False) seen = [] def fake_encode(x): seen.append(tuple(x.shape)) return torch.zeros((x.shape[0], 48, x.shape[2], 1, 1), device=x.device) encoder._adaptive_encode = fake_encode result = encoder.encode(torch.zeros(1, 3, 8, 8)) self.assertEqual(seen, [(1, 3, 1, 8, 8)]) self.assertEqual(tuple(result.shape), (1, 24, 1, 1, 1)) def test_downsample_pads_only_right_and_bottom(self): x = torch.arange(16, dtype=torch.float32).reshape(1, 1, 1, 4, 4) weight = torch.ones(1, 1, 3, 3, 3) params = {"w": weight, "b": torch.zeros(1), "time": 1, "space": 2} actual = _downsample(x, params) padded = F.pad(x, (0, 1, 0, 1, 0, 0), mode="reflect") expected = F.conv3d(padded, weight[:, :, -1:], params["b"], stride=(1, 2, 2)) torch.testing.assert_close(actual, expected) def test_hot_runtime_preserves_zero_to_one_images(self): runtime = H3HotRuntime.__new__(H3HotRuntime) image = torch.tensor([[[[0.0, 0.5, 1.0]]]]) converted = runtime._image_to_uint8_nhwc(image) self.assertEqual(converted.flatten().tolist(), [0, 127, 255]) class AttentionBackendContracts(unittest.TestCase): def test_hot_backends_include_benchmark_candidates(self): self.assertTrue({"sage2", "cudnn_sdpa", "ck_int8"}.issubset(AVAILABLE_BACKENDS)) def test_cudnn_backend_is_forced_without_fallback(self): q = torch.randn(1, 2, 3, 4) expected = torch.randn_like(q) context = unittest.mock.MagicMock() with ( patch("torch.nn.attention.sdpa_kernel", return_value=context) as kernel, patch("h3_blackwell_runtime.attention.functional.scaled_dot_product_attention", return_value=expected) as sdpa, ): actual = run_attention(q, q, q, backend="cudnn_sdpa", is_causal=False) self.assertIs(actual, expected) self.assertEqual(kernel.call_args.args[0], [torch.nn.attention.SDPBackend.CUDNN_ATTENTION]) sdpa.assert_called_once_with(q, q, q, is_causal=False) def test_comfy_kitchen_int8_backend_dispatches_hnd_tensors(self): q = torch.randn(1, 2, 3, 4) expected = torch.randn_like(q) kitchen = SimpleNamespace(int8_attention=unittest.mock.MagicMock(return_value=expected)) with patch.dict(sys.modules, {"comfy_kitchen": kitchen}): actual = run_attention(q, q, q, backend="ck_int8", is_causal=False) self.assertIs(actual, expected) kitchen.int8_attention.assert_called_once_with(q, q, q) def test_comfy_kitchen_int8_rejects_causal_attention(self): q = torch.randn(1, 2, 3, 4) with self.assertRaisesRegex(ValueError, "does not support causal"): run_attention(q, q, q, backend="ck_int8", is_causal=True) class Fl2vaVisionContracts(unittest.TestCase): def test_visual_rotary_coordinates_are_block_major(self): class CoordinateTable(nn.Module): def forward(self, length): return torch.arange(length, dtype=torch.float32).unsqueeze(1) vision = Qwen3VL32BVision.__new__(Qwen3VL32BVision) nn.Module.__init__(vision) vision.spatial_merge_size = 2 vision.rotary_pos_emb = CoordinateTable() coordinates = vision.rot_pos_emb(torch.tensor([[1, 4, 4]])).tolist() self.assertEqual(coordinates[:8], [ [0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0], [0.0, 2.0], [0.0, 3.0], [1.0, 2.0], [1.0, 3.0], ]) def test_sdpa_output_is_restored_to_token_major_layout(self): torch.manual_seed(7) sequence, heads, head_dim = 3, 2, 2 hidden = heads * head_dim qkv_weight = torch.randn(hidden * 3, hidden) qkv_bias = torch.randn(hidden * 3) proj_weight = torch.randn(hidden, hidden) proj_bias = torch.randn(hidden) module = _VisionAttention(qkv_weight, qkv_bias, proj_weight, proj_bias, num_heads=heads, head_dim=head_dim) x = torch.randn(sequence, hidden) cos = torch.ones(sequence, 1, head_dim) sin = torch.zeros(sequence, 1, head_dim // 2) actual = module(x, torch.tensor([0, sequence], dtype=torch.int32), (cos, sin, sin)) qkv = F.linear(x, qkv_weight, qkv_bias) query, key, value = qkv.reshape(sequence, 3, heads, head_dim).permute(1, 0, 2, 3).unbind(0) output = F.scaled_dot_product_attention( query.transpose(0, 1).unsqueeze(0), key.transpose(0, 1).unsqueeze(0), value.transpose(0, 1).unsqueeze(0), ) expected = F.linear(output.transpose(1, 2).reshape(sequence, hidden), proj_weight, proj_bias) torch.testing.assert_close(actual, expected) def test_vision_rope_uses_original_halves(self): q = torch.tensor([[[1.0, 2.0, 3.0, 4.0]]]) k = q + 4 cos = torch.full_like(q, 0.5) sin = torch.full_like(q[..., :2], 0.25) neg_sin = -sin actual_q, actual_k = _apply_rope_vision(q, k, (cos, sin, neg_sin)) def expected(x): return torch.cat((x[..., :2] * 0.5 + x[..., 2:] * -0.25, x[..., 2:] * 0.5 + x[..., :2] * 0.25), dim=-1) torch.testing.assert_close(actual_q, expected(q)) torch.testing.assert_close(actual_k, expected(k)) def test_mrope_uses_reference_section_boundaries(self): positions = torch.stack((torch.arange(8), torch.arange(8) + 10, torch.arange(8) + 20)) actual = mrope_freqs_cis(positions)[0] inv_freq = 1.0 / ( TEXT_ROPE_THETA ** (torch.arange(0, TEXT_HEAD_DIM, 2, dtype=torch.float32) / TEXT_HEAD_DIM) ) freqs = (inv_freq[None, :, None].expand(3, -1, 1) @ positions[:, None, :].float()).transpose(1, 2) interleaved = freqs[0].clone() for axis, offset in ((1, 1), (2, 2)): index = slice(offset, TEXT_ROPE_DIMS[axis] * 3, 3) interleaved[..., index] = freqs[axis, ..., index] expected = torch.cat((interleaved, interleaved), dim=-1).cos().unsqueeze(0) torch.testing.assert_close(actual, expected) def test_mergers_preserve_existing_block_major_order(self): x = torch.arange(4 * VISION_HIDDEN, dtype=torch.float32).reshape(4, VISION_HIDDEN) passthrough = lambda value, *args, **kwargs: value with ( patch("h3_blackwell_runtime.qwen3vl_vision.F.layer_norm", side_effect=passthrough), patch("h3_blackwell_runtime.qwen3vl_vision.F.linear", side_effect=passthrough), patch("h3_blackwell_runtime.qwen3vl_vision.F.gelu", side_effect=passthrough), ): main = _VisionPatchMerger(*(torch.empty(1) for _ in range(6)), merge_size=2, out_hidden_size=1, norm_dim=VISION_HIDDEN) deepstack = _VisionPatchMerger(*(torch.empty(1) for _ in range(6)), merge_size=2, out_hidden_size=1) torch.testing.assert_close(main(x), x.reshape(1, -1)) torch.testing.assert_close(deepstack(x), x.reshape(1, -1)) def test_configured_tokenizer_preserves_empty_prompt(self): calls = [] class RawTokenizer: def __call__(self, text, **kwargs): calls.append((text, kwargs)) return SimpleNamespace(input_ids=torch.empty((1, 0), dtype=torch.long)) self.assertEqual(_text_run_ids(SimpleNamespace(tokenizer=RawTokenizer()), ""), []) self.assertEqual(calls[0][0], "") self.assertFalse(calls[0][1]["add_special_tokens"]) class _FakeCheckpoint: def tensor(self, name, dtype=None): if name == "video_patch_proj.weight": value = torch.zeros(5376, 96) elif name == "video_patch_proj.bias": value = torch.zeros(5376) elif name == "audio_patch_proj.weight": value = torch.zeros(5376, 32) elif name == "audio_patch_proj.bias": value = torch.zeros(5376) else: value = torch.empty(0) return value.to(dtype=dtype) if dtype is not None else value class Fl2vaPackingContracts(unittest.TestCase): def test_each_keyframe_keeps_its_own_condition_segment(self): packer = H3PromptPacker(_FakeCheckpoint()) text = torch.zeros(1, 3, 5376) video = torch.zeros(1, 24, 2, 2, 2) audio = torch.zeros(1, 32, 2, 2) keyframes = [torch.zeros(1, 24, 1, 2, 2) for _ in range(2)] _, _, segments, _, _, _ = packer( text, video, audio, 0.5, cond_latents=keyframes, cond_frame_indices=[0, 21], frame_count=22, ) self.assertEqual(segments[1][:2], (3, 4)) self.assertEqual(segments[2][:2], (4, 5)) self.assertEqual(segments[1][2], segments[2][2]) def test_last_only_anchor_and_targets_share_reference_cursor(self): packer = H3PromptPacker(_FakeCheckpoint()) text = torch.zeros(1, 3, 5376) video = torch.zeros(1, 24, 2, 2, 2) audio = torch.zeros(1, 32, 2, 2) last = torch.zeros(1, 24, 1, 2, 2) _, _, _, positions, _, _ = packer( text, video, audio, 0.5, cond_latents=[last], cond_frame_indices=[21], frame_count=22, seed=1, ) expected_last_t = 3.0 + sum(_video_t_spans(2)) - FRAME_RESCALE self.assertTrue(math.isclose(float(positions[3, 0]), expected_last_t)) self.assertEqual(float(positions[4, 0]), 3.0) # target audio self.assertEqual(float(positions[8, 0]), 3.0) # target video if __name__ == "__main__": unittest.main()