Flux Controlnet Lora Fails To Load When Transformers Are Quantized.

[CyberVy]

Describe the bug

Failed to load Flux controlnet lora when the transformer is quantized to 4bit by bitsandbytes.
Maybe #10337 is a bit related to this issue.

Reproduction

import torch
from diffusers import FluxControlPipeline
pipe = FluxControlPipeline.from_pretrained("eramth/flux-4bit", torch_dtype=torch.float16).to("cuda")
pipe.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora")

Logs

/usr/local/lib/python3.11/dist-packages/diffusers/loaders/lora_pipeline.py in load_lora_weights(self, pretrained_model_name_or_path_or_dict, adapter_name, **kwargs)
   1549 
   1550         transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer
-> 1551         has_param_with_expanded_shape = self._maybe_expand_transformer_param_shape_or_error_(
   1552             transformer, transformer_lora_state_dict, transformer_norm_state_dict
   1553         )

/usr/local/lib/python3.11/dist-packages/diffusers/loaders/lora_pipeline.py in _maybe_expand_transformer_param_shape_or_error_(cls, transformer, lora_state_dict, norm_state_dict, prefix)
   2019 
   2020                     with torch.device("meta"):
-> 2021                         expanded_module = torch.nn.Linear(
   2022                             in_features, out_features, bias=bias, dtype=module_weight.dtype
   2023                         )

/usr/local/lib/python3.11/dist-packages/torch/nn/modules/linear.py in __init__(self, in_features, out_features, bias, device, dtype)
    103         self.in_features = in_features
    104         self.out_features = out_features
--> 105         self.weight = Parameter(
    106             torch.empty((out_features, in_features), **factory_kwargs)
    107         )

/usr/local/lib/python3.11/dist-packages/torch/nn/parameter.py in __new__(cls, data, requires_grad)
     44             # For ease of BC maintenance, keep this path for standard Tensor.
     45             # Eventually (tm), we should change the behavior for standard Tensor to match.
---> 46             return torch.Tensor._make_subclass(cls, data, requires_grad)
     47 
     48         # Path for custom tensors: set a flag on the instance to indicate parameter-ness.

RuntimeError: Only Tensors of floating point and complex dtype can require gradients

System Info

• 🤗 Diffusers version: 0.33.0.dev0
• Platform: Linux-6.1.85+-x86_64-with-glibc2.35
• Running on Google Colab?: Yes
• Python version: 3.11.11
• PyTorch version (GPU?): 2.5.1+cu124 (True)
• Flax version (CPU?/GPU?/TPU?): 0.10.4 (gpu)
• Jax version: 0.4.33
• JaxLib version: 0.4.33
• Huggingface_hub version: 0.28.1
• Transformers version: 4.48.3
• Accelerate version: 1.3.0
• PEFT version: 0.14.0
• Bitsandbytes version: 0.45.3
• Safetensors version: 0.5.3
• xFormers version: not installed
• Accelerator: Tesla T4, 15360 MiB

Who can help?

@hlky

[hlky]

Hi @CyberVy. With #10990 I am able to load eramth/flux-4bit with black-forest-labs/FLUX.1-Canny-dev-lora, can you confirm this fixes the issue on your end?

[CyberVy]

Hi @CyberVy. With #10990 I am able to load eramth/flux-4bit with black-forest-labs/FLUX.1-Canny-dev-lora, can you confirm this fixes the issue on your end?

@hlky
Great Job! Thank you very much. The issue is fixed now.

[CyberVy]

@hlky
That's very good for loading.
However, there is an another inference issue.
BTW, sorry, I am not very familar with bitsandbytes. Otherwise I will fix this issue myself. I am very appreciated for your help.

import torch
from diffusers import FluxControlPipeline
pipe = FluxControlPipeline.from_pretrained("eramth/flux-4bit", torch_dtype=torch.float16).to("cuda")
pipe.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora")
# The issue of load_lora_weights is fixed.

# There is a new issue about inference with lora.
from controlnet_aux import CannyDetector
from diffusers.utils import load_image
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,
    height=1024,
    width=1024,
    num_inference_steps=20,
    guidance_scale=3.5,
).images[0]

