| 215 | |
| 216 | |
| 217 | class PopMixBlock(torch.nn.Module): |
| 218 | def __init__(self, in_channels=None): |
| 219 | super().__init__() |
| 220 | self.mix_factor = torch.nn.Parameter(torch.Tensor([0.5])) |
| 221 | self.need_proj = in_channels is not None |
| 222 | if self.need_proj: |
| 223 | self.proj = torch.nn.Linear(in_channels, in_channels) |
| 224 | |
| 225 | def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs): |
| 226 | res_hidden_states = res_stack.pop() |
| 227 | alpha = torch.sigmoid(self.mix_factor) |
| 228 | hidden_states = alpha * res_hidden_states + (1 - alpha) * hidden_states |
| 229 | if self.need_proj: |
| 230 | hidden_states = hidden_states.permute(0, 2, 3, 1) |
| 231 | hidden_states = self.proj(hidden_states) |
| 232 | hidden_states = hidden_states.permute(0, 3, 1, 2) |
| 233 | res_hidden_states = res_stack.pop() |
| 234 | hidden_states = hidden_states + res_hidden_states |
| 235 | return hidden_states, time_emb, text_emb, res_stack |
| 236 | |
| 237 | |
| 238 | class SVDUNet(torch.nn.Module): |