| 21 | |
| 22 | class LoRALinearLayer(nn.Module): |
| 23 | def __init__(self, in_features, out_features, rank=4, network_alpha=None, device=None, dtype=None): |
| 24 | super().__init__() |
| 25 | |
| 26 | self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype) |
| 27 | self.up = nn.Linear(rank, out_features, bias=False, device=device, dtype=dtype) |
| 28 | # This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. |
| 29 | # See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning |
| 30 | self.network_alpha = network_alpha |
| 31 | self.rank = rank |
| 32 | self.out_features = out_features |
| 33 | self.in_features = in_features |
| 34 | |
| 35 | nn.init.normal_(self.down.weight, std=1 / rank) |
| 36 | nn.init.zeros_(self.up.weight) |
| 37 | |
| 38 | def forward(self, hidden_states): |
| 39 | orig_dtype = hidden_states.dtype |