Skip to content

Commit 766aa50

Browse files
patrickvonplatenpatrickvonplaten
authored
[LoRA Attn Processors] Refactor LoRA Attn Processors (huggingface#4765)
* [LoRA Attn] Refactor LoRA attn * correct for network alphas * fix more * fix more tests * fix more tests * Move below * Finish * better version * correct serialization format * fix * fix more * fix more * fix more * Apply suggestions from code review * Update src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py * deprecation * relax atol for slow test slighly * Finish tests * make style * make style
1 parent c4d2823 commit 766aa50

11 files changed

Lines changed: 546 additions & 652 deletions

File tree

‎src/diffusers/loaders.py‎

Lines changed: 66 additions & 136 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
import requests
2525
import safetensors
2626
import torch
27-
import torch.nn.functional as F
2827
from huggingface_hub import hf_hub_download, model_info
2928
from torch import nn
3029

@@ -231,15 +230,7 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict
231230
232231
"""
233232
from .models.attention_processor import (
234-
AttnAddedKVProcessor,
235-
AttnAddedKVProcessor2_0,
236233
CustomDiffusionAttnProcessor,
237-
LoRAAttnAddedKVProcessor,
238-
LoRAAttnProcessor,
239-
LoRAAttnProcessor2_0,
240-
LoRAXFormersAttnProcessor,
241-
SlicedAttnAddedKVProcessor,
242-
XFormersAttnProcessor,
243234
)
244235
from .models.lora import LoRACompatibleConv, LoRACompatibleLinear, LoRAConv2dLayer, LoRALinearLayer
245236

@@ -314,24 +305,14 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict
314305
state_dict = pretrained_model_name_or_path_or_dict
315306

316307
# fill attn processors
317-
attn_processors = {}
318-
non_attn_lora_layers = []
308+
lora_layers_list = []
319309

320310
is_lora = all(("lora" in k or k.endswith(".alpha")) for k in state_dict.keys())
321311
is_custom_diffusion = any("custom_diffusion" in k for k in state_dict.keys())
322312

323313
if is_lora:
324-
is_new_lora_format = all(
325-
key.startswith(self.unet_name) or key.startswith(self.text_encoder_name) for key in state_dict.keys()
326-
)
327-
if is_new_lora_format:
328-
# Strip the `"unet"` prefix.
329-
is_text_encoder_present = any(key.startswith(self.text_encoder_name) for key in state_dict.keys())
330-
if is_text_encoder_present:
331-
warn_message = "The state_dict contains LoRA params corresponding to the text encoder which are not being used here. To use both UNet and text encoder related LoRA params, use [`pipe.load_lora_weights()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraLoaderMixin.load_lora_weights)."
332-
warnings.warn(warn_message)
333-
unet_keys = [k for k in state_dict.keys() if k.startswith(self.unet_name)]
334-
state_dict = {k.replace(f"{self.unet_name}.", ""): v for k, v in state_dict.items() if k in unet_keys}
314+
# correct keys
315+
state_dict, network_alphas = self.convert_state_dict_legacy_attn_format(state_dict, network_alphas)
335316

336317
lora_grouped_dict = defaultdict(dict)
337318
mapped_network_alphas = {}
@@ -367,87 +348,38 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict
367348

368349
# Process non-attention layers, which don't have to_{k,v,q,out_proj}_lora layers
369350
# or add_{k,v,q,out_proj}_proj_lora layers.
370-
if "lora.down.weight" in value_dict:
371-
rank = value_dict["lora.down.weight"].shape[0]
372-
373-
if isinstance(attn_processor, LoRACompatibleConv):
374-
in_features = attn_processor.in_channels
375-
out_features = attn_processor.out_channels
376-
kernel_size = attn_processor.kernel_size
377-
378-
lora = LoRAConv2dLayer(
379-
in_features=in_features,
380-
out_features=out_features,
381-
rank=rank,
382-
kernel_size=kernel_size,
383-
stride=attn_processor.stride,
384-
padding=attn_processor.padding,
385-
network_alpha=mapped_network_alphas.get(key),
386-
)
387-
elif isinstance(attn_processor, LoRACompatibleLinear):
388-
lora = LoRALinearLayer(
389-
attn_processor.in_features,
390-
attn_processor.out_features,
391-
rank,
392-
mapped_network_alphas.get(key),
393-
)
394-
else:
395-
raise ValueError(f"Module {key} is not a LoRACompatibleConv or LoRACompatibleLinear module.")
396-
397-
value_dict = {k.replace("lora.", ""): v for k, v in value_dict.items()}
398-
lora.load_state_dict(value_dict)
399-
non_attn_lora_layers.append((attn_processor, lora))
351+
rank = value_dict["lora.down.weight"].shape[0]
352+
353+
if isinstance(attn_processor, LoRACompatibleConv):
354+
in_features = attn_processor.in_channels
355+
out_features = attn_processor.out_channels
356+
kernel_size = attn_processor.kernel_size
357+
358+
lora = LoRAConv2dLayer(
359+
in_features=in_features,
360+
out_features=out_features,
361+
rank=rank,
362+
kernel_size=kernel_size,
363+
stride=attn_processor.stride,
364+
padding=attn_processor.padding,
365+
network_alpha=mapped_network_alphas.get(key),
366+
)
367+
elif isinstance(attn_processor, LoRACompatibleLinear):
368+
lora = LoRALinearLayer(
369+
attn_processor.in_features,
370+
attn_processor.out_features,
371+
rank,
372+
mapped_network_alphas.get(key),
373+
)
400374
else:
401-
# To handle SDXL.
402-
rank_mapping = {}
403-
hidden_size_mapping = {}
404-
for projection_id in ["to_k", "to_q", "to_v", "to_out"]:
405-
rank = value_dict[f"{projection_id}_lora.down.weight"].shape[0]
406-
hidden_size = value_dict[f"{projection_id}_lora.up.weight"].shape[0]
407-
408-
rank_mapping.update({f"{projection_id}_lora.down.weight": rank})
409-
hidden_size_mapping.update({f"{projection_id}_lora.up.weight": hidden_size})
410-
411-
if isinstance(
412-
attn_processor, (AttnAddedKVProcessor, SlicedAttnAddedKVProcessor, AttnAddedKVProcessor2_0)
413-
):
414-
cross_attention_dim = value_dict["add_k_proj_lora.down.weight"].shape[1]
415-
attn_processor_class = LoRAAttnAddedKVProcessor
416-
else:
417-
cross_attention_dim = value_dict["to_k_lora.down.weight"].shape[1]
418-
if isinstance(attn_processor, (XFormersAttnProcessor, LoRAXFormersAttnProcessor)):
419-
attn_processor_class = LoRAXFormersAttnProcessor
420-
else:
421-
attn_processor_class = (
422-
LoRAAttnProcessor2_0
423-
if hasattr(F, "scaled_dot_product_attention")
424-
else LoRAAttnProcessor
425-
)
426-
427-
if attn_processor_class is not LoRAAttnAddedKVProcessor:
428-
attn_processors[key] = attn_processor_class(
429-
rank=rank_mapping.get("to_k_lora.down.weight"),
430-
hidden_size=hidden_size_mapping.get("to_k_lora.up.weight"),
431-
cross_attention_dim=cross_attention_dim,
432-
network_alpha=mapped_network_alphas.get(key),
433-
q_rank=rank_mapping.get("to_q_lora.down.weight"),
434-
q_hidden_size=hidden_size_mapping.get("to_q_lora.up.weight"),
435-
v_rank=rank_mapping.get("to_v_lora.down.weight"),
436-
v_hidden_size=hidden_size_mapping.get("to_v_lora.up.weight"),
437-
out_rank=rank_mapping.get("to_out_lora.down.weight"),
438-
out_hidden_size=hidden_size_mapping.get("to_out_lora.up.weight"),
439-
)
440-
else:
441-
attn_processors[key] = attn_processor_class(
442-
rank=rank_mapping.get("to_k_lora.down.weight", None),
443-
hidden_size=hidden_size_mapping.get("to_k_lora.up.weight", None),
444-
cross_attention_dim=cross_attention_dim,
445-
network_alpha=mapped_network_alphas.get(key),
446-
)
375+
raise ValueError(f"Module {key} is not a LoRACompatibleConv or LoRACompatibleLinear module.")
447376

448-
attn_processors[key].load_state_dict(value_dict)
377+
value_dict = {k.replace("lora.", ""): v for k, v in value_dict.items()}
378+
lora.load_state_dict(value_dict)
379+
lora_layers_list.append((attn_processor, lora))
449380

450381
elif is_custom_diffusion:
382+
attn_processors = {}
451383
custom_diffusion_grouped_dict = defaultdict(dict)
452384
for key, value in state_dict.items():
453385
if len(value) == 0:
@@ -475,22 +407,47 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict
475407
cross_attention_dim=cross_attention_dim,
476408
)
477409
attn_processors[key].load_state_dict(value_dict)
410+
411+
self.set_attn_processor(attn_processors)
478412
else:
479413
raise ValueError(
480414
f"{model_file} does not seem to be in the correct format expected by LoRA or Custom Diffusion training."
481415
)
482416

483417
# set correct dtype & device
484-
attn_processors = {k: v.to(device=self.device, dtype=self.dtype) for k, v in attn_processors.items()}
485-
non_attn_lora_layers = [(t, l.to(device=self.device, dtype=self.dtype)) for t, l in non_attn_lora_layers]
486-
487-
# set layers
488-
self.set_attn_processor(attn_processors)
418+
lora_layers_list = [(t, l.to(device=self.device, dtype=self.dtype)) for t, l in lora_layers_list]
489419

490-
# set ff layers
491-
for target_module, lora_layer in non_attn_lora_layers:
420+
# set lora layers
421+
for target_module, lora_layer in lora_layers_list:
492422
target_module.set_lora_layer(lora_layer)
493423

424+
def convert_state_dict_legacy_attn_format(self, state_dict, network_alphas):
425+
is_new_lora_format = all(
426+
key.startswith(self.unet_name) or key.startswith(self.text_encoder_name) for key in state_dict.keys()
427+
)
428+
if is_new_lora_format:
429+
# Strip the `"unet"` prefix.
430+
is_text_encoder_present = any(key.startswith(self.text_encoder_name) for key in state_dict.keys())
431+
if is_text_encoder_present:
432+
warn_message = "The state_dict contains LoRA params corresponding to the text encoder which are not being used here. To use both UNet and text encoder related LoRA params, use [`pipe.load_lora_weights()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraLoaderMixin.load_lora_weights)."
433+
logger.warn(warn_message)
434+
unet_keys = [k for k in state_dict.keys() if k.startswith(self.unet_name)]
435+
state_dict = {k.replace(f"{self.unet_name}.", ""): v for k, v in state_dict.items() if k in unet_keys}
436+
437+
# change processor format to 'pure' LoRACompatibleLinear format
438+
if any("processor" in k.split(".") for k in state_dict.keys()):
439+
440+
def format_to_lora_compatible(key):
441+
if "processor" not in key.split("."):
442+
return key
443+
return key.replace(".processor", "").replace("to_out_lora", "to_out.0.lora").replace("_lora", ".lora")
444+
445+
state_dict = {format_to_lora_compatible(k): v for k, v in state_dict.items()}
446+
447+
if network_alphas is not None:
448+
network_alphas = {format_to_lora_compatible(k): v for k, v in network_alphas.items()}
449+
return state_dict, network_alphas
450+
494451
def save_attn_procs(
495452
self,
496453
save_directory: Union[str, os.PathLike],
@@ -1748,36 +1705,9 @@ def unload_lora_weights(self):
17481705
>>> ...
17491706
```
17501707
"""
1751-
from .models.attention_processor import (
1752-
LORA_ATTENTION_PROCESSORS,
1753-
AttnProcessor,
1754-
AttnProcessor2_0,
1755-
LoRAAttnAddedKVProcessor,
1756-
LoRAAttnProcessor,
1757-
LoRAAttnProcessor2_0,
1758-
LoRAXFormersAttnProcessor,
1759-
XFormersAttnProcessor,
1760-
)
1761-
1762-
unet_attention_classes = {type(processor) for _, processor in self.unet.attn_processors.items()}
1763-
1764-
if unet_attention_classes.issubset(LORA_ATTENTION_PROCESSORS):
1765-
# Handle attention processors that are a mix of regular attention and AddedKV
1766-
# attention.
1767-
if len(unet_attention_classes) > 1 or LoRAAttnAddedKVProcessor in unet_attention_classes:
1768-
self.unet.set_default_attn_processor()
1769-
else:
1770-
regular_attention_classes = {
1771-
LoRAAttnProcessor: AttnProcessor,
1772-
LoRAAttnProcessor2_0: AttnProcessor2_0,
1773-
LoRAXFormersAttnProcessor: XFormersAttnProcessor,
1774-
}
1775-
[attention_proc_class] = unet_attention_classes
1776-
self.unet.set_attn_processor(regular_attention_classes[attention_proc_class]())
1777-
1778-
for _, module in self.unet.named_modules():
1779-
if hasattr(module, "set_lora_layer"):
1780-
module.set_lora_layer(None)
1708+
for _, module in self.unet.named_modules():
1709+
if hasattr(module, "set_lora_layer"):
1710+
module.set_lora_layer(None)
17811711

17821712
# Safe to call the following regardless of LoRA.
17831713
self._remove_text_encoder_monkey_patch()

0 commit comments

Comments
 (0)