File size: 3,572 Bytes
bf2a6c7 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """The key plan must line up with what ``comfy/ldm/minimax/model.py`` and the loaders expect."""
from slimdit import mapping
def _official_keys(n_blocks=2):
keys = []
for i in range(n_blocks):
p = f"transformer_blocks.{i}"
keys += [
f"{p}.adaln_proj.linear.bias",
f"{p}.adaln_proj.linear.weight",
f"{p}.attn.norm_k.weight",
f"{p}.attn.norm_q.weight",
f"{p}.attn.to_k.weight",
f"{p}.attn.to_out.0.weight",
f"{p}.attn.to_q.weight",
f"{p}.attn.to_v.weight",
f"{p}.ff.net.0.proj.weight",
f"{p}.ff.net.2.weight",
f"{p}.norm1.weight",
f"{p}.norm2.weight",
]
for i in range(2):
p = f"token_refiner.refiner_blocks.{i}"
keys += [
f"{p}.attn.norm_k.weight",
f"{p}.attn.norm_q.weight",
f"{p}.attn.to_k.weight",
f"{p}.attn.to_out.0.weight",
f"{p}.attn.to_q.weight",
f"{p}.attn.to_v.weight",
f"{p}.ff.net.0.proj.weight",
f"{p}.ff.net.2.weight",
f"{p}.norm1.weight",
f"{p}.norm2.weight",
]
keys += [
"audio_proj_in.bias",
"audio_proj_in.weight",
"audio_proj_out.bias",
"audio_proj_out.weight",
"context_embedder.bias",
"context_embedder.weight",
"norm_out.linear.bias",
"norm_out.linear.weight",
"norm_out.norm.weight",
"proj_in.bias",
"proj_in.weight",
"proj_out.bias",
"proj_out.weight",
"time_embedder.linear_1.bias",
"time_embedder.linear_1.weight",
"time_embedder.linear_2.bias",
"time_embedder.linear_2.weight",
"rope.inv_freq",
"token_refiner.final_norm.weight",
]
return keys
def test_every_source_resolves_and_blocks_are_counted():
keys = _official_keys(n_blocks=3)
plan = mapping.build_plan(keys)
assert mapping.block_count(keys) == 3
assert mapping.missing_sources(plan, keys) == []
def test_output_names_are_the_ones_the_loader_detects():
keys = _official_keys(n_blocks=1)
plan = mapping.build_plan(keys)
outs = {name for op in plan for name in op.out}
for expected in [
"blocks.0.attn.qkv_proj.weight",
"blocks.0.attn.out_proj.weight",
"blocks.0.mlp.fc1.weight",
"blocks.0.mlp.fc2.weight",
"blocks.0.attn.q_norm.weight",
"blocks.0.attn.k_norm.weight",
"blocks.0.norm1.weight",
"blocks.0.norm2.weight",
"blocks.0.adaln_proj.linear.weight",
"video_patch_proj.weight",
"final_layer.video_out.weight",
"final_layer.norm.weight",
"final_layer.adaln_proj.linear.weight",
"condition_proj.weight",
"token_refiner.blocks.0.attn.qkv_proj.weight",
]:
assert expected in outs, expected
def test_alignment_projection_keeps_the_pipeline_unquantized():
"""Drbaph's checkpoint quantizes exactly the four per-block linears; everything else is copied."""
plan = mapping.build_plan(_official_keys(n_blocks=1))
quantized = {op.out[0] for op in plan if op.kind == "quant"}
assert quantized == {
"blocks.0.attn.qkv_proj.weight",
"blocks.0.attn.out_proj.weight",
"blocks.0.mlp.fc1.weight",
"blocks.0.mlp.fc2.weight",
}
assert {op.out[0] for op in plan if op.kind == "curve"} == {
"blocks.0.adaln_proj.linear.weight",
"final_layer.adaln_proj.linear.weight",
}
|