From be46362e3e5648b1eacac3b7daf0849973022bd2 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 22:43:05 +0100 Subject: [PATCH] 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. --- pipelines/chroma/chroma_lora.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pipelines/chroma/chroma_lora.py b/pipelines/chroma/chroma_lora.py index 382bc6f3f..c92635948 100644 --- a/pipelines/chroma/chroma_lora.py +++ b/pipelines/chroma/chroma_lora.py @@ -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():