Spaces:
Runtime error
Runtime error
| # pytorch_diffusion + derived encoder decoder | |
| import math | |
| from urllib.request import proxy_bypass | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from einops import rearrange | |
| class VectorQuantizer(nn.Module): | |
| """ | |
| Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly | |
| avoids costly matrix multiplications and allows for post-hoc remapping of indices. | |
| """ | |
| # NOTE: due to a bug the beta term was applied to the wrong term. for | |
| # backwards compatibility we use the buggy version by default, but you can | |
| # specify legacy=False to fix it. | |
| def __init__(self, | |
| n_e, | |
| e_dim, | |
| beta, | |
| remap=None, | |
| unknown_index="random", | |
| sane_index_shape=False, | |
| legacy=True): | |
| super().__init__() | |
| self.n_e = n_e | |
| self.e_dim = e_dim | |
| self.beta = beta | |
| self.legacy = legacy | |
| self.embedding = nn.Embedding(self.n_e, self.e_dim) | |
| self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) | |
| self.remap = remap | |
| if self.remap is not None: | |
| self.register_buffer("used", torch.tensor(np.load(self.remap))) | |
| self.re_embed = self.used.shape[0] | |
| self.unknown_index = unknown_index # "random" or "extra" or integer | |
| if self.unknown_index == "extra": | |
| self.unknown_index = self.re_embed | |
| self.re_embed = self.re_embed + 1 | |
| print(f"Remapping {self.n_e} indices to {self.re_embed} indices. " | |
| f"Using {self.unknown_index} for unknown indices.") | |
| else: | |
| self.re_embed = n_e | |
| self.sane_index_shape = sane_index_shape | |
| def remap_to_used(self, inds): | |
| ishape = inds.shape | |
| assert len(ishape) > 1 | |
| inds = inds.reshape(ishape[0], -1) | |
| used = self.used.to(inds) | |
| match = (inds[:, :, None] == used[None, None, ...]).long() | |
| new = match.argmax(-1) | |
| unknown = match.sum(2) < 1 | |
| if self.unknown_index == "random": | |
| new[unknown] = torch.randint( | |
| 0, self.re_embed, | |
| size=new[unknown].shape).to(device=new.device) | |
| else: | |
| new[unknown] = self.unknown_index | |
| return new.reshape(ishape) | |
| def unmap_to_all(self, inds): | |
| ishape = inds.shape | |
| assert len(ishape) > 1 | |
| inds = inds.reshape(ishape[0], -1) | |
| used = self.used.to(inds) | |
| if self.re_embed > self.used.shape[0]: # extra token | |
| inds[inds >= self.used.shape[0]] = 0 # simply set to zero | |
| back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) | |
| return back.reshape(ishape) | |
| def forward(self, z, temp=None, rescale_logits=False, return_logits=False): | |
| assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel" | |
| assert rescale_logits == False, "Only for interface compatible with Gumbel" | |
| assert return_logits == False, "Only for interface compatible with Gumbel" | |
| # reshape z -> (batch, height, width, channel) and flatten | |
| z = rearrange(z, 'b c h w -> b h w c').contiguous() | |
| z_flattened = z.view(-1, self.e_dim) | |
| # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z | |
| d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \ | |
| torch.sum(self.embedding.weight**2, dim=1) - 2 * \ | |
| torch.einsum('bd,dn->bn', z_flattened, rearrange(self.embedding.weight, 'n d -> d n')) | |
| min_encoding_indices = torch.argmin(d, dim=1) | |
| z_q = self.embedding(min_encoding_indices).view(z.shape) | |
| perplexity = None | |
| min_encodings = None | |
| # compute loss for embedding | |
| if not self.legacy: | |
| loss = self.beta * torch.mean((z_q.detach()-z)**2) + \ | |
| torch.mean((z_q - z.detach()) ** 2) | |
| else: | |
| loss = torch.mean((z_q.detach()-z)**2) + self.beta * \ | |
| torch.mean((z_q - z.detach()) ** 2) | |
| # preserve gradients | |
| z_q = z + (z_q - z).detach() | |
| # reshape back to match original input shape | |
| z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous() | |
| if self.remap is not None: | |
| min_encoding_indices = min_encoding_indices.reshape( | |
| z.shape[0], -1) # add batch axis | |
| min_encoding_indices = self.remap_to_used(min_encoding_indices) | |
| min_encoding_indices = min_encoding_indices.reshape(-1, | |
| 1) # flatten | |
| if self.sane_index_shape: | |
| min_encoding_indices = min_encoding_indices.reshape( | |
| z_q.shape[0], z_q.shape[2], z_q.shape[3]) | |
| return z_q, loss, (perplexity, min_encodings, min_encoding_indices) | |
| def get_codebook_entry(self, indices, shape): | |
| # shape specifying (batch, height, width, channel) | |
| if self.remap is not None: | |
| indices = indices.reshape(shape[0], -1) # add batch axis | |
| indices = self.unmap_to_all(indices) | |
| indices = indices.reshape(-1) # flatten again | |
| # get quantized latent vectors | |
| z_q = self.embedding(indices) | |
| if shape is not None: | |
| z_q = z_q.view(shape) | |
| # reshape back to match original input shape | |
| z_q = z_q.permute(0, 3, 1, 2).contiguous() | |
| return z_q | |
| class VectorQuantizerTexture(nn.Module): | |
| """ | |
| Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly | |
| avoids costly matrix multiplications and allows for post-hoc remapping of indices. | |
| """ | |
| # NOTE: due to a bug the beta term was applied to the wrong term. for | |
| # backwards compatibility we use the buggy version by default, but you can | |
| # specify legacy=False to fix it. | |
| def __init__(self, | |
| n_e, | |
| e_dim, | |
| beta, | |
| remap=None, | |
| unknown_index="random", | |
| sane_index_shape=False, | |
| legacy=True): | |
| super().__init__() | |
| self.n_e = n_e | |
| self.e_dim = e_dim | |
| self.beta = beta | |
| self.legacy = legacy | |
| # TODO: decide number of embeddings | |
| self.embedding_list = nn.ModuleList( | |
| [nn.Embedding(self.n_e, self.e_dim) for i in range(18)]) | |
| for embedding in self.embedding_list: | |
| embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) | |
| self.remap = remap | |
| if self.remap is not None: | |
| self.register_buffer("used", torch.tensor(np.load(self.remap))) | |
| self.re_embed = self.used.shape[0] | |
| self.unknown_index = unknown_index # "random" or "extra" or integer | |
| if self.unknown_index == "extra": | |
| self.unknown_index = self.re_embed | |
| self.re_embed = self.re_embed + 1 | |
| print(f"Remapping {self.n_e} indices to {self.re_embed} indices. " | |
| f"Using {self.unknown_index} for unknown indices.") | |
| else: | |
| self.re_embed = n_e | |
| self.sane_index_shape = sane_index_shape | |
| def remap_to_used(self, inds): | |
| ishape = inds.shape | |
| assert len(ishape) > 1 | |
| inds = inds.reshape(ishape[0], -1) | |
| used = self.used.to(inds) | |
| match = (inds[:, :, None] == used[None, None, ...]).long() | |
| new = match.argmax(-1) | |
| unknown = match.sum(2) < 1 | |
| if self.unknown_index == "random": | |
| new[unknown] = torch.randint( | |
| 0, self.re_embed, | |
| size=new[unknown].shape).to(device=new.device) | |
| else: | |
| new[unknown] = self.unknown_index | |
| return new.reshape(ishape) | |
| def unmap_to_all(self, inds): | |
| ishape = inds.shape | |
| assert len(ishape) > 1 | |
| inds = inds.reshape(ishape[0], -1) | |
| used = self.used.to(inds) | |
| if self.re_embed > self.used.shape[0]: # extra token | |
| inds[inds >= self.used.shape[0]] = 0 # simply set to zero | |
| back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) | |
| return back.reshape(ishape) | |
| def forward(self, | |
| z, | |
| segm_map, | |
| temp=None, | |
| rescale_logits=False, | |
| return_logits=False): | |
| assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel" | |
| assert rescale_logits == False, "Only for interface compatible with Gumbel" | |
| assert return_logits == False, "Only for interface compatible with Gumbel" | |
| segm_map = F.interpolate(segm_map, size=z.size()[2:], mode='nearest') | |
| # reshape z -> (batch, height, width, channel) and flatten | |
| z = rearrange(z, 'b c h w -> b h w c').contiguous() | |
| z_flattened = z.view(-1, self.e_dim) | |
| # flatten segm_map (b, h, w) | |
| segm_map_flatten = segm_map.view(-1) | |
| z_q = torch.zeros_like(z_flattened) | |
| min_encoding_indices_list = [] | |
| min_encoding_indices_continual = torch.full( | |
| segm_map_flatten.size(), | |
| fill_value=-1, | |
| dtype=torch.long, | |
| device=segm_map_flatten.device) | |
| for codebook_idx in range(18): | |
| min_encoding_indices = torch.full( | |
| segm_map_flatten.size(), | |
| fill_value=-1, | |
| dtype=torch.long, | |
| device=segm_map_flatten.device) | |
| if torch.sum(segm_map_flatten == codebook_idx) > 0: | |
| z_selected = z_flattened[segm_map_flatten == codebook_idx] | |
| # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z | |
| d_selected = torch.sum( | |
| z_selected**2, dim=1, keepdim=True) + torch.sum( | |
| self.embedding_list[codebook_idx].weight**2, | |
| dim=1) - 2 * torch.einsum( | |
| 'bd,dn->bn', z_selected, | |
| rearrange(self.embedding_list[codebook_idx].weight, | |
| 'n d -> d n')) | |
| min_encoding_indices_selected = torch.argmin(d_selected, dim=1) | |
| z_q_selected = self.embedding_list[codebook_idx]( | |
| min_encoding_indices_selected) | |
| z_q[segm_map_flatten == codebook_idx] = z_q_selected | |
| min_encoding_indices[ | |
| segm_map_flatten == | |
| codebook_idx] = min_encoding_indices_selected | |
| min_encoding_indices_continual[ | |
| segm_map_flatten == | |
| codebook_idx] = min_encoding_indices_selected + 1024 * codebook_idx | |
| min_encoding_indices = min_encoding_indices.reshape( | |
| z.shape[0], z.shape[1], z.shape[2]) | |
| min_encoding_indices_list.append(min_encoding_indices) | |
| min_encoding_indices_continual = min_encoding_indices_continual.reshape( | |
| z.shape[0], z.shape[1], z.shape[2]) | |
| z_q = z_q.view(z.shape) | |
| perplexity = None | |
| # compute loss for embedding | |
| if not self.legacy: | |
| loss = self.beta * torch.mean((z_q.detach()-z)**2) + \ | |
| torch.mean((z_q - z.detach()) ** 2) | |
| else: | |
| loss = torch.mean((z_q.detach()-z)**2) + self.beta * \ | |
| torch.mean((z_q - z.detach()) ** 2) | |
| # preserve gradients | |
| z_q = z + (z_q - z).detach() | |
| # reshape back to match original input shape | |
| z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous() | |
| return z_q, loss, (perplexity, min_encoding_indices_continual, | |
| min_encoding_indices_list) | |
| def get_codebook_entry(self, indices_list, segm_map, shape): | |
| # flatten segm_map (b, h, w) | |
| segm_map = F.interpolate( | |
| segm_map, size=(shape[1], shape[2]), mode='nearest') | |
| segm_map_flatten = segm_map.view(-1) | |
| z_q = torch.zeros((shape[0] * shape[1] * shape[2]), | |
| self.e_dim).to(segm_map.device) | |
| for codebook_idx in range(18): | |
| if torch.sum(segm_map_flatten == codebook_idx) > 0: | |
| min_encoding_indices_selected = indices_list[ | |
| codebook_idx].view(-1)[segm_map_flatten == codebook_idx] | |
| z_q_selected = self.embedding_list[codebook_idx]( | |
| min_encoding_indices_selected) | |
| z_q[segm_map_flatten == codebook_idx] = z_q_selected | |
| z_q = z_q.view(shape) | |
| # reshape back to match original input shape | |
| z_q = z_q.permute(0, 3, 1, 2).contiguous() | |
| return z_q | |
| def sample_patches(inputs, patch_size=3, stride=1): | |
| """Extract sliding local patches from an input feature tensor. | |
| The sampled pathes are row-major. | |
| Args: | |
| inputs (Tensor): the input feature maps, shape: (n, c, h, w). | |
| patch_size (int): the spatial size of sampled patches. Default: 3. | |
| stride (int): the stride of sampling. Default: 1. | |
| Returns: | |
| patches (Tensor): extracted patches, shape: (n, c * patch_size * | |
| patch_size, n_patches). | |
| """ | |
| patches = F.unfold(inputs, (patch_size, patch_size), stride=stride) | |
| return patches | |
| class VectorQuantizerSpatialTextureAware(nn.Module): | |
| """ | |
| Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly | |
| avoids costly matrix multiplications and allows for post-hoc remapping of indices. | |
| """ | |
| # NOTE: due to a bug the beta term was applied to the wrong term. for | |
| # backwards compatibility we use the buggy version by default, but you can | |
| # specify legacy=False to fix it. | |
| def __init__(self, | |
| n_e, | |
| e_dim, | |
| beta, | |
| spatial_size, | |
| remap=None, | |
| unknown_index="random", | |
| sane_index_shape=False, | |
| legacy=True): | |
| super().__init__() | |
| self.n_e = n_e | |
| self.e_dim = e_dim * spatial_size * spatial_size | |
| self.beta = beta | |
| self.legacy = legacy | |
| self.spatial_size = spatial_size | |
| # TODO: decide number of embeddings | |
| self.embedding_list = nn.ModuleList( | |
| [nn.Embedding(self.n_e, self.e_dim) for i in range(18)]) | |
| for embedding in self.embedding_list: | |
| embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) | |
| self.remap = remap | |
| if self.remap is not None: | |
| self.register_buffer("used", torch.tensor(np.load(self.remap))) | |
| self.re_embed = self.used.shape[0] | |
| self.unknown_index = unknown_index # "random" or "extra" or integer | |
| if self.unknown_index == "extra": | |
| self.unknown_index = self.re_embed | |
| self.re_embed = self.re_embed + 1 | |
| print(f"Remapping {self.n_e} indices to {self.re_embed} indices. " | |
| f"Using {self.unknown_index} for unknown indices.") | |
| else: | |
| self.re_embed = n_e | |
| self.sane_index_shape = sane_index_shape | |
| def forward(self, | |
| z, | |
| segm_map, | |
| temp=None, | |
| rescale_logits=False, | |
| return_logits=False): | |
| assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel" | |
| assert rescale_logits == False, "Only for interface compatible with Gumbel" | |
| assert return_logits == False, "Only for interface compatible with Gumbel" | |
| segm_map = F.interpolate( | |
| segm_map, | |
| size=(z.size(2) // self.spatial_size, | |
| z.size(3) // self.spatial_size), | |
| mode='nearest') | |
| # reshape z -> (batch, height, width, channel) and flatten | |
| # z = rearrange(z, 'b c h w -> b h w c').contiguous() ? | |
| z_patches = sample_patches( | |
| z, patch_size=self.spatial_size, | |
| stride=self.spatial_size).permute(0, 2, 1) | |
| z_patches_flattened = z_patches.reshape(-1, self.e_dim) | |
| # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z | |
| # flatten segm_map (b, h, w) | |
| segm_map_flatten = segm_map.view(-1) | |
| z_q = torch.zeros_like(z_patches_flattened) | |
| min_encoding_indices_list = [] | |
| min_encoding_indices_continual = torch.full( | |
| segm_map_flatten.size(), | |
| fill_value=-1, | |
| dtype=torch.long, | |
| device=segm_map_flatten.device) | |
| for codebook_idx in range(18): | |
| min_encoding_indices = torch.full( | |
| segm_map_flatten.size(), | |
| fill_value=-1, | |
| dtype=torch.long, | |
| device=segm_map_flatten.device) | |
| if torch.sum(segm_map_flatten == codebook_idx) > 0: | |
| z_selected = z_patches_flattened[segm_map_flatten == | |
| codebook_idx] | |
| # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z | |
| d_selected = torch.sum( | |
| z_selected**2, dim=1, keepdim=True) + torch.sum( | |
| self.embedding_list[codebook_idx].weight**2, | |
| dim=1) - 2 * torch.einsum( | |
| 'bd,dn->bn', z_selected, | |
| rearrange(self.embedding_list[codebook_idx].weight, | |
| 'n d -> d n')) | |
| min_encoding_indices_selected = torch.argmin(d_selected, dim=1) | |
| z_q_selected = self.embedding_list[codebook_idx]( | |
| min_encoding_indices_selected) | |
| z_q[segm_map_flatten == codebook_idx] = z_q_selected | |
| min_encoding_indices[ | |
| segm_map_flatten == | |
| codebook_idx] = min_encoding_indices_selected | |
| min_encoding_indices_continual[ | |
| segm_map_flatten == | |
| codebook_idx] = min_encoding_indices_selected + self.n_e * codebook_idx | |
| min_encoding_indices = min_encoding_indices.reshape( | |
| z_patches.shape[0], segm_map.shape[2], segm_map.shape[3]) | |
| min_encoding_indices_list.append(min_encoding_indices) | |
| z_q = F.fold( | |
| z_q.view(z_patches.shape).permute(0, 2, 1), | |
| z.size()[2:], | |
| kernel_size=(self.spatial_size, self.spatial_size), | |
| stride=self.spatial_size) | |
| perplexity = None | |
| # compute loss for embedding | |
| if not self.legacy: | |
| loss = self.beta * torch.mean((z_q.detach()-z)**2) + \ | |
| torch.mean((z_q - z.detach()) ** 2) | |
| else: | |
| loss = torch.mean((z_q.detach()-z)**2) + self.beta * \ | |
| torch.mean((z_q - z.detach()) ** 2) | |
| # preserve gradients | |
| z_q = z + (z_q - z).detach() | |
| return z_q, loss, (perplexity, min_encoding_indices_continual, | |
| min_encoding_indices_list) | |
| def get_codebook_entry(self, indices_list, segm_map, shape): | |
| # flatten segm_map (b, h, w) | |
| segm_map = F.interpolate( | |
| segm_map, size=(shape[1], shape[2]), mode='nearest') | |
| segm_map_flatten = segm_map.view(-1) | |
| z_q = torch.zeros((shape[0] * shape[1] * shape[2]), | |
| self.e_dim).to(segm_map.device) | |
| for codebook_idx in range(18): | |
| if torch.sum(segm_map_flatten == codebook_idx) > 0: | |
| min_encoding_indices_selected = indices_list[ | |
| codebook_idx].view(-1)[segm_map_flatten == codebook_idx] | |
| z_q_selected = self.embedding_list[codebook_idx]( | |
| min_encoding_indices_selected) | |
| z_q[segm_map_flatten == codebook_idx] = z_q_selected | |
| z_q = F.fold( | |
| z_q.view(((shape[0], shape[1] * shape[2], | |
| self.e_dim))).permute(0, 2, 1), | |
| (shape[1] * self.spatial_size, shape[2] * self.spatial_size), | |
| kernel_size=(self.spatial_size, self.spatial_size), | |
| stride=self.spatial_size) | |
| return z_q | |
| def get_timestep_embedding(timesteps, embedding_dim): | |
| """ | |
| This matches the implementation in Denoising Diffusion Probabilistic Models: | |
| From Fairseq. | |
| Build sinusoidal embeddings. | |
| This matches the implementation in tensor2tensor, but differs slightly | |
| from the description in Section 3.5 of "Attention Is All You Need". | |
| """ | |
| assert len(timesteps.shape) == 1 | |
| half_dim = embedding_dim // 2 | |
| emb = math.log(10000) / (half_dim - 1) | |
| emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) | |
| emb = emb.to(device=timesteps.device) | |
| emb = timesteps.float()[:, None] * emb[None, :] | |
| emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) | |
| if embedding_dim % 2 == 1: # zero pad | |
| emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) | |
| return emb | |
| def nonlinearity(x): | |
| # swish | |
| return x * torch.sigmoid(x) | |
| def Normalize(in_channels): | |
| return torch.nn.GroupNorm( | |
| num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) | |
| class Upsample(nn.Module): | |
| def __init__(self, in_channels, with_conv): | |
| super().__init__() | |
| self.with_conv = with_conv | |
| if self.with_conv: | |
| self.conv = torch.nn.Conv2d( | |
| in_channels, in_channels, kernel_size=3, stride=1, padding=1) | |
| def forward(self, x): | |
| x = torch.nn.functional.interpolate( | |
| x, scale_factor=2.0, mode="nearest") | |
| if self.with_conv: | |
| x = self.conv(x) | |
| return x | |
| class Downsample(nn.Module): | |
| def __init__(self, in_channels, with_conv): | |
| super().__init__() | |
| self.with_conv = with_conv | |
| if self.with_conv: | |
| # no asymmetric padding in torch conv, must do it ourselves | |
| self.conv = torch.nn.Conv2d( | |
| in_channels, in_channels, kernel_size=3, stride=2, padding=0) | |
| def forward(self, x): | |
| if self.with_conv: | |
| pad = (0, 1, 0, 1) | |
| x = torch.nn.functional.pad(x, pad, mode="constant", value=0) | |
| x = self.conv(x) | |
| else: | |
| x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) | |
| return x | |
| class ResnetBlock(nn.Module): | |
| def __init__(self, | |
| *, | |
| in_channels, | |
| out_channels=None, | |
| conv_shortcut=False, | |
| dropout, | |
| temb_channels=512): | |
| super().__init__() | |
| self.in_channels = in_channels | |
| out_channels = in_channels if out_channels is None else out_channels | |
| self.out_channels = out_channels | |
| self.use_conv_shortcut = conv_shortcut | |
| self.norm1 = Normalize(in_channels) | |
| self.conv1 = torch.nn.Conv2d( | |
| in_channels, out_channels, kernel_size=3, stride=1, padding=1) | |
| if temb_channels > 0: | |
| self.temb_proj = torch.nn.Linear(temb_channels, out_channels) | |
| self.norm2 = Normalize(out_channels) | |
| self.dropout = torch.nn.Dropout(dropout) | |
| self.conv2 = torch.nn.Conv2d( | |
| out_channels, out_channels, kernel_size=3, stride=1, padding=1) | |
| if self.in_channels != self.out_channels: | |
| if self.use_conv_shortcut: | |
| self.conv_shortcut = torch.nn.Conv2d( | |
| in_channels, | |
| out_channels, | |
| kernel_size=3, | |
| stride=1, | |
| padding=1) | |
| else: | |
| self.nin_shortcut = torch.nn.Conv2d( | |
| in_channels, | |
| out_channels, | |
| kernel_size=1, | |
| stride=1, | |
| padding=0) | |
| def forward(self, x, temb): | |
| h = x | |
| h = self.norm1(h) | |
| h = nonlinearity(h) | |
| h = self.conv1(h) | |
| if temb is not None: | |
| h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None] | |
| h = self.norm2(h) | |
| h = nonlinearity(h) | |
| h = self.dropout(h) | |
| h = self.conv2(h) | |
| if self.in_channels != self.out_channels: | |
| if self.use_conv_shortcut: | |
| x = self.conv_shortcut(x) | |
| else: | |
| x = self.nin_shortcut(x) | |
| return x + h | |
| class AttnBlock(nn.Module): | |
| def __init__(self, in_channels): | |
| super().__init__() | |
| self.in_channels = in_channels | |
| self.norm = Normalize(in_channels) | |
| self.q = torch.nn.Conv2d( | |
| in_channels, in_channels, kernel_size=1, stride=1, padding=0) | |
| self.k = torch.nn.Conv2d( | |
| in_channels, in_channels, kernel_size=1, stride=1, padding=0) | |
| self.v = torch.nn.Conv2d( | |
| in_channels, in_channels, kernel_size=1, stride=1, padding=0) | |
| self.proj_out = torch.nn.Conv2d( | |
| in_channels, in_channels, kernel_size=1, stride=1, padding=0) | |
| def forward(self, x): | |
| h_ = x | |
| h_ = self.norm(h_) | |
| q = self.q(h_) | |
| k = self.k(h_) | |
| v = self.v(h_) | |
| # compute attention | |
| b, c, h, w = q.shape | |
| q = q.reshape(b, c, h * w) | |
| q = q.permute(0, 2, 1) # b,hw,c | |
| k = k.reshape(b, c, h * w) # b,c,hw | |
| w_ = torch.bmm(q, k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j] | |
| w_ = w_ * (int(c)**(-0.5)) | |
| w_ = torch.nn.functional.softmax(w_, dim=2) | |
| # attend to values | |
| v = v.reshape(b, c, h * w) | |
| w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q) | |
| h_ = torch.bmm( | |
| v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] | |
| h_ = h_.reshape(b, c, h, w) | |
| h_ = self.proj_out(h_) | |
| return x + h_ | |
| class Model(nn.Module): | |
| def __init__(self, | |
| *, | |
| ch, | |
| out_ch, | |
| ch_mult=(1, 2, 4, 8), | |
| num_res_blocks, | |
| attn_resolutions, | |
| dropout=0.0, | |
| resamp_with_conv=True, | |
| in_channels, | |
| resolution, | |
| use_timestep=True): | |
| super().__init__() | |
| self.ch = ch | |
| self.temb_ch = self.ch * 4 | |
| self.num_resolutions = len(ch_mult) | |
| self.num_res_blocks = num_res_blocks | |
| self.resolution = resolution | |
| self.in_channels = in_channels | |
| self.use_timestep = use_timestep | |
| if self.use_timestep: | |
| # timestep embedding | |
| self.temb = nn.Module() | |
| self.temb.dense = nn.ModuleList([ | |
| torch.nn.Linear(self.ch, self.temb_ch), | |
| torch.nn.Linear(self.temb_ch, self.temb_ch), | |
| ]) | |
| # downsampling | |
| self.conv_in = torch.nn.Conv2d( | |
| in_channels, self.ch, kernel_size=3, stride=1, padding=1) | |
| curr_res = resolution | |
| in_ch_mult = (1, ) + tuple(ch_mult) | |
| self.down = nn.ModuleList() | |
| for i_level in range(self.num_resolutions): | |
| block = nn.ModuleList() | |
| attn = nn.ModuleList() | |
| block_in = ch * in_ch_mult[i_level] | |
| block_out = ch * ch_mult[i_level] | |
| for i_block in range(self.num_res_blocks): | |
| block.append( | |
| ResnetBlock( | |
| in_channels=block_in, | |
| out_channels=block_out, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout)) | |
| block_in = block_out | |
| if curr_res in attn_resolutions: | |
| attn.append(AttnBlock(block_in)) | |
| down = nn.Module() | |
| down.block = block | |
| down.attn = attn | |
| if i_level != self.num_resolutions - 1: | |
| down.downsample = Downsample(block_in, resamp_with_conv) | |
| curr_res = curr_res // 2 | |
| self.down.append(down) | |
| # middle | |
| self.mid = nn.Module() | |
| self.mid.block_1 = ResnetBlock( | |
| in_channels=block_in, | |
| out_channels=block_in, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout) | |
| self.mid.attn_1 = AttnBlock(block_in) | |
| self.mid.block_2 = ResnetBlock( | |
| in_channels=block_in, | |
| out_channels=block_in, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout) | |
| # upsampling | |
| self.up = nn.ModuleList() | |
| for i_level in reversed(range(self.num_resolutions)): | |
| block = nn.ModuleList() | |
| attn = nn.ModuleList() | |
| block_out = ch * ch_mult[i_level] | |
| skip_in = ch * ch_mult[i_level] | |
| for i_block in range(self.num_res_blocks + 1): | |
| if i_block == self.num_res_blocks: | |
| skip_in = ch * in_ch_mult[i_level] | |
| block.append( | |
| ResnetBlock( | |
| in_channels=block_in + skip_in, | |
| out_channels=block_out, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout)) | |
| block_in = block_out | |
| if curr_res in attn_resolutions: | |
| attn.append(AttnBlock(block_in)) | |
| up = nn.Module() | |
| up.block = block | |
| up.attn = attn | |
| if i_level != 0: | |
| up.upsample = Upsample(block_in, resamp_with_conv) | |
| curr_res = curr_res * 2 | |
| self.up.insert(0, up) # prepend to get consistent order | |
| # end | |
| self.norm_out = Normalize(block_in) | |
| self.conv_out = torch.nn.Conv2d( | |
| block_in, out_ch, kernel_size=3, stride=1, padding=1) | |
| def forward(self, x, t=None): | |
| #assert x.shape[2] == x.shape[3] == self.resolution | |
| if self.use_timestep: | |
| # timestep embedding | |
| assert t is not None | |
| temb = get_timestep_embedding(t, self.ch) | |
| temb = self.temb.dense[0](temb) | |
| temb = nonlinearity(temb) | |
| temb = self.temb.dense[1](temb) | |
| else: | |
| temb = None | |
| # downsampling | |
| hs = [self.conv_in(x)] | |
| for i_level in range(self.num_resolutions): | |
| for i_block in range(self.num_res_blocks): | |
| h = self.down[i_level].block[i_block](hs[-1], temb) | |
| if len(self.down[i_level].attn) > 0: | |
| h = self.down[i_level].attn[i_block](h) | |
| hs.append(h) | |
| if i_level != self.num_resolutions - 1: | |
| hs.append(self.down[i_level].downsample(hs[-1])) | |
| # middle | |
| h = hs[-1] | |
| h = self.mid.block_1(h, temb) | |
| h = self.mid.attn_1(h) | |
| h = self.mid.block_2(h, temb) | |
| # upsampling | |
| for i_level in reversed(range(self.num_resolutions)): | |
| for i_block in range(self.num_res_blocks + 1): | |
| h = self.up[i_level].block[i_block](torch.cat([h, hs.pop()], | |
| dim=1), temb) | |
| if len(self.up[i_level].attn) > 0: | |
| h = self.up[i_level].attn[i_block](h) | |
| if i_level != 0: | |
| h = self.up[i_level].upsample(h) | |
| # end | |
| h = self.norm_out(h) | |
| h = nonlinearity(h) | |
| h = self.conv_out(h) | |
| return h | |
| class Encoder(nn.Module): | |
| def __init__(self, | |
| ch, | |
| num_res_blocks, | |
| attn_resolutions, | |
| in_channels, | |
| resolution, | |
| z_channels, | |
| ch_mult=(1, 2, 4, 8), | |
| dropout=0.0, | |
| resamp_with_conv=True, | |
| double_z=True): | |
| super().__init__() | |
| self.ch = ch | |
| self.temb_ch = 0 | |
| self.num_resolutions = len(ch_mult) | |
| self.num_res_blocks = num_res_blocks | |
| self.resolution = resolution | |
| self.in_channels = in_channels | |
| # downsampling | |
| self.conv_in = torch.nn.Conv2d( | |
| in_channels, self.ch, kernel_size=3, stride=1, padding=1) | |
| curr_res = resolution | |
| in_ch_mult = (1, ) + tuple(ch_mult) | |
| self.down = nn.ModuleList() | |
| for i_level in range(self.num_resolutions): | |
| block = nn.ModuleList() | |
| attn = nn.ModuleList() | |
| block_in = ch * in_ch_mult[i_level] | |
| block_out = ch * ch_mult[i_level] | |
| for i_block in range(self.num_res_blocks): | |
| block.append( | |
| ResnetBlock( | |
| in_channels=block_in, | |
| out_channels=block_out, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout)) | |
| block_in = block_out | |
| if curr_res in attn_resolutions: | |
| attn.append(AttnBlock(block_in)) | |
| down = nn.Module() | |
| down.block = block | |
| down.attn = attn | |
| if i_level != self.num_resolutions - 1: | |
| down.downsample = Downsample(block_in, resamp_with_conv) | |
| curr_res = curr_res // 2 | |
| self.down.append(down) | |
| # middle | |
| self.mid = nn.Module() | |
| self.mid.block_1 = ResnetBlock( | |
| in_channels=block_in, | |
| out_channels=block_in, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout) | |
| self.mid.attn_1 = AttnBlock(block_in) | |
| self.mid.block_2 = ResnetBlock( | |
| in_channels=block_in, | |
| out_channels=block_in, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout) | |
| # end | |
| self.norm_out = Normalize(block_in) | |
| self.conv_out = torch.nn.Conv2d( | |
| block_in, | |
| 2 * z_channels if double_z else z_channels, | |
| kernel_size=3, | |
| stride=1, | |
| padding=1) | |
| def forward(self, x): | |
| #assert x.shape[2] == x.shape[3] == self.resolution, "{}, {}, {}".format(x.shape[2], x.shape[3], self.resolution) | |
| # timestep embedding | |
| temb = None | |
| # downsampling | |
| hs = [self.conv_in(x)] | |
| for i_level in range(self.num_resolutions): | |
| for i_block in range(self.num_res_blocks): | |
| h = self.down[i_level].block[i_block](hs[-1], temb) | |
| if len(self.down[i_level].attn) > 0: | |
| h = self.down[i_level].attn[i_block](h) | |
| hs.append(h) | |
| if i_level != self.num_resolutions - 1: | |
| hs.append(self.down[i_level].downsample(hs[-1])) | |
| # middle | |
| h = hs[-1] | |
| h = self.mid.block_1(h, temb) | |
| h = self.mid.attn_1(h) | |
| h = self.mid.block_2(h, temb) | |
| # end | |
| h = self.norm_out(h) | |
| h = nonlinearity(h) | |
| h = self.conv_out(h) | |
| return h | |
| class Decoder(nn.Module): | |
| def __init__(self, | |
| in_channels, | |
| resolution, | |
| z_channels, | |
| ch, | |
| out_ch, | |
| num_res_blocks, | |
| attn_resolutions, | |
| ch_mult=(1, 2, 4, 8), | |
| dropout=0.0, | |
| resamp_with_conv=True, | |
| give_pre_end=False): | |
| super().__init__() | |
| self.ch = ch | |
| self.temb_ch = 0 | |
| self.num_resolutions = len(ch_mult) | |
| self.num_res_blocks = num_res_blocks | |
| self.resolution = resolution | |
| self.in_channels = in_channels | |
| self.give_pre_end = give_pre_end | |
| # compute in_ch_mult, block_in and curr_res at lowest res | |
| in_ch_mult = (1, ) + tuple(ch_mult) | |
| block_in = ch * ch_mult[self.num_resolutions - 1] | |
| curr_res = resolution // 2**(self.num_resolutions - 1) | |
| self.z_shape = (1, z_channels, curr_res, curr_res // 2) | |
| print("Working with z of shape {} = {} dimensions.".format( | |
| self.z_shape, np.prod(self.z_shape))) | |
| # z to block_in | |
| self.conv_in = torch.nn.Conv2d( | |
| z_channels, block_in, kernel_size=3, stride=1, padding=1) | |
| # middle | |
| self.mid = nn.Module() | |
| self.mid.block_1 = ResnetBlock( | |
| in_channels=block_in, | |
| out_channels=block_in, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout) | |
| self.mid.attn_1 = AttnBlock(block_in) | |
| self.mid.block_2 = ResnetBlock( | |
| in_channels=block_in, | |
| out_channels=block_in, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout) | |
| # upsampling | |
| self.up = nn.ModuleList() | |
| for i_level in reversed(range(self.num_resolutions)): | |
| block = nn.ModuleList() | |
| attn = nn.ModuleList() | |
| block_out = ch * ch_mult[i_level] | |
| for i_block in range(self.num_res_blocks + 1): | |
| block.append( | |
| ResnetBlock( | |
| in_channels=block_in, | |
| out_channels=block_out, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout)) | |
| block_in = block_out | |
| if curr_res in attn_resolutions: | |
| attn.append(AttnBlock(block_in)) | |
| up = nn.Module() | |
| up.block = block | |
| up.attn = attn | |
| if i_level != 0: | |
| up.upsample = Upsample(block_in, resamp_with_conv) | |
| curr_res = curr_res * 2 | |
| self.up.insert(0, up) # prepend to get consistent order | |
| # end | |
| self.norm_out = Normalize(block_in) | |
| self.conv_out = torch.nn.Conv2d( | |
| block_in, out_ch, kernel_size=3, stride=1, padding=1) | |
| def forward(self, z, bot_h=None): | |
| #assert z.shape[1:] == self.z_shape[1:] | |
| self.last_z_shape = z.shape | |
| # timestep embedding | |
| temb = None | |
| # z to block_in | |
| h = self.conv_in(z) | |
| # middle | |
| h = self.mid.block_1(h, temb) | |
| h = self.mid.attn_1(h) | |
| h = self.mid.block_2(h, temb) | |
| # upsampling | |
| for i_level in reversed(range(self.num_resolutions)): | |
| for i_block in range(self.num_res_blocks + 1): | |
| h = self.up[i_level].block[i_block](h, temb) | |
| if len(self.up[i_level].attn) > 0: | |
| h = self.up[i_level].attn[i_block](h) | |
| if i_level != 0: | |
| h = self.up[i_level].upsample(h) | |
| if i_level == 4 and bot_h is not None: | |
| h += bot_h | |
| # end | |
| if self.give_pre_end: | |
| return h | |
| h = self.norm_out(h) | |
| h = nonlinearity(h) | |
| h = self.conv_out(h) | |
| return h | |
| def get_feature_top(self, z): | |
| #assert z.shape[1:] == self.z_shape[1:] | |
| self.last_z_shape = z.shape | |
| # timestep embedding | |
| temb = None | |
| # z to block_in | |
| h = self.conv_in(z) | |
| # middle | |
| h = self.mid.block_1(h, temb) | |
| h = self.mid.attn_1(h) | |
| h = self.mid.block_2(h, temb) | |
| # upsampling | |
| for i_level in reversed(range(self.num_resolutions)): | |
| for i_block in range(self.num_res_blocks + 1): | |
| h = self.up[i_level].block[i_block](h, temb) | |
| if len(self.up[i_level].attn) > 0: | |
| h = self.up[i_level].attn[i_block](h) | |
| if i_level != 0: | |
| h = self.up[i_level].upsample(h) | |
| if i_level == 4: | |
| return h | |
| def get_feature_middle(self, z, mid_h): | |
| #assert z.shape[1:] == self.z_shape[1:] | |
| self.last_z_shape = z.shape | |
| # timestep embedding | |
| temb = None | |
| # z to block_in | |
| h = self.conv_in(z) | |
| # middle | |
| h = self.mid.block_1(h, temb) | |
| h = self.mid.attn_1(h) | |
| h = self.mid.block_2(h, temb) | |
| # upsampling | |
| for i_level in reversed(range(self.num_resolutions)): | |
| for i_block in range(self.num_res_blocks + 1): | |
| h = self.up[i_level].block[i_block](h, temb) | |
| if len(self.up[i_level].attn) > 0: | |
| h = self.up[i_level].attn[i_block](h) | |
| if i_level != 0: | |
| h = self.up[i_level].upsample(h) | |
| if i_level == 4: | |
| h += mid_h | |
| if i_level == 3: | |
| return h | |
| class DecoderRes(nn.Module): | |
| def __init__(self, | |
| in_channels, | |
| resolution, | |
| z_channels, | |
| ch, | |
| num_res_blocks, | |
| ch_mult=(1, 2, 4, 8), | |
| dropout=0.0, | |
| give_pre_end=False): | |
| super().__init__() | |
| self.ch = ch | |
| self.temb_ch = 0 | |
| self.num_resolutions = len(ch_mult) | |
| self.num_res_blocks = num_res_blocks | |
| self.resolution = resolution | |
| self.in_channels = in_channels | |
| self.give_pre_end = give_pre_end | |
| # compute in_ch_mult, block_in and curr_res at lowest res | |
| in_ch_mult = (1, ) + tuple(ch_mult) | |
| block_in = ch * ch_mult[self.num_resolutions - 1] | |
| curr_res = resolution // 2**(self.num_resolutions - 1) | |
| self.z_shape = (1, z_channels, curr_res, curr_res // 2) | |
| print("Working with z of shape {} = {} dimensions.".format( | |
| self.z_shape, np.prod(self.z_shape))) | |
| # z to block_in | |
| self.conv_in = torch.nn.Conv2d( | |
| z_channels, block_in, kernel_size=3, stride=1, padding=1) | |
| # middle | |
| self.mid = nn.Module() | |
| self.mid.block_1 = ResnetBlock( | |
| in_channels=block_in, | |
| out_channels=block_in, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout) | |
| self.mid.attn_1 = AttnBlock(block_in) | |
| self.mid.block_2 = ResnetBlock( | |
| in_channels=block_in, | |
| out_channels=block_in, | |
| temb_channels=self.temb_ch, | |
| dropout=dropout) | |
| def forward(self, z): | |
| #assert z.shape[1:] == self.z_shape[1:] | |
| self.last_z_shape = z.shape | |
| # timestep embedding | |
| temb = None | |
| # z to block_in | |
| h = self.conv_in(z) | |
| # middle | |
| h = self.mid.block_1(h, temb) | |
| h = self.mid.attn_1(h) | |
| h = self.mid.block_2(h, temb) | |
| return h | |
| # patch based discriminator | |
| class Discriminator(nn.Module): | |
| def __init__(self, nc, ndf, n_layers=3): | |
| super().__init__() | |
| layers = [ | |
| nn.Conv2d(nc, ndf, kernel_size=4, stride=2, padding=1), | |
| nn.LeakyReLU(0.2, True) | |
| ] | |
| ndf_mult = 1 | |
| ndf_mult_prev = 1 | |
| for n in range(1, | |
| n_layers): # gradually increase the number of filters | |
| ndf_mult_prev = ndf_mult | |
| ndf_mult = min(2**n, 8) | |
| layers += [ | |
| nn.Conv2d( | |
| ndf * ndf_mult_prev, | |
| ndf * ndf_mult, | |
| kernel_size=4, | |
| stride=2, | |
| padding=1, | |
| bias=False), | |
| nn.BatchNorm2d(ndf * ndf_mult), | |
| nn.LeakyReLU(0.2, True) | |
| ] | |
| ndf_mult_prev = ndf_mult | |
| ndf_mult = min(2**n_layers, 8) | |
| layers += [ | |
| nn.Conv2d( | |
| ndf * ndf_mult_prev, | |
| ndf * ndf_mult, | |
| kernel_size=4, | |
| stride=1, | |
| padding=1, | |
| bias=False), | |
| nn.BatchNorm2d(ndf * ndf_mult), | |
| nn.LeakyReLU(0.2, True) | |
| ] | |
| layers += [ | |
| nn.Conv2d(ndf * ndf_mult, 1, kernel_size=4, stride=1, padding=1) | |
| ] # output 1 channel prediction map | |
| self.main = nn.Sequential(*layers) | |
| def forward(self, x): | |
| return self.main(x) | |