lycoris-lora 2.1.0.dev7__tar.gz → 2.1.0.dev9__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (47) hide show
  1. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/PKG-INFO +1 -1
  2. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/__init__.py +2 -0
  3. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/kohya/__init__.py +45 -42
  4. lycoris_lora-2.1.0.dev9/lycoris/logging.py +36 -0
  5. lycoris_lora-2.1.0.dev9/lycoris/modules/boft.py +211 -0
  6. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/utils/__init__.py +8 -6
  7. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/wrapper.py +17 -17
  8. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris_lora.egg-info/PKG-INFO +1 -1
  9. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris_lora.egg-info/SOURCES.txt +2 -0
  10. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/setup.py +1 -1
  11. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/LICENSE.md +0 -0
  12. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/README.md +0 -0
  13. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/config.py +0 -0
  14. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/hcp/__init__.py +0 -0
  15. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/hcp/base.py +0 -0
  16. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/hcp/diag_oft.py +0 -0
  17. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/hcp/loha.py +0 -0
  18. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/hcp/lokr.py +0 -0
  19. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/kohya/model_utils.py +0 -0
  20. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/kohya/original_unet.py +0 -0
  21. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/kohya/sdxl_model_util.py +0 -0
  22. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/kohya/sdxl_original_unet.py +0 -0
  23. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/kohya/utils.py +0 -0
  24. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/__init__.py +0 -0
  25. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/attention.py +0 -0
  26. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/base.py +0 -0
  27. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/diag_oft.py +0 -0
  28. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/dylora.py +0 -0
  29. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/full.py +0 -0
  30. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/glokr.py +0 -0
  31. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/glora.py +0 -0
  32. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/hypernet/__init__.py +0 -0
  33. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/hypernet/generater.py +0 -0
  34. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/hypernet/text_encoder.py +0 -0
  35. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/ia3.py +0 -0
  36. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/locon.py +0 -0
  37. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/loha.py +0 -0
  38. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/lokr.py +0 -0
  39. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/modules/norms.py +0 -0
  40. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/utils/logger.py +0 -0
  41. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/utils/preset.py +0 -0
  42. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris/utils/xformers_utils.py +0 -0
  43. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris_lora.egg-info/dependency_links.txt +0 -0
  44. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris_lora.egg-info/not-zip-safe +0 -0
  45. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris_lora.egg-info/requires.txt +0 -0
  46. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/lycoris_lora.egg-info/top_level.txt +0 -0
  47. {lycoris_lora-2.1.0.dev7 → lycoris_lora-2.1.0.dev9}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: lycoris_lora
3
- Version: 2.1.0.dev7
3
+ Version: 2.1.0.dev9
4
4
  Summary: Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion
5
5
  Home-page: https://github.com/KohakuBlueleaf/LyCORIS
6
6
  Author: Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao
@@ -19,3 +19,5 @@ from .wrapper import (
19
19
  create_lycoris,
20
20
  create_lycoris_from_weights,
21
21
  )
22
+
23
+ from .logging import logger
@@ -4,13 +4,13 @@
4
4
  # https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
5
5
  # https://github.com/kohya-ss/sd-scripts/blob/main/networks/lora.py
6
6
 
7
- import math
8
- from warnings import warn
9
7
  import os
8
+ import re
10
9
  import sys
10
+ from typing import List
11
11
 
12
12
  sys.setrecursionlimit(10000)
13
- from typing import List
13
+
14
14
  import torch
15
15
  import torch.utils.checkpoint as checkpoint
16
16
 
@@ -25,11 +25,13 @@ from ..modules.glora import GLoRAModule
25
25
  from ..modules.norms import NormModule
26
26
  from ..modules.full import FullModule
27
27
  from ..modules.diag_oft import DiagOFTModule
28
+ from ..modules.boft import ButterflyOFTModule
28
29
  from ..modules import make_module
29
30
 
30
31
  from ..config import PRESET
31
32
  from ..utils.preset import read_preset
32
33
  from ..utils import get_module, str_bool
34
+ from ..logging import logger
33
35
 
34
36
 
