Args: x (torch.Tensor): In shape of (B, C, H, W)
(self, x)
| 153 | self.gamma = nn.Conv2d(c_dim, rank, kernel_size=1, bias=False) |
| 154 | |
| 155 | def forward(self, x): |
| 156 | """ |
| 157 | Args: |
| 158 | x (torch.Tensor): In shape of (B, C, H, W) |
| 159 | """ |
| 160 | w_out = self.W(x) |
| 161 | a_out = self.A(x) |
| 162 | |
| 163 | # inject conditioning into LoRA |
| 164 | c = self.data_provider.get_batch(a_out) |
| 165 | element_shift = self.beta(c) |
| 166 | element_scale = self.gamma(c) + 1 |
| 167 | a_cond = a_out * element_scale + element_shift |
| 168 | |
| 169 | b_out = self.B(a_cond) |
| 170 | |
| 171 | return w_out + b_out * self.lora_scale |