File ~/external-libraries/torch/utils/_contextlib.py:116, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    113 @functools.wraps(func)
    114 def decorate_context(*args, **kwargs):
    115     with ctx_factory():
--> 116         return func(*args, **kwargs)

File ~/external-libraries/diffusers/pipelines/flux/pipeline_flux_control.py:835, in FluxControlPipeline.__call__(self, prompt, prompt_2, control_image, height, width, num_inference_steps, sigmas, guidance_scale, num_images_per_prompt, generator, latents, prompt_embeds, pooled_prompt_embeds, output_type, return_dict, joint_attention_kwargs, callback_on_step_end, callback_on_step_end_tensor_inputs, max_sequence_length)
    832 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
    833 timestep = t.expand(latents.shape[0]).to(latents.dtype)
--> 835 noise_pred = self.transformer(
    836     hidden_states=latent_model_input,
    837     timestep=timestep / 1000,
    838     guidance=guidance,
    839     pooled_projections=pooled_prompt_embeds,
    840     encoder_hidden_states=prompt_embeds,
    841     txt_ids=text_ids,
    842     img_ids=latent_image_ids,
    843     joint_attention_kwargs=self.joint_attention_kwargs,
    844     return_dict=False,
    845 )[0]
    847 # compute the previous noisy sample x_t -> x_t-1
    848 latents_dtype = latents.dtype

File ~/external-libraries/torch/nn/modules/module.py:1736, in Module._wrapped_call_impl(self, *args, **kwargs)
   1734     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735 else:
-> 1736     return self._call_impl(*args, **kwargs)

File ~/external-libraries/torch/nn/modules/module.py:1747, in Module._call_impl(self, *args, **kwargs)
   1742 # If we don't have any hooks, we want to skip the rest of the logic in
   1743 # this function, and just call forward.
   1744 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1745         or _global_backward_pre_hooks or _global_backward_hooks
   1746         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747     return forward_call(*args, **kwargs)
   1749 result = None
   1750 called_always_called_hooks = set()

File ~/external-libraries/diffusers/models/transformers/transformer_flux.py:445, in FluxTransformer2DModel.forward(self, hidden_states, encoder_hidden_states, pooled_projections, timestep, img_ids, txt_ids, guidance, joint_attention_kwargs, controlnet_block_samples, controlnet_single_block_samples, return_dict, controlnet_blocks_repeat)
    440     if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
    441         logger.warning(
    442             "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
    443         )
--> 445 hidden_states = self.x_embedder(hidden_states)
    447 timestep = timestep.to(hidden_states.dtype) * 1000
    448 if guidance is not None:

File ~/external-libraries/torch/nn/modules/module.py:1736, in Module._wrapped_call_impl(self, *args, **kwargs)
   1734     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735 else:
-> 1736     return self._call_impl(*args, **kwargs)

File ~/external-libraries/torch/nn/modules/module.py:1747, in Module._call_impl(self, *args, **kwargs)
   1742 # If we don't have any hooks, we want to skip the rest of the logic in
   1743 # this function, and just call forward.
   1744 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1745         or _global_backward_pre_hooks or _global_backward_hooks
   1746         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747     return forward_call(*args, **kwargs)
   1749 result = None
   1750 called_always_called_hooks = set()

File ~/external-libraries/peft/tuners/lora/layer.py:621, in Linear.forward(self, x, *args, **kwargs)
    618 x = x.to(lora_A.weight.dtype)
    620 if not self.use_dora[active_adapter]:
--> 621     result = result + lora_B(lora_A(dropout(x))) * scaling
    622 else:
    623     if isinstance(dropout, nn.Identity) or not self.training:

File ~/external-libraries/torch/nn/modules/module.py:1736, in Module._wrapped_call_impl(self, *args, **kwargs)
   1734     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735 else:
-> 1736     return self._call_impl(*args, **kwargs)

