| |
| """Check every real checkpoint key/shape against native Comfy on the meta device.""" |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
|
|
| import convert_omniref_comfy as conv |
| from test_conversion import native_namespace |
|
|
|
|
| def audit(src): |
| c, source = conv.load_source(src) |
| expected = conv.expected_shapes(c) |
| assert set(source) == set(expected) |
| assert all(source[k].shape == shape for k, shape in expected.items()) |
| native = native_namespace() |
| receipts = [] |
| for rank in [0, 8]: |
| target = {} |
| for name, tensor in source.items(): |
| if ".attn.to_k." in name or ".attn.to_v." in name: |
| continue |
| shape = list(tensor.shape) |
| output_name = conv.renamed(name) |
| if ".attn.to_q." in name: |
| output_name = output_name.replace(".attn.to_q.", ".attn.qkv_proj.") |
| shape[0] *= 3 |
| if rank and name.startswith("time_embedder."): |
| continue |
| if rank and (name == "norm_out.linear.weight" or name.endswith(".adaln_proj.linear.weight")): |
| shape[-1] = rank |
| target[output_name] = torch.empty(shape, device="meta") |
| target["rope.inv_freq"] = torch.empty([c["rope_freq_dim"]], device="meta") |
| if rank: |
| target["adaln_t_table"] = torch.empty([1025, rank], device="meta") |
| metadata = conv.metadata_for(c, {"video_scheduler_shift": 12.0, "audio_scheduler_shift": 3.0}, rank) |
| detected = native["detect_unet_config"](target, "", metadata=metadata) |
| model = native["MiniMaxH3Model"](**detected, dtype=torch.float32, device="meta", operations=torch.nn) |
| result = model.load_state_dict(target, strict=True, assign=True) |
| assert not result.missing_keys and not result.unexpected_keys |
| assert all(target[k].shape == model.state_dict()[k].shape for k in target) |
| receipts.append({"source_tensors": len(source), "target_tensors": len(target), |
| "curve_rank": rank, "detected_config": detected, |
| "native_strict_meta_load": "passed", "gpu_inference": "not_run"}) |
| return receipts |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--src", required=True, type=Path) |
| args = parser.parse_args() |
| print(json.dumps(audit(args.src), indent=2)) |
|
|
