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

[sayakpaul]

What is the problem?

[christopher5106]

Unloading flux control lora does not work anymore.

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.unload_lora_weights()

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]

    image = pipe(
            ^^^^^
  File "venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.12/site-packages/diffusers/pipelines/flux/pipeline_flux_img2img.py", line 773, in __call__
    latents, latent_image_ids = self.prepare_latents(
                                ^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.12/site-packages/diffusers/pipelines/flux/pipeline_flux_img2img.py", line 565, in prepare_latents
    latents = self.scheduler.scale_noise(image_latents, timestep, noise)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.12/site-packages/diffusers/schedulers/scheduling_flow_match_euler_discrete.py", line 187, in scale_noise
    sample = sigma * noise + (1.0 - sigma) * sample
             ~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (32) must match the size of tensor b (16) at non-singleton dimension 1

[christopher5106]

When you asked me last month, it was working, in particular I double check now, unloading works well with commit 1b202c5 but no more with recent version of diffusers

[sayakpaul]

Could you follow the guidelines from https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux#note-about-unloadloraweights-when-using-flux-loras and check again?

[christopher5106]

There was an introduction of a parameter reset_to_overwritten_params=True inbetween my validation... So with this parameter it now works well.