| """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(): |
| |
| |
| @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) |
|
|