35
37
  network_module_dict = {
@@ -42,6 +44,7 @@ network_module_dict = {
42
44
  "glora": GLoRAModule,
43
45
  "full": FullModule,
44
46
  "diag-oft": DiagOFTModule,
47
+ "boft": ButterflyOFTModule,
45
48
  }
46
49
 
47
50
 
@@ -63,7 +66,7 @@ def create_network(
63
66
  or kwargs.get("use_tucker", False)
64
67
  )
65
68
  if "disable_conv_cp" in kwargs or "use_cp" in kwargs or "use_conv_cp" in kwargs:
66
- warn(
69
+ logger.warning(
67
70
  "disable_conv_cp and use_cp are deprecated. Please use use_tucker instead.",
68
71
  stacklevel=2,
69
72
  )
@@ -75,7 +78,7 @@ def create_network(
75
78
 
76
79
  if algo == "glora" and conv_dim > 0:
77
80
  conv_dim = 0
78
- print("Disable conv layer for GLoRA")
81
+ logger.info("Disable conv layer for GLoRA")
79
82
 
80
83
  preset = kwargs.get("preset", "full")
81
84
  if preset not in PRESET:
@@ -85,22 +88,20 @@ def create_network(
85
88
  assert preset is not None
86
89
  LycorisNetworkKohya.apply_preset(preset)
87
90
 
88
- print(f"Using rank adaptation algo: {algo}")
91
+ logger.info(f"Using rank adaptation algo: {algo}")
89
92
 
90
93
  if (
91
94
  (algo == "loha")
92
95
  and not kwargs.get("no_dim_warn", False)
93
96
  and (network_dim > 64 or conv_dim > 64)
94
97
  ):
95
- print("=" * 20 + "WARNING" + "=" * 20)
96
98
  warning_type = {"loha": "Hadamard Product representation"}
97
99
  warning_msg = f"""You are not supposed to use dim>64 (64*64 = 4096, it already has enough rank)\n
98
100
  in {warning_type[algo]}!\n
99
101
  Please consider use lower dim or disable this warning with --network_args no_dim_warn=True\n
100
102
  If you just want to use high dim {algo}, please consider use lower lr.
101
103
  """
102
- warn(warning_msg, stacklevel=2)
103
- print("=" * 20 + "WARNING" + "=" * 20)
104
+ logger.warning(warning_msg, stacklevel=2)
104
105
 
105
106
  if algo == "ia3":
106
107
  network = IA3Network(
@@ -240,7 +241,7 @@ def create_hypernetwork(
240
241
  or kwargs.get("use_tucker", False)
241
242
  )
242
243
  if "disable_conv_cp" in kwargs or "use_cp" in kwargs or "use_conv_cp" in kwargs:
243
- warn(
244
+ logger.warning(
244
245
  "disable_conv_cp and use_cp are deprecated. Please use use_tucker instead.",
245
246
  stacklevel=2,
246
247
  )
@@ -254,7 +255,7 @@ def create_hypernetwork(
254
255
  "locon": LoConModule,
255
256
  }[algo]
256
257
 
257
- print(f"Using rank adaptation algo: {algo}")
258
+ logger.info(f"Using rank adaptation algo: {algo}")
258
259
 
259
260
  return HyperDreamNetwork(
260
261
  text_encoder,
@@ -352,19 +353,19 @@ class LycorisNetworkKohya(LycorisNetwork):
352
353
 
353
354
  self.conv_lora_dim = int(conv_lora_dim)
354
355
  if self.conv_lora_dim and self.conv_lora_dim != self.lora_dim:
355
- print("Apply different lora dim for conv layer")
356
- print(f"Conv Dim: {conv_lora_dim}, Linear Dim: {lora_dim}")
356
+ logger.info("Apply different lora dim for conv layer")
357
+ logger.info(f"Conv Dim: {conv_lora_dim}, Linear Dim: {lora_dim}")
357
358
  elif self.conv_lora_dim == 0:
358
- print("Disable conv layer")
359
+ logger.info("Disable conv layer")
359
360
 
360
361
  self.alpha = alpha
361
362
  self.conv_alpha = float(conv_alpha)
362
363
  if self.conv_lora_dim and self.alpha != self.conv_alpha:
363
- print("Apply different alpha value for conv layer")
364
- print(f"Conv alpha: {conv_alpha}, Linear alpha: {alpha}")
364
+ logger.info("Apply different alpha value for conv layer")
365
+ logger.info(f"Conv alpha: {conv_alpha}, Linear alpha: {alpha}")
365
366
 
366
367
  if 1 >= dropout >= 0:
367
- print(f"Use Dropout value: {dropout}")
368
+ logger.info(f"Use Dropout value: {dropout}")
368
369
  self.dropout = dropout
369
370
  self.rank_dropout = rank_dropout
370
371
  self.module_dropout = module_dropout
@@ -465,7 +466,7 @@ class LycorisNetworkKohya(LycorisNetwork):
465
466
  target_replace_modules,
466
467
  target_replace_names=[],
467
468
  ) -> List:
468
- print("Create LyCORIS Module")
469
+ logger.info("Create LyCORIS Module")
469
470
  loras = []
470
471
  next_config = {}
471
472
  for name, module in root_module.named_modules():
@@ -482,7 +483,9 @@ class LycorisNetworkKohya(LycorisNetwork):
482
483
  ]
483
484
  )
484
485
  next_config = {}
485
- elif name in target_replace_names:
486
+ elif name in target_replace_names or any(
487
+ re.match(t, name) for t in target_replace_names
488
+ ):
486
489
  if name in self.NAME_ALGO_MAP:
487
490
  next_config = self.NAME_ALGO_MAP[name]
488
491
  algo = next_config.get("algo", network_module)
@@ -500,7 +503,7 @@ class LycorisNetworkKohya(LycorisNetwork):
500
503
  return loras
501
504
 
502
505
  if network_module == GLoRAModule:
503
- print("GLoRA enabled, only train transformer")
506
+ logger.info("GLoRA enabled, only train transformer")
504
507
  # only train transformer (for GLoRA)
505
508
  LycorisNetworkKohya.UNET_TARGET_REPLACE_MODULE = [
506
509
  "Transformer2DModel",
@@ -526,7 +529,7 @@ class LycorisNetworkKohya(LycorisNetwork):
526
529
  LycorisNetworkKohya.TEXT_ENCODER_TARGET_REPLACE_NAME,
527
530
  )
528
531
  )
529
- print(
532
+ logger.info(
530
533
  f"create LyCORIS for Text Encoder: {len(self.text_encoder_loras)} modules."
531
534
  )
532
535
 
@@ -536,14 +539,14 @@ class LycorisNetworkKohya(LycorisNetwork):
536
539
  LycorisNetworkKohya.UNET_TARGET_REPLACE_MODULE,
537
540
  LycorisNetworkKohya.UNET_TARGET_REPLACE_NAME,
538
541
  )
539
- print(f"create LyCORIS for U-Net: {len(self.unet_loras)} modules.")
542
+ logger.info(f"create LyCORIS for U-Net: {len(self.unet_loras)} modules.")
540
543
 
541
544
  algo_table = {}
542
545
  for lora in self.text_encoder_loras + self.unet_loras:
543
546
  algo_table[lora.__class__.__name__] = (
544
547
  algo_table.get(lora.__class__.__name__, 0) + 1
545
548
  )
546
- print(f"module type table: {algo_table}")
549
+ logger.info(f"module type table: {algo_table}")
547
550
 
548
551
  self.weights_sd = None
549
552
 
@@ -577,12 +580,12 @@ class LycorisNetworkKohya(LycorisNetwork):
577
580
  ), f"internal error: flag not set"
578
581
 
579
582
  if apply_text_encoder:
580
- print("enable LyCORIS for text encoder")
583
+ logger.info("enable LyCORIS for text encoder")
581
584
  else:
582
585
  self.text_encoder_loras = []
583
586
 
584
587
  if apply_unet:
585
- print("enable LyCORIS for U-Net")
588
+ logger.info("enable LyCORIS for U-Net")
586
589
  else:
587
590
  self.unet_loras = []
588
591
 
@@ -595,7 +598,7 @@ class LycorisNetworkKohya(LycorisNetwork):
595
598
  if self.weights_sd:
596
599
  # if some weights are not in state dict, it is ok because initial LoRA does nothing (lora_up is initialized by zeros)
597
600
  info = self.load_state_dict(self.weights_sd, False)
598
- print(f"weights are loaded: {info}")
601
+ logger.info(f"weights are loaded: {info}")
599
602
 
600
603
  # TODO refactor to common function with apply_to
601
604
  def merge_to(self, text_encoder, unet, weights_sd, dtype, device):
@@ -607,12 +610,12 @@ class LycorisNetworkKohya(LycorisNetwork):
607
610
  apply_unet = True
608
611
 
609
612
  if apply_text_encoder:
610
- print("enable LoRA for text encoder")
613
+ logger.info("enable LoRA for text encoder")
611
614
  else:
612
615
  self.text_encoder_loras = []
613
616
 
614
617
  if apply_unet:
615
- print("enable LoRA for U-Net")
618
+ logger.info("enable LoRA for U-Net")
616
619
  else:
617
620
  self.unet_loras = []
618
621
 
@@ -714,12 +717,12 @@ class HyperDreamNetwork(torch.nn.Module):
714
717
  self.alpha = alpha
715
718
 
716
719
  if 1 >= dropout >= 0:
717
- print(f"Use Dropout value: {dropout}")
720
+ logger.info(f"Use Dropout value: {dropout}")
718
721
  if network_module != LoConModule:
719
- print("HyperDreamBooth only support LoRA at this time")
722
+ logger.info("HyperDreamBooth only support LoRA at this time")
720
723
  raise NotImplementedError
721
724
  if lora_dim * (down_dim + up_dim) > 4096:
722
- print(
725
+ logger.info(
723
726
  "weight elements > 4096 (dim * (down_dim + up_dim)) is not recommended!"
724
727
  )
725
728
 
@@ -734,7 +737,7 @@ class HyperDreamNetwork(torch.nn.Module):
734
737
  target_replace_modules,
735
738
  target_replace_names=[],
736
739
  ) -> List[network_module]:
737
- print("Create LyCORIS Module")
740
+ logger.info("Create LyCORIS Module")
738
741
  loras = []
739
742
  for name, module in root_module.named_modules():
740
743
  if module.__class__.__name__ in target_replace_modules:
@@ -829,7 +832,7 @@ class HyperDreamNetwork(torch.nn.Module):
829
832
  LycorisNetworkKohya.TEXT_ENCODER_TARGET_REPLACE_MODULE,
830
833
  )
831
834
  )
