"""CPU regression checks against the pinned native ComfyUI consumer.""" import ast from contextlib import contextmanager import importlib.util import json import math from pathlib import Path import sys import tempfile from types import SimpleNamespace import unittest import torch from safetensors import safe_open from safetensors.torch import load_file, save_file HERE = Path(__file__).parent spec = importlib.util.spec_from_file_location("omniref_converter", HERE / "convert_omniref_comfy.py") conv = importlib.util.module_from_spec(spec) sys.modules[spec.name] = conv spec.loader.exec_module(conv) def native_namespace(): # Execute the actual pinned architecture, with ordinary CPU Linear/RMSNorm # and an explicit activation implementation in place of CUDA kitchen ops. @contextmanager def cast_weights(head, h, offloadable=True): yield head.weight.to(h), head.bias.to(h) def swiglu(linear, packed, activation): assert activation == "swiglu" gate, value = packed.chunk(2, dim=-1) return linear(torch.nn.functional.silu(gate) * value) namespace = {"torch": torch, "nn": torch.nn, "math": math, "ComfyAttention": lambda: None, "comfy": SimpleNamespace(ops=SimpleNamespace( CastBiasWeightContext=cast_weights, linear_input_act=swiglu))} tree = ast.parse((HERE / "upstream/comfy_model.py").read_text()) body = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef))] exec(compile(ast.Module(body=body, type_ignores=[]), "pinned-ComfyUI-model", "exec"), namespace) detection = ast.parse((HERE / "upstream/comfy_model_detection.py").read_text()) funcs = [n for n in detection.body if isinstance(n, ast.FunctionDef) and n.name in {"count_blocks", "detect_unet_config", "any_suffix_in"}] namespace["json"] = json exec(compile(ast.Module(body=funcs, type_ignores=[]), "pinned-ComfyUI-detection", "exec"), namespace) return namespace def fixture(path): config = {"hidden_size": 4, "num_attention_heads": 2, "attention_head_dim": 2, "ffn_dim": 6, "time_embed_dim": 3, "pdd_steps": 4, "in_channels": 2, "audio_in_channels": 2, "patch_size": [1, 2, 2], "text_dim": 5, "time_embed_hidden_dim": 4, "freq_dim": 4, "num_layers": 2, "num_refiner_layers": 1, "rope_theta": 10000.0, "rope_freq_dim": 2, "norm_eps": 1e-5, "qk_norm_eps": 1e-5, "final_norm_eps": 1e-5} g = torch.Generator().manual_seed(24) tensors = {k: (torch.randn(shape, generator=g) * 0.1).bfloat16() for k, shape in conv.expected_shapes(config).items()} names = sorted(tensors) mapping = {} for j in range(2): selected = {k: tensors[k] for k in names[j::2]} filename = f"part-{j}.safetensors" save_file(selected, str(path / filename)) mapping.update({k: filename for k in selected}) (path / "config.json").write_text(json.dumps(config)) (path / "diffusion_pytorch_model.safetensors.index.json").write_text(json.dumps({"weight_map": mapping})) return config, tensors class ConversionTests(unittest.TestCase): @classmethod def setUpClass(cls): torch.set_num_threads(2) cls.native = native_namespace() def convert(self, path, rank=0): c, original = fixture(path) _, source = conv.load_source(path) scratch = path / "scratch" scratch.mkdir() output, curve = conv.make_plan(c, source, scratch, rank=rank, grid=1025) contract = {"pdd_steps": 4, "grid_max_t": 0.999, "video_scheduler_shift": 12.0, "audio_scheduler_shift": 3.0} metadata = conv.metadata_for(c, contract, rank) target = path / "export.safetensors" receipt = conv.write_file(target, output, metadata) return c, original, load_file(str(target)), metadata, curve, receipt def test_native_detection_and_strict_dense_state_dict_load(self): with tempfile.TemporaryDirectory() as d: c, source, target, metadata, _, _ = self.convert(Path(d)) detected = self.native["detect_unet_config"](target, "", metadata=metadata) self.assertEqual(detected["pdd_steps"], 4) self.assertTrue(detected["gate_compress"]) self.assertEqual(detected["latents_dim"], 2) model = self.native["MiniMaxH3Model"](**detected, device="cpu", dtype=torch.float32, operations=torch.nn) model.load_state_dict(target, strict=True) qkv = model.blocks[0].attn.qkv_proj.weight expected = torch.cat([source[f"transformer_blocks.0.attn.to_{p}.weight"].float() for p in ("q", "k", "v")]) self.assertTrue(torch.equal(qkv, expected)) def test_swiglu_native_forward_matches_diffusers_value_first(self): with tempfile.TemporaryDirectory() as d: c, source, target, _, _, _ = self.convert(Path(d)) mlp = self.native["MLP"](c["hidden_size"], c["ffn_dim"], operations=torch.nn, dtype=torch.float32) mlp.load_state_dict({"fc1.weight": target["blocks.0.mlp.fc1.weight"], "fc2.weight": target["blocks.0.mlp.fc2.weight"]}, strict=True) x = torch.randn(7, c["hidden_size"]) value, gate = torch.nn.functional.linear(x, source["transformer_blocks.0.ff.net.0.proj.weight"].float()).chunk(2, dim=-1) expected = torch.nn.functional.linear(value * torch.nn.functional.silu(gate), source["transformer_blocks.0.ff.net.2.weight"].float()) torch.testing.assert_close(mlp(x), expected, rtol=0, atol=0) def test_pdd_native_consumer_recovers_weighted_full_heads(self): with tempfile.TemporaryDirectory() as d: c, source, target, _, _, _ = self.convert(Path(d)) h = torch.randn(3, c["hidden_size"]) for stream, key in [("video", "proj_out"), ("audio", "audio_proj_out")]: weight = target[f"final_layer.{stream}_out.weight"] bias = target[f"final_layer.{stream}_out.bias"] head = SimpleNamespace(weight=weight, bias=bias) original_w = source[key + ".weight"].float().reshape(4, -1, c["hidden_size"]) original_b = source[key + ".bias"].float().reshape(4, -1) for start, stop in [(0, 1), (0, 2), (1, 3), (3, 4)]: for shift in [3.0, 12.0]: grid = torch.linspace(1.0, 0.0, 5, dtype=torch.float64) dt = (1 - shift * grid / (1 + (shift - 1) * grid)).diff()[start:stop] normalized = (dt / dt.sum()).float() expected = torch.nn.functional.linear(h, torch.einsum("n,noi->oi", normalized, original_w[start:stop]), torch.einsum("n,no->o", normalized, original_b[start:stop])) actual = self.native["_pdd_head"](head, h, 4, start, stop, shift) torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) def test_curve_export_native_strict_load_and_holdout_error(self): with tempfile.TemporaryDirectory() as d: c, source, target, metadata, curve, _ = self.convert(Path(d), rank=3) detected = self.native["detect_unet_config"](target, "", metadata=metadata) self.assertEqual(detected["adaln_curve_grid"], 1025) self.assertEqual(detected["time_embed_dim"], 3) model = self.native["MiniMaxH3Model"](**detected, device="cpu", dtype=torch.float32, operations=torch.nn) model.load_state_dict(target, strict=True) self.assertFalse(any(k.startswith("time_embedder.") for k in target)) self.assertTrue(curve["approximation"]) self.assertLess(max(e["relative_l2"] for e in curve["holdout_modulation_samples"]), 1e-5) def test_reject_unexpected_tensor_before_writing(self): with tempfile.TemporaryDirectory() as d: path = Path(d) fixture(path) c, source = conv.load_source(path) source["optimizer.unknown"] = source["proj_in.weight"] with self.assertRaisesRegex(ValueError, "Unexpected source keys"): conv.make_plan(c, source, path) def test_output_hash_and_refuse_overwrite(self): import hashlib with tempfile.TemporaryDirectory() as d: path = Path(d) _, _, _, _, _, receipt = self.convert(path) self.assertEqual(receipt["sha256"], hashlib.sha256((path / "export.safetensors").read_bytes()).hexdigest()) with self.assertRaises(FileExistsError): conv.write_file(path / "export.safetensors", {}, {}) if __name__ == "__main__": unittest.main(verbosity=2)