Initial release

tools/onnx-export/gen_weight_map.py

@@ -0,0 +1,126 @@
+"""
+Generate weight_names.json sidecar for TRT adapter refit.
+
+Maps ONNX val_N initializer names to human-readable parameter names.
+Uses the deterministic linear_N numbering from dynamo decomposition.
+
+Can be run standalone (no model loading needed, just the ONNX file).
+"""
+import onnx
+import json
+import sys
+import os
+
+def build_weight_map(onnx_path):
+    m = onnx.load(onnx_path, load_external_data=False)
+    
+    # Build val_N -> MatMul node name mapping
+    val_to_node = {}
+    for node in m.graph.node:
+        if node.op_type == 'MatMul':
+            for inp in node.input:
+                if inp.startswith('val_'):
+                    val_to_node[inp] = node.name
+    
+    # Get all val_ weights sorted by N
+    val_inits = sorted(
+        [(i.name, list(i.dims)) for i in m.graph.initializer if i.name.startswith('val_')],
+        key=lambda x: int(x[0].split('_')[1])
+    )
+    
+    # Filter to only MatMul weights (skip small constants)
+    matmul_weights = []
+    for vname, vshape in val_inits:
+        if vname in val_to_node:
+            numel = 1
+            for d in vshape:
+                numel *= d
+            if numel > 1000:
+                matmul_weights.append((vname, vshape, val_to_node[vname]))
+    
+    # The first two MatMul nodes are proj_in and condition_embedder
+    # (before the per-layer linears start)
+    # node_MatMul_66 = proj_in.1.weight (Conv1d -> PatchEmbedLinear)
+    # node_MatMul_68 = condition_embedder.weight
+    
+    # After that, the per-layer linears follow a repeating pattern.
+    # Each layer has 11 linear projections in this order:
+    LAYER_PATTERN = [
+        # (param_suffix, expected_shapes)
+        ("self_attn.q_proj.weight",   None),
+        ("self_attn.k_proj.weight",   None),
+        ("self_attn.v_proj.weight",   None),
+        ("self_attn.o_proj.weight",   None),
+        ("cross_attn.q_proj.weight",  None),
+        ("cross_attn.k_proj.weight",  None),
+        ("cross_attn.v_proj.weight",  None),
+        ("cross_attn.o_proj.weight",  None),
+        ("mlp.gate_proj.weight",      None),
+        ("mlp.up_proj.weight",        None),
+        ("mlp.down_proj.weight",      None),
+    ]
+    
+    # After the last layer, there should be a proj_out linear
+    
+    rename_map = {}  # val_N -> param_name
+    
+    # Map the first two special cases
+    if len(matmul_weights) >= 2:
+        # proj_in
+        rename_map[matmul_weights[0][0]] = "dit.proj_in.1.linear.weight"
+        # condition_embedder  
+        rename_map[matmul_weights[1][0]] = "dit.condition_embedder.weight"
+    
+    # Map per-layer linears
+    layer_start = 2  # skip proj_in + condition_embedder
+    linears_per_layer = len(LAYER_PATTERN)
+    remaining = matmul_weights[layer_start:]
+    
+    # Detect number of layers from count
+    # Last entry might be proj_out
+    num_layers = len(remaining) // linears_per_layer
+    leftover = len(remaining) % linears_per_layer
+    
+    print(f"Total MatMul weights: {len(matmul_weights)}")
+    print(f"Layer weights: {len(remaining)} ({num_layers} layers * {linears_per_layer} + {leftover} extra)")
+    
+    for layer_idx in range(num_layers):
+        for proj_idx, (suffix, _) in enumerate(LAYER_PATTERN):
+            w_idx = layer_start + layer_idx * linears_per_layer + proj_idx
+            if w_idx < len(matmul_weights):
+                vname = matmul_weights[w_idx][0]
+                param_name = f"dit.layers.{layer_idx}.{suffix}"
+                rename_map[vname] = param_name
+    
+    # Map leftover (proj_out)
+    if leftover > 0:
+        proj_out_idx = layer_start + num_layers * linears_per_layer
+        if proj_out_idx < len(matmul_weights):
+            rename_map[matmul_weights[proj_out_idx][0]] = "dit.proj_out.1.inner.linear.weight"
+    
+    # Build both directions
+    forward_map = rename_map  # val_N -> param_name
+    reverse_map = {v: k for k, v in rename_map.items()}  # param_name -> val_N
+    
+    return {
+        "val_to_param": forward_map,
+        "param_to_val": reverse_map,
+    }
+
+
+if __name__ == "__main__":
+    onnx_path = sys.argv[1] if len(sys.argv) > 1 else r'D:\Ace-Step-Latest\hot-step-cpp\models\onnx\dit_acestep-v15-merge-sft-turbo-xl-ta-0.7.onnx'
+    
+    mapping = build_weight_map(onnx_path)
+    
+    # Print summary
+    print(f"\nMapped {len(mapping['val_to_param'])} weights")
+    print("\nFirst 15 mappings:")
+    for val_name, param_name in sorted(mapping['val_to_param'].items(), key=lambda x: int(x[0].split('_')[1]))[:15]:
+        print(f"  {val_name} -> {param_name}")
+    
+    # Save
+    out_path = onnx_path + ".weight_names.json"
+    with open(out_path, 'w') as f:
+        json.dump(mapping, f, indent=2)
+    print(f"\nSaved to {out_path}")