Flux Control Lora not unloaded correctly

[christopher5106]

Describe the bug

Hi,

There is a bug while switching pipeline from Flux dev when lora control has been loaded:

Reproduction

import torch
from controlnet_aux import CannyDetector
from diffusers import FluxControlPipeline
from diffusers.utils import load_image
from diffusers import FluxImg2ImgPipeline


model = "black-forest-labs/FLUX.1-dev"
pipe = FluxControlPipeline.from_pretrained(
 model, 
 torch_dtype=torch.bfloat16
).to("cuda")

pipe.load_lora_weights(
 "black-forest-labs/FLUX.1-Canny-dev-lora"
)

prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts."
control_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png")

processor = CannyDetector()
control_image = processor(control_image, low_threshold=50, high_threshold=200, detect_resolution=1024, image_resolution=1024)

image = pipe(
 prompt=prompt,
 control_image=control_image,
 num_inference_steps=50,
 guidance_scale=30.0,
).images[0]

# pipe = FluxImg2ImgPipeline.from_pretrained(
# "black-forest-labs/FLUX.1-dev", 
# torch_dtype=torch.bfloat16
# )

pipe = FluxImg2ImgPipeline.from_pipe(
 pipe,
 torch_dtype=torch.bfloat16
)

pipe = pipe.to("cuda")

init_image = load_image("https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg").resize((1024, 1024))
prompt = "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k"
image = pipe(
 prompt=prompt, 
 image=init_image, 
 num_inference_steps=28, 
 strength=0.5, 
 guidance_scale=2.5
).images[0]

Remplacing the from_pipe loading by standard loading shows the previous code should work

pipe = FluxImg2ImgPipeline.from_pretrained(
 "black-forest-labs/FLUX.1-dev", 
 torch_dtype=torch.bfloat16
)

Logs

No response

System Info

Ubuntu

Who can help?

@yiyixuxu @sayakpaul @DN6 @asomoza

[sayakpaul]

I think this is expected.

If you're doing a from_pipe on the Control pipeline the transformer has 128 input channels, which is not compatible with the Flux img2img pipeline.

We expand the transformer state dict to support loading the LoRA in question and we make explicit logging about that too.

So, we'd want to call pipe.unload_lora_weights() before calling from_pipe().

[christopher5106]

I tried with pipe.unload_lora_weights() before from_pipe and it still breaks

[sayakpaul]

Yeah looking into that now.

[sayakpaul]

@christopher5106 #10206 should probably solve the problem.

[christopher5106]

yes, it does work, when using unload_lora_weights before switch

[sayakpaul]

That is expected for the reasons I mentioned:
#10202 (comment)

Do you think this should be documented?

@yiyixuxu any thoughts?

[christopher5106]

@sayakpaul
just for my knowledge, what is the purpose of the torch.cat on single element array:

diffusers/src/diffusers/loaders/lora_conversion_utils.py

Line 792 in fafce0c

converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([sample_q])

[a-r-r-o-w]

It's a null-op here I think. It comes from copying the Flux conversion script which also has this:

diffusers/scripts/convert_flux_to_diffusers.py

Line 159 in fafce0c

converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([sample_q])

. It probably slipped through by mistake, or was serving a different purpose when it was added and eventually evolved to a single element list. Can be tackled in a refactor

[christopher5106]

@sayakpaul I thought it was solved. I don't know why but no more, something has change. Please reopen