832
- print(
835
+ logger.info(
833
836
  f"create LyCORIS for Text Encoder: "
834
837
  f"{len(self.text_encoder_loras)} modules."
835
838
  )
@@ -839,7 +842,7 @@ class HyperDreamNetwork(torch.nn.Module):
839
842
  unet,
840
843
  LycorisNetworkKohya.UNET_TARGET_REPLACE_MODULE,
841
844
  )
842
- print(f"create LyCORIS for U-Net: {len(self.unet_loras)} modules.")
845
+ logger.info(f"create LyCORIS for U-Net: {len(self.unet_loras)} modules.")
843
846
 
844
847
  self.loras: list[LoConModule] = self.text_encoder_loras + self.unet_loras
845
848
  self.img_weight_generater = ImgWeightGenerator(
@@ -965,12 +968,12 @@ class HyperDreamNetwork(torch.nn.Module):
965
968
  ), f"internal error: flag not set"
966
969
 
967
970
  if apply_text_encoder:
968
- print("enable LyCORIS for text encoder")
971
+ logger.info("enable LyCORIS for text encoder")
969
972
  else:
970
973
  self.text_encoder_loras = []
971
974
 
972
975
  if apply_unet:
973
- print("enable LyCORIS for U-Net")
976
+ logger.info("enable LyCORIS for U-Net")
974
977
  else:
