[SD3 branch] flux_extract_lora lora key not loaded|unused lora tensor

[fractal-fumbler]

Hello :)
So, after fixing #2206 i executed flux_extract_lora.py from sd3 branch and lora was saved successfully.
Command to extract lora

python flux_extract_lora.py --save_to pixelwave_flux1_dev_fp8_03_512.safetensors --model_org flux1-dev-fp8.safetensors --model_tuned pixelwave_flux1_dev_fp8_03.safetensors --device cpu --dim 512 --save_precision fp16

After i tried to apply extracted lora in generation proccess in ComfyUI or stable-diffusion.cpp there is an error info in log that lora isn't applied

ComfyUI - lora key not loaded

Using pytorch attention in VAE
Using pytorch attention in VAE
VAE load device: xpu:0, offload device: cpu, dtype: torch.bfloat16
Requested to load FluxClipModel_
loaded completely 9.5367431640625e+25 9319.4814453125 True
CLIP/text encoder model load device: cpu, offload device: cpu, current: cpu, dtype: torch.float16
gguf qtypes: F32 (471), Q6_K (304), F16 (5)
model weight dtype torch.bfloat16, manual cast: None
model_type FLUX
lora key not loaded: lora_unet_model_diffusion_model_double_blocks_0_img_attn_proj.alpha
lora key not loaded: lora_unet_model_diffusion_model_double_blocks_0_img_attn_proj.lora_down.weight
lora key not loaded: lora_unet_model_diffusion_model_double_blocks_0_img_attn_proj.lora_up.weight
lora key not loaded: lora_unet_model_diffusion_model_double_blocks_0_img_attn_qkv.alpha
lora key not loaded: lora_unet_model_diffusion_model_double_blocks_0_img_attn_qkv.lora_down.weight
lora key not loaded: lora_unet_model_diffusion_model_double_blocks_0_img_attn_qkv.lora_up.weight

full log of ComfyUI.log

stable-diffusion.cpp - unused lora tensor

[WARN ] lora.hpp:860  - unused lora tensor |lora.model_diffusion_model_model_diffusion_model_double_blocks_0_img_attn_proj.alpha|
 (f16): shape(1, 1, 1, 1)
[WARN ] lora.hpp:860  - unused lora tensor |lora.model_diffusion_model_model_diffusion_model_double_blocks_0_img_attn_proj.lora_down.weight|
 (f16): shape(3072, 512, 1, 1)
[WARN ] lora.hpp:860  - unused lora tensor |lora.model_diffusion_model_model_diffusion_model_double_blocks_0_img_attn_proj.lora_up.weight|
 (f16): shape(512, 3072, 1, 1)
[WARN ] lora.hpp:860  - unused lora tensor |lora.model_diffusion_model_model_diffusion_model_double_blocks_0_img_attn_qkv.alpha|
 (f16): shape(1, 1, 1, 1)
[WARN ] lora.hpp:860  - unused lora tensor |lora.model_diffusion_model_model_diffusion_model_double_blocks_0_img_attn_qkv.lora_down.weight|
 (f16): shape(3072, 512, 1, 1)
[WARN ] lora.hpp:860  - unused lora tensor |lora.model_diffusion_model_model_diffusion_model_double_blocks_0_img_attn_qkv.lora_up.weight|
 (f16): shape(512, 9216, 1, 1)
[WARN ] lora.hpp:860  - unused lora tensor |lora.model_diffusion_model_model_diffusion_model_double_blocks_0_img_mlp_0.alpha|
 (f16): shape(1, 1, 1, 1)
[WARN ] lora.hpp:860  - unused lora tensor |lora.model_diffusion_model_model_diffusion_model_double_blocks_0_img_mlp_0.lora_down.weight|
 (f16): shape(3072, 512, 1, 1)
