Flux quantized with lora

src/diffusers/loaders/lora_pipeline.py

@@ -18,6 +18,7 @@
 import torch
 from huggingface_hub.utils import validate_hf_hub_args
 
+from ..quantizers.bitsandbytes import dequantize_bnb_weight
 from ..utils import (
     USE_PEFT_BACKEND,
     deprecate,
@@ -1970,7 +1971,11 @@ def _maybe_expand_transformer_param_shape_or_error_(
         is_peft_loaded = getattr(transformer, "peft_config", None) is not None
         for name, module in transformer.named_modules():
             if isinstance(module, torch.nn.Linear):
-                module_weight = module.weight.data
+                module_weight = (
+                    dequantize_bnb_weight(module.weight, state=module.weight.quant_state).data
+                    if module.weight.__class__.__name__ == "Params4bit"
+                    else module.weight.data
+                )
                 module_bias = module.bias.data if module.bias is not None else None
                 bias = module_bias is not None
 
@@ -1994,7 +1999,7 @@ def _maybe_expand_transformer_param_shape_or_error_(
 
                 # TODO (sayakpaul): We still need to consider if the module we're expanding is
                 # quantized and handle it accordingly if that is the case.
-                module_out_features, module_in_features = module_weight.shape
+                module_out_features, module_in_features = module_weight_shape
                 debug_message = ""
                 if in_features > module_in_features:
                     debug_message += (
@@ -2028,7 +2033,7 @@ def _maybe_expand_transformer_param_shape_or_error_(
                     new_weight = torch.zeros_like(
                         expanded_module.weight.data, device=module_weight.device, dtype=module_weight.dtype
                     )
-                    slices = tuple(slice(0, dim) for dim in module_weight.shape)
+                    slices = tuple(slice(0, dim) for dim in module_weight_shape)
                     new_weight[slices] = module_weight
                     tmp_state_dict = {"weight": new_weight}
                     if module_bias is not None: