Spaces:
Running
Running
| from functools import partial | |
| import torch | |
| from torch import nn | |
| import torch.nn.functional as F | |
| from .Advanced_Network_Helpers import * | |
| class Unet(nn.Module): | |
| def __init__( | |
| self, | |
| dim, | |
| init_dim=None, | |
| out_dim=None, | |
| dim_mults=(1, 2, 4, 8), | |
| channels=3, | |
| with_time_emb=True, | |
| resnet_block_groups=8, | |
| use_convnext=True, | |
| convnext_mult=2, | |
| ): | |
| super().__init__() | |
| # determine dimensions | |
| self.channels = channels # since we are concatenating the images and the conditionings along the channel dimension | |
| init_dim = default(init_dim, dim // 3 * 2) | |
| self.init_conv = nn.Conv2d(self.channels * 2, init_dim, 7, padding=3) | |
| self.conditioning_init = nn.Conv2d(self.channels, init_dim, 7, padding=3) | |
| dims = [init_dim, *map(lambda m: dim * m, dim_mults)] | |
| in_out = list(zip(dims[:-1], dims[1:])) | |
| self.in_out = in_out | |
| if use_convnext: | |
| block_klass = partial(ConvNextBlock, mult=convnext_mult) | |
| else: | |
| block_klass = partial(ResnetBlock, groups=resnet_block_groups) | |
| # time embeddings | |
| if with_time_emb: | |
| time_dim = dim * 4 | |
| self.time_mlp = nn.Sequential( | |
| SinusoidalPositionEmbeddings(dim), | |
| nn.Linear(dim, time_dim), | |
| nn.GELU(), | |
| nn.Linear(time_dim, time_dim), | |
| ) | |
| else: | |
| time_dim = None | |
| self.time_mlp = None | |
| # layers | |
| self.downs = nn.ModuleList([]) | |
| self.ups = nn.ModuleList([]) | |
| self.conditioning_encoder = nn.ModuleList([]) | |
| num_resolutions = len(in_out) | |
| self.num_resolutions = num_resolutions | |
| # conditioning encoder | |
| for ind, (dim_in, dim_out) in enumerate(in_out): | |
| is_last = ind >= (num_resolutions - 1) | |
| self.conditioning_encoder.append( | |
| nn.ModuleList( | |
| [ | |
| block_klass(dim_in, dim_out), | |
| Residual(PreNorm(dim_out, LinearAttention(dim_out))), | |
| Downsample(dim_out) if not is_last else nn.Identity(), | |
| ] | |
| ) | |
| ) | |
| for ind, (dim_in, dim_out) in enumerate(in_out): | |
| is_last = ind >= (num_resolutions - 1) | |
| self.downs.append( | |
| nn.ModuleList( | |
| [ | |
| block_klass(dim_in, dim_out, time_emb_dim=time_dim), | |
| block_klass(dim_out, dim_out, time_emb_dim=time_dim), | |
| Residual(PreNorm(dim_out, LinearAttention(dim_out))), | |
| Downsample(dim_out) if not is_last else nn.Identity(), | |
| ] | |
| ) | |
| ) | |
| mid_dim = dims[-1] | |
| self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim) | |
| self.cross_attention_1 = Residual( | |
| PreNorm(mid_dim, LinearCrossAttention(mid_dim)) | |
| ) | |
| self.cross_attention_2 = Residual( | |
| PreNorm(mid_dim, LinearCrossAttention(mid_dim)) | |
| ) | |
| self.cross_attention_3 = Residual( | |
| PreNorm(mid_dim, LinearCrossAttention(mid_dim)) | |
| ) | |
| self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim) | |
| for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): | |
| is_last = ind >= (num_resolutions - 1) | |
| self.ups.append( | |
| nn.ModuleList( | |
| [ | |
| block_klass(dim_out * 2, dim_in, time_emb_dim=time_dim), | |
| block_klass(dim_in, dim_in, time_emb_dim=time_dim), | |
| Residual(PreNorm(dim_in, LinearAttention(dim_in))), | |
| Upsample(dim_in) if not is_last else nn.Identity(), | |
| ] | |
| ) | |
| ) | |
| out_dim = default(out_dim, channels) | |
| self.final_conv = nn.Sequential( | |
| block_klass(dim, dim), nn.Conv2d(dim, out_dim, 1) | |
| ) | |
| def forward(self, x, time, implicit_conditioning, explicit_conditioning): | |
| x = torch.cat((x, explicit_conditioning), dim=1) | |
| x = self.init_conv(x) | |
| conditioning = self.conditioning_init(implicit_conditioning) | |
| t = self.time_mlp(time) if exists(self.time_mlp) else None | |
| h = [] | |
| # conditioning encoder | |
| for block1, attn, downsample in self.conditioning_encoder: | |
| conditioning = block1(conditioning) | |
| conditioning = attn(conditioning) | |
| conditioning = downsample(conditioning) | |
| for block1, block2, attn, downsample in self.downs: | |
| x = block1(x, t) | |
| x = block2(x, t) | |
| x = attn(x) | |
| h.append(x) | |
| x = downsample(x) | |
| # reverse the c list | |
| # bottleneck | |
| x = self.cross_attention_1(x, conditioning) | |
| x = self.mid_block1(x, t) | |
| x = self.cross_attention_2(x, conditioning) | |
| x = self.mid_block2(x, t) | |
| x = self.cross_attention_3(x, conditioning) | |
| for block1, block2, attn, upsample in self.ups: | |
| x = torch.cat((x, h.pop()), dim=1) | |
| x = block1(x, t) | |
| x = block2(x, t) | |
| x = attn(x) | |
| x = upsample(x) | |
| return self.final_conv(x) | |