fix(chroma): prepend lora_transformer_ to LoKR slice_info keys

try_load_lokr renamed slice_info keys via static_rename but didn't add
the lora_transformer_ prefix that apply_static_rename adds to the groups
dict. slice_info.get(network_key) always returned None for fused targets,
so the loader fell back to NetworkModuleLokr instead of
NetworkModuleLokrSliceChunk on every fused-QKV / fused-linear1 adapter.

The full kron(w1, w2) was applied against split target modules, either
shape-mismatching at apply time or broadcasting wrong.

Caught by test_lokr_bfl_img_attn_qkv_slice_chunked. No real-world chroma
LoKR adapters on fused targets are known in the wild.
This commit is contained in:
CalamitousFelicitousness
2026-05-18 22:43:05 +01:00
parent 409a30f9e0
commit be46362e3e
+4 -1
View File
@@ -191,7 +191,10 @@ def try_load_lokr(name, network_on_disk, lora_scale):
groups = group_by_suffixes(state_dict, LOKR_SUFFIXES)
groups, slice_info = expand_chroma_fused_lokr(groups)
groups = apply_static_rename(groups, static_rename)
slice_info = {static_rename.get(k, k): v for k, v in slice_info.items()}
# Mirror apply_static_rename: the slice_info dict must use the same final
# network keys as ``groups`` (static rename applied + ``lora_transformer_``
# prefix) so the per-target lookup in the loop matches.
slice_info = {'lora_transformer_' + static_rename.get(k, k): v for k, v in slice_info.items()}
unmapped = 0
for network_key, w in groups.items():