From 329d69f5dff483d3a55d2aee359a886ae075c70a Mon Sep 17 00:00:00 2001
From: CalamitousFelicitousness <iowasovereign@gmail.com>
Date: Fri, 17 Jul 2026 04:52:27 +0100
Subject: [PATCH] fix(cli): match module paths the way the loader does

Kohya-style files carry an already-underscored base (lora_unet_layers_0_
mlp_gate_proj), which the loader resolves by comparing network_prefix +
path.replace('.', '_') against each module's stamped name, so both sides
are underscored and the file loads. The analyzer instead looked the base
up as a literal dotted module path, so every module of such a file was
reported unmatched: 76 files in a local collection, including 36 of 57
anima and 5 of 10 chroma.

Fall back to a stamped-name index when the direct lookup misses. Dotted
bases are unaffected.
---
 cli/lora-quant-fidelity.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/cli/lora-quant-fidelity.py b/cli/lora-quant-fidelity.py
index 1219d7189..8b8051d95 100644
--- a/cli/lora-quant-fidelity.py
+++ b/cli/lora-quant-fidelity.py
@@ -161,6 +161,17 @@ def map_lora_modules(lora_path, arch_mod):
     return mapped, census, chunked
 
 
+def stamp_index(paths):
+    """Map each module path to its stamped form, the way the loader matches.
+
+    The loader compares ``network_prefix + path.replace('.', '_')`` against each
+    module's stamped ``network_layer_name``, so kohya-style ``lora_unet_`` keys
+    (whose base arrives already underscored) resolve fine there. Matching on the
+    stamped form reproduces that and keeps dotted bases working unchanged.
+    """
+    return {p.replace('.', '_'): p for p in paths}
+
+
 def make_stub(shape, dtype=torch.bfloat16):
     """Minimal sd_module standing in for a bf16 repo weight: the module classes key off its type and shape."""
     if len(shape) == 2:
@@ -282,11 +293,14 @@ def main():
     pre_quantized = model_config.get('quantization_config') is not None
 
     quant_layers, bf16_repo = {}, None
+    quant_stamps, bf16_stamps = {}, {}
     if pre_quantized:
         rprint(f'model: "{model_dir}" pre-quantized={pre_quantized}')
         quant_layers = load_quantized_model(model_dir, arch=args.arch, class_name=model_config.get('_class_name'))
+        quant_stamps = stamp_index(quant_layers)
     else:
         bf16_repo = Bf16Repo(model_dir)
+        bf16_stamps = stamp_index(k[:-len('.weight')] for k in bf16_repo.weight_map if k.endswith('.weight'))
         if args.dtype is None:
             rprint('model is not quantized and no --dtype given: loras apply exactly, nothing to analyze')
             return 0
@@ -305,7 +319,7 @@ def main():
         for path in keys:
             entries = mapped[path]
             if pre_quantized:
-                layer = quant_layers.get(path)
+                layer = quant_layers.get(path) or quant_layers.get(quant_stamps.get(path.replace('.', '_'), ''))
                 if layer is None:
                     unmatched.append(path)
                     continue
@@ -319,6 +333,8 @@ def main():
                 sd_module = layer
             else:
                 W = bf16_repo.get(f'{path}.weight')
+                if W is None:
+                    W = bf16_repo.get(f'{bf16_stamps.get(path.replace(".", "_"), "")}.weight')
                 if W is None:
                     unmatched.append(path)
                     continue