[WARN ] lora.hpp:860  - unused lora tensor |lora.model_diffusion_model_model_diffusion_model_double_blocks_0_img_mlp_0.lora_up.weight|
 (f16): shape(512, 12288, 1, 1)

full log of stable-diffusion.cpp.log

pip list

Package                     Version
--------------------------- -----------------
accelerate                  1.10.1
aiohappyeyeballs            2.6.1
aiohttp                     3.12.15
aiosignal                   1.4.0
alembic                     1.16.5
annotated-types             0.7.0
appdirs                     1.4.4
attrs                       25.3.0
autocommand                 2.2.2
av                          15.1.0
cffi                        2.0.0
charset-normalizer          3.4.3
comfyui-embedded-docs       0.2.6
comfyui_frontend_package    1.26.13
comfyui_workflow_templates  0.1.86
contourpy                   1.3.3
CppHeaderParser             2.7.4
cryptography                46.0.1
cycler                      0.12.1
diffusers                   0.35.1
einops                      0.8.1
filelock                    3.19.1
fonttools                   4.60.0
frozenlist                  1.7.0
fsspec                      2025.9.0
gguf                        0.17.1
gmpy2                       2.2.1
greenlet                    3.2.4
hf-xet                      1.1.10
huggingface-hub             0.35.1
idna                        3.10
importlib_metadata          8.7.0
intel_extension_for_pytorch 2.8.10+git0933033
ipex_to_cuda                0.0.1
jaraco.collections          5.1.0
jaraco.context              6.0.1
jaraco.functools            4.1.0
jaraco.text                 4.0.0
Jinja2                      3.1.6
jsonschema                  4.25.1
jsonschema-specifications   2025.9.1
kiwisolver                  1.4.7
kornia                      0.8.1
kornia_rs                   0.1.9
lit                         20.1.0.dev0
Mako                        1.3.10
MarkupSafe                  3.0.2
matplotlib                  3.10.6
more-itertools              10.8.0
mpmath                      1.3.0
multidict                   6.6.4
networkx                    3.5
numpy                       2.3.3
packaging                   25.0
pandas                      2.3.1
pillow                      11.3.0
pip                         25.2
platformdirs                4.4.0
ply                         3.11
pooch                       1.8.2
propcache                   0.3.2
psutil                      7.1.0
py-cpuinfo                  9.0.0
pybind11                    3.0.1
pycparser                   2.23
pydantic                    2.11.9
pydantic_core               2.33.2
pydantic-settings           2.11.0
pyparsing                   3.2.5
python-dateutil             2.9.0
python-dotenv               1.1.1
pytz                        2025.2
PyYAML                      6.0.3
referencing                 0.36.2
regex                       2025.9.18
requests                    2.32.5
rpds-py                     0.27.1
safetensors                 0.6.2
scipy                       1.16.2
sentencepiece               0.2.1
setuptools                  80.9.0
six                         1.17.0
soundfile                   0.13.1
spandrel                    0.4.1
SQLAlchemy                  2.0.43
sympy                       1.14.0
TBB                         0.2
tokenizers                  0.22.1
toml                        0.10.2
torch                       2.8.0
torchaudio                  2.8.0a0
torchsde                    0.2.6
torchvision                 0.23.0a0
tqdm                        4.67.1
trampoline                  0.1.2
transformers                4.56.2
typing_extensions           4.15.0
typing-inspection           0.4.1
tzdata                      2025.2.post0
urllib3                     2.5.0
wheel                       0.45.1
yarl                        1.20.1
zipp                        3.23.0

Please, if you have spare time for this issue, tell can this issue with lora be resolved or it needed to be addressed to ComfyUI/stable-diffusion.cpp?

[kohya-ss]

The extracted LoRA should have the keys like lora_unet_double_blocks_0_img_attn_proj.lora_down.weight. Perhaps you are extracting LoRA from Diffusers format FLUX weights? Please use BFL format weights.

There may be a script in the Diffusers repository to convert to and from BFL.

