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.