File ~/external-libraries/torch/nn/modules/module.py:1747, in Module._call_impl(self, *args, **kwargs)
   1742 # If we don't have any hooks, we want to skip the rest of the logic in
   1743 # this function, and just call forward.
   1744 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1745         or _global_backward_pre_hooks or _global_backward_hooks
   1746         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747     return forward_call(*args, **kwargs)
   1749 result = None
   1750 called_always_called_hooks = set()

File ~/external-libraries/torch/nn/modules/linear.py:125, in Linear.forward(self, input)
    124 def forward(self, input: Tensor) -> Tensor:
--> 125     return F.linear(input, self.weight, self.bias)

File ~/external-libraries/torch/_prims_common/wrappers.py:273, in out_wrapper.<locals>._out_wrapper.<locals>._fn(out, *args, **kwargs)
    271     result = fn(*args, is_out=(out is not None), **kwargs)  # type: ignore[arg-type]
    272 else:
--> 273     result = fn(*args, **kwargs)
    274 assert (
    275     isinstance(result, TensorLike)
    276     and is_tensor
    277     or isinstance(result, Tuple)  # type: ignore[arg-type]
    278     and len(result) == len(out_names)  # type: ignore[arg-type]
    279 )
    280 if out is not None:
    281     # Naively you might expect this assert to be true, but
    282     # it's not:
   (...)
    295     # be a normal meta tensor, but this is perfectly
    296     # harmless.

File ~/external-libraries/torch/_decomp/decompositions.py:83, in type_casts.<locals>.inner(*args, **kwargs)
     80     else:
     81         return x
---> 83 r = f(*tree_map(increase_prec, args), **tree_map(increase_prec, kwargs))
     84 if compute_dtype_only:
     85     return r

File ~/external-libraries/torch/_decomp/decompositions.py:1525, in addmm(self, mat1, mat2, beta, alpha)
   1517     return out
   1519 # The output of aten.addmm is contiguous, we need to match this behavior in the decomposition.
   1520 # The original implementation 'beta * self + out' would return a strided tensor if `self` is strided.
   1521 # We thus use `out`, the output of torch.mm, which is always contiguous, as the first argument for addition.
   1522 # This is relying on TensorIterator's behavior that it takes higher precedence on the stride of first input.
   1523 # Alternative, we can write `(beta * self + out).contiguous()`, but it introduces another copy in some cases.
   1524 # This implementation is not ideal, and we should revisit this when we have a better solution.
-> 1525 return out + beta * self

File ~/external-libraries/torch/_prims_common/wrappers.py:273, in out_wrapper.<locals>._out_wrapper.<locals>._fn(out, *args, **kwargs)
    271     result = fn(*args, is_out=(out is not None), **kwargs)  # type: ignore[arg-type]
    272 else:
--> 273     result = fn(*args, **kwargs)
    274 assert (
    275     isinstance(result, TensorLike)
    276     and is_tensor
    277     or isinstance(result, Tuple)  # type: ignore[arg-type]
    278     and len(result) == len(out_names)  # type: ignore[arg-type]
    279 )
    280 if out is not None:
    281     # Naively you might expect this assert to be true, but
    282     # it's not:
   (...)
    295     # be a normal meta tensor, but this is perfectly
    296     # harmless.

File ~/external-libraries/torch/_prims_common/wrappers.py:141, in elementwise_type_promotion_wrapper.__call__.<locals>._fn(*args, **kwargs)
    134 promoted_args = {
    135     x: _maybe_convert_to_dtype(bound.arguments[x], compute_dtype)
    136     for x in self.type_promoting_arg_names  # type: ignore[union-attr]
    137     if x in bound.arguments.keys()
    138 }
    139 bound.arguments.update(promoted_args)
--> 141 result = fn(**bound.arguments)
    143 # Override the return_dtype if a dtype arg is present and not None
    144 if "dtype" in bound.arguments:

File ~/external-libraries/torch/_refs/__init__.py:1099, in add(a, b, alpha)
   1096     else:
   1097         b = b * alpha
-> 1099 output = prims.add(a, b)
   1100 return handle_noncontiguous_outputs([a, b], output)

