lycoris-lora 3.0.1.dev4__tar.gz → 3.0.1.dev5__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-3.0.1.dev4/lycoris_lora.egg-info → lycoris_lora-3.0.1.dev5}/PKG-INFO +1 -1
  2. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/kohya.py +20 -4
  3. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/wrapper.py +29 -7
  4. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5/lycoris_lora.egg-info}/PKG-INFO +1 -1
  5. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/setup.py +1 -1
  6. lycoris_lora-3.0.1.dev5/test/wrapper.py +290 -0
  7. lycoris_lora-3.0.1.dev4/test/wrapper.py +0 -128
  8. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/LICENSE.md +0 -0
  9. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/README.md +0 -0
  10. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/__init__.py +0 -0
  11. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/config.py +0 -0
  12. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/functional/__init__.py +0 -0
  13. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/functional/boft.py +0 -0
  14. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/functional/diag_oft.py +0 -0
  15. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/functional/general.py +0 -0
  16. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/functional/locon.py +0 -0
  17. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/functional/loha.py +0 -0
  18. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/functional/lokr.py +0 -0
  19. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/logging.py +0 -0
  20. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/__init__.py +0 -0
  21. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/base.py +0 -0
  22. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/boft.py +0 -0
  23. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/diag_oft.py +0 -0
  24. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/dylora.py +0 -0
  25. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/full.py +0 -0
  26. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/glora.py +0 -0
  27. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/ia3.py +0 -0
  28. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/locon.py +0 -0
  29. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/loha.py +0 -0
  30. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/lokr.py +0 -0
  31. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/modules/norms.py +0 -0
  32. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/utils/__init__.py +0 -0
  33. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/utils/bnb.py +0 -0
  34. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/utils/general.py +0 -0
  35. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/utils/logger.py +0 -0
  36. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/utils/preset.py +0 -0
  37. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris/utils/xformers_utils.py +0 -0
  38. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris_lora.egg-info/SOURCES.txt +0 -0
  39. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris_lora.egg-info/dependency_links.txt +0 -0
  40. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris_lora.egg-info/not-zip-safe +0 -0
  41. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris_lora.egg-info/requires.txt +0 -0
  42. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/lycoris_lora.egg-info/top_level.txt +0 -0
  43. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/setup.cfg +0 -0
  44. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/test/__init__.py +0 -0
  45. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/test/functional.py +0 -0
  46. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/test/kohya.py +0 -0
  47. {lycoris_lora-3.0.1.dev4 → lycoris_lora-3.0.1.dev5}/test/module.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: lycoris_lora
3
- Version: 3.0.1.dev4
3
+ Version: 3.0.1.dev5
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
@@ -1,7 +1,7 @@
1
1
  import os
2
2
  import re
3
3
  import logging
4
- from typing import List
4
+ from typing import Any, List
5
5
 
6
6
  import torch
7
7
 
@@ -404,7 +404,9 @@ class LycorisNetworkKohya(LycorisNetwork):
404
404
  next_config = {}
405
405
  for name, module in root_module.named_modules():
406
406
  module_name = module.__class__.__name__
407
- if module_name in target_replace_modules:
407
+ if module_name in target_replace_modules and not any(
408
+ re.match(t, name) for t in target_replace_names
409
+ ):
408
410
  if module_name in self.MODULE_ALGO_MAP:
409
411
  next_config = self.MODULE_ALGO_MAP[module_name]
410
412
  algo = next_config.get("algo", network_module)
@@ -419,8 +421,9 @@ class LycorisNetworkKohya(LycorisNetwork):
419
421
  elif name in target_replace_names or any(
420
422
  re.match(t, name) for t in target_replace_names
421
423
  ):
422
- if name in self.NAME_ALGO_MAP:
423
- next_config = self.NAME_ALGO_MAP[name]
424
+ conf_from_name = self.find_conf_for_name(name)
425
+ if conf_from_name is not None:
426
+ next_config = conf_from_name
424
427
  algo = next_config.get("algo", network_module)
425
428
  elif module_name in self.MODULE_ALGO_MAP:
426
429
  next_config = self.MODULE_ALGO_MAP[module_name]
