| """AutoencoderKL implementation compatible with diffusers weights.""" |
|
|
| |
| from dataclasses import dataclass |
| from typing import Optional, Tuple |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| @dataclass |
| class AutoencoderKLOutput: |
| sample: torch.Tensor |
|
|
|
|
| class AutoencoderConfig: |
| def __init__(self, **kwargs): |
| self.__dict__.update(kwargs) |
|
|
| def get(self, key, default=None): |
| return self.__dict__.get(key, default) |
|
|
| def __getattr__(self, name): |
| return self.__dict__.get(name) |
|
|
|
|
| def swish(x): |
| return x * torch.sigmoid(x) |
|
|
|
|
| class ResnetBlock2D(nn.Module): |
| def __init__(self, in_channels, out_channels=None, dropout=0.0, temb_channels=512, groups=32, eps=1e-6): |
| super().__init__() |
| out_channels = out_channels or in_channels |
| self.in_channels = in_channels |
| self.out_channels = out_channels |
|
|
| self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) |
| self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) |
| self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True) |
| self.dropout = nn.Dropout(dropout) |
| self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) |
|
|
| self.nonlinearity = swish |
|
|
| if self.in_channels != self.out_channels: |
| self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) |
| else: |
| self.conv_shortcut = None |
|
|
| def forward(self, input_tensor, temb=None): |
| hidden_states = input_tensor |
| hidden_states = self.norm1(hidden_states) |
| hidden_states = self.nonlinearity(hidden_states) |
| hidden_states = self.conv1(hidden_states) |
|
|
| hidden_states = self.norm2(hidden_states) |
| hidden_states = self.nonlinearity(hidden_states) |
| hidden_states = self.dropout(hidden_states) |
| hidden_states = self.conv2(hidden_states) |
|
|
| if self.conv_shortcut is not None: |
| input_tensor = self.conv_shortcut(input_tensor) |
|
|
| output_tensor = (input_tensor + hidden_states) / 1.0 |
| return output_tensor |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, in_channels, heads=1, dim_head=None, groups=32, eps=1e-6): |
| super().__init__() |
| self.heads = heads |
| self.in_channels = in_channels |
| self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) |
|
|
| self.to_q = nn.Linear(in_channels, in_channels) |
| self.to_k = nn.Linear(in_channels, in_channels) |
| self.to_v = nn.Linear(in_channels, in_channels) |
| self.to_out = nn.ModuleList([nn.Linear(in_channels, in_channels)]) |
|
|
| def forward(self, hidden_states): |
| b, c, h, w = hidden_states.shape |
| residual = hidden_states |
| hidden_states = self.group_norm(hidden_states) |
| hidden_states = hidden_states.view(b, c, -1).transpose(1, 2) |
|
|
| query = self.to_q(hidden_states) |
| key = self.to_k(hidden_states) |
| value = self.to_v(hidden_states) |
|
|
| import torch.nn.functional as F |
|
|
| hidden_states = F.scaled_dot_product_attention(query, key, value) |
|
|
| hidden_states = self.to_out[0](hidden_states) |
| hidden_states = hidden_states.transpose(1, 2).view(b, c, h, w) |
|
|
| return residual + hidden_states |
|
|
|
|
| class Downsample2D(nn.Module): |
| def __init__(self, channels, with_conv=True, out_channels=None, padding=1): |
| super().__init__() |
| out_channels = out_channels or channels |
| self.with_conv = with_conv |
| if with_conv: |
| self.conv = nn.Conv2d(channels, out_channels, kernel_size=3, stride=2, padding=padding) |
|
|
| def forward(self, hidden_states): |
| if self.with_conv: |
| return self.conv(hidden_states) |
| else: |
| return torch.nn.functional.avg_pool2d(hidden_states, kernel_size=2, stride=2) |
|
|
|
|
| class Upsample2D(nn.Module): |
| def __init__(self, channels, with_conv=True, out_channels=None): |
| super().__init__() |
| out_channels = out_channels or channels |
| self.with_conv = with_conv |
| if with_conv: |
| self.conv = nn.Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1) |
|
|
| def forward(self, hidden_states): |
| hidden_states = torch.nn.functional.interpolate(hidden_states, scale_factor=2.0, mode="nearest") |
| if self.with_conv: |
| hidden_states = self.conv(hidden_states) |
| return hidden_states |
|
|
|
|
| class DownEncoderBlock2D(nn.Module): |
| def __init__(self, in_channels, out_channels, num_layers=1, resnet_eps=1e-6, resnet_groups=32, add_downsample=True): |
| super().__init__() |
| resnets = [] |
| for i in range(num_layers): |
| in_c = in_channels if i == 0 else out_channels |
| resnets.append(ResnetBlock2D(in_c, out_channels, eps=resnet_eps, groups=resnet_groups)) |
| self.resnets = nn.ModuleList(resnets) |
|
|
| if add_downsample: |
| self.downsamplers = nn.ModuleList( |
| [Downsample2D(out_channels, with_conv=True, out_channels=out_channels, padding=0)] |
| ) |
| else: |
| self.downsamplers = None |
|
|
| def forward(self, hidden_states): |
| for resnet in self.resnets: |
| hidden_states = resnet(hidden_states) |
|
|
| if self.downsamplers is not None: |
| for downsampler in self.downsamplers: |
| pad = (0, 1, 0, 1) |
| hidden_states = torch.nn.functional.pad(hidden_states, pad, mode="constant", value=0) |
| hidden_states = downsampler(hidden_states) |
|
|
| return hidden_states |
|
|
|
|
| class UpDecoderBlock2D(nn.Module): |
| def __init__(self, in_channels, out_channels, num_layers=1, resnet_eps=1e-6, resnet_groups=32, add_upsample=True): |
| super().__init__() |
| resnets = [] |
| for i in range(num_layers): |
| in_c = in_channels if i == 0 else out_channels |
| resnets.append(ResnetBlock2D(in_c, out_channels, eps=resnet_eps, groups=resnet_groups)) |
| self.resnets = nn.ModuleList(resnets) |
|
|
| if add_upsample: |
| self.upsamplers = nn.ModuleList([Upsample2D(out_channels, with_conv=True, out_channels=out_channels)]) |
| else: |
| self.upsamplers = None |
|
|
| def forward(self, hidden_states): |
| for resnet in self.resnets: |
| hidden_states = resnet(hidden_states) |
|
|
| if self.upsamplers is not None: |
| for upsampler in self.upsamplers: |
| hidden_states = upsampler(hidden_states) |
|
|
| return hidden_states |
|
|
|
|
| class UNetMidBlock2D(nn.Module): |
| def __init__(self, in_channels, resnet_eps=1e-6, resnet_groups=32, attention_head_dim=None): |
| super().__init__() |
| self.resnets = nn.ModuleList( |
| [ |
| ResnetBlock2D(in_channels, in_channels, eps=resnet_eps, groups=resnet_groups), |
| ResnetBlock2D(in_channels, in_channels, eps=resnet_eps, groups=resnet_groups), |
| ] |
| ) |
| self.attentions = nn.ModuleList([Attention(in_channels, heads=1, groups=resnet_groups, eps=resnet_eps)]) |
|
|
| def forward(self, hidden_states): |
| hidden_states = self.resnets[0](hidden_states) |
| for attn in self.attentions: |
| hidden_states = attn(hidden_states) |
| hidden_states = self.resnets[1](hidden_states) |
| return hidden_states |
|
|
|
|
| class Encoder(nn.Module): |
| def __init__( |
| self, |
| in_channels=3, |
| out_channels=3, |
| block_out_channels=(64,), |
| layers_per_block=2, |
| norm_num_groups=32, |
| double_z=True, |
| ): |
| super().__init__() |
| self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) |
|
|
| self.down_blocks = nn.ModuleList([]) |
| output_channel = block_out_channels[0] |
| for i, block_out_channel in enumerate(block_out_channels): |
| input_channel = output_channel |
| output_channel = block_out_channel |
| is_final_block = i == len(block_out_channels) - 1 |
|
|
| block = DownEncoderBlock2D( |
| input_channel, |
| output_channel, |
| num_layers=layers_per_block, |
| resnet_groups=norm_num_groups, |
| add_downsample=not is_final_block, |
| ) |
| self.down_blocks.append(block) |
|
|
| self.mid_block = UNetMidBlock2D( |
| block_out_channels[-1], |
| resnet_groups=norm_num_groups, |
| ) |
|
|
| self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) |
| self.conv_act = nn.SiLU() |
|
|
| conv_out_channels = 2 * out_channels if double_z else out_channels |
| self.conv_out = nn.Conv2d(block_out_channels[-1], conv_out_channels, 3, padding=1) |
|
|
| def forward(self, x): |
| x = self.conv_in(x) |
| for block in self.down_blocks: |
| x = block(x) |
| x = self.mid_block(x) |
| x = self.conv_norm_out(x) |
| x = self.conv_act(x) |
| x = self.conv_out(x) |
| return x |
|
|
|
|
| class Decoder(nn.Module): |
| def __init__( |
| self, |
| in_channels=3, |
| out_channels=3, |
| block_out_channels=(64,), |
| layers_per_block=2, |
| norm_num_groups=32, |
| ): |
| super().__init__() |
| self.conv_in = nn.Conv2d(in_channels, block_out_channels[-1], kernel_size=3, stride=1, padding=1) |
|
|
| self.mid_block = UNetMidBlock2D( |
| block_out_channels[-1], |
| resnet_groups=norm_num_groups, |
| ) |
|
|
| self.up_blocks = nn.ModuleList([]) |
| reversed_block_out_channels = list(reversed(block_out_channels)) |
| output_channel = reversed_block_out_channels[0] |
|
|
| for i, block_out_channel in enumerate(reversed_block_out_channels): |
| input_channel = output_channel |
| output_channel = block_out_channel |
| is_final_block = i == len(block_out_channels) - 1 |
| block = UpDecoderBlock2D( |
| input_channel, |
| output_channel, |
| num_layers=layers_per_block + 1, |
| resnet_groups=norm_num_groups, |
| add_upsample=not is_final_block, |
| ) |
| self.up_blocks.append(block) |
|
|
| self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) |
| self.conv_act = nn.SiLU() |
| self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, padding=1) |
|
|
| def forward(self, x): |
| x = self.conv_in(x) |
| x = self.mid_block(x) |
| for block in self.up_blocks: |
| x = block(x) |
| x = self.conv_norm_out(x) |
| x = self.conv_act(x) |
| x = self.conv_out(x) |
| return x |
|
|
|
|
| class AutoencoderKL(nn.Module): |
| def __init__( |
| self, |
| in_channels: int = 3, |
| out_channels: int = 3, |
| down_block_types: Tuple[str] = ("DownEncoderBlock2D",), |
| up_block_types: Tuple[str] = ("UpDecoderBlock2D",), |
| block_out_channels: Tuple[int] = (64,), |
| layers_per_block: int = 1, |
| act_fn: str = "silu", |
| latent_channels: int = 4, |
| norm_num_groups: int = 32, |
| sample_size: int = 32, |
| scaling_factor: float = 0.18215, |
| shift_factor: Optional[float] = None, |
| force_upcast: bool = True, |
| use_quant_conv: bool = True, |
| use_post_quant_conv: bool = True, |
| mid_block_add_attention: bool = True, |
| **kwargs, |
| ): |
| super().__init__() |
| self.config = AutoencoderConfig( |
| in_channels=in_channels, |
| out_channels=out_channels, |
| block_out_channels=block_out_channels, |
| layers_per_block=layers_per_block, |
| latent_channels=latent_channels, |
| scaling_factor=scaling_factor, |
| shift_factor=shift_factor, |
| ) |
|
|
| self.encoder = Encoder( |
| in_channels=in_channels, |
| out_channels=latent_channels, |
| block_out_channels=block_out_channels, |
| layers_per_block=layers_per_block, |
| norm_num_groups=norm_num_groups, |
| double_z=True, |
| ) |
|
|
| self.decoder = Decoder( |
| in_channels=latent_channels, |
| out_channels=out_channels, |
| block_out_channels=block_out_channels, |
| layers_per_block=layers_per_block, |
| norm_num_groups=norm_num_groups, |
| ) |
|
|
| self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None |
| self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None |
|
|
| @property |
| def dtype(self): |
| return next(self.parameters()).dtype |
|
|
| def decode(self, z: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput: |
| if self.post_quant_conv is not None: |
| z = self.post_quant_conv(z) |
|
|
| dec = self.decoder(z) |
|
|
| if not return_dict: |
| return (dec,) |
|
|
| return AutoencoderKLOutput(sample=dec) |
|
|