975
978
  self.unet_loras = []
976
979
 
@@ -1122,7 +1125,7 @@ class IA3Network(torch.nn.Module):
1122
1125
  target_replace_names=[],
1123
1126
  target_train_input=[],
1124
1127
  ) -> List[IA3Module]:
1125
- print("Create LyCORIS Module")
1128
+ logger.info("Create LyCORIS Module")
1126
1129
  loras = []
1127
1130
  for name, module in root_module.named_modules():
1128
1131
  if module.__class__.__name__ in target_replace_modules:
@@ -1171,7 +1174,7 @@ class IA3Network(torch.nn.Module):
1171
1174
  IA3Network.TRAIN_INPUT,
1172
1175
  )
1173
1176
  )
1174
- print(
1177
+ logger.info(
1175
1178
  f"create LyCORIS for Text Encoder: {len(self.text_encoder_loras)} modules."
1176
1179
  )
1177
1180
 
@@ -1182,7 +1185,7 @@ class IA3Network(torch.nn.Module):
1182
1185
  IA3Network.UNET_TARGET_REPLACE_NAME,
1183
1186
  IA3Network.TRAIN_INPUT,
1184
1187
  )
1185
- print(f"create LyCORIS for U-Net: {len(self.unet_loras)} modules.")
1188
+ logger.info(f"create LyCORIS for U-Net: {len(self.unet_loras)} modules.")
1186
1189
 
1187
1190
  self.weights_sd = None
1188
1191
 
@@ -1235,12 +1238,12 @@ class IA3Network(torch.nn.Module):
1235
1238
  ), f"internal error: flag not set"
1236
1239
 
1237
1240
  if apply_text_encoder:
1238
- print("enable LyCORIS for text encoder")
1241
+ logger.info("enable LyCORIS for text encoder")
1239
1242
  else:
1240
1243
  self.text_encoder_loras = []
1241
1244
 
1242
1245
  if apply_unet:
1243
- print("enable LyCORIS for U-Net")
1246
+ logger.info("enable LyCORIS for U-Net")
1244
1247
  else:
1245
1248
  self.unet_loras = []
1246
1249
 
@@ -1251,7 +1254,7 @@ class IA3Network(torch.nn.Module):
1251
1254
  if self.weights_sd:
1252
1255
  # if some weights are not in state dict, it is ok because initial LoRA does nothing (lora_up is initialized by zeros)
1253
1256
  info = self.load_state_dict(self.weights_sd, False)
1254
- print(f"weights are loaded: {info}")
1257
+ logger.info(f"weights are loaded: {info}")
1255
1258
 