File ~/external-libraries/torch/_ops.py:716, in OpOverload.__call__(self, *args, **kwargs)
    715 def __call__(self, /, *args, **kwargs):
--> 716     return self._op(*args, **kwargs)

File ~/external-libraries/torch/_library/fake_impl.py:93, in construct_meta_kernel.<locals>.meta_kernel(*args, **kwargs)
     83     raise RuntimeError(
     84         f"Attempted to call get_ctx() for the meta implementation "
     85         f"for {qualname} (implemented at {source})"
   (...)
     89         f"behavior."
     90     )
     92 with set_ctx_getter(error_on_ctx):
---> 93     return fake_impl_holder.kernel(*args, **kwargs)

File ~/external-libraries/torch/_library/utils.py:20, in Kernel.__call__(self, *args, **kwargs)
     19 def __call__(self, *args, **kwargs):
---> 20     return self.func(*args, **kwargs)

File ~/external-libraries/torch/library.py:1151, in _check_pystubs_once.<locals>.inner(*args, **kwargs)
   1149 nonlocal checked
   1150 if checked:
-> 1151     return func(*args, **kwargs)
   1153 op = torch._library.utils.lookup_op(qualname)
   1154 if op._defined_in_python:

File ~/external-libraries/torch/_library/custom_ops.py:614, in CustomOpDef._register_to_dispatcher.<locals>.fake_impl(*args, **kwargs)
    607         return None
    608     raise RuntimeError(
    609         f"There was no fake impl registered for {self}. "
    610         f"This is necessary for torch.compile/export/fx tracing to work. "
    611         f"Please use `{self._init_fn.__name__}.register_fake` to add an "
    612         f"fake impl."
    613     )
--> 614 return self._abstract_fn(*args, **kwargs)

File ~/external-libraries/torch/_prims/__init__.py:402, in _prim_elementwise_meta(type_promotion, args_with_fixed_dtypes, *args)
    399 if args_with_fixed_dtypes is not None:
    400     args_ = list(args_with_fixed_dtypes) + args_
--> 402 utils.check_same_device(*args_, allow_cpu_scalar_tensors=True)
    403 utils.check_same_shape(*args_, allow_cpu_scalar_tensors=True)
    405 l2p_perm = utils.compute_elementwise_output_logical_to_physical_perm(*args_)

File ~/external-libraries/torch/_prims_common/__init__.py:742, in check_same_device(allow_cpu_scalar_tensors, *args)
    734     if device != arg.device:
    735         msg = (
    736             "Tensor on device "
    737             + str(arg.device)
   (...)
    740             + "!"
    741         )
--> 742         raise RuntimeError(msg)
    743 else:
    744     msg = (
    745         "Unexpected type when checking for same device, " + str(type(arg)) + "!"
    746     )

RuntimeError: Tensor on device meta is not on the expected device cuda:0!

[hlky]

Sorry about that! Inference is working on my end now

[CyberVy]

Nice fix! Now it works for me now.

[CyberVy]

@hlky

Thanks again!
But there is an another issue about unloading.
It seems that the transformer is not back to the initial state after unloading lora weights, no matter I use delete_adapters or unload_lora_weights.

import torch
from diffusers import FluxControlPipeline
pipe = FluxControlPipeline.from_pretrained("eramth/flux-4bit", torch_dtype=torch.float16).to("cuda")
pipe.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora")


from diffusers import FluxPipeline
pipe.unload_lora_weights()
components = pipe.components
p = FluxPipeline(**components)
image = p(
    prompt="a cute cat",
    height=1024,
    width=1024,
    num_inference_steps=20,
    guidance_scale=3.5,
).images[0]

File ~/external-libraries/torch/utils/_contextlib.py:116, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    113 @functools.wraps(func)
    114 def decorate_context(*args, **kwargs):
    115     with ctx_factory():
--> 116         return func(*args, **kwargs)

