lycoris-lora 2.3.0.dev9__tar.gz → 2.3.0.dev11__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 (54) hide show
  1. {lycoris_lora-2.3.0.dev9/lycoris_lora.egg-info → lycoris_lora-2.3.0.dev11}/PKG-INFO +1 -1
  2. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/README.md +9 -6
  3. lycoris_lora-2.3.0.dev11/lycoris/functional/__init__.py +4 -0
  4. lycoris_lora-2.3.0.dev11/lycoris/functional/general.py +8 -0
  5. lycoris_lora-2.3.0.dev11/lycoris/functional/locon.py +109 -0
  6. lycoris_lora-2.3.0.dev11/lycoris/functional/loha.py +153 -0
  7. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/lilora.py +1 -1
  8. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/locon.py +5 -4
  9. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/loha.py +25 -97
  10. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/lokr.py +1 -1
  11. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/wrapper.py +2 -0
  12. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11/lycoris_lora.egg-info}/PKG-INFO +1 -1
  13. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris_lora.egg-info/SOURCES.txt +4 -0
  14. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/setup.py +1 -1
  15. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/LICENSE.md +0 -0
  16. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/__init__.py +0 -0
  17. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/config.py +0 -0
  18. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/hcp/__init__.py +0 -0
  19. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/hcp/base.py +0 -0
  20. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/hcp/diag_oft.py +0 -0
  21. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/hcp/loha.py +0 -0
  22. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/hcp/lokr.py +0 -0
  23. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/kohya/__init__.py +0 -0
  24. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/kohya/model_utils.py +0 -0
  25. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/kohya/original_unet.py +0 -0
  26. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/kohya/sdxl_model_util.py +0 -0
  27. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/kohya/sdxl_original_unet.py +0 -0
  28. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/kohya/test_hyperdream.py +0 -0
  29. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/kohya/utils.py +0 -0
  30. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/logging.py +0 -0
  31. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/__init__.py +0 -0
  32. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/attention.py +0 -0
  33. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/base.py +0 -0
  34. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/boft.py +0 -0
  35. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/diag_oft.py +0 -0
  36. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/dylora.py +0 -0
  37. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/full.py +0 -0
  38. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/glokr.py +0 -0
  39. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/glora.py +0 -0
  40. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/hypernet/__init__.py +0 -0
  41. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/hypernet/generater.py +0 -0
  42. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/hypernet/text_encoder.py +0 -0
  43. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/ia3.py +0 -0
  44. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/modules/norms.py +0 -0
  45. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/utils/__init__.py +0 -0
  46. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/utils/bnb.py +0 -0
  47. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/utils/logger.py +0 -0
  48. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/utils/preset.py +0 -0
  49. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris/utils/xformers_utils.py +0 -0
  50. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris_lora.egg-info/dependency_links.txt +0 -0
  51. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris_lora.egg-info/not-zip-safe +0 -0
  52. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris_lora.egg-info/requires.txt +0 -0
  53. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/lycoris_lora.egg-info/top_level.txt +0 -0
  54. {lycoris_lora-2.3.0.dev9 → lycoris_lora-2.3.0.dev11}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: lycoris_lora
3
- Version: 2.3.0.dev9
3
+ Version: 2.3.0.dev11
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,4 +1,4 @@
1
- # LyCORIS - Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion.
1
+ # LyCORIS - Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion. (ICLR'24)
2
2
 
3
3
  ![banner image](docs/images/banner.png)
4
4
 
@@ -10,7 +10,7 @@ This project originated from LoCon (see archive branch).
10
10
 