[fractal-fumbler]

Extracting from FLUX.1 Dev 03
civitai page of Full Model fp8 (11.08 GB))

will check for BFL script. TY

[fractal-fumbler]

So i did some research and found out that both pixelwave_flux1_dev_fp8_03_512.safetensors and flux1-dev-fp8.safetensors are indeed have diffusers format, e.g. having "model.diffusion_model" in dict keys. That is why flux_extract_lora.py successfully executed from my humble point of view.

So i've found flux1 model in black lab forest format and tried flux_extract_lora.py which get me an error with tuned model

Traceback (most recent call last):
  File "flux_extract_lora.py", line 219, in <module>
    svd(**vars(args))
    ~~~^^^^^^^^^^^^^^
  File "flux_extract_lora.py", line 86, in svd
    value_t = f_tuned.get_tensor(key)
  File "/library/safetensors_utils.py", line 158, in get_tensor
    raise KeyError(f"Tensor '{key}' not found in the file")
KeyError: "Tensor 'double_blocks.0.img_attn.proj.weight' not found in the file"

here is a little script for those who will encounter same problem to check for keys in dict for BLF or diffusers format:

from safetensors.torch import load_file
lora_path = "pixelwave_flux1_dev_fp8_03_512.safetensors"
original_dict = load_file(lora_path)
print(original_dict.keys())

@kohya-ss, thanks for your help!
sorry for bothering, is there a sense to remove keys "model.diffusion_model" from flux model?

[kohya-ss]

Indeed, if both original and tuned have prefix for the keys, the extraction will succeed.

Unfortunately, I don't have time to test it with weights that have prefixes, but I wrote code in this branch that removes the prefix when extracting: feat-extract-flux-lora-from-weights-with-key-prefix. If you don't mind, could you test it?

https://github.com/kohya-ss/sd-scripts/tree/feat-extract-flux-lora-from-weights-with-key-prefix

[fractal-fumbler]

test case:

1. org model
   1.1. flux1-dev -- BFL format
   1.2. flux1-dev -- Diffusers format
2. tuned model
   2.1. tuned flux model -- -- Diffusers format

a) when extracting LoRA 1.2 (org Diffusers) + 2.1 (tuned Diffusers) = success. LoRA is usable in ComfyUI and generated picture is different from picture, that is generated without LoRA applied.

dict of successful LoRA

b) when extracting LoRA 1.1 (org BFL) + 2.1 (tuned Diffusers) = error

with memsafeopen

2025-09-29 21:47:57 INFO     Using memory efficient safe_open                                flux_extract_lora_new.py:67
  0%|                                                                                           | 0/304 [00:00<?, ?it/s]
Traceback (most recent call last):                                                                                      
  File "flux_extract_lora_new.py", line 92, in svd                                         
    value_t = f_tuned.get_tensor(original_key)  # the tuned model may not have this key                                 
  File "library/safetensors_utils.py", line 158, in get_tensor
    raise KeyError(f"Tensor '{key}' not found in the file")
KeyError: "Tensor 'double_blocks.0.img_attn.proj.weight' not found in the file"

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "flux_extract_lora_new.py", line 229, in <module>
    svd(**vars(args))
    ~~~^^^^^^^^^^^^^^
  File "flux_extract_lora_new.py", line 94, in svd
    value_t = f_tuned.get_tensor(key)
  File "library/safetensors_utils.py", line 158, in get_tensor
    raise KeyError(f"Tensor '{key}' not found in the file")
KeyError: "Tensor 'double_blocks.0.img_attn.proj.weight' not found in the file"

without memsafeopen

  0%|                                                                                           | 0/304 [00:00<?, ?it/s]
Traceback (most recent call last):                                                                                      
  File "/tmp/yay/extract_lora_resize/flux_extract_lora_new.py", line 92, in svd                                         
    value_t = f_tuned.get_tensor(original_key)  # the tuned model may not have this key                                 