@@ -492,6 +495,19 @@ class LycorisNetworkKohya(LycorisNetwork):
492
495
  ), f"duplicated lora name: {lora.lora_name}"
493
496
  names.add(lora.lora_name)
494
497
 
498
+ def find_conf_for_name(
499
+ self,
500
+ name: str,
501
+ ) -> dict[str, Any]:
502
+ if name in self.NAME_ALGO_MAP.keys():
503
+ return self.NAME_ALGO_MAP[name]
504
+
505
+ for key, value in self.NAME_ALGO_MAP.items():
506
+ if re.match(key, name):
507
+ return value
508
+
509
+ return None
510
+
495
511
  def load_weights(self, file):
496
512
  if os.path.splitext(file)[1] == ".safetensors":
497
513
  from safetensors.torch import load_file, safe_open
@@ -1,9 +1,9 @@
1
1
  # General LyCORIS wrapper based on kohya-ss/sd-scripts' style
2
2
  import os
3
- import regex as re
3
+ import re
4
4
  import logging
5
5
 
6
- from typing import List
6
+ from typing import Any, List
7
7
 
8
8
  import torch
9
9
  import torch.nn as nn
@@ -367,7 +367,9 @@ class LycorisNetwork(torch.nn.Module):
367
367
  next_config = {}
368
368
  for name, module in root_module.named_modules():
369
369
  module_name = module.__class__.__name__
370
- if module_name in target_replace_modules:
370
+ if module_name in target_replace_modules and not any(
371
+ re.match(t, name) for t in target_replace_names
372
+ ):
371
373
  if module_name in self.MODULE_ALGO_MAP:
372
374
  next_config = self.MODULE_ALGO_MAP[module_name]
373
375
  algo = next_config.get("algo", network_module)
@@ -382,8 +384,9 @@ class LycorisNetwork(torch.nn.Module):
382
384
  elif name in target_replace_names or any(
383
385
  re.match(t, name) for t in target_replace_names
384
386
  ):
385
- if name in self.NAME_ALGO_MAP:
386
- next_config = self.NAME_ALGO_MAP[name]
387
+ conf_from_name = self.find_conf_for_name(name)
388
+ if conf_from_name is not None:
389
+ next_config = conf_from_name
387
390
  algo = next_config.get("algo", network_module)
388
391
  elif module_name in self.MODULE_ALGO_MAP:
389
392
  next_config = self.MODULE_ALGO_MAP[module_name]
@@ -401,8 +404,14 @@ class LycorisNetwork(torch.nn.Module):
401
404
  self.loras = create_modules(
402
405
  LycorisNetwork.LORA_PREFIX,
403
406
  module,
404
- LycorisNetwork.TARGET_REPLACE_MODULE,
405
- LycorisNetwork.TARGET_REPLACE_NAME,
407
+ list(set([
408
+ *LycorisNetwork.TARGET_REPLACE_MODULE,
409
+ *LycorisNetwork.MODULE_ALGO_MAP.keys(),
410
+ ])),
411
+ list(set([
412
+ *LycorisNetwork.TARGET_REPLACE_NAME,
413
+ *LycorisNetwork.NAME_ALGO_MAP.keys(),
414
+ ])),
406
415
  )
407
416
  logger.info(f"create LyCORIS: {len(self.loras)} modules.")
408
417
 
@@ -421,6 +430,19 @@ class LycorisNetwork(torch.nn.Module):
421
430
  ), f"duplicated lora name: {lora.lora_name}"
422
431
  names.add(lora.lora_name)
423
432
 
433
+ def find_conf_for_name(
434
+ self,
435
+ name: str,
436
+ ) -> dict[str, Any]:
437
+ if name in self.NAME_ALGO_MAP.keys():
438
+ return self.NAME_ALGO_MAP[name]
439
+
440
+ for key, value in self.NAME_ALGO_MAP.items():
441
+ if re.match(key, name):
442
+ return value
443
+
444
+ return None
445
+
424
446
  def set_multiplier(self, multiplier):
425
447
  self.multiplier = multiplier