1256
1259
  def enable_gradient_checkpointing(self):
1257
1260
  # not supported
@@ -0,0 +1,36 @@
1
+ import sys
2
+ import copy
3
+ import logging
4
+
5
+
6
+ class ColoredFormatter(logging.Formatter):
7
+ COLORS = {
8
+ "DEBUG": "\033[0;36m", # CYAN
9
+ "INFO": "\033[0;32m", # GREEN
10
+ "WARNING": "\033[0;33m", # YELLOW
11
+ "ERROR": "\033[0;31m", # RED
12
+ "CRITICAL": "\033[0;37;41m", # WHITE ON RED
13
+ "RESET": "\033[0m", # RESET COLOR
14
+ }
15
+
16
+ def format(self, record):
17
+ colored_record = copy.copy(record)
18
+ levelname = colored_record.levelname
19
+ seq = self.COLORS.get(levelname, self.COLORS["RESET"])
20
+ colored_record.levelname = f"{seq}{levelname}{self.COLORS['RESET']}"
21
+ return super().format(colored_record)
22
+
23
+
24
+ logger = logging.getLogger("LyCORIS")
25
+ logger.propagate = False
26
+ logger.setLevel(logging.INFO)
27
+
28
+
29
+ if not logger.handlers:
30
+ handler = logging.StreamHandler(sys.stdout)
31
+ handler.setFormatter(
32
+ ColoredFormatter(
33
+ "%(asctime)s|[%(name)s]-%(levelname)s: %(message)s", "%Y-%m-%d %H:%M:%S"
34
+ )
35
+ )
36
+ logger.addHandler(handler)
@@ -0,0 +1,211 @@
1
+ from functools import cache
2
+ from math import log2, ceil
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+ import torch.nn.functional as F
7
+
8
+ from einops import rearrange
9
+
10
+ from .base import ModuleCustomSD
11
+ from lycoris.logging import logger
12
+
13
+
14
+ @cache
15
+ def log_butterfly_factorize(dim, factor, result):
16
+ logger.info(
17
+ f"Use BOFT({int(log2(result[1]))}, {result[0]//2})"
18
+ f" (equivalent to factor={result[0]}) "
19
+ f"for {dim=} and {factor=}"
20
+ )
21
+
22
+
23
+ def butterfly_factor(dimension: int, factor: int = -1) -> tuple[int, int]:
24
+ """
25
+ m = 2k
26
+ n = 2**p
27
+ m*n = dim
28
+ """
29
+
30
+ # Find the first solution and check if it is even doable
31
+ m = n = 0
32
+ while m <= factor:
33
+ m += 2
34
+ while dimension % m != 0 and m < dimension:
35
+ m += 2
36
+ if m > factor:
37
+ break
38
+ if sum(int(i) for i in f"{dimension//m:b}") == 1:
39
+ n = dimension // m
40
+
41
+ if n == 0:
42
+ raise ValueError(
43
+ f"It is impossible to decompose {dimension} with factor {factor} under BOFT constrains."
44
+ )
45
+
46
+ log_butterfly_factorize(dimension, factor, (dimension // n, n))
47
+ return dimension // n, n
48
+
49
+
50
+ class ButterflyOFTModule(ModuleCustomSD):
51
+ """
52
+ modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
53
+ """
54
+
55
+ def __init__(
56
+ self,
57
+ lora_name,
58
+ org_module: nn.Module,
59
+ multiplier=1.0,
60
+ lora_dim=4,
61
+ alpha=1,
62
+ dropout=0.0,
63
+ rank_dropout=0.0,
64
+ module_dropout=0.0,
65
+ use_tucker=False,
66
+ use_scalar=False,
67
+ rank_dropout_scale=False,
68
+ constrain=0,
69
+ rescaled=False,
70
+ **kwargs,
71
+ ):
72
+ """if alpha == 0 or None, alpha is rank (no scaling)."""
73
+ super().__init__()
74
+ self.lora_name = lora_name
75
+
76
+ if isinstance(org_module, nn.Linear):
77
+ self.op = F.linear
78
+ self.dim = org_module.out_features
79
+ self.kw_dict = {}
80
+ elif isinstance(org_module, nn.Conv2d):
81
+ self.op = F.conv2d
82
+ self.dim = org_module.out_channels
83
+ self.kw_dict = {
84
+ "stride": org_module.stride,
85
+ "padding": org_module.padding,
86
+ "dilation": org_module.dilation,
87
+ "groups": org_module.groups,
88
+ }
89
+ else:
90
+ raise NotImplementedError
91
+
92
+ out_dim = org_module.weight.shape[0]
93
+ b, m_exp = butterfly_factor(out_dim, lora_dim)
94
+ self.block_size = b
95
+ self.block_num = m_exp
96
+ #BOFT(m, b)
97
+ self.boft_b = b
98
+ self.boft_m = sum(int(i) for i in f"{m_exp-1:b}") + 1
99
+ # block_num > block_size
100
+ self.rescaled = rescaled
101
+ self.constrain = constrain * out_dim
102
+ self.oft_blocks = nn.Parameter(
103
+ torch.zeros(self.boft_m, self.block_num, self.block_size, self.block_size)
104
+ )
105
+ if rescaled:
106
+ self.rescale = nn.Parameter(torch.ones(self.block_num, self.block_size, 1))
107
+
108
+ self.rank_dropout = rank_dropout
109
+ self.rank_dropout_scale = rank_dropout_scale
110
+ self.module_dropout = module_dropout
111
+
112
+ self.multiplier = multiplier
113
+ self.org_module = [org_module]
114
+ self.org_forward = self.org_module[0].forward
115
+
116
+ @property
117
+ def I(self):
118
+ return torch.eye(self.block_size, device=next(self.parameters()).device)
119
+
120
+ def apply_to(self, **kwargs):
121
+ self.org_forward = self.org_module[0].forward
122
+ self.org_module[0].forward = self.forward
123
+
124
+ def restore(self):
125
+ self.org_module[0].forward = self.org_forward
126
+
127
+ def merge_to(self, multiplier=1.0):
128
+ weight = self.make_weight(scale=multiplier)
129
+ self.org_module[0].weight.data.copy_(weight)
130
+
131
+ def custom_state_dict(self):
132
+ return {"oft_diag": self.get_r()}
133
+
134
+ def get_r(self):
135
+ I = self.I
136
+ # for Q = -Q^T
137
+ q = self.oft_blocks - self.oft_blocks.transpose(-1, -2)
138
+ normed_q = q
139
+ # Diag OFT style constrain
140
+ if self.constrain > 0:
141
+ q_norm = torch.norm(q) + 1e-8
142
+ if q_norm > self.constrain:
143
+ normed_q = q * self.constrain / q_norm
144
+ # use float() to prevent unsupported type
145
+ r = (I + normed_q) @ (I - normed_q).float().inverse()
146
+ return r
147
+
148
+ def make_weight(self, scale=1, device=None):
149
+ if self.rank_dropout and self.training:
150
+ drop = torch.rand(self.dim, device=device) < self.rank_dropout
151
+ drop = drop.to(self.oft_blocks.dtype)
152
+ if self.rank_dropout_scale:
153
+ drop /= drop.mean()
154
+ else:
155
+ drop = 1
156
+
157
+ m = self.boft_m
158
+ b = self.boft_b
159
+ r_b = b // 2
160
+ r = self.get_r()
161
+ inp = self.org_module[0].weight.to(device, dtype=r.dtype)
162
+
163
+ for i in range(m):
164
+ bi = r[i] # b_num, b_size, b_size
165
+ if i == 0:
166
+ # Apply multiplier/scale and rescale into first weight
167
+ bi = bi * scale + (1 - scale) * self.I
168
+ if self.rescaled:
169
+ bi = bi * self.rescale
170
+ inp = rearrange(inp, "(c g k) ... -> (c k g) ...", g=2, k=2**i * r_b)
171
+ inp = rearrange(inp, "(d b) ... -> d b ...", b=b)
172
+ inp = torch.einsum("b i j, b j ... -> b i ...", bi, inp)
173
+ inp = rearrange(inp, "d b ... -> (d b) ...")
174
+ inp = rearrange(inp, "(c k g) ... -> (c g k) ...", g=2, k=2**i * r_b)
175
+
176
+ weight = inp
177
+ return weight*drop
178
+
179
+ @torch.no_grad()
180
+ def apply_max_norm(self, max_norm, device=None):
181
+ orig_norm = self.oft_blocks.to(device).norm()
182
+ norm = torch.clamp(orig_norm, max_norm / 2)
183
+ desired = torch.clamp(norm, max=max_norm)
184
+ ratio = desired / norm
185
+
186
+ scaled = ratio.item() != 1.0
187
+ if scaled:
188
+ self.oft_blocks *= ratio
189
+
190
+ return scaled, orig_norm * ratio
191
+
192
+ def forward(self, x):
193
+ if self.module_dropout and self.training:
194
+ if torch.rand(1) < self.module_dropout:
195
+ return self.org_forward(x)
196
+ scale = self.multiplier
197
+
198
+ w = self.make_weight(scale, x.device)
199
+ kw_dict = self.kw_dict | {"weight": w, "bias": self.org_module[0].bias}
200
+ return self.op(x, **kw_dict)
201
+
202
+
203
+ if __name__ == '__main__':
204
+ dim = 320
205
+ test_layer = nn.Linear(dim, dim)
206
+ test_layer.weight = nn.Parameter(torch.eye(dim))
207
+ test_boft = ButterflyOFTModule('test', test_layer, lora_dim=16)
208
+ test_boft.oft_blocks = nn.Parameter(torch.randn_like(test_boft.oft_blocks))
209
+
210
+ print(test_boft.make_weight(1, 'cpu').shape)
211
+ print(torch.count_nonzero(test_boft.make_weight(1, 'cpu')), dim*dim)
@@ -230,12 +230,14 @@ def extract_diff(
230
230
  extract_c.detach().cpu().contiguous().half()
231
231
  )
232
232
  diff = (
233
- root_weight
234
- - torch.einsum(
235
- "i j k l, j r, p i -> p r k l",
236
- extract_c,
237
- extract_a.flatten(1, -1),
238
- extract_b.flatten(1, -1),
233
+ (
234
+ root_weight
235
+ - torch.einsum(
236
+ "i j k l, j r, p i -> p r k l",
237
+ extract_c,
238
+ extract_a.flatten(1, -1),
239
+ extract_b.flatten(1, -1),
240
+ )
239
241
  )
240
242
  .detach()
241
243
  .cpu()
@@ -7,7 +7,6 @@ import sys
7
7
 
8
8
  sys.setrecursionlimit(10000)
9
9
  from typing import List
10
- from warnings import warn
11
10
 
12
11
  import torch
13
12
 
@@ -20,11 +19,13 @@ from .modules.glora import GLoRAModule
20
19
  from .modules.norms import NormModule
21
20
  from .modules.full import FullModule
22
21
  from .modules.diag_oft import DiagOFTModule
22
+ from .modules.boft import ButterflyOFTModule
23
23
  from .modules import make_module
24
24
 
25
25
  from .config import PRESET
26
26
  from .utils.preset import read_preset
27
27
  from .utils import get_module, str_bool
28
+ from .logging import logger
28
29
 
29
30
 
30
31
  network_module_dict = {
@@ -36,6 +37,7 @@ network_module_dict = {
36
37
  "glora": GLoRAModule,
37
38
  "full": FullModule,
38
39
  "diag-oft": DiagOFTModule,
40
+ "boft": ButterflyOFTModule,
39
41
  }
40
42
 
41
43
 
@@ -55,7 +57,7 @@ def create_lycoris(module, multiplier, linear_dim, linear_alpha, **kwargs):
55
57
  or kwargs.get("use_tucker", False)
56
58
  )
57
59
  if "disable_conv_cp" in kwargs or "use_cp" in kwargs or "use_conv_cp" in kwargs:
58
- warn(
60
+ logger.warning(
59
61
  "disable_conv_cp and use_cp are deprecated. Please use use_tucker instead.",
60
62
  stacklevel=2,
61
63
  )
@@ -67,7 +69,7 @@ def create_lycoris(module, multiplier, linear_dim, linear_alpha, **kwargs):
67
69
 
68
70
  if algo == "glora" and conv_dim > 0:
69
71
  conv_dim = 0
70
- print("Disable conv layer for GLoRA")
72
+ logger.info("Disable conv layer for GLoRA")
71
73
 
72
74
  preset = kwargs.get("preset", "full")
73
75
  if preset not in PRESET:
@@ -77,14 +79,13 @@ def create_lycoris(module, multiplier, linear_dim, linear_alpha, **kwargs):
77
79
  assert preset is not None
78
80
  LycorisNetwork.apply_preset(preset)
79
81
 
80
- print(f"Using rank adaptation algo: {algo}")
82
+ logger.info(f"Using rank adaptation algo: {algo}")
81
83
 
82
84
  if (
83
85
  (algo == "loha")
84
86
  and not kwargs.get("no_dim_warn", False)
85
87
  and (linear_dim > 64 or conv_dim > 64)
86
88
  ):
87
- print("=" * 20 + "WARNING" + "=" * 20)
88
89
  warning_type = {"loha": "Hadamard Product representation"}
89
90
  warning_msg = (
90
91
  "You are not supposed to use dim>64 (64*64 = 4096, it already has enough rank)\n"
@@ -92,8 +93,7 @@ def create_lycoris(module, multiplier, linear_dim, linear_alpha, **kwargs):
92
93
  "Please consider use lower dim or disable this warning with --network_args no_dim_warn=True\n"
93
94
  f"If you just want to use high dim {algo}, please consider use lower lr."
94
95
  )
95
- warn(warning_msg, stacklevel=2)
96
- print("=" * 20 + "WARNING" + "=" * 20)
96
+ logger.warning(warning_msg, stacklevel=2)
97
97
 
98
98
  network = LycorisNetwork(
99
99
  module,
@@ -230,19 +230,19 @@ class LycorisNetwork(torch.nn.Module):
230
230
 
231
231
  self.conv_lora_dim = int(conv_lora_dim)
232
232
  if self.conv_lora_dim and self.conv_lora_dim != self.lora_dim:
233
- print("Apply different lora dim for conv layer")
234
- print(f"Conv Dim: {conv_lora_dim}, Linear Dim: {lora_dim}")
233
+ logger.info("Apply different lora dim for conv layer")
234
+ logger.info(f"Conv Dim: {conv_lora_dim}, Linear Dim: {lora_dim}")
235
235
  elif self.conv_lora_dim == 0:
236
- print("Disable conv layer")
236
+ logger.info("Disable conv layer")
237
237
 
238
238
  self.alpha = alpha
239
239
  self.conv_alpha = float(conv_alpha)
240
240
  if self.conv_lora_dim and self.alpha != self.conv_alpha:
241
- print("Apply different alpha value for conv layer")
242
- print(f"Conv alpha: {conv_alpha}, Linear alpha: {alpha}")
241
+ logger.info("Apply different alpha value for conv layer")
242
+ logger.info(f"Conv alpha: {conv_alpha}, Linear alpha: {alpha}")
243
243
 
244
244
  if 1 >= dropout >= 0:
245
- print(f"Use Dropout value: {dropout}")
245
+ logger.info(f"Use Dropout value: {dropout}")
246
246
  self.dropout = dropout
247
247
  self.rank_dropout = rank_dropout
248
248
  self.module_dropout = module_dropout
@@ -341,7 +341,7 @@ class LycorisNetwork(torch.nn.Module):
341
341
  target_replace_modules,
342
342
  target_replace_names=[],
343
343
  ) -> List:
344
- print("Create LyCORIS Module")
344
+ logger.info("Create LyCORIS Module")
345
345
  loras = []
346
346
  next_config = {}
347
347
  for name, module in root_module.named_modules():
@@ -383,14 +383,14 @@ class LycorisNetwork(torch.nn.Module):
383
383
  LycorisNetwork.TARGET_REPLACE_MODULE,
384
384
  LycorisNetwork.TARGET_REPLACE_NAME,
385
385
  )
386
- print(f"create LyCORIS: {len(self.loras)} modules.")
386
+ logger.info(f"create LyCORIS: {len(self.loras)} modules.")
387
387
 
388
388
  algo_table = {}
389
389
  for lora in self.loras:
390
390
  algo_table[lora.__class__.__name__] = (
391
391
  algo_table.get(lora.__class__.__name__, 0) + 1
392
392
  )
393
- print(f"module type table: {algo_table}")
393
+ logger.info(f"module type table: {algo_table}")
394
394
 
395
395
  self.weights_sd = None
396
396
 
@@ -430,7 +430,7 @@ class LycorisNetwork(torch.nn.Module):
430
430
  if self.weights_sd:
431
431
  # if some weights are not in state dict, it is ok because initial LoRA does nothing (lora_up is initialized by zeros)
432
432
  info = self.load_state_dict(self.weights_sd, False)
433
- print(f"weights are loaded: {info}")
433
+ logger.info(f"weights are loaded: {info}")
434
434
 
435
435
  def is_mergeable(self):
436
436
  return True
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: lycoris-lora
3
- Version: 2.1.0.dev7
3
+ Version: 2.1.0.dev9
4
4
  Summary: Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion
5
5
  Home-page: https://github.com/KohakuBlueleaf/LyCORIS
6
6
  Author: Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao
@@ -3,6 +3,7 @@ README.md
3
3
  setup.py
4
4
  lycoris/__init__.py
5
5
  lycoris/config.py
6
+ lycoris/logging.py
6
7
  lycoris/wrapper.py
7
8
  lycoris/hcp/__init__.py
8
9
  lycoris/hcp/base.py
@@ -18,6 +19,7 @@ lycoris/kohya/utils.py
18
19
  lycoris/modules/__init__.py
19
20
  lycoris/modules/attention.py
20
21
  lycoris/modules/base.py
22
+ lycoris/modules/boft.py
21
23
  lycoris/modules/diag_oft.py
22
24
  lycoris/modules/dylora.py
23
25
  lycoris/modules/full.py
@@ -4,7 +4,7 @@ from setuptools import setup, find_packages
4
4
  setup(
5
5
  name="lycoris_lora",
6
6
  packages=find_packages(),
7
- version="2.1.0.dev7",
7
+ version="2.1.0.dev9",
8
8
  url="https://github.com/KohakuBlueleaf/LyCORIS",
9
9
  description="Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion",
10
10
  author="Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao",