Update app.py

app.py

@@ -13,8 +13,75 @@ with open('loras.json', 'r') as f:
 # Initialize the base model
 base_model = "black-forest-labs/FLUX.1-dev"
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
+original_load_lora = copy.deepcopy(pipe.load_lora_into_transformer)
 pipe.to("cuda")
 
+def load_lora_into_transformer_patched(cls, state_dict, transformer, adapter_name=None, alpha=None, _pipeline=None):
+    from peft import LoraConfig, inject_adapter_in_model, set_peft_model_state_dict
+
+    keys = list(state_dict.keys())
+
+    transformer_keys = [k for k in keys if k.startswith(cls.transformer_name)]
+    state_dict = {
+        k.replace(f"{cls.transformer_name}.", ""): v for k, v in state_dict.items() if k in transformer_keys
+    }
+
+    if len(state_dict.keys()) > 0:
+        # check with first key if is not in peft format
+        first_key = next(iter(state_dict.keys()))
+        if "lora_A" not in first_key:
+            state_dict = convert_unet_state_dict_to_peft(state_dict)
+
+        if adapter_name in getattr(transformer, "peft_config", {}):
+            raise ValueError(
+                f"Adapter name {adapter_name} already in use in the transformer - please select a new adapter name."
+            )
+
+        rank = {}
+        for key, val in state_dict.items():
+            if "lora_B" in key:
+                rank[key] = val.shape[1]
+
+        lora_config_kwargs = get_peft_kwargs(rank, network_alpha_dict=None, peft_state_dict=state_dict)
+        if "use_dora" in lora_config_kwargs:
+            if lora_config_kwargs["use_dora"] and is_peft_version("<", "0.9.0"):
+                raise ValueError(
+                    "You need `peft` 0.9.0 at least to use DoRA-enabled LoRAs. Please upgrade your installation of `peft`."
+                )
+            else:
+                lora_config_kwargs.pop("use_dora")
+
+        
+        lora_config_kwargs["lora_alpha"] = 32
+        lora_config = LoraConfig(**lora_config_kwargs)
+
+        # adapter_name
+        if adapter_name is None:
+            adapter_name = get_adapter_name(transformer)
+
+        # In case the pipeline has been already offloaded to CPU - temporarily remove the hooks
+        # otherwise loading LoRA weights will lead to an error
+        is_model_cpu_offload, is_sequential_cpu_offload = cls._optionally_disable_offloading(_pipeline)
+
+        inject_adapter_in_model(lora_config, transformer, adapter_name=adapter_name)
+        incompatible_keys = set_peft_model_state_dict(transformer, state_dict, adapter_name)
+
+        if incompatible_keys is not None:
+            # check only for unexpected keys
+            unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
+            if unexpected_keys:
+                logger.warning(
+                    f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
+                    f" {unexpected_keys}. "
+                )
+
+        # Offload back.
+        if is_model_cpu_offload:
+            _pipeline.enable_model_cpu_offload()
+        elif is_sequential_cpu_offload:
+            _pipeline.enable_sequential_cpu_offload()
+        # Unsafe code />
+
 def update_selection(evt: gr.SelectData):
     selected_lora = loras[evt.index]
     new_placeholder = f"Type a prompt for {selected_lora['title']}"
@@ -41,6 +108,11 @@ def run_lora(prompt, cfg_scale, steps, selected_index, seed, width, height, lora
     else:
         pipe.load_lora_weights(lora_path)
 
+    if "custom_alpha" in selected_lora:
+        pipe.load_lora_into_transformer = load_lora_into_transformer_patched
+    else:
+        pipe.load_lora_into_transformer = original_load_lora
+        
     # Set random seed for reproducibility
     generator = torch.Generator(device="cuda").manual_seed(seed)