multimodalart HF staff commited on
Commit
c59400c
1 Parent(s): 71c1eef

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +72 -0
app.py CHANGED
@@ -13,8 +13,75 @@ with open('loras.json', 'r') as f:
13
  # Initialize the base model
14
  base_model = "black-forest-labs/FLUX.1-dev"
15
  pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
 
16
  pipe.to("cuda")
17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18
  def update_selection(evt: gr.SelectData):
19
  selected_lora = loras[evt.index]
20
  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
41
  else:
42
  pipe.load_lora_weights(lora_path)
43
 
 
 
 
 
 
44
  # Set random seed for reproducibility
45
  generator = torch.Generator(device="cuda").manual_seed(seed)
46
 
 
13
  # Initialize the base model
14
  base_model = "black-forest-labs/FLUX.1-dev"
15
  pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
16
+ original_load_lora = copy.deepcopy(pipe.load_lora_into_transformer)
17
  pipe.to("cuda")
18
 
19
+ def load_lora_into_transformer_patched(cls, state_dict, transformer, adapter_name=None, alpha=None, _pipeline=None):
20
+ from peft import LoraConfig, inject_adapter_in_model, set_peft_model_state_dict
21
+
22
+ keys = list(state_dict.keys())
23
+
24
+ transformer_keys = [k for k in keys if k.startswith(cls.transformer_name)]
25
+ state_dict = {
26
+ k.replace(f"{cls.transformer_name}.", ""): v for k, v in state_dict.items() if k in transformer_keys
27
+ }
28
+
29
+ if len(state_dict.keys()) > 0:
30
+ # check with first key if is not in peft format
31
+ first_key = next(iter(state_dict.keys()))
32
+ if "lora_A" not in first_key:
33
+ state_dict = convert_unet_state_dict_to_peft(state_dict)
34
+
35
+ if adapter_name in getattr(transformer, "peft_config", {}):
36
+ raise ValueError(
37
+ f"Adapter name {adapter_name} already in use in the transformer - please select a new adapter name."
38
+ )
39
+
40
+ rank = {}
41
+ for key, val in state_dict.items():
42
+ if "lora_B" in key:
43
+ rank[key] = val.shape[1]
44
+
45
+ lora_config_kwargs = get_peft_kwargs(rank, network_alpha_dict=None, peft_state_dict=state_dict)
46
+ if "use_dora" in lora_config_kwargs:
47
+ if lora_config_kwargs["use_dora"] and is_peft_version("<", "0.9.0"):
48
+ raise ValueError(
49
+ "You need `peft` 0.9.0 at least to use DoRA-enabled LoRAs. Please upgrade your installation of `peft`."
50
+ )
51
+ else:
52
+ lora_config_kwargs.pop("use_dora")
53
+
54
+
55
+ lora_config_kwargs["lora_alpha"] = 32
56
+ lora_config = LoraConfig(**lora_config_kwargs)
57
+
58
+ # adapter_name
59
+ if adapter_name is None:
60
+ adapter_name = get_adapter_name(transformer)
61
+
62
+ # In case the pipeline has been already offloaded to CPU - temporarily remove the hooks
63
+ # otherwise loading LoRA weights will lead to an error
64
+ is_model_cpu_offload, is_sequential_cpu_offload = cls._optionally_disable_offloading(_pipeline)
65
+
66
+ inject_adapter_in_model(lora_config, transformer, adapter_name=adapter_name)
67
+ incompatible_keys = set_peft_model_state_dict(transformer, state_dict, adapter_name)
68
+
69
+ if incompatible_keys is not None:
70
+ # check only for unexpected keys
71
+ unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
72
+ if unexpected_keys:
73
+ logger.warning(
74
+ f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
75
+ f" {unexpected_keys}. "
76
+ )
77
+
78
+ # Offload back.
79
+ if is_model_cpu_offload:
80
+ _pipeline.enable_model_cpu_offload()
81
+ elif is_sequential_cpu_offload:
82
+ _pipeline.enable_sequential_cpu_offload()
83
+ # Unsafe code />
84
+
85
  def update_selection(evt: gr.SelectData):
86
  selected_lora = loras[evt.index]
87
  new_placeholder = f"Type a prompt for {selected_lora['title']}"
 
108
  else:
109
  pipe.load_lora_weights(lora_path)
110
 
111
+ if "custom_alpha" in selected_lora:
112
+ pipe.load_lora_into_transformer = load_lora_into_transformer_patched
113
+ else:
114
+ pipe.load_lora_into_transformer = original_load_lora
115
+
116
  # Set random seed for reproducibility
117
  generator = torch.Generator(device="cuda").manual_seed(seed)
118