# from einops._torch_specific import allow_ops_in_compiled_graph
# allow_ops_in_compiled_graph()
import einops
import torch
import torch as th
import torch.nn as nn
from einops import rearrange, repeat
from sgm.modules.diffusionmodules.util import (
avg_pool_nd,
checkpoint,
conv_nd,
linear,
normalization,
timestep_embedding,
zero_module,
)
from sgm.modules.diffusionmodules.openaimodel import Downsample, Upsample, UNetModel, Timestep, \
TimestepEmbedSequential, ResBlock, AttentionBlock, TimestepBlock
from sgm.modules.attention import SpatialTransformer, MemoryEfficientCrossAttention, CrossAttention
from sgm.util import default, log_txt_as_img, exists, instantiate_from_config
import re
import torch
from functools import partial
try:
import xformers
import xformers.ops
XFORMERS_IS_AVAILBLE = True
except:
XFORMERS_IS_AVAILBLE = False
# dummy replace
def convert_module_to_f16(x):
pass
def convert_module_to_f32(x):
pass
class ZeroConv(nn.Module):
def __init__(self, label_nc, norm_nc, mask=False):
super().__init__()
self.zero_conv = zero_module(conv_nd(2, label_nc, norm_nc, 1, 1, 0))
self.mask = mask
def forward(self, c, h, h_ori=None):
# with torch.cuda.amp.autocast(enabled=False, dtype=torch.float32):
if not self.mask:
h = h + self.zero_conv(c)
else:
h = h + self.zero_conv(c) * torch.zeros_like(h)
if h_ori is not None:
h = th.cat([h_ori, h], dim=1)
return h
class ZeroSFT(nn.Module):
def __init__(self, label_nc, norm_nc, concat_channels=0, norm=True, mask=False):
super().__init__()
# param_free_norm_type = str(parsed.group(1))
ks = 3
pw = ks // 2
self.norm = norm
if self.norm:
self.param_free_norm = normalization(norm_nc + concat_channels)
else:
self.param_free_norm = nn.Identity()
nhidden = 128
self.mlp_shared = nn.Sequential(
nn.Conv2d(label_nc, nhidden, kernel_size=ks, padding=pw),
nn.SiLU()
)
self.zero_mul = zero_module(nn.Conv2d(nhidden, norm_nc + concat_channels, kernel_size=ks, padding=pw))
self.zero_add = zero_module(nn.Conv2d(nhidden, norm_nc + concat_channels, kernel_size=ks, padding=pw))
# self.zero_mul = nn.Conv2d(nhidden, norm_nc + concat_channels, kernel_size=ks, padding=pw)
# self.zero_add = nn.Conv2d(nhidden, norm_nc + concat_channels, kernel_size=ks, padding=pw)
self.zero_conv = zero_module(conv_nd(2, label_nc, norm_nc, 1, 1, 0))
self.pre_concat = bool(concat_channels != 0)
self.mask = mask
def forward(self, c, h, h_ori=None, control_scale=1):
assert self.mask is False
if h_ori is not None and self.pre_concat:
h_raw = th.cat([h_ori, h], dim=1)
else:
h_raw = h
if self.mask:
h = h + self.zero_conv(c) * torch.zeros_like(h)
else:
h = h + self.zero_conv(c)
if h_ori is not None and self.pre_concat:
h = th.cat([h_ori, h], dim=1)
actv = self.mlp_shared(c)
gamma = self.zero_mul(actv)
beta = self.zero_add(actv)
if self.mask:
gamma = gamma * torch.zeros_like(gamma)
beta = beta * torch.zeros_like(beta)
h = self.param_free_norm(h) * (gamma + 1) + beta
if h_ori is not None and not self.pre_concat:
h = th.cat([h_ori, h], dim=1)
return h * control_scale + h_raw * (1 - control_scale)
class ZeroCrossAttn(nn.Module):
ATTENTION_MODES = {
"softmax": CrossAttention, # vanilla attention
"softmax-xformers": MemoryEfficientCrossAttention
}
def __init__(self, context_dim, query_dim, zero_out=True, mask=False):
super().__init__()
attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILBLE else "softmax"
assert attn_mode in self.ATTENTION_MODES
attn_cls = self.ATTENTION_MODES[attn_mode]
self.attn = attn_cls(query_dim=query_dim, context_dim=context_dim, heads=query_dim//64, dim_head=64)
self.norm1 = normalization(query_dim)
self.norm2 = normalization(context_dim)
self.mask = mask
# if zero_out:
# # for p in self.attn.to_out.parameters():
# # p.detach().zero_()
# self.attn.to_out = zero_module(self.attn.to_out)
def forward(self, context, x, control_scale=1):
assert self.mask is False
x_in = x
x = self.norm1(x)
context = self.norm2(context)
b, c, h, w = x.shape
x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
context = rearrange(context, 'b c h w -> b (h w) c').contiguous()
x = self.attn(x, context)
x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
if self.mask:
x = x * torch.zeros_like(x)
x = x_in + x * control_scale
return x
class GLVControl(nn.Module):
def __init__(
self,
in_channels,
model_channels,
out_channels,
num_res_blocks,
attention_resolutions,
dropout=0,
channel_mult=(1, 2, 4, 8),
conv_resample=True,
dims=2,
num_classes=None,
use_checkpoint=False,
use_fp16=False,
num_heads=-1,
num_head_channels=-1,
num_heads_upsample=-1,
use_scale_shift_norm=False,
resblock_updown=False,
use_new_attention_order=False,
use_spatial_transformer=False, # custom transformer support
transformer_depth=1, # custom transformer support
context_dim=None, # custom transformer support
n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
legacy=True,
disable_self_attentions=None,
num_attention_blocks=None,
disable_middle_self_attn=False,
use_linear_in_transformer=False,
spatial_transformer_attn_type="softmax",
adm_in_channels=None,
use_fairscale_checkpoint=False,
offload_to_cpu=False,
transformer_depth_middle=None,
input_upscale=1,
):
super().__init__()
from omegaconf.listconfig import ListConfig
if use_spatial_transformer:
assert (
context_dim is not None
), "Fool!! You forgot to include the dimension of your cross-attention conditioning..."
if context_dim is not None:
assert (
use_spatial_transformer
), "Fool!! You forgot to use the spatial transformer for your cross-attention conditioning..."
if type(context_dim) == ListConfig:
context_dim = list(context_dim)
if num_heads_upsample == -1:
num_heads_upsample = num_heads
if num_heads == -1:
assert (
num_head_channels != -1
), "Either num_heads or num_head_channels has to be set"
if num_head_channels == -1:
assert (
num_heads != -1
), "Either num_heads or num_head_channels has to be set"
self.in_channels = in_channels
self.model_channels = model_channels
self.out_channels = out_channels
if isinstance(transformer_depth, int):
transformer_depth = len(channel_mult) * [transformer_depth]
elif isinstance(transformer_depth, ListConfig):
transformer_depth = list(transformer_depth)
transformer_depth_middle = default(
transformer_depth_middle, transformer_depth[-1]
)
if isinstance(num_res_blocks, int):
self.num_res_blocks = len(channel_mult) * [num_res_blocks]
else:
if len(num_res_blocks) != len(channel_mult):
raise ValueError(
"provide num_res_blocks either as an int (globally constant) or "
"as a list/tuple (per-level) with the same length as channel_mult"
)
self.num_res_blocks = num_res_blocks
# self.num_res_blocks = num_res_blocks
if disable_self_attentions is not None:
# should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
assert len(disable_self_attentions) == len(channel_mult)
if num_attention_blocks is not None:
assert len(num_attention_blocks) == len(self.num_res_blocks)
assert all(
map(
lambda i: self.num_res_blocks[i] >= num_attention_blocks[i],
range(len(num_attention_blocks)),
)
)
print(
f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
f"attention will still not be set."
) # todo: convert to warning
self.attention_resolutions = attention_resolutions
self.dropout = dropout
self.channel_mult = channel_mult
self.conv_resample = conv_resample
self.num_classes = num_classes
self.use_checkpoint = use_checkpoint
if use_fp16:
print("WARNING: use_fp16 was dropped and has no effect anymore.")
# self.dtype = th.float16 if use_fp16 else th.float32
self.num_heads = num_heads
self.num_head_channels = num_head_channels
self.num_heads_upsample = num_heads_upsample
self.predict_codebook_ids = n_embed is not None
assert use_fairscale_checkpoint != use_checkpoint or not (
use_checkpoint or use_fairscale_checkpoint
)
self.use_fairscale_checkpoint = False
checkpoint_wrapper_fn = (
partial(checkpoint_wrapper, offload_to_cpu=offload_to_cpu)
if self.use_fairscale_checkpoint
else lambda x: x
)
time_embed_dim = model_channels * 4
self.time_embed = checkpoint_wrapper_fn(
nn.Sequential(
linear(model_channels, time_embed_dim),
nn.SiLU(),
linear(time_embed_dim, time_embed_dim),
)
)
if self.num_classes is not None:
if isinstance(self.num_classes, int):
self.label_emb = nn.Embedding(num_classes, time_embed_dim)
elif self.num_classes == "continuous":
print("setting up linear c_adm embedding layer")
self.label_emb = nn.Linear(1, time_embed_dim)
elif self.num_classes == "timestep":
self.label_emb = checkpoint_wrapper_fn(
nn.Sequential(
Timestep(model_channels),
nn.Sequential(
linear(model_channels, time_embed_dim),
nn.SiLU(),
linear(time_embed_dim, time_embed_dim),
),
)
)
elif self.num_classes == "sequential":
assert adm_in_channels is not None
self.label_emb = nn.Sequential(
nn.Sequential(
linear(adm_in_channels, time_embed_dim),
nn.SiLU(),
linear(time_embed_dim, time_embed_dim),
)
)
else:
raise ValueError()
self.input_blocks = nn.ModuleList(
[
TimestepEmbedSequential(
conv_nd(dims, in_channels, model_channels, 3, padding=1)
)
]
)
self._feature_size = model_channels
input_block_chans = [model_channels]
ch = model_channels
ds = 1
for level, mult in enumerate(channel_mult):
for nr in range(self.num_res_blocks[level]):
layers = [
checkpoint_wrapper_fn(
ResBlock(
ch,
time_embed_dim,
dropout,
out_channels=mult * model_channels,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
)
)
]
ch = mult * model_channels
if ds in attention_resolutions:
if num_head_channels == -1:
dim_head = ch // num_heads
else:
num_heads = ch // num_head_channels
dim_head = num_head_channels
if legacy:
# num_heads = 1
dim_head = (
ch // num_heads
if use_spatial_transformer
else num_head_channels
)
if exists(disable_self_attentions):
disabled_sa = disable_self_attentions[level]
else:
disabled_sa = False
if (
not exists(num_attention_blocks)
or nr < num_attention_blocks[level]
):
layers.append(
checkpoint_wrapper_fn(
AttentionBlock(
ch,
use_checkpoint=use_checkpoint,
num_heads=num_heads,
num_head_channels=dim_head,
use_new_attention_order=use_new_attention_order,
)
)
if not use_spatial_transformer
else checkpoint_wrapper_fn(
SpatialTransformer(
ch,
num_heads,
dim_head,
depth=transformer_depth[level],
context_dim=context_dim,
disable_self_attn=disabled_sa,
use_linear=use_linear_in_transformer,
attn_type=spatial_transformer_attn_type,
use_checkpoint=use_checkpoint,
)
)
)
self.input_blocks.append(TimestepEmbedSequential(*layers))
self._feature_size += ch
input_block_chans.append(ch)
if level != len(channel_mult) - 1:
out_ch = ch
self.input_blocks.append(
TimestepEmbedSequential(
checkpoint_wrapper_fn(
ResBlock(
ch,
time_embed_dim,
dropout,
out_channels=out_ch,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
down=True,
)
)
if resblock_updown
else Downsample(
ch, conv_resample, dims=dims, out_channels=out_ch
)
)
)
ch = out_ch
input_block_chans.append(ch)
ds *= 2
self._feature_size += ch
if num_head_channels == -1:
dim_head = ch // num_heads
else:
num_heads = ch // num_head_channels
dim_head = num_head_channels
if legacy:
# num_heads = 1
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
self.middle_block = TimestepEmbedSequential(
checkpoint_wrapper_fn(
ResBlock(
ch,
time_embed_dim,
dropout,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
)
),
checkpoint_wrapper_fn(
AttentionBlock(
ch,
use_checkpoint=use_checkpoint,
num_heads=num_heads,
num_head_channels=dim_head,
use_new_attention_order=use_new_attention_order,
)
)
if not use_spatial_transformer
else checkpoint_wrapper_fn(
SpatialTransformer( # always uses a self-attn
ch,
num_heads,
dim_head,
depth=transformer_depth_middle,
context_dim=context_dim,
disable_self_attn=disable_middle_self_attn,
use_linear=use_linear_in_transformer,
attn_type=spatial_transformer_attn_type,
use_checkpoint=use_checkpoint,
)
),
checkpoint_wrapper_fn(
ResBlock(
ch,
time_embed_dim,
dropout,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
)
),
)
self.input_upscale = input_upscale
self.input_hint_block = TimestepEmbedSequential(
zero_module(conv_nd(dims, in_channels, model_channels, 3, padding=1))
)
def convert_to_fp16(self):
"""
Convert the torso of the model to float16.
"""
self.input_blocks.apply(convert_module_to_f16)
self.middle_block.apply(convert_module_to_f16)
def convert_to_fp32(self):
"""
Convert the torso of the model to float32.
"""
self.input_blocks.apply(convert_module_to_f32)
self.middle_block.apply(convert_module_to_f32)
def forward(self, x, timesteps, xt, context=None, y=None, **kwargs):
# with torch.cuda.amp.autocast(enabled=False, dtype=torch.float32):
# x = x.to(torch.float32)
# timesteps = timesteps.to(torch.float32)
# xt = xt.to(torch.float32)
# context = context.to(torch.float32)
# y = y.to(torch.float32)
# print(x.dtype)
xt, context, y = xt.to(x.dtype), context.to(x.dtype), y.to(x.dtype)
if self.input_upscale != 1:
x = nn.functional.interpolate(x, scale_factor=self.input_upscale, mode='bilinear', antialias=True)
assert (y is not None) == (
self.num_classes is not None
), "must specify y if and only if the model is class-conditional"
hs = []
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
# import pdb
# pdb.set_trace()
emb = self.time_embed(t_emb)
if self.num_classes is not None:
assert y.shape[0] == xt.shape[0]
emb = emb + self.label_emb(y)
guided_hint = self.input_hint_block(x, emb, context)
# h = x.type(self.dtype)
h = xt
for module in self.input_blocks:
if guided_hint is not None:
h = module(h, emb, context)
h += guided_hint
guided_hint = None
else:
h = module(h, emb, context)
hs.append(h)
# print(module)
# print(h.shape)
h = self.middle_block(h, emb, context)
hs.append(h)
return hs
class LightGLVUNet(UNetModel):
def __init__(self, mode='', project_type='ZeroSFT', project_channel_scale=1,
*args, **kwargs):
super().__init__(*args, **kwargs)
if mode == 'XL-base':
cond_output_channels = [320] * 4 + [640] * 3 + [1280] * 3
project_channels = [160] * 4 + [320] * 3 + [640] * 3
concat_channels = [320] * 2 + [640] * 3 + [1280] * 4 + [0]
cross_attn_insert_idx = [6, 3]
self.progressive_mask_nums = [0, 3, 7, 11]
elif mode == 'XL-refine':
cond_output_channels = [384] * 4 + [768] * 3 + [1536] * 6
project_channels = [192] * 4 + [384] * 3 + [768] * 6
concat_channels = [384] * 2 + [768] * 3 + [1536] * 7 + [0]
cross_attn_insert_idx = [9, 6, 3]
self.progressive_mask_nums = [0, 3, 6, 10, 14]
else:
raise NotImplementedError
project_channels = [int(c * project_channel_scale) for c in project_channels]
self.project_modules = nn.ModuleList()
for i in range(len(cond_output_channels)):
# if i == len(cond_output_channels) - 1:
# _project_type = 'ZeroCrossAttn'
# else:
# _project_type = project_type
_project_type = project_type
if _project_type == 'ZeroSFT':
self.project_modules.append(ZeroSFT(project_channels[i], cond_output_channels[i],
concat_channels=concat_channels[i]))
elif _project_type == 'ZeroCrossAttn':
self.project_modules.append(ZeroCrossAttn(cond_output_channels[i], project_channels[i]))
else:
raise NotImplementedError
for i in cross_attn_insert_idx:
self.project_modules.insert(i, ZeroCrossAttn(cond_output_channels[i], concat_channels[i]))
# print(self.project_modules[i])
def step_progressive_mask(self):
if len(self.progressive_mask_nums) > 0:
mask_num = self.progressive_mask_nums.pop()
for i in range(len(self.project_modules)):
if i < mask_num:
self.project_modules[i].mask = True
else:
self.project_modules[i].mask = False
return
# print(f'step_progressive_mask, current masked layers: {mask_num}')
else:
return
# print('step_progressive_mask, no more masked layers')
# for i in range(len(self.project_modules)):
# print(self.project_modules[i].mask)
def forward(self, x, timesteps=None, context=None, y=None, control=None, control_scale=1, **kwargs):
"""
Apply the model to an input batch.
:param x: an [N x C x ...] Tensor of inputs.
:param timesteps: a 1-D batch of timesteps.
:param context: conditioning plugged in via crossattn
:param y: an [N] Tensor of labels, if class-conditional.
:return: an [N x C x ...] Tensor of outputs.
"""
assert (y is not None) == (
self.num_classes is not None
), "must specify y if and only if the model is class-conditional"
hs = []
_dtype = control[0].dtype
x, context, y = x.to(_dtype), context.to(_dtype), y.to(_dtype)
with torch.no_grad():
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
emb = self.time_embed(t_emb)
if self.num_classes is not None:
assert y.shape[0] == x.shape[0]
emb = emb + self.label_emb(y)
# h = x.type(self.dtype)
h = x
for module in self.input_blocks:
h = module(h, emb, context)
hs.append(h)
adapter_idx = len(self.project_modules) - 1
control_idx = len(control) - 1
h = self.middle_block(h, emb, context)
h = self.project_modules[adapter_idx](control[control_idx], h, control_scale=control_scale)
adapter_idx -= 1
control_idx -= 1
for i, module in enumerate(self.output_blocks):
_h = hs.pop()
h = self.project_modules[adapter_idx](control[control_idx], _h, h, control_scale=control_scale)
adapter_idx -= 1
# h = th.cat([h, _h], dim=1)
if len(module) == 3:
assert isinstance(module[2], Upsample)
for layer in module[:2]:
if isinstance(layer, TimestepBlock):
h = layer(h, emb)
elif isinstance(layer, SpatialTransformer):
h = layer(h, context)
else:
h = layer(h)
# print('cross_attn_here')
h = self.project_modules[adapter_idx](control[control_idx], h, control_scale=control_scale)
adapter_idx -= 1
h = module[2](h)
else:
h = module(h, emb, context)
control_idx -= 1
# print(module)
# print(h.shape)
h = h.type(x.dtype)
if self.predict_codebook_ids:
assert False, "not supported anymore. what the f*** are you doing?"
else:
return self.out(h)
if __name__ == '__main__':
from omegaconf import OmegaConf
# refiner
# opt = OmegaConf.load('../../options/train/debug_p2_xl.yaml')
#
# model = instantiate_from_config(opt.model.params.control_stage_config)
# hint = model(torch.randn([1, 4, 64, 64]), torch.randn([1]), torch.randn([1, 4, 64, 64]))
# hint = [h.cuda() for h in hint]
# print(sum(map(lambda hint: hint.numel(), model.parameters())))
#
# unet = instantiate_from_config(opt.model.params.network_config)
# unet = unet.cuda()
#
# _output = unet(torch.randn([1, 4, 64, 64]).cuda(), torch.randn([1]).cuda(), torch.randn([1, 77, 1280]).cuda(),
# torch.randn([1, 2560]).cuda(), hint)
# print(sum(map(lambda _output: _output.numel(), unet.parameters())))
# base
with torch.no_grad():
opt = OmegaConf.load('../../options/dev/SUPIR_tmp.yaml')
model = instantiate_from_config(opt.model.params.control_stage_config)
model = model.cuda()
hint = model(torch.randn([1, 4, 64, 64]).cuda(), torch.randn([1]).cuda(), torch.randn([1, 4, 64, 64]).cuda(), torch.randn([1, 77, 2048]).cuda(),
torch.randn([1, 2816]).cuda())
for h in hint:
print(h.shape)
#
unet = instantiate_from_config(opt.model.params.network_config)
unet = unet.cuda()
_output = unet(torch.randn([1, 4, 64, 64]).cuda(), torch.randn([1]).cuda(), torch.randn([1, 77, 2048]).cuda(),
torch.randn([1, 2816]).cuda(), hint)
# model = instantiate_from_config(opt.model.params.control_stage_config)
# model = model.cuda()
# # hint = model(torch.randn([1, 4, 64, 64]), torch.randn([1]), torch.randn([1, 4, 64, 64]))
# hint = model(torch.randn([1, 4, 64, 64]).cuda(), torch.randn([1]).cuda(), torch.randn([1, 4, 64, 64]).cuda(), torch.randn([1, 77, 1280]).cuda(),
# torch.randn([1, 2560]).cuda())
# # hint = [h.cuda() for h in hint]
#
# for h in hint:
# print(h.shape)
#
# unet = instantiate_from_config(opt.model.params.network_config)
# unet = unet.cuda()
# _output = unet(torch.randn([1, 4, 64, 64]).cuda(), torch.randn([1]).cuda(), torch.randn([1, 77, 1280]).cuda(),
# torch.randn([1, 2560]).cuda(), hint)