File ~/external-libraries/diffusers/pipelines/flux/pipeline_flux.py:950, in FluxPipeline.__call__(self, prompt, prompt_2, negative_prompt, negative_prompt_2, true_cfg_scale, height, width, num_inference_steps, sigmas, guidance_scale, num_images_per_prompt, generator, latents, prompt_embeds, pooled_prompt_embeds, ip_adapter_image, ip_adapter_image_embeds, negative_ip_adapter_image, negative_ip_adapter_image_embeds, negative_prompt_embeds, negative_pooled_prompt_embeds, output_type, return_dict, joint_attention_kwargs, callback_on_step_end, callback_on_step_end_tensor_inputs, max_sequence_length)
    948 # compute the previous noisy sample x_t -> x_t-1
    949 latents_dtype = latents.dtype
--> 950 latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
    952 if latents.dtype != latents_dtype:
    953     if torch.backends.mps.is_available():
    954         # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272

File ~/external-libraries/diffusers/schedulers/scheduling_flow_match_euler_discrete.py:433, in FlowMatchEulerDiscreteScheduler.step(self, model_output, timestep, sample, s_churn, s_tmin, s_tmax, s_noise, generator, return_dict)
    430 sigma = self.sigmas[self.step_index]
    431 sigma_next = self.sigmas[self.step_index + 1]
--> 433 prev_sample = sample + (sigma_next - sigma) * model_output
    435 # Cast sample back to model compatible dtype
    436 prev_sample = prev_sample.to(model_output.dtype)

RuntimeError: The size of tensor a (128) must match the size of tensor b (64) at non-singleton dimension 2

[hlky]

When unloading a Flux lora like black-forest-labs/FLUX.1-Canny-dev-lora that expands the shapes we have to use pipe.unload_lora_weights(reset_to_overwritten_params=True)

cc @sayakpaul should this be the default?

[sayakpaul]

I think we make it sufficiently clear in the docs how unload_lora_weights() should be used for special cases like this. So, I think fine as is.

The docstrings are also updated:

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 1878 in 748cb0f

Args:

[CyberVy]

@hlky @sayakpaul

Thanks for your clarification.

I just noticed FluxLoraLoaderMixin.unload_lora_weights has an optional argument.
Yes, it works well with reset_to_overwritten_params=True.

However this method unloads all loaded LoRA adapters, unlike LoraBaseMixin.delete_adapters which allows unloading one or multiple specific LoRA adapters.

Actually, not only LoraBaseMixin.delete_adapters, but also all other LoRA unloading/disable related methods that are not overridden in FluxLoraLoaderMixin don't work properly if the LoRA expands the shape. For example, LoraBaseMixin.disable_lora is one of them also.

Should we rewrite those methods in FluxLoraLoaderMixin?

[sayakpaul]

However this method unloads all loaded LoRA adapters, unlike LoraBaseMixin.delete_adapters which allows unloading one or multiple specific LoRA adapters.

That is what it is supposed to do and we make that sufficiently clear in the docstrings and also our guides.

For example, LoraBaseMixin.disable_lora is one of them also.

Should we rewrite those methods in FluxLoraLoaderMixin?

Thanks for flagging. Would you like to maybe open a PR to help us fix that?

[CyberVy]

That is what it is supposed to do and we make that sufficiently clear in the docstrings and also our guides.

Yes, I know unloading all LoRA is what it is supposed to do. What I mean is sometime we don't want to unload all LoRA, so maybe we should fix the methods which can unload the specified LoRA like LoraBaseMixin.delete_adapters.

Thanks for flagging. Would you like to maybe open a PR to help us fix that?

Of course. :)

[sayakpaul]

Yes, I know unloading all LoRA is what it is supposed to do. What I mean is sometime we don't want to unload all LoRA, so maybe we should fix the methods which can unload the specified LoRA like LoraBaseMixin.delete_adapters.

We did consider that at some point with @BenjaminBossan. We considered that users could completely unload a specific LoRA from a model (delete_adapters() doesn't completely unload). See the discussions here:
#9325 (comment). We ultimately decided not to pursue the feature as it's better to keep the LoRA layers when dealing with multiple LoRA layers. I realize this is also diverging a bit from the original issue, so, maybe we could switch to a different one.

Of course. :)

Thanks much. Please let us know when we can help.