File size: 2,421 Bytes
2f48fc1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | #!/usr/bin/env python3
"""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))
|