| 639 | |
| 640 | |
| 641 | class LocalCondProj(nn.Module): |
| 642 | |
| 643 | def __init__(self, input_dim, out_dim, cond_dims): |
| 644 | super().__init__() |
| 645 | self.proj_in=(nn.Linear(input_dim, out_dim, bias=False)) |
| 646 | self.ff=FeedForwardBlock(out_dim, out_dim, cond_dims) |
| 647 | |
| 648 | |
| 649 | def forward(self, x, noise_cond): |
| 650 | x = rearrange(x, 'b c h w -> b h w c') |
| 651 | x=self.proj_in(x) |
| 652 | x=self.ff(x, noise_cond) |
| 653 | x = rearrange(x, 'b h w c -> b c h w') |
| 654 | |
| 655 | return x |