426
448
  for lora in self.loras:
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: lycoris_lora
3
- Version: 3.0.1.dev4
3
+ Version: 3.0.1.dev5
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
@@ -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="3.0.1.dev4",
7
+ version="3.0.1.dev5",
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",
@@ -0,0 +1,290 @@
1
+ import unittest
2
+ import re
3
+
4
+ from itertools import product
5
+ from parameterized import parameterized
6
+
7
+ import torch
8
+ import torch.nn as nn
9
+
10
+ from lycoris import create_lycoris, create_lycoris_from_weights, LycorisNetwork
11
+
12
+
13
+ class TestNetwork(nn.Module):
14
+ def __init__(self, dim):
15
+ super().__init__()
16
+ self.conv3d = nn.Conv3d(dim, dim, (3, 3, 3), 1, 1)
17
+ self.conv2d = nn.Conv2d(dim, dim, (3, 3), 1, 1)
18
+ self.group_norm = nn.GroupNorm(4, dim)
19
+ self.conv1d = nn.Conv1d(dim, dim, 3, 1, 1)
20
+ self.layer_norm = nn.LayerNorm(dim)
21
+ self.linear = nn.Linear(dim, dim)
22
+
23
+ def forward(self, x):
24
+ h = self.conv3d(x)
25
+ h = h.flatten(2, 3)
26
+ h = self.group_norm(h)
27
+ h = self.conv2d(h)
28
+ h = h.flatten(2, 3)
29
+ h = self.conv1d(h)
30
+ h = h.transpose(-1, -2)
31
+ h = self.layer_norm(h)
32
+ h = self.linear(h)
33
+ return h
34
+
35
+
36
+ algos: list[str] = [
37
+ "lora",
38
+ "loha",
39
+ "lokr",
40
+ "full",
41
+ "diag-oft",
42
+ "boft",
43
+ "glora",
44
+ # "ia3",
45
+ ]
46
+ device_and_dtype = [
47
+ (torch.device("cpu"), torch.float32),
48
+ ]
49
+ weight_decompose = [
50
+ False,
51
+ True,
52
+ ]
53
+ use_tucker = [
54
+ False,
55
+ True,
56
+ ]
57
+ use_scalar = [
58
+ False,
59
+ True,
60
+ ]
61
+
62
+ if torch.cuda.is_available():
63
+ device_and_dtype.append((torch.device("cuda"), torch.float32))
64
+ device_and_dtype.append((torch.device("cuda"), torch.float16))
65
+ device_and_dtype.append((torch.device("cuda"), torch.bfloat16))
66
+
67
+ if torch.backends.mps.is_available():
68
+ device_and_dtype.append((torch.device("mps"), torch.float32))
69
+
70
+
71
+ patch_forward_param_list = list(
72
+ product(
73
+ algos,
74
+ device_and_dtype,
75
+ weight_decompose,
76
+ use_tucker,
77
+ use_scalar,
78
+ )
79
+ )
80
+
81
+
82
+ class FunLinearNetwork(nn.Module):
83
+ def __init__(self, dim):
84
+ super().__init__()
85
+ self.layer_norm = nn.LayerNorm(dim)
86
+ self.linear = nn.Linear(dim, dim)
87
+ self.act = nn.ReLU()
88
+
89
+ def forward(self, x):
90
+ x = self.layer_norm(x)
91
+ x = self.linear(x)
92
+ x = self.act(x)
93
+ return x
94
+
95
+
96
+ class FunConvNetwork(nn.Module):
97
+ def __init__(self, dim):
98
+ super().__init__()
99
+ self.layer_norm = nn.LayerNorm(dim)
100
+ self.conv = nn.Conv1d(dim, dim, 3, 1, 1)
101
+ self.act = nn.ReLU()
102
+
103
+ def forward(self, x):
104
+ x = self.layer_norm(x)
105
+ x = self.conv(x)
106
+ x = self.act(x)
107
+ return x
108
+
109
+
110
+ class NamedSequential(nn.Module):
111
+ def __init__(self, name, *layers):
112
+ super().__init__()
113
+ self.name = name
114
+ self.layers = nn.ModuleList(layers)
115
+
116
+ def forward(self, x):
117
+ for layer in self.layers:
118
+ x = layer(x)
119
+ return x
120
+
121
+
122
+ class TestNetwork2(nn.Module):
123
+ def __init__(self, dim):
124
+ super().__init__()
125
+ self.linear_layers = NamedSequential(
126
+ "linear_layers",
127
+ FunLinearNetwork(dim),
128
+ FunLinearNetwork(dim),
129
+ FunLinearNetwork(dim),
130
+ )
131
+ self.conv_layers = NamedSequential(
132
+ "conv_layers",
133
+ FunConvNetwork(dim),
134
+ FunConvNetwork(dim),
135
+ FunConvNetwork(dim),
136
+ )
137
+
138
+ def forward(self, x):
139
+ x = self.linear_layers(x)
140
+ x = self.conv_layers(x)
141
+ return x
142
+
143
+ def named_modules(self, memo=None, prefix=''):
144
+ # Call the parent class named_modules method
145
+ for layer_name, layer in super().named_modules(memo, prefix):
146
+ yield layer_name, layer
147
+
148
+
149
+ class LycorisWrapperTests(unittest.TestCase):
150
+ @parameterized.expand(patch_forward_param_list)
151
+ def test_lycoris_wrapper(self, algo, device_dtype, wd, tucker, scalar):
152
+ device, dtype = device_dtype
153
+ print(
154
+ f"{algo: <18}",
155
+ f"device={str(device): <5}",
156
+ f"dtype={str(dtype): <15}",
157
+ f"wd={str(wd): <6}",
158
+ f"tucker={str(tucker): <6}",
159
+ f"scalar={str(scalar): <6}",
160
+ sep="|| ",
161
+ )
162
+ test_net = TestNetwork(16).to(device, dtype)
163
+ test_lycoris: LycorisNetwork = create_lycoris(
164
+ test_net,
165
+ 1,
166
+ algo=algo,
167
+ linear_dim=4,
168
+ linear_alpha=2.0,
169
+ conv_dim=4,
170
+ conv_alpha=2.0,
171
+ dropout=0.0,
172
+ rank_dropout=0.0,
173
+ weight_decompose=wd,
174
+ use_tucker=tucker,
175
+ use_scalar=scalar,
176
+ train_norm=True,
177
+ )
178
+ test_lycoris.apply_to()
179
+ test_lycoris.to(device, dtype)
180
+
181
+ test_input = torch.randn(1, 16, 8, 8, 8).to(device, dtype)
182
+ test_output = test_net(test_input)
183
+ test_lycoris.restore()
184
+
185
+ state_dict = test_lycoris.state_dict()
186
+ test_lycoris_from_weights: LycorisNetwork
187
+ test_lycoris_from_weights, _ = create_lycoris_from_weights(
188
+ 1, None, test_net, state_dict
189
+ )
190
+ test_lycoris_from_weights.apply_to()
191
+ test_lycoris_from_weights.to(device, dtype)
192
+ test_output_from_weights = test_net(test_input)
193
+
194
+ test_lycoris_from_weights.load_state_dict(test_lycoris.state_dict())
195
+
196
+ self.assertTrue(len(test_lycoris.loras) == len(test_lycoris_from_weights.loras))
197
+ self.assertTrue(torch.allclose(test_output, test_output_from_weights))
198
+
199
+ def test_lycoris_wrapper_regex_named_modules(
200
+ self,
201
+ device_dtype=(torch.device("cpu"), torch.float32),
202
+ ):
203
+ device, dtype = device_dtype
204
+
205
+ # Define the preset configuration with regex
206
+ preset = {
207
+ "name_algo_map": {
208
+ "linear_layers.layers.[0-1]..*": {
209
+ "algo": "lokr",
210
+ "factor": 4,
211
+ "linear_dim":1000000,
212
+ "linear_alpha":1,
213
+ "full_matrix": True,
214
+ },
215
+ "linear_layers.layers.2..*": {
216
+ "algo": "lora", "dim": 8,
217
+ },
218
+ "conv_layers.layers.0..*": {
219
+ "algo": "locon", "dim": 8,
220
+ },
221
+ "conv_layers.layers.[1-2]..*": {
222
+ "algo": "glora", "dim": 128,
223
+ }
224
+ }
225
+ }
226
+
227
+ # Apply the preset to the LycorisNetwork class
228
+ LycorisNetwork.apply_preset(preset)
229
+
230
+ # Create the test network and the Lycoris network wrapper
231
+ test_net = TestNetwork2(16).to(device, dtype)
232
+ test_lycoris: LycorisNetwork = create_lycoris(
233
+ test_net,
234
+ multiplier=1,
235
+ linear_dim=16,
236
+ linear_alpha=1,
237
+ train_norm=True,
238
+ )
239
+ test_lycoris.apply_to()
240
+ test_lycoris.to(device, dtype)
241
+
242
+ # Assertions to verify the correct modules were created and configured
243
+ assert len(test_lycoris.loras) == 12
244
+
245
+ lora_names = sorted([ lora.lora_name for lora in test_lycoris.loras ])
246
+ expected = sorted([
247
+ "lycoris_linear_layers_layers_0_layer_norm",
248
+ "lycoris_linear_layers_layers_0_linear",
249
+ "lycoris_linear_layers_layers_1_layer_norm",
250
+ "lycoris_linear_layers_layers_1_linear",
251
+ "lycoris_linear_layers_layers_2_layer_norm",
252
+ "lycoris_linear_layers_layers_2_linear",
253
+ "lycoris_conv_layers_layers_0_layer_norm",
254
+ "lycoris_conv_layers_layers_0_conv",
255
+ "lycoris_conv_layers_layers_1_layer_norm",
256
+ "lycoris_conv_layers_layers_1_conv",
257
+ "lycoris_conv_layers_layers_2_layer_norm",
258
+ "lycoris_conv_layers_layers_2_conv",
259
+ ])
260
+ self.assertEqual(lora_names, expected)
261
+
262
+ for lora in test_lycoris.loras:
263
+ if "lycoris_linear_layers_layers_0" in lora.lora_name or \
264
+ "lycoris_linear_layers_layers_1" in lora.lora_name:
265
+ self.assertEqual(lora.dim, 16)
266
+ if 'norm' in lora.lora_name:
267
+ self.assertEqual(lora.__class__.__name__, "NormModule")
268
+ else:
269
+ self.assertEqual(lora.__class__.__name__, "LokrModule")
270
+ self.assertEqual(lora.multiplier, 1)
271
+ self.assertEqual(lora.full_matrix, True)
272
+ elif "lycoris_linear_layers_layers_2" in lora.lora_name:
273
+ self.assertEqual(lora.dim, 16)
274
+ if 'norm' in lora.lora_name:
275
+ self.assertEqual(lora.__class__.__name__, "NormModule")
276
+ else:
277
+ self.assertEqual(lora.__class__.__name__, "LoConModule")
278
+ elif "lycoris_conv_layers_layers_0" in lora.lora_name:
279
+ self.assertEqual(lora.dim, 16)
280
+ if 'norm' in lora.lora_name:
281
+ self.assertEqual(lora.__class__.__name__, "NormModule")
282
+ else:
283
+ self.assertEqual(lora.__class__.__name__, "LoConModule")
284
+ elif "lycoris_linear_layers_layers_1" in lora.lora_name or \
285
+ "lycoris_linear_layers_layers_2" in lora.lora_name:
286
+ self.assertEqual(lora.dim, 128)
287
+ if 'norm' in lora.lora_name:
288
+ self.assertEqual(lora.__class__.__name__, "NormModule")
289
+ else:
290
+ self.assertEqual(lora.__class__.__name__, "GLoRAModule")
@@ -1,128 +0,0 @@
1
- import unittest
2
- from itertools import product
3
- from parameterized import parameterized
4
-
5
- import torch
6
- import torch.nn as nn
7
-
8
- from lycoris import create_lycoris, create_lycoris_from_weights, LycorisNetwork
9
-
10
-
11
- class TestNetwork(nn.Module):
12
- def __init__(self, dim):
13
- super().__init__()
14
- self.conv3d = nn.Conv3d(dim, dim, (3, 3, 3), 1, 1)
15
- self.conv2d = nn.Conv2d(dim, dim, (3, 3), 1, 1)
16
- self.group_norm = nn.GroupNorm(4, dim)
17
- self.conv1d = nn.Conv1d(dim, dim, 3, 1, 1)
18
- self.layer_norm = nn.LayerNorm(dim)
19
- self.linear = nn.Linear(dim, dim)
20
-
21
- def forward(self, x):
22
- h = self.conv3d(x)
23
- h = h.flatten(2, 3)
24
- h = self.group_norm(h)
25
- h = self.conv2d(h)
26
- h = h.flatten(2, 3)
27
- h = self.conv1d(h)
28
- h = h.transpose(-1, -2)
29
- h = self.layer_norm(h)
30
- h = self.linear(h)
31
- return h
32
-
33
-
34
- algos: list[str] = [
35
- "lora",
36
- "loha",
37
- "lokr",
38
- "full",
39
- "diag-oft",
40
- "boft",
41
- "glora",
42
- # "ia3",
43
- ]
44
- device_and_dtype = [
45
- (torch.device("cpu"), torch.float32),
46
- ]
47
- weight_decompose = [
48
- False,
49
- True,
50
- ]
51
- use_tucker = [
52
- False,
53
- True,
54
- ]
55
- use_scalar = [
56
- False,
57
- True,
58
- ]
59
-
60
- if torch.cuda.is_available():
61
- device_and_dtype.append((torch.device("cuda"), torch.float32))
62
- device_and_dtype.append((torch.device("cuda"), torch.float16))
63
- device_and_dtype.append((torch.device("cuda"), torch.bfloat16))
64
-
65
- if torch.backends.mps.is_available():
66
- device_and_dtype.append((torch.device("mps"), torch.float32))
67
-
68
-
69
- patch_forward_param_list = list(
70
- product(
71
- algos,
72
- device_and_dtype,
73
- weight_decompose,
74
- use_tucker,
75
- use_scalar,
76
- )
77
- )
78
-
79
-
80
- class LycorisWrapperTests(unittest.TestCase):
81
- @parameterized.expand(patch_forward_param_list)
82
- def test_lycoris_wrapper(self, algo, device_dtype, wd, tucker, scalar):
83
- device, dtype = device_dtype
84
- print(
85
- f"{algo: <18}",
86
- f"device={str(device): <5}",
87
- f"dtype={str(dtype): <15}",
88
- f"wd={str(wd): <6}",
89
- f"tucker={str(tucker): <6}",
90
- f"scalar={str(scalar): <6}",
91
- sep="|| ",
92
- )
93
- test_net = TestNetwork(16).to(device, dtype)
94
- test_lycoris: LycorisNetwork = create_lycoris(
95
- test_net,
96
- 1,
97
- algo=algo,
98
- linear_dim=4,
99
- linear_alpha=2.0,
100
- conv_dim=4,
101
- conv_alpha=2.0,
102
- dropout=0.0,
103
- rank_dropout=0.0,
104
- weight_decompose=wd,
105
- use_tucker=tucker,
106
- use_scalar=scalar,
107
- train_norm=True,
108
- )
109
- test_lycoris.apply_to()
110
- test_lycoris.to(device, dtype)
111
-
112
- test_input = torch.randn(1, 16, 8, 8, 8).to(device, dtype)
113
- test_output = test_net(test_input)
114
- test_lycoris.restore()
115
-
116
- state_dict = test_lycoris.state_dict()
117
- test_lycoris_from_weights: LycorisNetwork
118
- test_lycoris_from_weights, _ = create_lycoris_from_weights(
119
- 1, None, test_net, state_dict
120
- )
121
- test_lycoris_from_weights.apply_to()
122
- test_lycoris_from_weights.to(device, dtype)
123
- test_output_from_weights = test_net(test_input)
124
-
125
- test_lycoris_from_weights.load_state_dict(test_lycoris.state_dict())
126
-
127
- self.assertTrue(len(test_lycoris.loras) == len(test_lycoris_from_weights.loras))
128
- self.assertTrue(torch.allclose(test_output, test_output_from_weights))