390 lines
13 KiB
Python
390 lines
13 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from einops import repeat
|
|
from torch.utils.checkpoint import checkpoint
|
|
|
|
from scepter.modules.model.backbone.autoencoder.ae_utils import (
|
|
XFORMERS_IS_AVAILBLE, AttnBlock, Downsample, MemoryEfficientAttention,
|
|
Normalize, ResnetBlock, Upsample, nonlinearity)
|
|
from scepter.modules.model.base_model import BaseModel
|
|
from scepter.modules.model.registry import BACKBONES
|
|
from scepter.modules.utils.config import dict_to_yaml
|
|
|
|
|
|
@BACKBONES.register_class()
|
|
class Encoder(BaseModel):
|
|
para_dict = {
|
|
'CH': {
|
|
'value': 128,
|
|
'description': ''
|
|
},
|
|
'NUM_RES_BLOCKS': {
|
|
'value': 2,
|
|
'description': ''
|
|
},
|
|
'IN_CHANNELS': {
|
|
'value': 3,
|
|
'description': ''
|
|
},
|
|
'ATTN_RESOLUTIONS': {
|
|
'value': [],
|
|
'description': ''
|
|
},
|
|
'CH_MULT': {
|
|
'value': [1, 2, 4, 4],
|
|
'description': ''
|
|
},
|
|
'Z_CHANNELS': {
|
|
'value': 4,
|
|
'description': ''
|
|
},
|
|
'DOUBLE_Z': {
|
|
'value': True,
|
|
'description': ''
|
|
},
|
|
'DROPOUT': {
|
|
'value': 0.0,
|
|
'description': ''
|
|
},
|
|
'RESAMP_WITH_CONV': {
|
|
'value': True,
|
|
'description': ''
|
|
}
|
|
}
|
|
para_dict.update(BaseModel.para_dict)
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super().__init__(cfg, logger=logger)
|
|
self.use_checkpoint = cfg.get('USE_CHECKPOINT', False)
|
|
self.ch = cfg.CH
|
|
self.out_ch = cfg.OUT_CH
|
|
self.num_res_blocks = cfg.NUM_RES_BLOCKS
|
|
self.in_channels = cfg.IN_CHANNELS
|
|
self.attn_resolutions = cfg.ATTN_RESOLUTIONS
|
|
self.ch_mult = tuple(cfg.get('CH_MULT', [1, 2, 4, 8]))
|
|
self.z_channels = cfg.Z_CHANNELS
|
|
self.double_z = cfg.get('DOUBLE_Z', True)
|
|
self.dropout = cfg.get('DROPOUT', 0.0)
|
|
self.resamp_with_conv = cfg.get('RESAMP_WITH_CONV', True)
|
|
self.temb_ch = 0
|
|
self.construct_model()
|
|
|
|
def construct_model(self):
|
|
self.num_resolutions = len(self.ch_mult)
|
|
self.logger.info(
|
|
f'AE Module XFORMERS_IS_AVAILBLE : {XFORMERS_IS_AVAILBLE}')
|
|
AttentionBuilder = MemoryEfficientAttention if XFORMERS_IS_AVAILBLE else AttnBlock
|
|
|
|
self.conv_in = torch.nn.Conv2d(self.in_channels,
|
|
self.ch,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1)
|
|
|
|
curr_res = 2**(self.num_resolutions - 1)
|
|
in_ch_mult = (1, ) + tuple(self.ch_mult)
|
|
self.in_ch_mult = in_ch_mult
|
|
self.down = nn.ModuleList()
|
|
for i_level in range(self.num_resolutions):
|
|
block = nn.ModuleList()
|
|
attn = nn.ModuleList()
|
|
block_in = self.ch * in_ch_mult[i_level]
|
|
block_out = self.ch * self.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=self.dropout))
|
|
block_in = block_out
|
|
if curr_res in self.attn_resolutions:
|
|
attn.append(AttentionBuilder(block_in))
|
|
down = nn.Module()
|
|
down.block = block
|
|
down.attn = attn
|
|
if i_level != self.num_resolutions - 1:
|
|
down.downsample = Downsample(block_in, self.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=self.dropout)
|
|
self.mid.attn_1 = AttentionBuilder(block_in)
|
|
self.mid.block_2 = ResnetBlock(in_channels=block_in,
|
|
out_channels=block_in,
|
|
temb_channels=self.temb_ch,
|
|
dropout=self.dropout)
|
|
|
|
# end
|
|
self.norm_out = Normalize(block_in)
|
|
self.conv_out = torch.nn.Conv2d(
|
|
block_in,
|
|
2 * self.z_channels if self.double_z else self.z_channels,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1)
|
|
|
|
def forward_ori(self, x):
|
|
# 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
|
|
|
|
def forward(self, x):
|
|
if self.use_checkpoint:
|
|
return checkpoint(self.forward_ori, x)
|
|
else:
|
|
return self.forward_ori(x)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('BACKBONE',
|
|
__class__.__name__,
|
|
Encoder.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@BACKBONES.register_class()
|
|
class Decoder(BaseModel):
|
|
para_dict = {
|
|
'CH': {
|
|
'value': 128,
|
|
'description': ''
|
|
},
|
|
'OUT_CH': {
|
|
'value': 3,
|
|
'description': ''
|
|
},
|
|
'NUM_RES_BLOCKS': {
|
|
'value': 2,
|
|
'description': ''
|
|
},
|
|
'IN_CHANNELS': {
|
|
'value': 3,
|
|
'description': ''
|
|
},
|
|
'ATTN_RESOLUTIONS': {
|
|
'value': [],
|
|
'description': ''
|
|
},
|
|
'CH_MULT': {
|
|
'value': [1, 2, 4, 4],
|
|
'description': ''
|
|
},
|
|
'Z_CHANNELS': {
|
|
'value': 4,
|
|
'description': ''
|
|
},
|
|
'DROPOUT': {
|
|
'value': 0.0,
|
|
'description': ''
|
|
},
|
|
'RESAMP_WITH_CONV': {
|
|
'value': True,
|
|
'description': ''
|
|
},
|
|
'GIVE_PRE_END': {
|
|
'value': False,
|
|
'description': ''
|
|
},
|
|
'TANH_OUT': {
|
|
'value': False,
|
|
'description': ''
|
|
}
|
|
}
|
|
para_dict.update(BaseModel.para_dict)
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super().__init__(cfg, logger=logger)
|
|
self.use_checkpoint = cfg.get('USE_CHECKPOINT', False)
|
|
self.ch = cfg.CH
|
|
self.out_ch = cfg.OUT_CH
|
|
self.num_res_blocks = cfg.NUM_RES_BLOCKS
|
|
self.in_channels = cfg.IN_CHANNELS
|
|
self.attn_resolutions = cfg.ATTN_RESOLUTIONS
|
|
self.ch_mult = tuple(cfg.get('CH_MULT', [1, 2, 4, 8]))
|
|
self.z_channels = cfg.Z_CHANNELS
|
|
self.dropout = cfg.get('DROPOUT', 0.0)
|
|
self.resamp_with_conv = cfg.get('RESAMP_WITH_CONV', True)
|
|
self.give_pre_end = cfg.get('GIVE_PRE_END', False)
|
|
self.tanh_out = cfg.get('TANH_OUT', False)
|
|
self.temb_ch = 0
|
|
self.construct_model()
|
|
|
|
def construct_model(self):
|
|
self.num_resolutions = len(self.ch_mult)
|
|
AttentionBuilder = MemoryEfficientAttention if XFORMERS_IS_AVAILBLE else AttnBlock
|
|
|
|
# compute in_ch_mult, block_in and curr_res at lowest res
|
|
block_in = self.ch * self.ch_mult[self.num_resolutions - 1]
|
|
self.block_in = block_in
|
|
curr_res = 1
|
|
# z to block_in
|
|
self.conv_in = torch.nn.Conv2d(self.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=self.dropout)
|
|
self.mid.attn_1 = AttentionBuilder(block_in)
|
|
self.mid.block_2 = ResnetBlock(in_channels=block_in,
|
|
out_channels=block_in,
|
|
temb_channels=self.temb_ch,
|
|
dropout=self.dropout)
|
|
|
|
# upsampling
|
|
self.up = nn.ModuleList()
|
|
for i_level in reversed(range(self.num_resolutions)):
|
|
block = nn.ModuleList()
|
|
attn = nn.ModuleList()
|
|
block_out = self.ch * self.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=self.dropout))
|
|
block_in = block_out
|
|
if curr_res in self.attn_resolutions:
|
|
attn.append(AttentionBuilder(block_in))
|
|
up = nn.Module()
|
|
up.block = block
|
|
up.attn = attn
|
|
if i_level != 0:
|
|
up.upsample = Upsample(block_in, self.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,
|
|
self.out_ch,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1)
|
|
|
|
def mid_upsclae_transform(self, h, temb):
|
|
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)
|
|
|
|
return h
|
|
|
|
def forward(self, z, cond=None):
|
|
# timestep embedding
|
|
temb = None
|
|
|
|
h = self.conv_in(z)
|
|
|
|
# middle
|
|
if not self.use_checkpoint:
|
|
h = self.mid_upsclae_transform(h, temb)
|
|
else:
|
|
h = checkpoint(self.mid_upsclae_transform, h, temb)
|
|
|
|
# end
|
|
if self.give_pre_end:
|
|
return h
|
|
|
|
h = self.norm_out(h)
|
|
h = nonlinearity(h)
|
|
h = self.conv_out(h)
|
|
if self.tanh_out:
|
|
h = torch.tanh(h)
|
|
return h
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('BACKBONE',
|
|
__class__.__name__,
|
|
Decoder.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@BACKBONES.register_class()
|
|
class RDecoder(Decoder):
|
|
def construct_model(self):
|
|
super().construct_model()
|
|
self.resize_level = nn.Sequential(
|
|
nn.Linear(self.block_in, self.block_in),
|
|
nn.SiLU(),
|
|
nn.Linear(self.block_in, self.block_in),
|
|
)
|
|
|
|
def forward(self, z, rembed=None):
|
|
# timestep embedding
|
|
temb = None
|
|
h = self.conv_in(z)
|
|
bs, channel, hdim, wdim = h.size()
|
|
if rembed is not None:
|
|
rembed = self.resize_level(rembed)
|
|
rembed = repeat(rembed, 'b e-> b e hd wd', hd=hdim, wd=wdim)
|
|
h = h + rembed
|
|
|
|
# middle
|
|
if not self.use_checkpoint:
|
|
h = self.mid_upsclae_transform(h, temb)
|
|
else:
|
|
h = checkpoint(self.mid_upsclae_transform, h, temb)
|
|
|
|
# end
|
|
if self.give_pre_end:
|
|
return h
|
|
|
|
h = self.norm_out(h)
|
|
h = nonlinearity(h)
|
|
h = self.conv_out(h)
|
|
if self.tanh_out:
|
|
h = torch.tanh(h)
|
|
return h
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('BACKBONE',
|
|
__class__.__name__,
|
|
Decoder.para_dict,
|
|
set_name=True)
|