| 50 | |
| 51 | class LoRAConv2dLayer(nn.Module): |
| 52 | def __init__( |
| 53 | self, in_features, out_features, rank=4, kernel_size=(1, 1), stride=(1, 1), padding=0, network_alpha=None |
| 54 | ): |
| 55 | super().__init__() |
| 56 | |
| 57 | self.down = nn.Conv2d(in_features, rank, kernel_size=kernel_size, stride=stride, padding=padding, bias=False) |
| 58 | # according to the official kohya_ss trainer kernel_size are always fixed for the up layer |
| 59 | # # see: https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L129 |
| 60 | self.up = nn.Conv2d(rank, out_features, kernel_size=(1, 1), stride=(1, 1), bias=False) |
| 61 | |
| 62 | # This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. |
| 63 | # See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning |
| 64 | self.network_alpha = network_alpha |
| 65 | self.rank = rank |
| 66 | |
| 67 | nn.init.normal_(self.down.weight, std=1 / rank) |
| 68 | nn.init.zeros_(self.up.weight) |
| 69 | |
| 70 | def forward(self, hidden_states): |
| 71 | orig_dtype = hidden_states.dtype |