fixes

Author: sayakpaul

src/diffusers/loaders/lora_conversion_utils.py

@@ -673,6 +673,10 @@ def _convert_bfl_flux_control_lora_to_diffusers(original_state_dict):
     inner_dim = 3072
     mlp_ratio = 4.0
 
+    for k in original_state_dict:
+        if "bias" in k and "img_in" in k:
+            print(f"{k=}")
+
     def swap_scale_shift(weight):
         shift, scale = weight.chunk(2, dim=0)
         new_weight = torch.cat([scale, shift], dim=0)
@@ -750,7 +754,7 @@ def swap_scale_shift(weight):
     for i in range(num_layers):
         block_prefix = f"transformer_blocks.{i}."
 
-        for lora_key, lora_key in zip(["lora_A", "lora_B"], ["lora_A", "lora_B"]):
+        for lora_key in ["lora_A", "lora_B"]:
             # norms
             converted_state_dict[f"{block_prefix}norm1.linear.{lora_key}.weight"] = original_state_dict.pop(
                 f"double_blocks.{i}.img_mod.lin.{lora_key}.weight"

src/diffusers/loaders/lora_pipeline.py

@@ -2318,7 +2318,7 @@ def _maybe_expand_transformer_param_shape_or_error_(
 
                 lora_A_weight_name = f"{name}.lora_A.weight"
                 lora_B_weight_name = f"{name}.lora_B.weight"
-                lora_B_bias_name = f"{name}.lora_B.bias"
+                # lora_B_bias_name = f"{name}.lora_B.bias"
 
                 if lora_A_weight_name not in state_dict.keys():
                     continue
@@ -2352,24 +2352,15 @@ def _maybe_expand_transformer_param_shape_or_error_(
                 expanded_module = torch.nn.Linear(
                     in_features, out_features, bias=bias, device=module_weight.device, dtype=module_weight.dtype
                 )
-
+                # Only weights are expanded and biases are not.
                 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)
                 new_weight[slices] = module_weight
                 expanded_module.weight.data.copy_(new_weight)
-
-                bias_present_for_lora_B = lora_B_bias_name in state_dict
-                if bias_present_for_lora_B:
-                    new_bias_shape = state_dict[lora_B_bias_name].shape
-                    if bias and module_bias.shape < new_bias_shape:
-                        new_bias = torch.zeros_like(
-                            expanded_module.bias.data, device=module_bias.device, dtype=module_bias.dtype
-                        )
-                        slices = tuple(slice(0, dim) for dim in module_bias.shape)
-                        new_bias[slices] = module_bias
-                        expanded_module.bias.data.copy_(new_bias)
+                if module_bias is not None:
+                    expanded_module.bias.data.copy_(module_bias)
 
                 setattr(parent_module, current_module_name, expanded_module)