Spaces:
Build error
Build error
| # ------------------------------------------------------------------------------------ | |
| # Modified from VQGAN (https://github.com/CompVis/taming-transformers) | |
| # Copyright (c) 2020 Patrick Esser and Robin Rombach and Björn Ommer. All Rights Reserved. | |
| # ------------------------------------------------------------------------------------ | |
| import torch | |
| import torch.nn as nn | |
| from typing import List, Tuple, Optional | |
| from einops import rearrange | |
| from omegaconf import OmegaConf | |
| from .layers import Encoder, Decoder | |
| class VectorQuantizer(nn.Module): | |
| """ | |
| Simplified VectorQuantizer in the original VQGAN repository | |
| by removing unncessary modules for sampling | |
| """ | |
| def __init__(self, dim: int, n_embed: int, beta: float) -> None: | |
| super().__init__() | |
| self.n_embed = n_embed | |
| self.dim = dim | |
| self.beta = beta | |
| self.embedding = nn.Embedding(self.n_embed, self.dim) | |
| self.embedding.weight.data.uniform_(-1.0 / self.n_embed, 1.0 / self.n_embed) | |
| def forward(self, | |
| z: torch.FloatTensor) -> Tuple[torch.FloatTensor, torch.LongTensor]: | |
| z = rearrange(z, 'b c h w -> b h w c').contiguous() # [B,C,H,W] -> [B,H,W,C] | |
| z_flattened = z.view(-1, self.dim) | |
| 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) | |
| return z_q, min_encoding_indices | |
| def get_codebook_entry(self, | |
| indices: torch.LongTensor, | |
| shape: Optional[List[int]] = None) -> torch.FloatTensor: | |
| z_q = self.embedding(indices) | |
| if shape is not None: | |
| z_q = z_q.view(shape) | |
| z_q = z_q.permute(0, 3, 1, 2).contiguous() | |
| return z_q | |
| class VQGAN(nn.Module): | |
| def __init__(self, n_embed: int, embed_dim: int, hparams: OmegaConf) -> None: | |
| super().__init__() | |
| self.encoder = Encoder(**hparams) | |
| self.decoder = Decoder(**hparams) | |
| self.quantize = VectorQuantizer(dim=embed_dim, n_embed=n_embed, beta=0.25) | |
| self.quant_conv = torch.nn.Conv2d(hparams.z_channels, embed_dim, 1) | |
| self.post_quant_conv = torch.nn.Conv2d(embed_dim, hparams.z_channels, 1) | |
| self.latent_dim = hparams.attn_resolutions[0] | |
| def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: | |
| quant = self.encode(x) | |
| dec = self.decode(quant) | |
| return dec | |
| def encode(self, x: torch.FloatTensor) -> torch.FloatTensor: | |
| h = self.encoder(x) | |
| h = self.quant_conv(h) | |
| quant = self.quantize(h)[0] | |
| quant = rearrange(quant, 'b h w c -> b c h w').contiguous() | |
| return quant | |
| def decode(self, quant: torch.FloatTensor) -> torch.FloatTensor: | |
| quant = self.post_quant_conv(quant) | |
| dec = self.decoder(quant) | |
| return dec | |
| def decode_code(self, code: torch.LongTensor) -> torch.FloatTensor: | |
| quant = self.quantize.get_codebook_entry(code) | |
| quant = quant.permute(0, 3, 1, 2) | |
| dec = self.decode(quant) | |
| return dec | |
| def get_codes(self, x: torch.FloatTensor) -> torch.LongTensor: | |
| h = self.encoder(x) | |
| h = self.quant_conv(h) | |
| codes = self.quantize(h)[1].view(x.shape[0], self.latent_dim ** 2) | |
| return codes | |
| def from_ckpt(self, path: str, strict: bool = True) -> None: | |
| ckpt = torch.load(path, map_location='cpu')['state_dict'] | |
| self.load_state_dict(ckpt, strict=strict) | |
| print(f'{path} successfully restored..') | |