# https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py # https://github.com/JingyunLiang/SwinIR/blob/main/models/network_swinir.py#L812 import copy import math from collections import namedtuple from contextlib import contextmanager, nullcontext from functools import partial, wraps from pathlib import Path from random import random from einops import rearrange, repeat, reduce, pack, unpack import torch import torch.nn.functional as F import torchvision.transforms as T from torch import einsum, nn from beartype.typing import List, Union from beartype import beartype from tqdm.auto import tqdm from pdb import set_trace as st # helper functions, from: # https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py def exists(val): return val is not None def identity(t, *args, **kwargs): return t def divisible_by(numer, denom): return (numer % denom) == 0 def first(arr, d=None): if len(arr) == 0: return d return arr[0] def maybe(fn): @wraps(fn) def inner(x): if not exists(x): return x return fn(x) return inner def once(fn): called = False @wraps(fn) def inner(x): nonlocal called if called: return called = True return fn(x) return inner print_once = once(print) def default(val, d): if exists(val): return val return d() if callable(d) else d def compact(input_dict): return {key: value for key, value in input_dict.items() if exists(value)} def maybe_transform_dict_key(input_dict, key, fn): if key not in input_dict: return input_dict copied_dict = input_dict.copy() copied_dict[key] = fn(copied_dict[key]) return copied_dict def cast_uint8_images_to_float(images): if not images.dtype == torch.uint8: return images return images / 255 def module_device(module): return next(module.parameters()).device def zero_init_(m): nn.init.zeros_(m.weight) if exists(m.bias): nn.init.zeros_(m.bias) def eval_decorator(fn): def inner(model, *args, **kwargs): was_training = model.training model.eval() out = fn(model, *args, **kwargs) model.train(was_training) return out return inner def pad_tuple_to_length(t, length, fillvalue=None): remain_length = length - len(t) if remain_length <= 0: return t return (*t, *((fillvalue, ) * remain_length)) # helper classes class Identity(nn.Module): def __init__(self, *args, **kwargs): super().__init__() def forward(self, x, *args, **kwargs): return x # tensor helpers def log(t, eps: float = 1e-12): return torch.log(t.clamp(min=eps)) def l2norm(t): return F.normalize(t, dim=-1) def right_pad_dims_to(x, t): padding_dims = x.ndim - t.ndim if padding_dims <= 0: return t return t.view(*t.shape, *((1, ) * padding_dims)) def masked_mean(t, *, dim, mask=None): if not exists(mask): return t.mean(dim=dim) denom = mask.sum(dim=dim, keepdim=True) mask = rearrange(mask, 'b n -> b n 1') masked_t = t.masked_fill(~mask, 0.) return masked_t.sum(dim=dim) / denom.clamp(min=1e-5) def resize_image_to(image, target_image_size, clamp_range=None, mode='nearest'): orig_image_size = image.shape[-1] if orig_image_size == target_image_size: return image out = F.interpolate(image, target_image_size, mode=mode) if exists(clamp_range): out = out.clamp(*clamp_range) return out def calc_all_frame_dims(downsample_factors: List[int], frames): if not exists(frames): return (tuple(), ) * len(downsample_factors) all_frame_dims = [] for divisor in downsample_factors: assert divisible_by(frames, divisor) all_frame_dims.append((frames // divisor, )) return all_frame_dims def safe_get_tuple_index(tup, index, default=None): if len(tup) <= index: return default return tup[index] # image normalization functions # ddpms expect images to be in the range of -1 to 1 def normalize_neg_one_to_one(img): return img * 2 - 1 def unnormalize_zero_to_one(normed_img): return (normed_img + 1) * 0.5 # def Upsample(dim, dim_out=None): # dim_out = default(dim_out, dim) # return nn.Sequential(nn.Upsample(scale_factor=2, mode='nearest'), # nn.Conv2d(dim, dim_out, 3, padding=1)) class PixelShuffleUpsample(nn.Module): """ code shared by @MalumaDev at DALLE2-pytorch for addressing checkboard artifacts https://arxiv.org/ftp/arxiv/papers/1707/1707.02937.pdf """ def __init__(self, dim, dim_out=None): super().__init__() dim_out = default(dim_out, dim) conv = nn.Conv2d(dim, dim_out * 4, 1) self.net = nn.Sequential(conv, nn.SiLU(), nn.PixelShuffle(2)) self.init_conv_(conv) def init_conv_(self, conv): o, i, h, w = conv.weight.shape conv_weight = torch.empty(o // 4, i, h, w) nn.init.kaiming_uniform_(conv_weight) conv_weight = repeat(conv_weight, 'o ... -> (o 4) ...') conv.weight.data.copy_(conv_weight) nn.init.zeros_(conv.bias.data) def forward(self, x): return self.net(x) class ResidualBlock(nn.Module): def __init__(self, dim_in, dim_out, dim_inter=None, use_norm=True, norm_layer=nn.BatchNorm2d, bias=False): super().__init__() if dim_inter is None: dim_inter = dim_out if use_norm: self.conv = nn.Sequential( norm_layer(dim_in), nn.ReLU(True), nn.Conv2d(dim_in, dim_inter, 3, 1, 1, bias=bias, padding_mode='reflect'), norm_layer(dim_inter), nn.ReLU(True), nn.Conv2d(dim_inter, dim_out, 3, 1, 1, bias=bias, padding_mode='reflect'), ) else: self.conv = nn.Sequential( nn.ReLU(True), nn.Conv2d(dim_in, dim_inter, 3, 1, 1), nn.ReLU(True), nn.Conv2d(dim_inter, dim_out, 3, 1, 1), ) self.short_cut = None if dim_in != dim_out: self.short_cut = nn.Conv2d(dim_in, dim_out, 1, 1) def forward(self, feats): feats_out = self.conv(feats) if self.short_cut is not None: feats_out = self.short_cut(feats) + feats_out else: feats_out = feats_out + feats return feats_out class Upsample(nn.Sequential): """Upsample module. Args: scale (int): Scale factor. Supported scales: 2^n and 3. num_feat (int): Channel number of intermediate features. """ def __init__(self, scale, num_feat): m = [] if (scale & (scale - 1)) == 0: # scale = 2^n for _ in range(int(math.log(scale, 2))): m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1)) m.append(nn.PixelShuffle(2)) elif scale == 3: m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1)) m.append(nn.PixelShuffle(3)) else: raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.') super(Upsample, self).__init__(*m) class PixelUnshuffleUpsample(nn.Module): def __init__(self, output_dim, num_feat=128, num_out_ch=3, sr_ratio=4, *args, **kwargs) -> None: super().__init__() self.conv_after_body = nn.Conv2d(output_dim, output_dim, 3, 1, 1) self.conv_before_upsample = nn.Sequential( nn.Conv2d(output_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)) self.upsample = Upsample(sr_ratio, num_feat) # 4 time SR self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) def forward(self, x, input_skip_connection=True, *args, **kwargs): # x = self.conv_first(x) if input_skip_connection: x = self.conv_after_body(x) + x else: x = self.conv_after_body(x) x = self.conv_before_upsample(x) x = self.conv_last(self.upsample(x)) return x class Conv3x3TriplaneTransformation(nn.Module): # used in the final layer before triplane def __init__(self, input_dim, output_dim) -> None: super().__init__() self.conv_after_unpachify = nn.Sequential( nn.Conv2d(input_dim, output_dim, 3, 1, 1), nn.LeakyReLU(inplace=True) ) self.conv_before_rendering = nn.Sequential( nn.Conv2d(output_dim, output_dim, 3, 1, 1), nn.LeakyReLU(inplace=True)) def forward(self, unpachified_latent): latent = self.conv_after_unpachify(unpachified_latent) # no residual connections here latent = self.conv_before_rendering(latent) + latent return latent # https://github.com/JingyunLiang/SwinIR/blob/6545850fbf8df298df73d81f3e8cba638787c8bd/models/network_swinir.py#L750 class NearestConvSR(nn.Module): """ code shared by @MalumaDev at DALLE2-pytorch for addressing checkboard artifacts https://arxiv.org/ftp/arxiv/papers/1707/1707.02937.pdf """ def __init__(self, output_dim, num_feat=128, num_out_ch=3, sr_ratio=4, *args, **kwargs) -> None: super().__init__() self.upscale = sr_ratio self.conv_after_body = nn.Conv2d(output_dim, output_dim, 3, 1, 1) self.conv_before_upsample = nn.Sequential(nn.Conv2d(output_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)) self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) if self.upscale == 4: self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1) self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) def forward(self, x, *args, **kwargs): # x = self.conv_first(x) x = self.conv_after_body(x) + x x = self.conv_before_upsample(x) x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest'))) if self.upscale == 4: x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest'))) x = self.conv_last(self.lrelu(self.conv_hr(x))) return x # https://github.com/yumingj/C2-Matching/blob/fa171ca6707c6f16a5d04194ce866ea70bb21d2b/mmsr/models/archs/ref_restoration_arch.py#L65 class NearestConvSR_Residual(NearestConvSR): # learn residual + normalize def __init__(self, output_dim, num_feat=128, num_out_ch=3, sr_ratio=4, *args, **kwargs) -> None: super().__init__(output_dim, num_feat, num_out_ch, sr_ratio, *args, **kwargs) # self.mean = torch.Tensor((0.485, 0.456, 0.406)).view(1,3,1,1) # imagenet mean self.act = nn.Tanh() def forward(self, x, base_x, *args, **kwargs): # base_x: low-resolution 3D rendering, for residual addition # self.mean = self.mean.type_as(x) # x = super().forward(x).clamp(-1,1) x = super().forward(x) x = self.act(x) # residual normalize to [-1,1] scale = x.shape[-1] // base_x.shape[-1] # 2 or 4 x = x + F.interpolate(base_x, None, scale, 'bilinear', False) # add residual; [-1,1] range # return x + 2 * self.mean return x class UpsampleOneStep(nn.Sequential): """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle) Used in lightweight SR to save parameters. Args: scale (int): Scale factor. Supported scales: 2^n and 3. num_feat (int): Channel number of intermediate features. """ def __init__(self, scale, num_feat, num_out_ch, input_resolution=None): self.num_feat = num_feat self.input_resolution = input_resolution m = [] m.append(nn.Conv2d(num_feat, (scale ** 2) * num_out_ch, 3, 1, 1)) m.append(nn.PixelShuffle(scale)) super(UpsampleOneStep, self).__init__(*m) def flops(self): H, W = self.input_resolution flops = H * W * self.num_feat * 3 * 9 return flops # class PixelShuffledDirect(nn.Module):