11
11
  [![Discord!](https://i.imgur.com/A8tOvFS.jpg)](https://discord.gg/VtTFKrj9gJ)
12
12
 
13
- **If you want to check more in-depth experiment results and discussions for LyCORIS, you can check our [paper](https://arxiv.org/abs/2309.14859)**
13
+ **If you want to check more in-depth experiment results and discussions for LyCORIS, you can check our [paper](https://openreview.net/forum?id=wfzXa8e783)**
14
14
 
15
15
  ## Algorithm Overview
16
16
 
@@ -20,10 +20,10 @@ Please check [List of Implemented Algorithms](docs/Algo-List.md) and [Guidelines
20
20
 
21
21
  A simple comparison of some of these methods are provided below (to be taken with a grain of salt)
22
22
 
23
- | | Full | LoRA | LoHa | LoKr low factor | LoKr high factor |
23
+ | | Full | LoRA | LoHa | LoKr low factor | LoKr high factor $^+$ |
24
24
  | --------------------- | ---- | ---- | ---- | --------------- | ---------------- |
25
25
  | Fidelity | ★ | ● | ▲ | ◉ | ▲ |
26
- | Flexibility$^*$ | ★ | ● | ◉ | ▲ | ●$^†$ |
26
+ | Flexibility $^*$ | ★ | ● | ◉ | ▲ | ● $^†$ |
27
27
  | Diversity | ▲ | ◉ | ★ | ● | ★ |
28
28
  | Size | ▲ | ● | ● | ● | ★ |
29
29
  | Training Speed Linear | ★ | ● | ● | ★ | ★ |
@@ -32,8 +32,11 @@ A simple comparison of some of these methods are provided below (to be taken wit
32
32
  ★ > ◉ > ● > ▲
33
33
  [> means better and smaller size is better]
34
34
 
35
- $^*$ Flexibility means anything related to generating images not similar to those in the training set, and combination of multiple concepts, whether they are trained together or not
36
- $^†$ It may become more difficult to switch base model or combine multiple concepts in this situation
35
+ $^+$ Usually we take `factor <= 0.5 * sqrt(dim)` as low factor and `factor >= sqrt(dim` as high factor. For example, factor<=8 for SD1.x/SD2.x/SDXL can be seen as low factor, and, factor>=16 can be seen as high factor. <br>
36
+ $^*$ Flexibility means anything related to generating images not similar to those in the training set, and combination of multiple concepts, whether they are trained together or not <br>
37
+ $^†$ It may become more difficult to switch base model or combine multiple concepts in this situation <br>
38
+
39
+ **The actual performance may vary depending on the datasets, tasks, and hyperparameters used. It is recommended to experiment with different settings to achieve optimal results.**
37
40
 
38
41
  ## Usage
39
42
 
@@ -0,0 +1,4 @@
1
+ from .locon import (
2
+ lora_diff_weight,
3
+ lora_bypass_forward_diff,
4
+ )
@@ -0,0 +1,8 @@
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.nn.functional as F
4
+
5
+
6
+ def rebuild_tucker(t, wa, wb):
7
+ rebuild2 = torch.einsum("i j ..., i p, j r -> p r ...", t, wa, wb)
8
+ return rebuild2
@@ -0,0 +1,109 @@
1
+ import math
2
+
3
+ import torch
4
+ import torch.nn as nn
5
+ import torch.nn.functional as F
6
+
7
+ from .general import rebuild_tucker
8
+
9
+
10
+ def lora_weight_gen(org_weight, rank, tucker=True):
11
+ """### lora_weight_gen
12
+
13
+ Args:
14
+ out_dim (int): output dimension
15
+ in_dim (int): input dimension
16
+ rank (int): low rank
17
+
18
+ Returns:
19
+ torch.Tensor: down proj weight
20
+ torch.Tensor: up proj weight
21
+ torch.Tensor(optional): mid proj weight
22
+ """
23
+ out_dim, in_dim, *k = org_weight.shape
24
+ if k and tucker:
25
+ down = torch.empty(rank, in_dim, *(1 for _ in k))
26
+ up = torch.empty(out_dim, rank, *(1 for _ in k))
27
+ mid = torch.empty(rank, rank, *k)
28
+ nn.init.kaiming_uniform_(down, a=math.sqrt(5))
29
+ nn.init.constant_(up, 0)
30
+ nn.init.kaiming_uniform_(mid, a=math.sqrt(5))
31
+ return down, up, mid
32
+ else:
33
+ down = torch.empty(rank, in_dim)
34
+ up = torch.empty(out_dim, rank)
35
+ nn.init.kaiming_uniform_(down, a=math.sqrt(5))
36
+ nn.init.constant_(up, 0)
37
+ return down, up
38
+
39
+
40
+ def lora_diff_weight(d, u, m=None, gamma=1.0):
41
+ """### lora_diff_weight
42
+
43
+ Get ΔW = BA, where BA is low rank decomposition
44
+
45
+ Args:
46
+ d (torch.Tensor): weight of down proj linear/conv layer
47
+ u (torch.Tensor): weight of up proj linear/conv layer
48
+ m (torch.Tensor, optional): weight of mid proj linear/conv layer, \
49
+ for tucker deomposition
50
+ gamma (float, optional): scale factor, normally alpha/rank here
51
+
52
+ Returns:
53
+ torch.Tensor: ΔW
54
+ """
55
+ assert d.size(0) == u.size(1)
56
+ R, I, *k = d.shape
57
+ O, R, *_ = u.shape
58
+ u = u * gamma
59
+
60
+ if m is None:
61
+ result = u.reshape(-1, u.size(1)) @ d.reshape(d.size(0), -1)
62
+ else:
63
+ R, R, *k = m.shape
64
+ u = u.reshape(u.size(0), -1).transpose(0, 1)
65
+ d = d.reshape(d.size(0), -1)
66
+ result = rebuild_tucker(m, u, d)
67
+ return result.reshape(O, I, *k)
68
+
69
+
70
+ FUNC_LIST = [None, None, F.linear, F.conv1d, F.conv2d, F.conv3d]
71
+
72
+
73
+ def lora_bypass_forward_diff(x, d, u, m=None, gamma=1.0, extra_args={}):
74
+ """### lora_bypass_forward_diff
75
+
76
+ Args:
77
+ x (torch.Tensor): input tensor
78
+ d (torch.Tensor): weight of down proj linear/conv layer
79
+ u (torch.Tensor): weight of up proj linear/conv layer
80
+ m (torch.Tensor, optional): weight of mid proj linear/conv layer, \
81
+ for tucker deomposition
82
+ gamma (float, optional): scale factor, normally alpha/rank here
83
+
84
+ Returns:
85
+ torch.Tensor: output tensor
86
+ """
87
+ if m is not None:
88
+ down = FUNC_LIST[d.dim()](x, d)
89
+ mid = FUNC_LIST[d.dim()](down, m, **extra_args)
90
+ up = FUNC_LIST[d.dim()](mid, u)
91
+ else:
92
+ down = FUNC_LIST[d.dim()](x, d, **extra_args)
93
+ up = FUNC_LIST[d.dim()](down, u)
94
+ return up * gamma
95
+
96
+
97
+ if __name__ == "__main__":
98
+ w = torch.randn(32, 32, 3, 3, 3)
99
+ d, u, m = lora_weight_gen(w, 4)
100
+ u = u+0.01
101
+ extra_args = {"padding": 1}
102
+
103
+ x = torch.randn(1, 32, 8, 8, 8)
104
+ y = FUNC_LIST[d.dim()](x, w, **extra_args)
105
+ diff_w = lora_diff_weight(d, u, m, 1)
106
+ diff_y = lora_bypass_forward_diff(x, d, u, m, 1, extra_args)
107
+
108
+ print(F.mse_loss(y, y + diff_y))
109
+ print(F.mse_loss(w, w + diff_w))
@@ -0,0 +1,153 @@
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.nn.functional as F
4
+
5
+
6
+ class HadaWeight(torch.autograd.Function):
7
+ @staticmethod
8
+ def forward(ctx, w1d, w1u, w2d, w2u, scale=torch.tensor(1)):
9
+ ctx.save_for_backward(w1d, w1u, w2d, w2u, scale)
10
+ diff_weight = ((w1u @ w1d) * (w2u @ w2d)) * scale
11
+ return diff_weight
12
+
13
+ @staticmethod
14
+ def backward(ctx, grad_out):
15
+ (w1d, w1u, w2d, w2u, scale) = ctx.saved_tensors
16
+ grad_out = grad_out * scale
17
+ temp = grad_out * (w2u @ w2d)
18
+ grad_w1d = temp @ w1d.T
19
+ grad_w1u = w1u.T @ temp
20
+
21
+ temp = grad_out * (w1u @ w1d)
22
+ grad_w2d = temp @ w2d.T
23
+ grad_w2u = w2u.T @ temp
24
+
25
+ del temp
26
+ return grad_w1d, grad_w1u, grad_w2d, grad_w2u, None
27
+
28
+
29
+ class HadaWeightTucker(torch.autograd.Function):
30
+ @staticmethod
31
+ def forward(ctx, t1, w1d, w1u, t2, w2d, w2u, scale=torch.tensor(1)):
32
+ ctx.save_for_backward(t1, w1d, w1u, t2, w2d, w2u, scale)
33
+
34
+ rebuild1 = torch.einsum("i j ..., j r, i p -> p r ...", t1, w1d, w1u)
35
+ rebuild2 = torch.einsum("i j ..., j r, i p -> p r ...", t2, w2d, w2u)
36
+
37
+ return rebuild1 * rebuild2 * scale
38
+
39
+ @staticmethod
40
+ def backward(ctx, grad_out):
41
+ (t1, w1d, w1u, t2, w2d, w2u, scale) = ctx.saved_tensors
42
+ grad_out = grad_out * scale
43
+
44
+ temp = torch.einsum("i j ..., j r -> i r ...", t2, w2d)
45
+ rebuild = torch.einsum("i j ..., i r -> r j ...", temp, w2u)
46
+
47
+ grad_w = rebuild * grad_out
48
+ del rebuild
49
+
50
+ grad_w1d = torch.einsum("r j ..., i j ... -> r i", temp, grad_w)
51
+ grad_temp = torch.einsum("i j ..., i r -> r j ...", grad_w, w1u.T)
52
+ del grad_w, temp
53
+
54
+ grad_w1u = torch.einsum("i r ..., i j ... -> r j", t1, grad_temp)
55
+ grad_t1 = torch.einsum("i j ..., j r -> i r ...", grad_temp, w1d.T)
56
+ del grad_temp
57
+
58
+ temp = torch.einsum("i j ..., j r -> i r ...", t1, w1d)
59
+ rebuild = torch.einsum("i j ..., i r -> r j ...", temp, w1u)
60
+
61
+ grad_w = rebuild * grad_out
62
+ del rebuild
63
+
64
+ grad_w2d = torch.einsum("r j ..., i j ... -> r i", temp, grad_w)
65
+ grad_temp = torch.einsum("i j ..., i r -> r j ...", grad_w, w2u.T)
66
+ del grad_w, temp
67
+
68
+ grad_w2u = torch.einsum("i r ..., i j ... -> r j", t2, grad_temp)
69
+ grad_t2 = torch.einsum("i j ..., j r -> i r ...", grad_temp, w2d.T)
70
+ del grad_temp
71
+ return grad_t1, grad_w1d, grad_w1u, grad_t2, grad_w2d, grad_w2u, None
72
+
73
+
74
+ def make_weight(w1d, w1u, w2d, w2u, scale):
75
+ return HadaWeight.apply(w1d, w1u, w2d, w2u, scale)
76
+
77
+
78
+ def make_weight_tucker(t1, w1d, w1u, t2, w2d, w2u, scale):
79
+ return HadaWeightTucker.apply(t1, w1d, w1u, t2, w2d, w2u, scale)
80
+
81
+
82
+ def loha_diff_weight(w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0):
83
+ """### loha_diff_weight
84
+
85
+ Get ΔW = BA, where BA is low rank decomposition
86
+
87
+ Args:
88
+ w1d, w2d (torch.Tensor): weight of down proj linear/conv layer
89
+ w1u, w2u (torch.Tensor): weight of up proj linear/conv layer
90
+ gamma (float, optional): scale factor, normally alpha/rank here
91
+
92
+ Returns:
93
+ torch.Tensor: ΔW
94
+ """
95
+ if t1 is not None and t2 is not None:
96
+ assert w1d.size(0) == t1.size(1)
97
+ assert w1u.size(0) == t1.size(0)
98
+ assert w2d.size(0) == t2.size(1)
99
+ assert w2u.size(0) == t2.size(0)
100
+ R, I = w1d.shape
101
+ R, O = w1u.shape
102
+ R, R, *k = t1.shape
103
+ result = make_weight_tucker(t1, w1d, w1u, t2, w2d, w2u, gamma)
104
+ else:
105
+ assert w1d.size(0) == w1u.size(1)
106
+ assert w2d.size(0) == w2u.size(1)
107
+ R, I, *k = w1d.shape
108
+ O, R, *_ = w1u.shape
109
+ w1d = w1d.reshape(w1d.size(0), -1)
110
+ w1u = w1u.reshape(-1, w1u.size(1))
111
+ w2d = w2d.reshape(w2d.size(0), -1)
112
+ w2u = w1u.reshape(-1, w2u.size(1))
113
+ result = make_weight(w1d, w1u, w2d, w2u, gamma)
114
+
115
+ result = result.reshape(O, I, *k)
116
+ return result
117
+
118
+
119
+ FUNC_LIST = [None, None, F.linear, F.conv1d, F.conv2d, F.conv3d]
120
+
121
+
122
+ def loha_bypass_forward_diff(
123
+ x, w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0, extra_args={}
124
+ ):
125
+ """### loha_bypass_forward_diff
126
+
127
+ Args:
128
+ x (torch.Tensor): input tensor
129
+ w1d, w2d (torch.Tensor): weight of up proj linear/conv layer
130
+ w1u, w2u (torch.Tensor): weight of down proj linear/conv layer
131
+ gamma (float, optional): scale factor, normally alpha/rank here
132
+
133
+ Returns:
134
+ torch.Tensor: output tensor
135
+ """
136
+ diff_w = loha_diff_weight(w1d, w1u, w2d, w2u, t1, t2, gamma)
137
+ return FUNC_LIST[w1d.dim() if t1 is None else t1.dim()](x, diff_w, **extra_args)
138
+
139
+
140
+ if __name__ == "__main__":
141
+ w = torch.randn(128, 128, 3, 3, 3)
142
+ a = torch.randn(16, 128) * 0.01
143
+ b = torch.randn(16, 128) * 0.01
144
+ m = torch.randn(16, 16, 3, 3, 3) * 0.01
145
+ extra_args = {"padding": 1}
146
+
147
+ x = torch.randn(1, 128, 8, 8, 8)
148
+ y = FUNC_LIST[w.dim()](x, w, **extra_args)
149
+ diff_w = loha_diff_weight(a, b, a, b, m, m, 1)
150
+ diff_y = loha_bypass_forward_diff(x, a, b, a, b, m, m, 1, extra_args)
151
+
152
+ print(F.mse_loss(y, y + diff_y))
153
+ print(F.mse_loss(w, w + diff_w))
@@ -157,7 +157,7 @@ class LiLoConModule(ModuleCustomSD):
157
157
  if "scalar" in key:
158
158
  del missing_keys[missing_keys.index(key)]
159
159
  if isinstance(self.scalar, nn.Parameter):
160
- self.scalar.copy_(torch.ones_like(self.scalar))
160
+ self.scalar.data.copy_(torch.ones_like(self.scalar))
161
161
  else:
162
162
  self.scalar = torch.ones_like(self.scalar)
163
163
 
@@ -6,6 +6,7 @@ import torch.nn as nn
6
6
  import torch.nn.functional as F
7
7
 
8
8
  from .base import LycorisBaseModule
9
+ from ..functional.general import rebuild_tucker
9
10
  from ..logging import logger
10
11
  from ..utils.bnb import LinearNF4
11
12
 
@@ -149,7 +150,7 @@ class LoConModule(LycorisBaseModule):
149
150
  if "scalar" in key:
150
151
  del missing_keys[missing_keys.index(key)]
151
152
  if isinstance(self.scalar, nn.Parameter):
152
- self.scalar.copy_(torch.ones_like(self.scalar))
153
+ self.scalar.data.copy_(torch.ones_like(self.scalar))
153
154
  else:
154
155
  self.scalar = torch.ones_like(self.scalar)
155
156
 
@@ -158,9 +159,9 @@ class LoConModule(LycorisBaseModule):
158
159
  wb = self.lora_down.weight.to(device)
159
160
  if self.tucker:
160
161
  t = self.lora_mid.weight.to(device)
161
- wa = wa.squeeze(-1, -2)
162
- wb = wb.squeeze(-1, -2)
163
- weight = torch.einsum("i j k l, p i, j r -> p r k l", t, wa, wb)
162
+ wa = wa.reshape(wa.size(0), -1).transpose(0, 1)
163
+ wb = wb.reshape(wb.size(0), -1)
164
+ weight = rebuild_tucker(t, wa, wb)
164
165
  else:
165
166
  weight = wa.view(wa.size(0), -1) @ wb.view(wb.size(0), -1)
166
167
 
@@ -5,85 +5,10 @@ import torch.nn as nn
5
5
  import torch.nn.functional as F
6
6
 
7
7
  from .base import LycorisBaseModule
8
+ from ..functional.loha import loha_diff_weight
8
9
  from ..utils.bnb import LinearNF4
9
10
 
10
11
 
11
- class HadaWeight(torch.autograd.Function):
12
- @staticmethod
13
- def forward(ctx, w1a, w1b, w2a, w2b, scale=torch.tensor(1)):
14
- ctx.save_for_backward(w1a, w1b, w2a, w2b, scale)
15
- diff_weight = ((w1a @ w1b) * (w2a @ w2b)) * scale
16
- return diff_weight
17
-
18
- @staticmethod
19
- def backward(ctx, grad_out):
20
- (w1a, w1b, w2a, w2b, scale) = ctx.saved_tensors
21
- grad_out = grad_out * scale
22
- temp = grad_out * (w2a @ w2b)
23
- grad_w1a = temp @ w1b.T
24
- grad_w1b = w1a.T @ temp
25
-
26
- temp = grad_out * (w1a @ w1b)
27
- grad_w2a = temp @ w2b.T
28
- grad_w2b = w2a.T @ temp
29
-
30
- del temp
31
- return grad_w1a, grad_w1b, grad_w2a, grad_w2b, None
32
-
33
-
34
- class HadaWeightCP(torch.autograd.Function):
35
- @staticmethod
36
- def forward(ctx, t1, w1a, w1b, t2, w2a, w2b, scale=torch.tensor(1)):
37
- ctx.save_for_backward(t1, w1a, w1b, t2, w2a, w2b, scale)
38
-
39
- rebuild1 = torch.einsum("i j k l, j r, i p -> p r k l", t1, w1b, w1a)
40
- rebuild2 = torch.einsum("i j k l, j r, i p -> p r k l", t2, w2b, w2a)
41
-
42
- return rebuild1 * rebuild2 * scale
43
-
44
- @staticmethod
45
- def backward(ctx, grad_out):
46
- (t1, w1a, w1b, t2, w2a, w2b, scale) = ctx.saved_tensors
47
- grad_out = grad_out * scale
48
-
49
- temp = torch.einsum("i j k l, j r -> i r k l", t2, w2b)
50
- rebuild = torch.einsum("i j k l, i r -> r j k l", temp, w2a)
51
-
52
- grad_w = rebuild * grad_out
53
- del rebuild
54
-
55
- grad_w1a = torch.einsum("r j k l, i j k l -> r i", temp, grad_w)
56
- grad_temp = torch.einsum("i j k l, i r -> r j k l", grad_w, w1a.T)
57
- del grad_w, temp
58
-
59
- grad_w1b = torch.einsum("i r k l, i j k l -> r j", t1, grad_temp)
60
- grad_t1 = torch.einsum("i j k l, j r -> i r k l", grad_temp, w1b.T)
61
- del grad_temp
62
-
63
- temp = torch.einsum("i j k l, j r -> i r k l", t1, w1b)
64
- rebuild = torch.einsum("i j k l, i r -> r j k l", temp, w1a)
65
-
66
- grad_w = rebuild * grad_out
67
- del rebuild
68
-
69
- grad_w2a = torch.einsum("r j k l, i j k l -> r i", temp, grad_w)
70
- grad_temp = torch.einsum("i j k l, i r -> r j k l", grad_w, w2a.T)
71
- del grad_w, temp
72
-
73
- grad_w2b = torch.einsum("i r k l, i j k l -> r j", t2, grad_temp)
74
- grad_t2 = torch.einsum("i j k l, j r -> i r k l", grad_temp, w2b.T)
75
- del grad_temp
76
- return grad_t1, grad_w1a, grad_w1b, grad_t2, grad_w2a, grad_w2b, None
77
-
78
-
79
- def make_weight(w1a, w1b, w2a, w2b, scale):
80
- return HadaWeight.apply(w1a, w1b, w2a, w2b, scale)
81
-
82
-
83
- def make_weight_tucker(t1, w1a, w1b, t2, w2a, w2b, scale):
84
- return HadaWeightCP.apply(t1, w1a, w1b, t2, w2a, w2b, scale)
85
-
86
-
87
12
  class LohaModule(LycorisBaseModule):
88
13
  support_module = {
89
14
  "linear",
@@ -214,28 +139,30 @@ class LohaModule(LycorisBaseModule):
214
139
  if "scalar" in key:
215
140
  del missing_keys[missing_keys.index(key)]
216
141
  if isinstance(self.scalar, nn.Parameter):
217
- self.scalar.copy_(torch.ones_like(self.scalar))
142
+ self.scalar.data.copy_(torch.ones_like(self.scalar))
218
143
  else:
219
144
  self.scalar = torch.ones_like(self.scalar)
220
145
 
221
146
  def get_weight(self, shape):
222
147
  if self.tucker:
223
- weight = make_weight_tucker(
224
- self.hada_t1,
225
- self.hada_w1_a,
148
+ weight = loha_diff_weight(
226
149
  self.hada_w1_b,
227
- self.hada_t2,
228
- self.hada_w2_a,
150
+ self.hada_w1_a,
229
151
  self.hada_w2_b,
230
- scale=torch.tensor(self.scale),
152
+ self.hada_w2_a,
153
+ self.hada_t1,
154
+ self.hada_t2,
155
+ gamma=torch.tensor(self.scale),
231
156
  )
232
157
  else:
233
- weight = make_weight(
234
- self.hada_w1_a,
158
+ weight = loha_diff_weight(
235
159
  self.hada_w1_b,
236
- self.hada_w2_a,
160
+ self.hada_w1_a,
237
161
  self.hada_w2_b,
238
- scale=torch.tensor(self.scale),
162
+ self.hada_w2_a,
163
+ None,
164
+ None,
165
+ gamma=torch.tensor(self.scale),
239
166
  )
240
167
  if shape is not None:
241
168
  weight = weight.reshape(shape)
@@ -339,20 +266,21 @@ class LohaModule(LycorisBaseModule):
339
266
 
340
267
 
341
268
  if __name__ == "__main__":
342
- base = nn.Linear(128, 128).cuda()
343
- lokr = LohaModule("test", base, 1, 4, 1, weight_decompose=True).cuda()
269
+ device = torch.device("mps")
270
+ base = nn.Linear(128, 128).to(device)
271
+ lokr = LohaModule("test", base, 1, 4, 1, weight_decompose=True).to(device)
344
272
  print(lokr)
345
- test_input = torch.randn(1, 128).cuda()
273
+ test_input = torch.randn(1, 128).to(device)
346
274
  test_output = lokr(test_input)
347
275
  print(test_output.shape)
348
276
 
349
- base_4bit = LinearNF4(128, 128)
350
- base_4bit.load_state_dict(base.state_dict())
351
- base_4bit.cuda()
352
- qlocon = LohaModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
353
- print(qlocon)
354
- test_input = torch.randn(1, 128).cuda()
355
- test_output = qlocon(test_input)
277
+ # base_4bit = LinearNF4(128, 128)
278
+ # base_4bit.load_state_dict(base.state_dict())
279
+ # base_4bit.to(device)
280
+ # qlocon = LohaModule("test", base_4bit, 1, 4, 1, weight_decompose=False).to(device)
281
+ # print(qlocon)
282
+ # test_input = torch.randn(1, 128).to(device)
283
+ # test_output = qlocon(test_input)
356
284
  print(test_output.shape)
357
285
 
358
286
  base = nn.Conv2d(128, 128, 3, 1, 1)
@@ -280,7 +280,7 @@ class LokrModule(LycorisBaseModule):
280
280
  if "scalar" in key:
281
281
  del missing_keys[missing_keys.index(key)]
282
282
  if isinstance(self.scalar, nn.Parameter):
283
- self.scalar.copy_(torch.ones_like(self.scalar))
283
+ self.scalar.data.copy_(torch.ones_like(self.scalar))
284
284
  else:
285
285
  self.scalar = torch.ones_like(self.scalar)
286
286
 
@@ -331,6 +331,8 @@ class LycorisNetwork(torch.nn.Module):
331
331
  loras = {}
332
332
  lora_names = []
333
333
  for name, module in root_module.named_modules():
334
+ if module is root_module:
335
+ continue
334
336
  module_name = module.__class__.__name__
335
337
  if module_name in self.MODULE_ALGO_MAP:
336
338
  next_config = self.MODULE_ALGO_MAP[module_name]
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: lycoris_lora
3
- Version: 2.3.0.dev9
3
+ Version: 2.3.0.dev11
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
@@ -5,6 +5,10 @@ lycoris/__init__.py
5
5
  lycoris/config.py
6
6
  lycoris/logging.py
7
7
  lycoris/wrapper.py
8
+ lycoris/functional/__init__.py
9
+ lycoris/functional/general.py
10
+ lycoris/functional/locon.py
11
+ lycoris/functional/loha.py
8
12
  lycoris/hcp/__init__.py
9
13
  lycoris/hcp/base.py
10
14
  lycoris/hcp/diag_oft.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.3.0.dev9",
7
+ version="2.3.0.dev11",
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",