safetensors_rust.SafetensorError: File does not contain tensor double_blocks.0.img_attn.proj.weight                     
                                                                                                                        
During handling of the above exception, another exception occurred:                                                     
                                                                                                                        
Traceback (most recent call last):                                                                                      
  File "/tmp/yay/extract_lora_resize/flux_extract_lora_new.py", line 229, in <module>                                   
    svd(**vars(args))                                                                                                   
    ~~~^^^^^^^^^^^^^^                                                                                                   
  File "/tmp/yay/extract_lora_resize/flux_extract_lora_new.py", line 94, in svd                                         
    value_t = f_tuned.get_tensor(key)                                                                                   
safetensors_rust.SafetensorError: File does not contain tensor double_blocks.0.img_attn.proj.weight  

both tuned models, that i've tried, have double_blocks.0.img_attn.proj.weight in their dicts
pixelwave tuned model
trubass tuned model

my guess that script logic should be:

1. Check if model_org has key 'model.diffusion_model' in dict.
   1.1. If yes, then remove re.sub(r'model.diffusion_model_', '', key) from model_org
   1.2 If no, then don't do anything to model_org
2. Check if model_tuned has key 'model.diffusion_model' in dict.
   2.1. If yes, then remove re.sub(r'model.diffusion_model_', '', key)
   2.2 If no, then don't do anything to model_tuned.
3. Proceed with extracting LoRA from flux model.

here is a code snippet, if you don't mind my little experiments with grok

    with open_fn(model_org) as f_org:
        # filter keys
        org_keys = []
        for key in f_org.keys():
            if not ("single_block" in key or "double_block" in key):
                continue
            if ".bias" in key:
                continue
            if "norm" in key:
                continue
            org_keys.append(key)

        # Check and rename keys for model_org
        has_diffusion_model_org = any(key.startswith("model.diffusion_model.") for key in org_keys)
        if has_diffusion_model_org:
            logger.info("Detected 'model.diffusion_model.' prefix in model_org keys. Renaming keys to remove this prefix.")
            org_keys = [re.sub(r'^model\.diffusion_model\.', '', key) for key in org_keys]
        else:
            logger.info("No 'model.diffusion_model.' prefix detected in model_org keys.")

        with open_fn(model_tuned) as f_tuned:
            # filter keys for model_tuned
            tuned_keys = []
            for key in f_tuned.keys():
                if not ("single_block" in key or "double_block" in key):
                    continue
                if ".bias" in key:
                    continue
                if "norm" in key:
                    continue
                tuned_keys.append(key)

            # Check and rename keys for model_tuned
            has_diffusion_model_tuned = any(key.startswith("model.diffusion_model.") for key in tuned_keys)
            if has_diffusion_model_tuned:
                logger.info("Detected 'model.diffusion_model.' prefix in model_tuned keys. Renaming keys to remove this prefix.")
                tuned_keys = [re.sub(r'^model\.diffusion_model\.', '', key) for key in tuned_keys]
            else:
                logger.info("No 'model.diffusion_model.' prefix detected in model_tuned keys.")

            # Ensure keys match between models
            if set(org_keys) != set(tuned_keys):
                logger.warning("Key mismatch between model_org and model_tuned after renaming. Using common keys.")
                keys = list(set(org_keys).intersection(tuned_keys))
            else:
                keys = org_keys

            for key in tqdm(keys):
                # get tensors and calculate difference
                org_key = key if not has_diffusion_model_org else f"model.diffusion_model.{key}"
                tuned_key = key if not has_diffusion_model_tuned else f"model.diffusion_model.{key}"
                value_o = f_org.get_tensor(org_key)
                value_t = f_tuned.get_tensor(tuned_key)
                mat = value_t.to(calc_dtype) - value_o.to(calc_dtype)
                del value_o, value_t

in both cases LoRA extraction was successful :)