Error in loading flux2 klein lora adapter

[saeedkhanehgir]

Hi,

I train my lora with DiffSynth repo for flux 2 klein model. When I want to inference with below code, I get warning and I think this problem is important.

inference code :

import torch
from diffusers import Flux2KleinPipeline
from glob import glob
from PIL import Image 

device = "cuda"
dtype = torch.bfloat16

model_path = "./weights/flux2-klein-9b"
pipe = Flux2KleinPipeline.from_pretrained(model_path, torch_dtype=dtype)
pipe.load_lora_weights(
    "./FLUX.2-klein-9B_lora",  
    weight_name="epoch-5.safetensors"  
)
pipe.fuse_lora(lora_scale=1.0)

pipe.enable_sequential_cpu_offload()  # save some VRAM by offloading the model to CPU

prompt = """You are an expert image adder."""

images_path = glob("./images/*")
for image_path in images_path:
    image  = Image.open(image_path)
    image_name = image_path.split('/')[-1]
    image = pipe(
        prompt=prompt,
        image = image,
        height=1024,
        width=1024,
        guidance_scale=1.0,
        num_inference_steps=4,
        generator=torch.Generator(device=device).manual_seed(0)
    ).images[0]
    image.save(f"./output/{image_name}")

warning

No LoRA keys associated to Flux2Transformer2DModel found with the prefix='transformer'. This is safe to ignore if LoRA state dict didn't originally have any Flux2Transformer2DModel related params. You can also try specifying prefix=None to resolve the warning.

Can someone help me ?

[rehan243]

Hey, saw that warning about No LoRA keys associated to Flux2Transformer2DModel in your output. Yeah, this is tricky — it usually pops up when the LoRA weights you're loading don't match the expected structure for the Flux2KleinPipeline transformer. Honestly, we've hit this kind of mismatch before when fine-tuning on custom datasets with different repos, and it took us a few tries to debug. The thing is, DiffSynth might be structuring the state dict in a way that diffusers can't directly map to the transformer keys.

I'd suggest explicitly checking the state dict keys in your LoRA file to see what's actually there. You can dump them with something like this to debug:

import safetensors.torch
state_dict = safetensors.torch.load_file("./FLUX.2-klein-9B_lora/epoch-5.safetensors")
for k in state_dict.keys():
    print(k)

Also, try setting prefix=None in load_lora_weights() as the warning suggests — sometimes that lets it fall back to a more generic key matching. Out of curiosity, did you train this LoRA with the exact same base model version as flux2-klein-9b in your inference setup?

[asomoza]

Hi @saeedkhanehgir @rehan243 this is because we don't know all the formats beforehand, most trainers change the key naming at their discretion so we need to wait until the loras comes out and people open an issue like this one for us to integrate them.

But for this we need a sfw lora uploaded to the hf hub, this is because it's a lot more work for us to download them and test them locally and it's also safer.

[rehan243]

Hey, gotcha. Yeah, that version issue... We've had similar struggles with LoRA mismatches. The docs don't mention this, but it's worth checking if the LoRA keys are aligned with the expected format for Flux2KleinPipeline.

[asomoza]

@rehan243 yeah, it helps and its a very useful advice, no issues there.