first commit fixed

app.py

@@ -24,17 +24,12 @@ if not os.path.exists(realesr_model_path):
     os.system(
         "wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -O experiments/pretrained_models/RealESRGAN_x4plus.pth")
 
-pmrf_model_path = 'blind_face_restoration_pmrf.ckpt'
-
 # background enhancer with RealESRGAN
 model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
 half = True if torch.cuda.is_available() else False
 upsampler = RealESRGANer(scale=4, model_path=realesr_model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
 
-pmrf = MMSERectifiedFlow.load_from_checkpoint('./blind_face_restoration_pmrf.ckpt',
-                                              mmse_model_arch='swinir_L',
-                                              mmse_model_ckpt_path=None,
-                                              map_location='cpu').to(device)
+pmrf = MMSERectifiedFlow.from_pretrained('ohayonguy/PMRF_blind_face_image_restoration').to(device)
 
 os.makedirs('output', exist_ok=True)
 

arch/hourglass/axial_rope.py

@@ -0,0 +1,113 @@
+"""k-diffusion transformer diffusion models, version 2.
+Codes adopted from https://github.com/crowsonkb/k-diffusion
+"""
+
+import math
+
+import torch
+import torch._dynamo
+from torch import nn
+
+from . import flags
+
+if flags.get_use_compile():
+    torch._dynamo.config.suppress_errors = True
+
+
+def rotate_half(x):
+    x1, x2 = x[..., 0::2], x[..., 1::2]
+    x = torch.stack((-x2, x1), dim=-1)
+    *shape, d, r = x.shape
+    return x.view(*shape, d * r)
+
+
+@flags.compile_wrap
+def apply_rotary_emb(freqs, t, start_index=0, scale=1.0):
+    freqs = freqs.to(t)
+    rot_dim = freqs.shape[-1]
+    end_index = start_index + rot_dim
+    assert rot_dim <= t.shape[-1], f"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}"
+    t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:]
+    t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)
+    return torch.cat((t_left, t, t_right), dim=-1)
+
+
+def centers(start, stop, num, dtype=None, device=None):
+    edges = torch.linspace(start, stop, num + 1, dtype=dtype, device=device)
+    return (edges[:-1] + edges[1:]) / 2
+
+
+def make_grid(h_pos, w_pos):
+    grid = torch.stack(torch.meshgrid(h_pos, w_pos, indexing='ij'), dim=-1)
+    h, w, d = grid.shape
+    return grid.view(h * w, d)
+
+
+def bounding_box(h, w, pixel_aspect_ratio=1.0):
+    # Adjusted dimensions
+    w_adj = w
+    h_adj = h * pixel_aspect_ratio
+
+    # Adjusted aspect ratio
+    ar_adj = w_adj / h_adj
+
+    # Determine bounding box based on the adjusted aspect ratio
+    y_min, y_max, x_min, x_max = -1.0, 1.0, -1.0, 1.0
+    if ar_adj > 1:
+        y_min, y_max = -1 / ar_adj, 1 / ar_adj
+    elif ar_adj < 1:
+        x_min, x_max = -ar_adj, ar_adj
+
+    return y_min, y_max, x_min, x_max
+
+
+def make_axial_pos(h, w, pixel_aspect_ratio=1.0, align_corners=False, dtype=None, device=None):
+    y_min, y_max, x_min, x_max = bounding_box(h, w, pixel_aspect_ratio)
+    if align_corners:
+        h_pos = torch.linspace(y_min, y_max, h, dtype=dtype, device=device)
+        w_pos = torch.linspace(x_min, x_max, w, dtype=dtype, device=device)
+    else:
+        h_pos = centers(y_min, y_max, h, dtype=dtype, device=device)
+        w_pos = centers(x_min, x_max, w, dtype=dtype, device=device)
+    return make_grid(h_pos, w_pos)
+
+
+def freqs_pixel(max_freq=10.0):
+    def init(shape):
+        freqs = torch.linspace(1.0, max_freq / 2, shape[-1]) * math.pi
+        return freqs.log().expand(shape)
+    return init
+
+
+def freqs_pixel_log(max_freq=10.0):
+    def init(shape):
+        log_min = math.log(math.pi)
+        log_max = math.log(max_freq * math.pi / 2)
+        return torch.linspace(log_min, log_max, shape[-1]).expand(shape)
+    return init
+
+
+class AxialRoPE(nn.Module):
+    def __init__(self, dim, n_heads, start_index=0, freqs_init=freqs_pixel_log(max_freq=10.0)):
+        super().__init__()
+        self.n_heads = n_heads
+        self.start_index = start_index
+        log_freqs = freqs_init((n_heads, dim // 4))
+        self.freqs_h = nn.Parameter(log_freqs.clone())
+        self.freqs_w = nn.Parameter(log_freqs.clone())
+
+    def extra_repr(self):
+        dim = (self.freqs_h.shape[-1] + self.freqs_w.shape[-1]) * 2
+        return f"dim={dim}, n_heads={self.n_heads}, start_index={self.start_index}"
+
+    def get_freqs(self, pos):
+        if pos.shape[-1] != 2:
+            raise ValueError("input shape must be (..., 2)")
+        freqs_h = pos[..., None, None, 0] * self.freqs_h.exp()
+        freqs_w = pos[..., None, None, 1] * self.freqs_w.exp()
+        freqs = torch.cat((freqs_h, freqs_w), dim=-1).repeat_interleave(2, dim=-1)
+        return freqs.transpose(-2, -3)
+
+    def forward(self, x, pos):
+        freqs = self.get_freqs(pos)
+        return apply_rotary_emb(freqs, x, self.start_index)

arch/hourglass/flags.py

@@ -0,0 +1,60 @@
+"""k-diffusion transformer diffusion models, version 2.
+Codes adopted from https://github.com/crowsonkb/k-diffusion
+"""
+
+from contextlib import contextmanager
+from functools import update_wrapper
+import os
+import threading
+
+import torch
+
+
+def get_use_compile():
+    return os.environ.get("K_DIFFUSION_USE_COMPILE", "1") == "1"
+
+
+def get_use_flash_attention_2():
+    return os.environ.get("K_DIFFUSION_USE_FLASH_2", "1") == "1"
+
+
+state = threading.local()
+state.checkpointing = False
+
+
+@contextmanager
+def checkpointing(enable=True):
+    try:
+        old_checkpointing, state.checkpointing = state.checkpointing, enable
+        yield
+    finally:
+        state.checkpointing = old_checkpointing
+
+
+def get_checkpointing():
+    return getattr(state, "checkpointing", False)
+
+
+class compile_wrap:
+    def __init__(self, function, *args, **kwargs):
+        self.function = function
+        self.args = args
+        self.kwargs = kwargs
+        self._compiled_function = None
+        update_wrapper(self, function)
+
+    @property
+    def compiled_function(self):
+        if self._compiled_function is not None:
+            return self._compiled_function
+        if get_use_compile():
+            try:
+                self._compiled_function = torch.compile(self.function, *self.args, **self.kwargs)
+            except RuntimeError:
+                self._compiled_function = self.function
+        else:
+            self._compiled_function = self.function
+        return self._compiled_function
+
+    def __call__(self, *args, **kwargs):
+        return self.compiled_function(*args, **kwargs)

arch/hourglass/flops.py

@@ -0,0 +1,58 @@
+"""k-diffusion transformer diffusion models, version 2.
+Codes adopted from https://github.com/crowsonkb/k-diffusion
+"""
+
+from contextlib import contextmanager
+import math
+import threading
+
+
+state = threading.local()
+state.flop_counter = None
+
+
+@contextmanager
+def flop_counter(enable=True):
+    try:
+        old_flop_counter = state.flop_counter
+        state.flop_counter = FlopCounter() if enable else None
+        yield state.flop_counter
+    finally:
+        state.flop_counter = old_flop_counter
+
+
+class FlopCounter:
+    def __init__(self):
+        self.ops = []
+
+    def op(self, op, *args, **kwargs):
+        self.ops.append((op, args, kwargs))
+
+    @property
+    def flops(self):
+        flops = 0
+        for op, args, kwargs in self.ops:
+            flops += op(*args, **kwargs)
+        return flops
+
+
+def op(op, *args, **kwargs):
+    if getattr(state, "flop_counter", None):
+        state.flop_counter.op(op, *args, **kwargs)
+
+
+def op_linear(x, weight):
+    return math.prod(x) * weight[0]
+
+
+def op_attention(q, k, v):
+    *b, s_q, d_q = q
+    *b, s_k, d_k = k
+    *b, s_v, d_v = v
+    return math.prod(b) * s_q * s_k * (d_q + d_v)
+
+
+def op_natten(q, k, v, kernel_size):
+    *q_rest, d_q = q
+    *_, d_v = v
+    return math.prod(q_rest) * (d_q + d_v) * kernel_size**2

arch/hourglass/image_transformer_v2.py

@@ -0,0 +1,772 @@
+"""k-diffusion transformer diffusion models, version 2.
+Codes adopted from https://github.com/crowsonkb/k-diffusion
+"""
+
+from dataclasses import dataclass
+from functools import lru_cache, reduce
+import math
+from typing import Union
+
+from einops import rearrange
+import torch
+from torch import nn
+import torch._dynamo
+from torch.nn import functional as F
+
+from . import flags, flops
+from .axial_rope import make_axial_pos
+
+
+try:
+    import natten
+except ImportError:
+    natten = None
+
+try:
+    import flash_attn
+except ImportError:
+    flash_attn = None
+
+
+if flags.get_use_compile():
+    torch._dynamo.config.cache_size_limit = max(64, torch._dynamo.config.cache_size_limit)
+    torch._dynamo.config.suppress_errors = True
+
+
+# Helpers
+
+def zero_init(layer):
+    nn.init.zeros_(layer.weight)
+    if layer.bias is not None:
+        nn.init.zeros_(layer.bias)
+    return layer
+
+
+def checkpoint(function, *args, **kwargs):
+    if flags.get_checkpointing():
+        kwargs.setdefault("use_reentrant", True)
+        return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
+    else:
+        return function(*args, **kwargs)
+
+
+def downscale_pos(pos):
+    pos = rearrange(pos, "... (h nh) (w nw) e -> ... h w (nh nw) e", nh=2, nw=2)
+    return torch.mean(pos, dim=-2)
+
+
+# Param tags
+
+def tag_param(param, tag):
+    if not hasattr(param, "_tags"):
+        param._tags = set([tag])
+    else:
+        param._tags.add(tag)
+    return param
+
+
+def tag_module(module, tag):
+    for param in module.parameters():
+        tag_param(param, tag)
+    return module
+
+
+def apply_wd(module):
+    for name, param in module.named_parameters():
+        if name.endswith("weight"):
+            tag_param(param, "wd")
+    return module
+
+
+def filter_params(function, module):
+    for param in module.parameters():
+        tags = getattr(param, "_tags", set())
+        if function(tags):
+            yield param
+
+
+# Kernels
+
+@flags.compile_wrap
+def linear_geglu(x, weight, bias=None):
+    x = x @ weight.mT
+    if bias is not None:
+        x = x + bias
+    x, gate = x.chunk(2, dim=-1)
+    return x * F.gelu(gate)
+
+
+@flags.compile_wrap
+def rms_norm(x, scale, eps):
+    dtype = reduce(torch.promote_types, (x.dtype, scale.dtype, torch.float32))
+    mean_sq = torch.mean(x.to(dtype)**2, dim=-1, keepdim=True)
+    scale = scale.to(dtype) * torch.rsqrt(mean_sq + eps)
+    return x * scale.to(x.dtype)
+
+
+@flags.compile_wrap
+def scale_for_cosine_sim(q, k, scale, eps):
+    dtype = reduce(torch.promote_types, (q.dtype, k.dtype, scale.dtype, torch.float32))
+    sum_sq_q = torch.sum(q.to(dtype)**2, dim=-1, keepdim=True)
+    sum_sq_k = torch.sum(k.to(dtype)**2, dim=-1, keepdim=True)
+    sqrt_scale = torch.sqrt(scale.to(dtype))
+    scale_q = sqrt_scale * torch.rsqrt(sum_sq_q + eps)
+    scale_k = sqrt_scale * torch.rsqrt(sum_sq_k + eps)
+    return q * scale_q.to(q.dtype), k * scale_k.to(k.dtype)
+
+
+@flags.compile_wrap
+def scale_for_cosine_sim_qkv(qkv, scale, eps):
+    q, k, v = qkv.unbind(2)
+    q, k = scale_for_cosine_sim(q, k, scale[:, None], eps)
+    return torch.stack((q, k, v), dim=2)
+
+
+# Layers
+
+class Linear(nn.Linear):
+    def forward(self, x):
+        flops.op(flops.op_linear, x.shape, self.weight.shape)
+        return super().forward(x)
+
+
+class LinearGEGLU(nn.Linear):
+    def __init__(self, in_features, out_features, bias=True):
+        super().__init__(in_features, out_features * 2, bias=bias)
+        self.out_features = out_features
+
+    def forward(self, x):
+        flops.op(flops.op_linear, x.shape, self.weight.shape)
+        return linear_geglu(x, self.weight, self.bias)
+
+
+class FourierFeatures(nn.Module):
+    def __init__(self, in_features, out_features, std=1.):
+        super().__init__()
+        assert out_features % 2 == 0
+        self.register_buffer('weight', torch.randn([out_features // 2, in_features]) * std)
+
+    def forward(self, input):
+        f = 2 * math.pi * input @ self.weight.T
+        return torch.cat([f.cos(), f.sin()], dim=-1)
+
+class RMSNorm(nn.Module):
+    def __init__(self, shape, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.scale = nn.Parameter(torch.ones(shape))
+
+    def extra_repr(self):
+        return f"shape={tuple(self.scale.shape)}, eps={self.eps}"
+
+    def forward(self, x):
+        return rms_norm(x, self.scale, self.eps)
+
+
+class AdaRMSNorm(nn.Module):
+    def __init__(self, features, cond_features, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.linear = apply_wd(zero_init(Linear(cond_features, features, bias=False)))
+        tag_module(self.linear, "mapping")
+
+    def extra_repr(self):
+        return f"eps={self.eps},"
+
+    def forward(self, x, cond):
+        return rms_norm(x, self.linear(cond)[:, None, None, :] + 1, self.eps)
+
+
+# Rotary position embeddings
+
+@flags.compile_wrap
+def apply_rotary_emb(x, theta, conj=False):
+    out_dtype = x.dtype
+    dtype = reduce(torch.promote_types, (x.dtype, theta.dtype, torch.float32))
+    d = theta.shape[-1]
+    assert d * 2 <= x.shape[-1]
+    x1, x2, x3 = x[..., :d], x[..., d : d * 2], x[..., d * 2 :]
+    x1, x2, theta = x1.to(dtype), x2.to(dtype), theta.to(dtype)
+    cos, sin = torch.cos(theta), torch.sin(theta)
+    sin = -sin if conj else sin
+    y1 = x1 * cos - x2 * sin
+    y2 = x2 * cos + x1 * sin
+    y1, y2 = y1.to(out_dtype), y2.to(out_dtype)
+    return torch.cat((y1, y2, x3), dim=-1)
+
+
+@flags.compile_wrap
+def _apply_rotary_emb_inplace(x, theta, conj):
+    dtype = reduce(torch.promote_types, (x.dtype, theta.dtype, torch.float32))
+    d = theta.shape[-1]
+    assert d * 2 <= x.shape[-1]
+    x1, x2 = x[..., :d], x[..., d : d * 2]
+    x1_, x2_, theta = x1.to(dtype), x2.to(dtype), theta.to(dtype)
+    cos, sin = torch.cos(theta), torch.sin(theta)
+    sin = -sin if conj else sin
+    y1 = x1_ * cos - x2_ * sin
+    y2 = x2_ * cos + x1_ * sin
+    x1.copy_(y1)
+    x2.copy_(y2)
+
+
+class ApplyRotaryEmbeddingInplace(torch.autograd.Function):
+    @staticmethod
+    def forward(x, theta, conj):
+        _apply_rotary_emb_inplace(x, theta, conj=conj)
+        return x
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        _, theta, conj = inputs
+        ctx.save_for_backward(theta)
+        ctx.conj = conj
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        theta, = ctx.saved_tensors
+        _apply_rotary_emb_inplace(grad_output, theta, conj=not ctx.conj)
+        return grad_output, None, None
+
+
+def apply_rotary_emb_(x, theta):
+    return ApplyRotaryEmbeddingInplace.apply(x, theta, False)
+
+
+class AxialRoPE(nn.Module):
+    def __init__(self, dim, n_heads):
+        super().__init__()
+        log_min = math.log(math.pi)
+        log_max = math.log(10.0 * math.pi)
+        freqs = torch.linspace(log_min, log_max, n_heads * dim // 4 + 1)[:-1].exp()
+        self.register_buffer("freqs", freqs.view(dim // 4, n_heads).T.contiguous())
+
+    def extra_repr(self):
+        return f"dim={self.freqs.shape[1] * 4}, n_heads={self.freqs.shape[0]}"
+
+    def forward(self, pos):
+        theta_h = pos[..., None, 0:1] * self.freqs.to(pos.dtype)
+        theta_w = pos[..., None, 1:2] * self.freqs.to(pos.dtype)
+        return torch.cat((theta_h, theta_w), dim=-1)
+
+
+# Shifted window attention
+
+def window(window_size, x):
+    *b, h, w, c = x.shape
+    x = torch.reshape(
+        x,
+        (*b, h // window_size, window_size, w // window_size, window_size, c),
+    )
+    x = torch.permute(
+        x,
+        (*range(len(b)), -5, -3, -4, -2, -1),
+    )
+    return x
+
+
+def unwindow(x):
+    *b, h, w, wh, ww, c = x.shape
+    x = torch.permute(x, (*range(len(b)), -5, -3, -4, -2, -1))
+    x = torch.reshape(x, (*b, h * wh, w * ww, c))
+    return x
+
+
+def shifted_window(window_size, window_shift, x):
+    x = torch.roll(x, shifts=(window_shift, window_shift), dims=(-2, -3))
+    windows = window(window_size, x)
+    return windows
+
+
+def shifted_unwindow(window_shift, x):
+    x = unwindow(x)
+    x = torch.roll(x, shifts=(-window_shift, -window_shift), dims=(-2, -3))
+    return x
+
+
+@lru_cache
+def make_shifted_window_masks(n_h_w, n_w_w, w_h, w_w, shift, device=None):
+    ph_coords = torch.arange(n_h_w, device=device)
+    pw_coords = torch.arange(n_w_w, device=device)
+    h_coords = torch.arange(w_h, device=device)
+    w_coords = torch.arange(w_w, device=device)
+    patch_h, patch_w, q_h, q_w, k_h, k_w = torch.meshgrid(
+        ph_coords,
+        pw_coords,
+        h_coords,
+        w_coords,
+        h_coords,
+        w_coords,
+        indexing="ij",
+    )
+    is_top_patch = patch_h == 0
+    is_left_patch = patch_w == 0
+    q_above_shift = q_h < shift
+    k_above_shift = k_h < shift
+    q_left_of_shift = q_w < shift
+    k_left_of_shift = k_w < shift
+    m_corner = (
+        is_left_patch
+        & is_top_patch
+        & (q_left_of_shift == k_left_of_shift)
+        & (q_above_shift == k_above_shift)
+    )
+    m_left = is_left_patch & ~is_top_patch & (q_left_of_shift == k_left_of_shift)
+    m_top = ~is_left_patch & is_top_patch & (q_above_shift == k_above_shift)
+    m_rest = ~is_left_patch & ~is_top_patch
+    m = m_corner | m_left | m_top | m_rest
+    return m
+
+
+def apply_window_attention(window_size, window_shift, q, k, v, scale=None):
+    # prep windows and masks
+    q_windows = shifted_window(window_size, window_shift, q)
+    k_windows = shifted_window(window_size, window_shift, k)
+    v_windows = shifted_window(window_size, window_shift, v)
+    b, heads, h, w, wh, ww, d_head = q_windows.shape
+    mask = make_shifted_window_masks(h, w, wh, ww, window_shift, device=q.device)
+    q_seqs = torch.reshape(q_windows, (b, heads, h, w, wh * ww, d_head))
+    k_seqs = torch.reshape(k_windows, (b, heads, h, w, wh * ww, d_head))
+    v_seqs = torch.reshape(v_windows, (b, heads, h, w, wh * ww, d_head))
+    mask = torch.reshape(mask, (h, w, wh * ww, wh * ww))
+
+    # do the attention here
+    flops.op(flops.op_attention, q_seqs.shape, k_seqs.shape, v_seqs.shape)
+    qkv = F.scaled_dot_product_attention(q_seqs, k_seqs, v_seqs, mask, scale=scale)
+
+    # unwindow
+    qkv = torch.reshape(qkv, (b, heads, h, w, wh, ww, d_head))
+    return shifted_unwindow(window_shift, qkv)
+
+
+# Transformer layers
+
+
+def use_flash_2(x):
+    if not flags.get_use_flash_attention_2():
+        return False
+    if flash_attn is None:
+        return False
+    if x.device.type != "cuda":
+        return False
+    if x.dtype not in (torch.float16, torch.bfloat16):
+        return False
+    return True
+
+
+class SelfAttentionBlock(nn.Module):
+    def __init__(self, d_model, d_head, cond_features, dropout=0.0):
+        super().__init__()
+        self.d_head = d_head
+        self.n_heads = d_model // d_head
+        self.norm = AdaRMSNorm(d_model, cond_features)
+        self.qkv_proj = apply_wd(Linear(d_model, d_model * 3, bias=False))
+        self.scale = nn.Parameter(torch.full([self.n_heads], 10.0))
+        self.pos_emb = AxialRoPE(d_head // 2, self.n_heads)
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = apply_wd(zero_init(Linear(d_model, d_model, bias=False)))
+
+    def extra_repr(self):
+        return f"d_head={self.d_head},"
+
+    def forward(self, x, pos, cond):
+        skip = x
+        x = self.norm(x, cond)
+        qkv = self.qkv_proj(x)
+        pos = rearrange(pos, "... h w e -> ... (h w) e").to(qkv.dtype)
+        theta = self.pos_emb(pos)
+        if use_flash_2(qkv):
+            qkv = rearrange(qkv, "n h w (t nh e) -> n (h w) t nh e", t=3, e=self.d_head)
+            qkv = scale_for_cosine_sim_qkv(qkv, self.scale, 1e-6)
+            theta = torch.stack((theta, theta, torch.zeros_like(theta)), dim=-3)
+            qkv = apply_rotary_emb_(qkv, theta)
+            flops_shape = qkv.shape[-5], qkv.shape[-2], qkv.shape[-4], qkv.shape[-1]
+            flops.op(flops.op_attention, flops_shape, flops_shape, flops_shape)
+            x = flash_attn.flash_attn_qkvpacked_func(qkv, softmax_scale=1.0)
+            x = rearrange(x, "n (h w) nh e -> n h w (nh e)", h=skip.shape[-3], w=skip.shape[-2])
+        else:
+            q, k, v = rearrange(qkv, "n h w (t nh e) -> t n nh (h w) e", t=3, e=self.d_head)
+            q, k = scale_for_cosine_sim(q, k, self.scale[:, None, None], 1e-6)
+            theta = theta.movedim(-2, -3)
+            q = apply_rotary_emb_(q, theta)
+            k = apply_rotary_emb_(k, theta)
+            flops.op(flops.op_attention, q.shape, k.shape, v.shape)
+            x = F.scaled_dot_product_attention(q, k, v, scale=1.0)
+            x = rearrange(x, "n nh (h w) e -> n h w (nh e)", h=skip.shape[-3], w=skip.shape[-2])
+        x = self.dropout(x)
+        x = self.out_proj(x)
+        return x + skip
+
+
+class NeighborhoodSelfAttentionBlock(nn.Module):
+    def __init__(self, d_model, d_head, cond_features, kernel_size, dropout=0.0):
+        super().__init__()
+        self.d_head = d_head
+        self.n_heads = d_model // d_head
+        self.kernel_size = kernel_size
+        self.norm = AdaRMSNorm(d_model, cond_features)
+        self.qkv_proj = apply_wd(Linear(d_model, d_model * 3, bias=False))
+        self.scale = nn.Parameter(torch.full([self.n_heads], 10.0))
+        self.pos_emb = AxialRoPE(d_head // 2, self.n_heads)
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = apply_wd(zero_init(Linear(d_model, d_model, bias=False)))
+
+    def extra_repr(self):
+        return f"d_head={self.d_head}, kernel_size={self.kernel_size}"
+
+    def forward(self, x, pos, cond):
+        skip = x
+        x = self.norm(x, cond)
+        qkv = self.qkv_proj(x)
+        if natten is None:
+            raise ModuleNotFoundError("natten is required for neighborhood attention")
+        if natten.has_fused_na():
+            q, k, v = rearrange(qkv, "n h w (t nh e) -> t n h w nh e", t=3, e=self.d_head)
+            q, k = scale_for_cosine_sim(q, k, self.scale[:, None], 1e-6)
+            theta = self.pos_emb(pos)
+            q = apply_rotary_emb_(q, theta)
+            k = apply_rotary_emb_(k, theta)
+            flops.op(flops.op_natten, q.shape, k.shape, v.shape, self.kernel_size)
+            x = natten.functional.na2d(q, k, v, self.kernel_size, scale=1.0)
+            x = rearrange(x, "n h w nh e -> n h w (nh e)")
+        else:
+            q, k, v = rearrange(qkv, "n h w (t nh e) -> t n nh h w e", t=3, e=self.d_head)
+            q, k = scale_for_cosine_sim(q, k, self.scale[:, None, None, None], 1e-6)
+            theta = self.pos_emb(pos).movedim(-2, -4)
+            q = apply_rotary_emb_(q, theta)
+            k = apply_rotary_emb_(k, theta)
+            flops.op(flops.op_natten, q.shape, k.shape, v.shape, self.kernel_size)
+            qk = natten.functional.na2d_qk(q, k, self.kernel_size)
+            a = torch.softmax(qk, dim=-1).to(v.dtype)
+            x = natten.functional.na2d_av(a, v, self.kernel_size)
+            x = rearrange(x, "n nh h w e -> n h w (nh e)")
+        x = self.dropout(x)
+        x = self.out_proj(x)
+        return x + skip
+
+
+class ShiftedWindowSelfAttentionBlock(nn.Module):
+    def __init__(self, d_model, d_head, cond_features, window_size, window_shift, dropout=0.0):
+        super().__init__()
+        self.d_head = d_head
+        self.n_heads = d_model // d_head
+        self.window_size = window_size
+        self.window_shift = window_shift
+        self.norm = AdaRMSNorm(d_model, cond_features)
+        self.qkv_proj = apply_wd(Linear(d_model, d_model * 3, bias=False))
+        self.scale = nn.Parameter(torch.full([self.n_heads], 10.0))
+        self.pos_emb = AxialRoPE(d_head // 2, self.n_heads)
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = apply_wd(zero_init(Linear(d_model, d_model, bias=False)))
+
+    def extra_repr(self):
+        return f"d_head={self.d_head}, window_size={self.window_size}, window_shift={self.window_shift}"
+
+    def forward(self, x, pos, cond):
+        skip = x
+        x = self.norm(x, cond)
+        qkv = self.qkv_proj(x)
+        q, k, v = rearrange(qkv, "n h w (t nh e) -> t n nh h w e", t=3, e=self.d_head)
+        q, k = scale_for_cosine_sim(q, k, self.scale[:, None, None, None], 1e-6)
+        theta = self.pos_emb(pos).movedim(-2, -4)
+        q = apply_rotary_emb_(q, theta)
+        k = apply_rotary_emb_(k, theta)
+        x = apply_window_attention(self.window_size, self.window_shift, q, k, v, scale=1.0)
+        x = rearrange(x, "n nh h w e -> n h w (nh e)")
+        x = self.dropout(x)
+        x = self.out_proj(x)
+        return x + skip
+
+
+class FeedForwardBlock(nn.Module):
+    def __init__(self, d_model, d_ff, cond_features, dropout=0.0):
+        super().__init__()
+        self.norm = AdaRMSNorm(d_model, cond_features)
+        self.up_proj = apply_wd(LinearGEGLU(d_model, d_ff, bias=False))
+        self.dropout = nn.Dropout(dropout)
+        self.down_proj = apply_wd(zero_init(Linear(d_ff, d_model, bias=False)))
+
+    def forward(self, x, cond):
+        skip = x
+        x = self.norm(x, cond)
+        x = self.up_proj(x)
+        x = self.dropout(x)
+        x = self.down_proj(x)
+        return x + skip
+
+
+class GlobalTransformerLayer(nn.Module):
+    def __init__(self, d_model, d_ff, d_head, cond_features, dropout=0.0):
+        super().__init__()
+        self.self_attn = SelfAttentionBlock(d_model, d_head, cond_features, dropout=dropout)
+        self.ff = FeedForwardBlock(d_model, d_ff, cond_features, dropout=dropout)
+
+    def forward(self, x, pos, cond):
+        x = checkpoint(self.self_attn, x, pos, cond)
+        x = checkpoint(self.ff, x, cond)
+        return x
+
+
+class NeighborhoodTransformerLayer(nn.Module):
+    def __init__(self, d_model, d_ff, d_head, cond_features, kernel_size, dropout=0.0):
+        super().__init__()
+        self.self_attn = NeighborhoodSelfAttentionBlock(d_model, d_head, cond_features, kernel_size, dropout=dropout)
+        self.ff = FeedForwardBlock(d_model, d_ff, cond_features, dropout=dropout)
+
+    def forward(self, x, pos, cond):
+        x = checkpoint(self.self_attn, x, pos, cond)
+        x = checkpoint(self.ff, x, cond)
+        return x
+
+
+class ShiftedWindowTransformerLayer(nn.Module):
+    def __init__(self, d_model, d_ff, d_head, cond_features, window_size, index, dropout=0.0):
+        super().__init__()
+        window_shift = window_size // 2 if index % 2 == 1 else 0
+        self.self_attn = ShiftedWindowSelfAttentionBlock(d_model, d_head, cond_features, window_size, window_shift, dropout=dropout)
+        self.ff = FeedForwardBlock(d_model, d_ff, cond_features, dropout=dropout)
+
+    def forward(self, x, pos, cond):
+        x = checkpoint(self.self_attn, x, pos, cond)
+        x = checkpoint(self.ff, x, cond)
+        return x
+
+
+class NoAttentionTransformerLayer(nn.Module):
+    def __init__(self, d_model, d_ff, cond_features, dropout=0.0):
+        super().__init__()
+        self.ff = FeedForwardBlock(d_model, d_ff, cond_features, dropout=dropout)
+
+    def forward(self, x, pos, cond):
+        x = checkpoint(self.ff, x, cond)
+        return x
+
+
+class Level(nn.ModuleList):
+    def forward(self, x, *args, **kwargs):
+        for layer in self:
+            x = layer(x, *args, **kwargs)
+        return x
+
+
+# Mapping network
+
+class MappingFeedForwardBlock(nn.Module):
+    def __init__(self, d_model, d_ff, dropout=0.0):
+        super().__init__()
+        self.norm = RMSNorm(d_model)
+        self.up_proj = apply_wd(LinearGEGLU(d_model, d_ff, bias=False))
+        self.dropout = nn.Dropout(dropout)
+        self.down_proj = apply_wd(zero_init(Linear(d_ff, d_model, bias=False)))
+
+    def forward(self, x):
+        skip = x
+        x = self.norm(x)
+        x = self.up_proj(x)
+        x = self.dropout(x)
+        x = self.down_proj(x)
+        return x + skip
+
+
+class MappingNetwork(nn.Module):
+    def __init__(self, n_layers, d_model, d_ff, dropout=0.0):
+        super().__init__()
+        self.in_norm = RMSNorm(d_model)
+        self.blocks = nn.ModuleList([MappingFeedForwardBlock(d_model, d_ff, dropout=dropout) for _ in range(n_layers)])
+        self.out_norm = RMSNorm(d_model)
+
+    def forward(self, x):
+        x = self.in_norm(x)
+        for block in self.blocks:
+            x = block(x)
+        x = self.out_norm(x)
+        return x
+
+
+# Token merging and splitting
+
+class TokenMerge(nn.Module):
+    def __init__(self, in_features, out_features, patch_size=(2, 2)):
+        super().__init__()
+        self.h = patch_size[0]
+        self.w = patch_size[1]
+        self.proj = apply_wd(Linear(in_features * self.h * self.w, out_features, bias=False))
+
+    def forward(self, x):
+        x = rearrange(x, "... (h nh) (w nw) e -> ... h w (nh nw e)", nh=self.h, nw=self.w)
+        return self.proj(x)
+
+
+class TokenSplitWithoutSkip(nn.Module):
+    def __init__(self, in_features, out_features, patch_size=(2, 2)):
+        super().__init__()
+        self.h = patch_size[0]
+        self.w = patch_size[1]
+        self.proj = apply_wd(Linear(in_features, out_features * self.h * self.w, bias=False))
+
+    def forward(self, x):
+        x = self.proj(x)
+        return rearrange(x, "... h w (nh nw e) -> ... (h nh) (w nw) e", nh=self.h, nw=self.w)
+
+
+class TokenSplit(nn.Module):
+    def __init__(self, in_features, out_features, patch_size=(2, 2)):
+        super().__init__()
+        self.h = patch_size[0]
+        self.w = patch_size[1]
+        self.proj = apply_wd(Linear(in_features, out_features * self.h * self.w, bias=False))
+        self.fac = nn.Parameter(torch.ones(1) * 0.5)
+
+    def forward(self, x, skip):
+        x = self.proj(x)
+        x = rearrange(x, "... h w (nh nw e) -> ... (h nh) (w nw) e", nh=self.h, nw=self.w)
+        return torch.lerp(skip, x, self.fac.to(x.dtype))
+
+
+# Configuration
+
+@dataclass
+class GlobalAttentionSpec:
+    d_head: int
+
+
+@dataclass
+class NeighborhoodAttentionSpec:
+    d_head: int
+    kernel_size: int
+
+
+@dataclass
+class ShiftedWindowAttentionSpec:
+    d_head: int
+    window_size: int
+
+
+@dataclass
+class NoAttentionSpec:
+    pass
+
+
+@dataclass
+class LevelSpec:
+    depth: int
+    width: int
+    d_ff: int
+    self_attn: Union[GlobalAttentionSpec, NeighborhoodAttentionSpec, ShiftedWindowAttentionSpec, NoAttentionSpec]
+    dropout: float
+
+
+@dataclass
+class MappingSpec:
+    depth: int
+    width: int
+    d_ff: int
+    dropout: float
+
+
+# Model class
+
+class ImageTransformerDenoiserModelV2(nn.Module):
+    def __init__(self, levels, mapping, in_channels, out_channels, patch_size, num_classes=0, mapping_cond_dim=0, degradation_params_dim=None):
+        super().__init__()
+        self.num_classes = num_classes
+        self.patch_in = TokenMerge(in_channels, levels[0].width, patch_size)
+        self.mapping_width = mapping.width
+        self.time_emb = FourierFeatures(1, mapping.width)
+        self.time_in_proj = Linear(mapping.width, mapping.width, bias=False)
+        self.aug_emb = FourierFeatures(9, mapping.width)
+        self.aug_in_proj = Linear(mapping.width, mapping.width, bias=False)
+        self.degradation_proj = Linear(degradation_params_dim, mapping.width, bias=False) if degradation_params_dim else None
+        self.class_emb = nn.Embedding(num_classes, mapping.width) if num_classes else None
+        self.mapping_cond_in_proj = Linear(mapping_cond_dim, mapping.width, bias=False) if mapping_cond_dim else None
+        self.mapping = tag_module(MappingNetwork(mapping.depth, mapping.width, mapping.d_ff, dropout=mapping.dropout), "mapping")
+
+        self.down_levels, self.up_levels = nn.ModuleList(), nn.ModuleList()
+        for i, spec in enumerate(levels):
+            if isinstance(spec.self_attn, GlobalAttentionSpec):
+                layer_factory = lambda _: GlobalTransformerLayer(spec.width, spec.d_ff, spec.self_attn.d_head, mapping.width, dropout=spec.dropout)
+            elif isinstance(spec.self_attn, NeighborhoodAttentionSpec):
+                layer_factory = lambda _: NeighborhoodTransformerLayer(spec.width, spec.d_ff, spec.self_attn.d_head, mapping.width, spec.self_attn.kernel_size, dropout=spec.dropout)
+            elif isinstance(spec.self_attn, ShiftedWindowAttentionSpec):
+                layer_factory = lambda i: ShiftedWindowTransformerLayer(spec.width, spec.d_ff, spec.self_attn.d_head, mapping.width, spec.self_attn.window_size, i, dropout=spec.dropout)
+            elif isinstance(spec.self_attn, NoAttentionSpec):
+                layer_factory = lambda _: NoAttentionTransformerLayer(spec.width, spec.d_ff, mapping.width, dropout=spec.dropout)
+            else:
+                raise ValueError(f"unsupported self attention spec {spec.self_attn}")
+
+            if i < len(levels) - 1:
+                self.down_levels.append(Level([layer_factory(i) for i in range(spec.depth)]))
+                self.up_levels.append(Level([layer_factory(i + spec.depth) for i in range(spec.depth)]))
+            else:
+                self.mid_level = Level([layer_factory(i) for i in range(spec.depth)])
+
+        self.merges = nn.ModuleList([TokenMerge(spec_1.width, spec_2.width) for spec_1, spec_2 in zip(levels[:-1], levels[1:])])
+        self.splits = nn.ModuleList([TokenSplit(spec_2.width, spec_1.width) for spec_1, spec_2 in zip(levels[:-1], levels[1:])])
+
+        self.out_norm = RMSNorm(levels[0].width)
+        self.patch_out = TokenSplitWithoutSkip(levels[0].width, out_channels, patch_size)
+        nn.init.zeros_(self.patch_out.proj.weight)
+
+    def param_groups(self, base_lr=5e-4, mapping_lr_scale=1 / 3):
+        wd = filter_params(lambda tags: "wd" in tags and "mapping" not in tags, self)
+        no_wd = filter_params(lambda tags: "wd" not in tags and "mapping" not in tags, self)
+        mapping_wd = filter_params(lambda tags: "wd" in tags and "mapping" in tags, self)
+        mapping_no_wd = filter_params(lambda tags: "wd" not in tags and "mapping" in tags, self)
+        groups = [
+            {"params": list(wd), "lr": base_lr},
+            {"params": list(no_wd), "lr": base_lr, "weight_decay": 0.0},
+            {"params": list(mapping_wd), "lr": base_lr * mapping_lr_scale},
+            {"params": list(mapping_no_wd), "lr": base_lr * mapping_lr_scale, "weight_decay": 0.0}
+        ]
+        return groups
+
+    def forward(self, x, sigma=None, aug_cond=None, class_cond=None, mapping_cond=None, degradation_params=None):
+        # Patching
+        x = x.movedim(-3, -1)
+        x = self.patch_in(x)
+        # TODO: pixel aspect ratio for nonsquare patches
+        pos = make_axial_pos(x.shape[-3], x.shape[-2], device=x.device).view(x.shape[-3], x.shape[-2], 2)
+
+        # Mapping network
+        if class_cond is None and self.class_emb is not None:
+            raise ValueError("class_cond must be specified if num_classes > 0")
+        if mapping_cond is None and self.mapping_cond_in_proj is not None:
+            raise ValueError("mapping_cond must be specified if mapping_cond_dim > 0")
+
+        # c_noise = torch.log(sigma) / 4
+        # c_noise = (sigma * 2.0 - 1.0)
+        # c_noise = sigma * 2 - 1
+        if sigma is not None:
+            time_emb = self.time_in_proj(self.time_emb(sigma[..., None]))
+        else:
+            time_emb = self.time_in_proj(torch.ones(1, 1, device=x.device, dtype=x.dtype).expand(x.shape[0], self.mapping_width))
+        # time_emb = self.time_in_proj(sigma[..., None])
+
+        aug_cond = x.new_zeros([x.shape[0], 9]) if aug_cond is None else aug_cond
+        aug_emb = self.aug_in_proj(self.aug_emb(aug_cond))
+        class_emb = self.class_emb(class_cond) if self.class_emb is not None else 0
+        mapping_emb = self.mapping_cond_in_proj(mapping_cond) if self.mapping_cond_in_proj is not None else 0
+        degradation_emb = self.degradation_proj(degradation_params) if degradation_params is not None else 0
+        cond = self.mapping(time_emb + aug_emb + class_emb + mapping_emb + degradation_emb)
+
+        # Hourglass transformer
+        skips, poses = [], []
+        for down_level, merge in zip(self.down_levels, self.merges):
+            x = down_level(x, pos, cond)
+            skips.append(x)
+            poses.append(pos)
+            x = merge(x)
+            pos = downscale_pos(pos)
+
+        x = self.mid_level(x, pos, cond)
+
+        for up_level, split, skip, pos in reversed(list(zip(self.up_levels, self.splits, skips, poses))):
+            x = split(x, skip)
+            x = up_level(x, pos, cond)
+
+        # Unpatching
+        x = self.out_norm(x)
+        x = self.patch_out(x)
+        x = x.movedim(-1, -3)
+
+        return x

arch/swinir/swinir.py

@@ -0,0 +1,904 @@
+# -----------------------------------------------------------------------------------
+# SwinIR: Image Restoration Using Swin Transformer, https://arxiv.org/abs/2108.10257
+# Originally Written by Ze Liu, Modified by Jingyun Liang.
+# -----------------------------------------------------------------------------------
+# Borrowed from DifFace (https://github.com/zsyOAOA/DifFace/blob/master/models/swinir.py)
+
+import math
+from typing import Set
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        # coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        # Fix: Pass indexing="ij" to avoid warning
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
+
+    def flops(self, N):
+        # calculate flops for 1 window with token length of N
+        flops = 0
+        # qkv = self.qkv(x)
+        flops += N * self.dim * 3 * self.dim
+        # attn = (q @ k.transpose(-2, -1))
+        flops += self.num_heads * N * (self.dim // self.num_heads) * N
+        #  x = (attn @ v)
+        flops += self.num_heads * N * N * (self.dim // self.num_heads)
+        # x = self.proj(x)
+        flops += N * self.dim * self.dim
+        return flops
+
+
+class SwinTransformerBlock(nn.Module):
+    r""" Swin Transformer Block.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        if min(self.input_resolution) <= self.window_size:
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if self.shift_size > 0:
+            attn_mask = self.calculate_mask(self.input_resolution)
+        else:
+            attn_mask = None
+
+        self.register_buffer("attn_mask", attn_mask)
+
+    def calculate_mask(self, x_size):
+        # calculate attention mask for SW-MSA
+        H, W = x_size
+        img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
+        h_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        w_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+
+        return attn_mask
+
+    def forward(self, x, x_size):
+        H, W = x_size
+        B, L, C = x.shape
+        # assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+        else:
+            shifted_x = x
+
+        # partition windows
+        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
+        if self.input_resolution == x_size:
+            attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
+        else:
+            attn_windows = self.attn(x_windows, mask=self.calculate_mask(x_size).to(x.device))
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
+               f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
+
+    def flops(self):
+        flops = 0
+        H, W = self.input_resolution
+        # norm1
+        flops += self.dim * H * W
+        # W-MSA/SW-MSA
+        nW = H * W / self.window_size / self.window_size
+        flops += nW * self.attn.flops(self.window_size * self.window_size)
+        # mlp
+        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
+        # norm2
+        flops += self.dim * H * W
+        return flops
+
+
+class PatchMerging(nn.Module):
+    r""" Patch Merging Layer.
+
+    Args:
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x):
+        """
+        x: B, H*W, C
+        """
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
+
+        x = x.view(B, H, W, C)
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"input_resolution={self.input_resolution}, dim={self.dim}"
+
+    def flops(self):
+        H, W = self.input_resolution
+        flops = H * W * self.dim
+        flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
+        return flops
+
+
+class BasicLayer(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False):
+
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
+                                 num_heads=num_heads, window_size=window_size,
+                                 shift_size=0 if (i % 2 == 0) else window_size // 2,
+                                 mlp_ratio=mlp_ratio,
+                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                 drop=drop, attn_drop=attn_drop,
+                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                 norm_layer=norm_layer)
+            for i in range(depth)])
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x, x_size):
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x, x_size)
+            else:
+                x = blk(x, x_size)
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
+
+    def flops(self):
+        flops = 0
+        for blk in self.blocks:
+            flops += blk.flops()
+        if self.downsample is not None:
+            flops += self.downsample.flops()
+        return flops
+
+
+class RSTB(nn.Module):
+    """Residual Swin Transformer Block (RSTB).
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+        img_size: Input image size.
+        patch_size: Patch size.
+        resi_connection: The convolutional block before residual connection.
+    """
+
+    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
+                 img_size=224, patch_size=4, resi_connection='1conv'):
+        super(RSTB, self).__init__()
+
+        self.dim = dim
+        self.input_resolution = input_resolution
+
+        self.residual_group = BasicLayer(dim=dim,
+                                         input_resolution=input_resolution,
+                                         depth=depth,
+                                         num_heads=num_heads,
+                                         window_size=window_size,
+                                         mlp_ratio=mlp_ratio,
+                                         qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                         drop=drop, attn_drop=attn_drop,
+                                         drop_path=drop_path,
+                                         norm_layer=norm_layer,
+                                         downsample=downsample,
+                                         use_checkpoint=use_checkpoint)
+
+        if resi_connection == '1conv':
+            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
+        elif resi_connection == '3conv':
+            # to save parameters and memory
+            self.conv = nn.Sequential(nn.Conv2d(dim, dim // 4, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                                      nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),
+                                      nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                                      nn.Conv2d(dim // 4, dim, 3, 1, 1))
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim,
+            norm_layer=None)
+
+        self.patch_unembed = PatchUnEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim,
+            norm_layer=None)
+
+    def forward(self, x, x_size):
+        return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))) + x
+
+    def flops(self):
+        flops = 0
+        flops += self.residual_group.flops()
+        H, W = self.input_resolution
+        flops += H * W * self.dim * self.dim * 9
+        flops += self.patch_embed.flops()
+        flops += self.patch_unembed.flops()
+
+        return flops
+
+
+class PatchEmbed(nn.Module):
+    r""" Image to Patch Embedding
+
+    Args:
+        img_size (int): Image size.  Default: 224.
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        x = x.flatten(2).transpose(1, 2)  # B Ph*Pw C
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+    def flops(self):
+        flops = 0
+        H, W = self.img_size
+        if self.norm is not None:
+            flops += H * W * self.embed_dim
+        return flops
+
+
+class PatchUnEmbed(nn.Module):
+    r""" Image to Patch Unembedding
+
+    Args:
+        img_size (int): Image size.  Default: 224.
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+    def forward(self, x, x_size):
+        B, HW, C = x.shape
+        x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1])  # B Ph*Pw C
+        return x
+
+    def flops(self):
+        flops = 0
+        return flops
+
+
+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 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 SwinIR(nn.Module):
+    r""" SwinIR
+        A PyTorch impl of : `SwinIR: Image Restoration Using Swin Transformer`, based on Swin Transformer.
+
+    Args:
+        img_size (int | tuple(int)): Input image size. Default 64
+        patch_size (int | tuple(int)): Patch size. Default: 1
+        in_chans (int): Number of input image channels. Default: 3
+        embed_dim (int): Patch embedding dimension. Default: 96
+        depths (tuple(int)): Depth of each Swin Transformer layer.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        window_size (int): Window size. Default: 7
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
+        sf: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction
+        img_range: Image range. 1. or 255.
+        upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None
+        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
+    """
+
+    def __init__(
+            self,
+            img_size=64,
+            patch_size=1,
+            in_chans=3,
+            num_out_ch=3,
+            embed_dim=96,
+            depths=[6, 6, 6, 6],
+            num_heads=[6, 6, 6, 6],
+            window_size=7,
+            mlp_ratio=4.,
+            qkv_bias=True,
+            qk_scale=None,
+            drop_rate=0.,
+            attn_drop_rate=0.,
+            drop_path_rate=0.1,
+            norm_layer=nn.LayerNorm,
+            ape=False,
+            patch_norm=True,
+            use_checkpoint=False,
+            sf=4,
+            img_range=1.,
+            upsampler='',
+            resi_connection='1conv',
+            unshuffle=False,
+            unshuffle_scale=None,
+            hq_key: str = "jpg",
+            lq_key: str = "hint",
+            learning_rate: float = None,
+            weight_decay: float = None
+    ) -> "SwinIR":
+        super(SwinIR, self).__init__()
+        num_in_ch = in_chans * (unshuffle_scale ** 2) if unshuffle else in_chans
+        num_feat = 64
+        self.img_range = img_range
+        if in_chans == 3:
+            rgb_mean = (0.4488, 0.4371, 0.4040)
+            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
+        else:
+            self.mean = torch.zeros(1, 1, 1, 1)
+        self.upscale = sf
+        self.upsampler = upsampler
+        self.window_size = window_size
+        self.unshuffle_scale = unshuffle_scale
+        self.unshuffle = unshuffle
+
+        #####################################################################################################
+        ################################### 1, shallow feature extraction ###################################
+        if unshuffle:
+            assert unshuffle_scale is not None
+            self.conv_first = nn.Sequential(
+                nn.PixelUnshuffle(sf),
+                nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1),
+            )
+        else:
+            self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)
+
+        #####################################################################################################
+        ################################### 2, deep feature extraction ######################################
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.num_features = embed_dim
+        self.mlp_ratio = mlp_ratio
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None
+        )
+        num_patches = self.patch_embed.num_patches
+        patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution
+
+        # merge non-overlapping patches into image
+        self.patch_unembed = PatchUnEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None
+        )
+
+        # absolute position embedding
+        if self.ape:
+            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
+
+        # build Residual Swin Transformer blocks (RSTB)
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = RSTB(
+                dim=embed_dim,
+                input_resolution=(patches_resolution[0], patches_resolution[1]),
+                depth=depths[i_layer],
+                num_heads=num_heads[i_layer],
+                window_size=window_size,
+                mlp_ratio=self.mlp_ratio,
+                qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate,
+                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],  # no impact on SR results
+                norm_layer=norm_layer,
+                downsample=None,
+                use_checkpoint=use_checkpoint,
+                img_size=img_size,
+                patch_size=patch_size,
+                resi_connection=resi_connection
+            )
+            self.layers.append(layer)
+        self.norm = norm_layer(self.num_features)
+
+        # build the last conv layer in deep feature extraction
+        if resi_connection == '1conv':
+            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
+        elif resi_connection == '3conv':
+            # to save parameters and memory
+            self.conv_after_body = nn.Sequential(
+                nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),
+                nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),
+                nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1)
+            )
+
+        #####################################################################################################
+        ################################ 3, high quality image reconstruction ################################
+        if self.upsampler == 'pixelshuffle':
+            # for classical SR
+            self.conv_before_upsample = nn.Sequential(
+                nn.Conv2d(embed_dim, num_feat, 3, 1, 1),
+                nn.LeakyReLU(inplace=True)
+            )
+            self.upsample = Upsample(sf, num_feat)
+            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+        elif self.upsampler == 'pixelshuffledirect':
+            # for lightweight SR (to save parameters)
+            self.upsample = UpsampleOneStep(
+                sf, embed_dim, num_out_ch,
+                (patches_resolution[0], patches_resolution[1])
+            )
+        elif self.upsampler == 'nearest+conv':
+            # for real-world SR (less artifacts)
+            self.conv_before_upsample = nn.Sequential(
+                nn.Conv2d(embed_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)
+            elif self.upscale == 8:
+                self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+                self.conv_up3 = 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)
+        else:
+            # for image denoising and JPEG compression artifact reduction
+            self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m: nn.Module) -> None:
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    # TODO: What's this ?
+    @torch.jit.ignore
+    def no_weight_decay(self) -> Set[str]:
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self) -> Set[str]:
+        return {'relative_position_bias_table'}
+
+    def check_image_size(self, x: torch.Tensor) -> torch.Tensor:
+        _, _, h, w = x.size()
+        mod_pad_h = (self.window_size - h % self.window_size) % self.window_size
+        mod_pad_w = (self.window_size - w % self.window_size) % self.window_size
+        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect')
+        return x
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        x_size = (x.shape[2], x.shape[3])
+        x = self.patch_embed(x)
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+
+        for layer in self.layers:
+            x = layer(x, x_size)
+
+        x = self.norm(x)  # B L C
+        x = self.patch_unembed(x, x_size)
+
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        H, W = x.shape[2:]
+        x = self.check_image_size(x)
+
+        self.mean = self.mean.type_as(x)
+        x = (x - self.mean) * self.img_range
+
+        if self.upsampler == 'pixelshuffle':
+            # for classical SR
+            x = self.conv_first(x)
+            x = self.conv_after_body(self.forward_features(x)) + x
+            x = self.conv_before_upsample(x)
+            x = self.conv_last(self.upsample(x))
+        elif self.upsampler == 'pixelshuffledirect':
+            # for lightweight SR
+            x = self.conv_first(x)
+            x = self.conv_after_body(self.forward_features(x)) + x
+            x = self.upsample(x)
+        elif self.upsampler == 'nearest+conv':
+            # for real-world SR
+            x = self.conv_first(x)
+            x = self.conv_after_body(self.forward_features(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')))
+            elif self.upscale == 8:
+                x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
+                x = self.lrelu(self.conv_up3(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
+            x = self.conv_last(self.lrelu(self.conv_hr(x)))
+        else:
+            # for image denoising and JPEG compression artifact reduction
+            x_first = self.conv_first(x)
+            res = self.conv_after_body(self.forward_features(x_first)) + x_first
+            x = x + self.conv_last(res)
+
+        x = x / self.img_range + self.mean
+
+        return x[:, :, :H * self.upscale, :W * self.upscale]
+
+    def flops(self) -> int:
+        flops = 0
+        H, W = self.patches_resolution
+        flops += H * W * 3 * self.embed_dim * 9
+        flops += self.patch_embed.flops()
+        for i, layer in enumerate(self.layers):
+            flops += layer.flops()
+        flops += H * W * 3 * self.embed_dim * self.embed_dim
+        flops += self.upsample.flops()
+        return flops

packages.txt

@@ -0,0 +1,3 @@
+ffmpeg
+libsm6
+libxext6

requirements.txt

@@ -0,0 +1,21 @@
+torch==2.2.2
+facexlib==0.2.5
+realesrgan==0.2.5
+numpy
+opencv-python
+torchvision
+pytorch-lightning==2.4.0
+scipy
+tqdm
+lmdb
+pyyaml
+basicsr==1.4.2
+yapf
+dctorch
+einops
+torch-ema==0.3
+huggingface_hub==0.24.5
+natten==0.17.1
+wandb
+timm
+huggingface_hub==0.24.5

utils/basicsr_custom.py

@@ -0,0 +1,954 @@
+# https://github.com/XPixelGroup/BasicSR/blob/master/basicsr/data/degradations.py
+# Copyright (c) OpenMMLab. All rights reserved.
+# https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
+
+import math
+import random
+import re
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import List, Dict
+from typing import Mapping, Any
+from typing import Optional, Union
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+from scipy import special
+from scipy.stats import multivariate_normal
+from torch import Tensor
+# from torchvision.transforms.functional_tensor import rgb_to_grayscale
+from torchvision.transforms._functional_tensor import rgb_to_grayscale
+
+
+# -------------------------------------------------------------------- #
+# --------------------------- blur kernels --------------------------- #
+# -------------------------------------------------------------------- #
+
+
+# --------------------------- util functions --------------------------- #
+def sigma_matrix2(sig_x, sig_y, theta):
+    """Calculate the rotated sigma matrix (two dimensional matrix).
+
+    Args:
+        sig_x (float):
+        sig_y (float):
+        theta (float): Radian measurement.
+
+    Returns:
+        ndarray: Rotated sigma matrix.
+    """
+    d_matrix = np.array([[sig_x ** 2, 0], [0, sig_y ** 2]])
+    u_matrix = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])
+    return np.dot(u_matrix, np.dot(d_matrix, u_matrix.T))
+
+
+def mesh_grid(kernel_size):
+    """Generate the mesh grid, centering at zero.
+
+    Args:
+        kernel_size (int):
+
+    Returns:
+        xy (ndarray): with the shape (kernel_size, kernel_size, 2)
+        xx (ndarray): with the shape (kernel_size, kernel_size)
+        yy (ndarray): with the shape (kernel_size, kernel_size)
+    """
+    ax = np.arange(-kernel_size // 2 + 1., kernel_size // 2 + 1.)
+    xx, yy = np.meshgrid(ax, ax)
+    xy = np.hstack((xx.reshape((kernel_size * kernel_size, 1)), yy.reshape(kernel_size * kernel_size,
+                                                                           1))).reshape(kernel_size, kernel_size, 2)
+    return xy, xx, yy
+
+
+def pdf2(sigma_matrix, grid):
+    """Calculate PDF of the bivariate Gaussian distribution.
+
+    Args:
+        sigma_matrix (ndarray): with the shape (2, 2)
+        grid (ndarray): generated by :func:`mesh_grid`,
+            with the shape (K, K, 2), K is the kernel size.
+
+    Returns:
+        kernel (ndarrray): un-normalized kernel.
+    """
+    inverse_sigma = np.linalg.inv(sigma_matrix)
+    kernel = np.exp(-0.5 * np.sum(np.dot(grid, inverse_sigma) * grid, 2))
+    return kernel
+
+
+def cdf2(d_matrix, grid):
+    """Calculate the CDF of the standard bivariate Gaussian distribution.
+        Used in skewed Gaussian distribution.
+
+    Args:
+        d_matrix (ndarrasy): skew matrix.
+        grid (ndarray): generated by :func:`mesh_grid`,
+            with the shape (K, K, 2), K is the kernel size.
+
+    Returns:
+        cdf (ndarray): skewed cdf.
+    """
+    rv = multivariate_normal([0, 0], [[1, 0], [0, 1]])
+    grid = np.dot(grid, d_matrix)
+    cdf = rv.cdf(grid)
+    return cdf
+
+
+def bivariate_Gaussian(kernel_size, sig_x, sig_y, theta, grid=None, isotropic=True):
+    """Generate a bivariate isotropic or anisotropic Gaussian kernel.
+
+    In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.
+
+    Args:
+        kernel_size (int):
+        sig_x (float):
+        sig_y (float):
+        theta (float): Radian measurement.
+        grid (ndarray, optional): generated by :func:`mesh_grid`,
+            with the shape (K, K, 2), K is the kernel size. Default: None
+        isotropic (bool):
+
+    Returns:
+        kernel (ndarray): normalized kernel.
+    """
+    if grid is None:
+        grid, _, _ = mesh_grid(kernel_size)
+    if isotropic:
+        sigma_matrix = np.array([[sig_x ** 2, 0], [0, sig_x ** 2]])
+    else:
+        sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)
+    kernel = pdf2(sigma_matrix, grid)
+    kernel = kernel / np.sum(kernel)
+    return kernel
+
+
+def bivariate_generalized_Gaussian(kernel_size, sig_x, sig_y, theta, beta, grid=None, isotropic=True):
+    """Generate a bivariate generalized Gaussian kernel.
+
+    ``Paper: Parameter Estimation For Multivariate Generalized Gaussian Distributions``
+
+    In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.
+
+    Args:
+        kernel_size (int):
+        sig_x (float):
+        sig_y (float):
+        theta (float): Radian measurement.
+        beta (float): shape parameter, beta = 1 is the normal distribution.
+        grid (ndarray, optional): generated by :func:`mesh_grid`,
+            with the shape (K, K, 2), K is the kernel size. Default: None
+
+    Returns:
+        kernel (ndarray): normalized kernel.
+    """
+    if grid is None:
+        grid, _, _ = mesh_grid(kernel_size)
+    if isotropic:
+        sigma_matrix = np.array([[sig_x ** 2, 0], [0, sig_x ** 2]])
+    else:
+        sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)
+    inverse_sigma = np.linalg.inv(sigma_matrix)
+    kernel = np.exp(-0.5 * np.power(np.sum(np.dot(grid, inverse_sigma) * grid, 2), beta))
+    kernel = kernel / np.sum(kernel)
+    return kernel
+
+
+def bivariate_plateau(kernel_size, sig_x, sig_y, theta, beta, grid=None, isotropic=True):
+    """Generate a plateau-like anisotropic kernel.
+
+    1 / (1+x^(beta))
+
+    Reference: https://stats.stackexchange.com/questions/203629/is-there-a-plateau-shaped-distribution
+
+    In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.
+
+    Args:
+        kernel_size (int):
+        sig_x (float):
+        sig_y (float):
+        theta (float): Radian measurement.
+        beta (float): shape parameter, beta = 1 is the normal distribution.
+        grid (ndarray, optional): generated by :func:`mesh_grid`,
+            with the shape (K, K, 2), K is the kernel size. Default: None
+
+    Returns:
+        kernel (ndarray): normalized kernel.
+    """
+    if grid is None:
+        grid, _, _ = mesh_grid(kernel_size)
+    if isotropic:
+        sigma_matrix = np.array([[sig_x ** 2, 0], [0, sig_x ** 2]])
+    else:
+        sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)
+    inverse_sigma = np.linalg.inv(sigma_matrix)
+    kernel = np.reciprocal(np.power(np.sum(np.dot(grid, inverse_sigma) * grid, 2), beta) + 1)
+    kernel = kernel / np.sum(kernel)
+    return kernel
+
+
+def random_bivariate_Gaussian(kernel_size,
+                              sigma_x_range,
+                              sigma_y_range,
+                              rotation_range,
+                              noise_range=None,
+                              isotropic=True):
+    """Randomly generate bivariate isotropic or anisotropic Gaussian kernels.
+
+    In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.
+
+    Args:
+        kernel_size (int):
+        sigma_x_range (tuple): [0.6, 5]
+        sigma_y_range (tuple): [0.6, 5]
+        rotation range (tuple): [-math.pi, math.pi]
+        noise_range(tuple, optional): multiplicative kernel noise,
+            [0.75, 1.25]. Default: None
+
+    Returns:
+        kernel (ndarray):
+    """
+    assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
+    assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
+    sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
+    if isotropic is False:
+        assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'
+        assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'
+        sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])
+        rotation = np.random.uniform(rotation_range[0], rotation_range[1])
+    else:
+        sigma_y = sigma_x
+        rotation = 0
+
+    kernel = bivariate_Gaussian(kernel_size, sigma_x, sigma_y, rotation, isotropic=isotropic)
+
+    # add multiplicative noise
+    if noise_range is not None:
+        assert noise_range[0] < noise_range[1], 'Wrong noise range.'
+        noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)
+        kernel = kernel * noise
+    kernel = kernel / np.sum(kernel)
+    return kernel
+
+
+def random_bivariate_generalized_Gaussian(kernel_size,
+                                          sigma_x_range,
+                                          sigma_y_range,
+                                          rotation_range,
+                                          beta_range,
+                                          noise_range=None,
+                                          isotropic=True):
+    """Randomly generate bivariate generalized Gaussian kernels.
+
+    In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.
+
+    Args:
+        kernel_size (int):
+        sigma_x_range (tuple): [0.6, 5]
+        sigma_y_range (tuple): [0.6, 5]
+        rotation range (tuple): [-math.pi, math.pi]
+        beta_range (tuple): [0.5, 8]
+        noise_range(tuple, optional): multiplicative kernel noise,
+            [0.75, 1.25]. Default: None
+
+    Returns:
+        kernel (ndarray):
+    """
+    assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
+    assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
+    sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
+    if isotropic is False:
+        assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'
+        assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'
+        sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])
+        rotation = np.random.uniform(rotation_range[0], rotation_range[1])
+    else:
+        sigma_y = sigma_x
+        rotation = 0
+
+    # assume beta_range[0] < 1 < beta_range[1]
+    if np.random.uniform() < 0.5:
+        beta = np.random.uniform(beta_range[0], 1)
+    else:
+        beta = np.random.uniform(1, beta_range[1])
+
+    kernel = bivariate_generalized_Gaussian(kernel_size, sigma_x, sigma_y, rotation, beta, isotropic=isotropic)
+
+    # add multiplicative noise
+    if noise_range is not None:
+        assert noise_range[0] < noise_range[1], 'Wrong noise range.'
+        noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)
+        kernel = kernel * noise
+    kernel = kernel / np.sum(kernel)
+    return kernel
+
+
+def random_bivariate_plateau(kernel_size,
+                             sigma_x_range,
+                             sigma_y_range,
+                             rotation_range,
+                             beta_range,
+                             noise_range=None,
+                             isotropic=True):
+    """Randomly generate bivariate plateau kernels.
+
+    In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.
+
+    Args:
+        kernel_size (int):
+        sigma_x_range (tuple): [0.6, 5]
+        sigma_y_range (tuple): [0.6, 5]
+        rotation range (tuple): [-math.pi/2, math.pi/2]
+        beta_range (tuple): [1, 4]
+        noise_range(tuple, optional): multiplicative kernel noise,
+            [0.75, 1.25]. Default: None
+
+    Returns:
+        kernel (ndarray):
+    """
+    assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
+    assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
+    sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
+    if isotropic is False:
+        assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'
+        assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'
+        sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])
+        rotation = np.random.uniform(rotation_range[0], rotation_range[1])
+    else:
+        sigma_y = sigma_x
+        rotation = 0
+
+    # TODO: this may be not proper
+    if np.random.uniform() < 0.5:
+        beta = np.random.uniform(beta_range[0], 1)
+    else:
+        beta = np.random.uniform(1, beta_range[1])
+
+    kernel = bivariate_plateau(kernel_size, sigma_x, sigma_y, rotation, beta, isotropic=isotropic)
+    # add multiplicative noise
+    if noise_range is not None:
+        assert noise_range[0] < noise_range[1], 'Wrong noise range.'
+        noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)
+        kernel = kernel * noise
+    kernel = kernel / np.sum(kernel)
+
+    return kernel
+
+
+def random_mixed_kernels(kernel_list,
+                         kernel_prob,
+                         kernel_size=21,
+                         sigma_x_range=(0.6, 5),
+                         sigma_y_range=(0.6, 5),
+                         rotation_range=(-math.pi, math.pi),
+                         betag_range=(0.5, 8),
+                         betap_range=(0.5, 8),
+                         noise_range=None):
+    """Randomly generate mixed kernels.
+
+    Args:
+        kernel_list (tuple): a list name of kernel types,
+            support ['iso', 'aniso', 'skew', 'generalized', 'plateau_iso',
+            'plateau_aniso']
+        kernel_prob (tuple): corresponding kernel probability for each
+            kernel type
+        kernel_size (int):
+        sigma_x_range (tuple): [0.6, 5]
+        sigma_y_range (tuple): [0.6, 5]
+        rotation range (tuple): [-math.pi, math.pi]
+        beta_range (tuple): [0.5, 8]
+        noise_range(tuple, optional): multiplicative kernel noise,
+            [0.75, 1.25]. Default: None
+
+    Returns:
+        kernel (ndarray):
+    """
+    kernel_type = random.choices(kernel_list, kernel_prob)[0]
+    if kernel_type == 'iso':
+        kernel = random_bivariate_Gaussian(
+            kernel_size, sigma_x_range, sigma_y_range, rotation_range, noise_range=noise_range, isotropic=True)
+    elif kernel_type == 'aniso':
+        kernel = random_bivariate_Gaussian(
+            kernel_size, sigma_x_range, sigma_y_range, rotation_range, noise_range=noise_range, isotropic=False)
+    elif kernel_type == 'generalized_iso':
+        kernel = random_bivariate_generalized_Gaussian(
+            kernel_size,
+            sigma_x_range,
+            sigma_y_range,
+            rotation_range,
+            betag_range,
+            noise_range=noise_range,
+            isotropic=True)
+    elif kernel_type == 'generalized_aniso':
+        kernel = random_bivariate_generalized_Gaussian(
+            kernel_size,
+            sigma_x_range,
+            sigma_y_range,
+            rotation_range,
+            betag_range,
+            noise_range=noise_range,
+            isotropic=False)
+    elif kernel_type == 'plateau_iso':
+        kernel = random_bivariate_plateau(
+            kernel_size, sigma_x_range, sigma_y_range, rotation_range, betap_range, noise_range=None, isotropic=True)
+    elif kernel_type == 'plateau_aniso':
+        kernel = random_bivariate_plateau(
+            kernel_size, sigma_x_range, sigma_y_range, rotation_range, betap_range, noise_range=None, isotropic=False)
+    return kernel
+
+
+np.seterr(divide='ignore', invalid='ignore')
+
+
+def circular_lowpass_kernel(cutoff, kernel_size, pad_to=0):
+    """2D sinc filter
+
+    Reference: https://dsp.stackexchange.com/questions/58301/2-d-circularly-symmetric-low-pass-filter
+
+    Args:
+        cutoff (float): cutoff frequency in radians (pi is max)
+        kernel_size (int): horizontal and vertical size, must be odd.
+        pad_to (int): pad kernel size to desired size, must be odd or zero.
+    """
+    assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
+    kernel = np.fromfunction(
+        lambda x, y: cutoff * special.j1(cutoff * np.sqrt(
+            (x - (kernel_size - 1) / 2) ** 2 + (y - (kernel_size - 1) / 2) ** 2)) / (2 * np.pi * np.sqrt(
+            (x - (kernel_size - 1) / 2) ** 2 + (y - (kernel_size - 1) / 2) ** 2)), [kernel_size, kernel_size])
+    kernel[(kernel_size - 1) // 2, (kernel_size - 1) // 2] = cutoff ** 2 / (4 * np.pi)
+    kernel = kernel / np.sum(kernel)
+    if pad_to > kernel_size:
+        pad_size = (pad_to - kernel_size) // 2
+        kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
+    return kernel
+
+
+# ------------------------------------------------------------- #
+# --------------------------- noise --------------------------- #
+# ------------------------------------------------------------- #
+
+# ----------------------- Gaussian Noise ----------------------- #
+
+def instantiate_from_config(config: Mapping[str, Any]) -> Any:
+    if not "target" in config:
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+class BaseStorageBackend(metaclass=ABCMeta):
+    """Abstract class of storage backends.
+
+    All backends need to implement two apis: ``get()`` and ``get_text()``.
+    ``get()`` reads the file as a byte stream and ``get_text()`` reads the file
+    as texts.
+    """
+
+    @property
+    def name(self) -> str:
+        return self.__class__.__name__
+
+    @abstractmethod
+    def get(self, filepath: str) -> bytes:
+        pass
+
+
+class PetrelBackend(BaseStorageBackend):
+    """Petrel storage backend (for internal use).
+
+    PetrelBackend supports reading and writing data to multiple clusters.
+    If the file path contains the cluster name, PetrelBackend will read data
+    from specified cluster or write data to it. Otherwise, PetrelBackend will
+    access the default cluster.
+
+    Args:
+        path_mapping (dict, optional): Path mapping dict from local path to
+            Petrel path. When ``path_mapping={'src': 'dst'}``, ``src`` in
+            ``filepath`` will be replaced by ``dst``. Default: None.
+        enable_mc (bool, optional): Whether to enable memcached support.
+            Default: True.
+        conf_path (str, optional): Config path of Petrel client. Default: None.
+            `New in version 1.7.1`.
+
+    Examples:
+        >>> filepath1 = 's3://path/of/file'
+        >>> filepath2 = 'cluster-name:s3://path/of/file'
+        >>> client = PetrelBackend()
+        >>> client.get(filepath1)  # get data from default cluster
+        >>> client.get(filepath2)  # get data from 'cluster-name' cluster
+    """
+
+    def __init__(self,
+                 path_mapping: Optional[dict] = None,
+                 enable_mc: bool = False,
+                 conf_path: str = None):
+        try:
+            from petrel_client import client
+        except ImportError:
+            raise ImportError('Please install petrel_client to enable '
+                              'PetrelBackend.')
+
+        self._client = client.Client(conf_path=conf_path, enable_mc=enable_mc)
+        assert isinstance(path_mapping, dict) or path_mapping is None
+        self.path_mapping = path_mapping
+
+    def _map_path(self, filepath: Union[str, Path]) -> str:
+        """Map ``filepath`` to a string path whose prefix will be replaced by
+        :attr:`self.path_mapping`.
+
+        Args:
+            filepath (str): Path to be mapped.
+        """
+        filepath = str(filepath)
+        if self.path_mapping is not None:
+            for k, v in self.path_mapping.items():
+                filepath = filepath.replace(k, v, 1)
+        return filepath
+
+    def _format_path(self, filepath: str) -> str:
+        """Convert a ``filepath`` to standard format of petrel oss.
+
+        If the ``filepath`` is concatenated by ``os.path.join``, in a Windows
+        environment, the ``filepath`` will be the format of
+        's3://bucket_name\\image.jpg'. By invoking :meth:`_format_path`, the
+        above ``filepath`` will be converted to 's3://bucket_name/image.jpg'.
+
+        Args:
+            filepath (str): Path to be formatted.
+        """
+        return re.sub(r'\\+', '/', filepath)
+
+    def get(self, filepath: Union[str, Path]) -> bytes:
+        """Read data from a given ``filepath`` with 'rb' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+
+        Returns:
+            bytes: The loaded bytes.
+        """
+        filepath = self._map_path(filepath)
+        filepath = self._format_path(filepath)
+        value = self._client.Get(filepath)
+        return value
+
+
+class HardDiskBackend(BaseStorageBackend):
+    """Raw hard disks storage backend."""
+
+    def get(self, filepath: Union[str, Path]) -> bytes:
+        """Read data from a given ``filepath`` with 'rb' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+
+        Returns:
+            bytes: Expected bytes object.
+        """
+        with open(filepath, 'rb') as f:
+            value_buf = f.read()
+        return value_buf
+
+
+def generate_gaussian_noise(img, sigma=10, gray_noise=False):
+    """Generate Gaussian noise.
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        sigma (float): Noise scale (measured in range 255). Default: 10.
+
+    Returns:
+        (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    if gray_noise:
+        noise = np.float32(np.random.randn(*(img.shape[0:2]))) * sigma / 255.
+        noise = np.expand_dims(noise, axis=2).repeat(3, axis=2)
+    else:
+        noise = np.float32(np.random.randn(*(img.shape))) * sigma / 255.
+    return noise
+
+
+def add_gaussian_noise(img, sigma=10, clip=True, rounds=False, gray_noise=False):
+    """Add Gaussian noise.
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        sigma (float): Noise scale (measured in range 255). Default: 10.
+
+    Returns:
+        (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    noise = generate_gaussian_noise(img, sigma, gray_noise)
+    out = img + noise
+    if clip and rounds:
+        out = np.clip((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = np.clip(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+def generate_gaussian_noise_pt(img, sigma=10, gray_noise=0):
+    """Add Gaussian noise (PyTorch version).
+
+    Args:
+        img (Tensor): Shape (b, c, h, w), range[0, 1], float32.
+        scale (float | Tensor): Noise scale. Default: 1.0.
+
+    Returns:
+        (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
+            float32.
+    """
+    b, _, h, w = img.size()
+    if not isinstance(sigma, (float, int)):
+        sigma = sigma.view(img.size(0), 1, 1, 1)
+    if isinstance(gray_noise, (float, int)):
+        cal_gray_noise = gray_noise > 0
+    else:
+        gray_noise = gray_noise.view(b, 1, 1, 1)
+        cal_gray_noise = torch.sum(gray_noise) > 0
+
+    if cal_gray_noise:
+        noise_gray = torch.randn(*img.size()[2:4], dtype=img.dtype, device=img.device) * sigma / 255.
+        noise_gray = noise_gray.view(b, 1, h, w)
+
+    # always calculate color noise
+    noise = torch.randn(*img.size(), dtype=img.dtype, device=img.device) * sigma / 255.
+
+    if cal_gray_noise:
+        noise = noise * (1 - gray_noise) + noise_gray * gray_noise
+    return noise
+
+
+def add_gaussian_noise_pt(img, sigma=10, gray_noise=0, clip=True, rounds=False):
+    """Add Gaussian noise (PyTorch version).
+
+    Args:
+        img (Tensor): Shape (b, c, h, w), range[0, 1], float32.
+        scale (float | Tensor): Noise scale. Default: 1.0.
+
+    Returns:
+        (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
+            float32.
+    """
+    noise = generate_gaussian_noise_pt(img, sigma, gray_noise)
+    out = img + noise
+    if clip and rounds:
+        out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = torch.clamp(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+# ----------------------- Random Gaussian Noise ----------------------- #
+def random_generate_gaussian_noise(img, sigma_range=(0, 10), gray_prob=0):
+    sigma = np.random.uniform(sigma_range[0], sigma_range[1])
+    if np.random.uniform() < gray_prob:
+        gray_noise = True
+    else:
+        gray_noise = False
+    return generate_gaussian_noise(img, sigma, gray_noise)
+
+
+def random_add_gaussian_noise(img, sigma_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
+    noise = random_generate_gaussian_noise(img, sigma_range, gray_prob)
+    out = img + noise
+    if clip and rounds:
+        out = np.clip((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = np.clip(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+def random_generate_gaussian_noise_pt(img, sigma_range=(0, 10), gray_prob=0):
+    sigma = torch.rand(
+        img.size(0), dtype=img.dtype, device=img.device) * (sigma_range[1] - sigma_range[0]) + sigma_range[0]
+    gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=img.device)
+    gray_noise = (gray_noise < gray_prob).float()
+    return generate_gaussian_noise_pt(img, sigma, gray_noise)
+
+
+def random_add_gaussian_noise_pt(img, sigma_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
+    noise = random_generate_gaussian_noise_pt(img, sigma_range, gray_prob)
+    out = img + noise
+    if clip and rounds:
+        out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = torch.clamp(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+# ----------------------- Poisson (Shot) Noise ----------------------- #
+
+
+def generate_poisson_noise(img, scale=1.0, gray_noise=False):
+    """Generate poisson noise.
+
+    Reference: https://github.com/scikit-image/scikit-image/blob/main/skimage/util/noise.py#L37-L219
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        scale (float): Noise scale. Default: 1.0.
+        gray_noise (bool): Whether generate gray noise. Default: False.
+
+    Returns:
+        (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    if gray_noise:
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    # round and clip image for counting vals correctly
+    img = np.clip((img * 255.0).round(), 0, 255) / 255.
+    vals = len(np.unique(img))
+    vals = 2 ** np.ceil(np.log2(vals))
+    out = np.float32(np.random.poisson(img * vals) / float(vals))
+    noise = out - img
+    if gray_noise:
+        noise = np.repeat(noise[:, :, np.newaxis], 3, axis=2)
+    return noise * scale
+
+
+def add_poisson_noise(img, scale=1.0, clip=True, rounds=False, gray_noise=False):
+    """Add poisson noise.
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        scale (float): Noise scale. Default: 1.0.
+        gray_noise (bool): Whether generate gray noise. Default: False.
+
+    Returns:
+        (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    noise = generate_poisson_noise(img, scale, gray_noise)
+    out = img + noise
+    if clip and rounds:
+        out = np.clip((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = np.clip(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+def generate_poisson_noise_pt(img, scale=1.0, gray_noise=0):
+    """Generate a batch of poisson noise (PyTorch version)
+
+    Args:
+        img (Tensor): Input image, shape (b, c, h, w), range [0, 1], float32.
+        scale (float | Tensor): Noise scale. Number or Tensor with shape (b).
+            Default: 1.0.
+        gray_noise (float | Tensor): 0-1 number or Tensor with shape (b).
+            0 for False, 1 for True. Default: 0.
+
+    Returns:
+        (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
+            float32.
+    """
+    b, _, h, w = img.size()
+    if isinstance(gray_noise, (float, int)):
+        cal_gray_noise = gray_noise > 0
+    else:
+        gray_noise = gray_noise.view(b, 1, 1, 1)
+        cal_gray_noise = torch.sum(gray_noise) > 0
+    if cal_gray_noise:
+        img_gray = rgb_to_grayscale(img, num_output_channels=1)
+        # round and clip image for counting vals correctly
+        img_gray = torch.clamp((img_gray * 255.0).round(), 0, 255) / 255.
+        # use for-loop to get the unique values for each sample
+        vals_list = [len(torch.unique(img_gray[i, :, :, :])) for i in range(b)]
+        vals_list = [2 ** np.ceil(np.log2(vals)) for vals in vals_list]
+        vals = img_gray.new_tensor(vals_list).view(b, 1, 1, 1)
+        out = torch.poisson(img_gray * vals) / vals
+        noise_gray = out - img_gray
+        noise_gray = noise_gray.expand(b, 3, h, w)
+
+    # always calculate color noise
+    # round and clip image for counting vals correctly
+    img = torch.clamp((img * 255.0).round(), 0, 255) / 255.
+    # use for-loop to get the unique values for each sample
+    vals_list = [len(torch.unique(img[i, :, :, :])) for i in range(b)]
+    vals_list = [2 ** np.ceil(np.log2(vals)) for vals in vals_list]
+    vals = img.new_tensor(vals_list).view(b, 1, 1, 1)
+    out = torch.poisson(img * vals) / vals
+    noise = out - img
+    if cal_gray_noise:
+        noise = noise * (1 - gray_noise) + noise_gray * gray_noise
+    if not isinstance(scale, (float, int)):
+        scale = scale.view(b, 1, 1, 1)
+    return noise * scale
+
+
+def add_poisson_noise_pt(img, scale=1.0, clip=True, rounds=False, gray_noise=0):
+    """Add poisson noise to a batch of images (PyTorch version).
+
+    Args:
+        img (Tensor): Input image, shape (b, c, h, w), range [0, 1], float32.
+        scale (float | Tensor): Noise scale. Number or Tensor with shape (b).
+            Default: 1.0.
+        gray_noise (float | Tensor): 0-1 number or Tensor with shape (b).
+            0 for False, 1 for True. Default: 0.
+
+    Returns:
+        (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
+            float32.
+    """
+    noise = generate_poisson_noise_pt(img, scale, gray_noise)
+    out = img + noise
+    if clip and rounds:
+        out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = torch.clamp(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+# ----------------------- Random Poisson (Shot) Noise ----------------------- #
+
+
+def random_generate_poisson_noise(img, scale_range=(0, 1.0), gray_prob=0):
+    scale = np.random.uniform(scale_range[0], scale_range[1])
+    if np.random.uniform() < gray_prob:
+        gray_noise = True
+    else:
+        gray_noise = False
+    return generate_poisson_noise(img, scale, gray_noise)
+
+
+def random_add_poisson_noise(img, scale_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
+    noise = random_generate_poisson_noise(img, scale_range, gray_prob)
+    out = img + noise
+    if clip and rounds:
+        out = np.clip((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = np.clip(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+def random_generate_poisson_noise_pt(img, scale_range=(0, 1.0), gray_prob=0):
+    scale = torch.rand(
+        img.size(0), dtype=img.dtype, device=img.device) * (scale_range[1] - scale_range[0]) + scale_range[0]
+    gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=img.device)
+    gray_noise = (gray_noise < gray_prob).float()
+    return generate_poisson_noise_pt(img, scale, gray_noise)
+
+
+def random_add_poisson_noise_pt(img, scale_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
+    noise = random_generate_poisson_noise_pt(img, scale_range, gray_prob)
+    out = img + noise
+    if clip and rounds:
+        out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = torch.clamp(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+# ------------------------------------------------------------------------ #
+# --------------------------- JPEG compression --------------------------- #
+# ------------------------------------------------------------------------ #
+
+
+def add_jpg_compression(img, quality=90):
+    """Add JPG compression artifacts.
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        quality (float): JPG compression quality. 0 for lowest quality, 100 for
+            best quality. Default: 90.
+
+    Returns:
+        (Numpy array): Returned image after JPG, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    img = np.clip(img, 0, 1)
+    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
+    _, encimg = cv2.imencode('.jpg', img * 255., encode_param)
+    img = np.float32(cv2.imdecode(encimg, 1)) / 255.
+    return img
+
+
+def random_add_jpg_compression(img, quality_range=(90, 100)):
+    """Randomly add JPG compression artifacts.
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        quality_range (tuple[float] | list[float]): JPG compression quality
+            range. 0 for lowest quality, 100 for best quality.
+            Default: (90, 100).
+
+    Returns:
+        (Numpy array): Returned image after JPG, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    quality = np.random.uniform(quality_range[0], quality_range[1])
+    return add_jpg_compression(img, int(quality))
+
+
+def load_file_list(file_list_path: str) -> List[Dict[str, str]]:
+    files = []
+    with open(file_list_path, "r") as fin:
+        for line in fin:
+            path = line.strip()
+            if path:
+                files.append({"image_path": path, "prompt": ""})
+    return files
+
+
+# https://github.com/openai/guided-diffusion/blob/main/guided_diffusion/image_datasets.py
+def center_crop_arr(pil_image, image_size):
+    # We are not on a new enough PIL to support the `reducing_gap`
+    # argument, which uses BOX downsampling at powers of two first.
+    # Thus, we do it by hand to improve downsample quality.
+    while min(*pil_image.size) >= 2 * image_size:
+        pil_image = pil_image.resize(
+            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
+        )
+
+    scale = image_size / min(*pil_image.size)
+    pil_image = pil_image.resize(
+        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+    )
+
+    arr = np.array(pil_image)
+    crop_y = (arr.shape[0] - image_size) // 2
+    crop_x = (arr.shape[1] - image_size) // 2
+    return arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size]
+
+
+# https://github.com/openai/guided-diffusion/blob/main/guided_diffusion/image_datasets.py
+def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0):
+    min_smaller_dim_size = math.ceil(image_size / max_crop_frac)
+    max_smaller_dim_size = math.ceil(image_size / min_crop_frac)
+    smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1)
+
+    # We are not on a new enough PIL to support the `reducing_gap`
+    # argument, which uses BOX downsampling at powers of two first.
+    # Thus, we do it by hand to improve downsample quality.
+    while min(*pil_image.size) >= 2 * smaller_dim_size:
+        pil_image = pil_image.resize(
+            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
+        )
+
+    scale = smaller_dim_size / min(*pil_image.size)
+    pil_image = pil_image.resize(
+        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+    )
+
+    arr = np.array(pil_image)
+    crop_y = random.randrange(arr.shape[0] - image_size + 1)
+    crop_x = random.randrange(arr.shape[1] - image_size + 1)
+    return arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size]

utils/create_arch.py

@@ -0,0 +1,143 @@
+from arch.hourglass import image_transformer_v2 as itv2
+from arch.hourglass.image_transformer_v2 import ImageTransformerDenoiserModelV2
+from arch.swinir.swinir import SwinIR
+
+
+def create_arch(arch, condition_channels=0):
+    # arch should be, e.g., swinir_XL, or hdit_XL
+    arch_name, arch_size = arch.split('_')
+    arch_config = arch_configs[arch_name][arch_size].copy()
+    arch_config['in_channels'] += condition_channels
+    return arch_name_to_object[arch_name](**arch_config)
+
+
+arch_configs = {
+    'hdit': {
+        "ImageNet256Sp4": {
+            'in_channels': 3,
+            'out_channels': 3,
+            'widths': [256, 512, 1024],
+            'depths': [2, 2, 8],
+            'patch_size': [4, 4],
+            'self_attns': [
+                {"type": "neighborhood", "d_head": 64, "kernel_size": 7},
+                {"type": "neighborhood", "d_head": 64, "kernel_size": 7},
+                {"type": "global", "d_head": 64}
+            ],
+            'mapping_depth': 2,
+            'mapping_width': 768,
+            'dropout_rate': [0, 0, 0],
+            'mapping_dropout_rate': 0.0
+        },
+        "XL2": {
+            'in_channels': 3,
+            'out_channels': 3,
+            'widths': [384, 768],
+            'depths': [2, 11],
+            'patch_size': [4, 4],
+            'self_attns': [
+                {"type": "neighborhood", "d_head": 64, "kernel_size": 7},
+                {"type": "global", "d_head": 64}
+            ],
+            'mapping_depth': 2,
+            'mapping_width': 768,
+            'dropout_rate': [0, 0],
+            'mapping_dropout_rate': 0.0
+        }
+
+    },
+    'swinir': {
+        "M": {
+            'in_channels': 3,
+            'out_channels': 3,
+            'embed_dim': 120,
+            'depths': [6, 6, 6, 6, 6],
+            'num_heads': [6, 6, 6, 6, 6],
+            'resi_connection': '1conv',
+            'sf': 8
+
+        },
+        "L": {
+            'in_channels': 3,
+            'out_channels': 3,
+            'embed_dim': 180,
+            'depths': [6, 6, 6, 6, 6, 6, 6, 6],
+            'num_heads': [6, 6, 6, 6, 6, 6, 6, 6],
+            'resi_connection': '1conv',
+            'sf': 8
+        },
+    },
+}
+
+
+def create_swinir_model(in_channels, out_channels, embed_dim, depths, num_heads, resi_connection,
+                        sf):
+    return SwinIR(
+        img_size=64,
+        patch_size=1,
+        in_chans=in_channels,
+        num_out_ch=out_channels,
+        embed_dim=embed_dim,
+        depths=depths,
+        num_heads=num_heads,
+        window_size=8,
+        mlp_ratio=2,
+        sf=sf,
+        img_range=1.0,
+        upsampler="nearest+conv",
+        resi_connection=resi_connection,
+        unshuffle=True,
+        unshuffle_scale=8
+    )
+
+
+def create_hdit_model(widths,
+                      depths,
+                      self_attns,
+                      dropout_rate,
+                      mapping_depth,
+                      mapping_width,
+                      mapping_dropout_rate,
+                      in_channels,
+                      out_channels,
+                      patch_size
+                      ):
+    assert len(widths) == len(depths)
+    assert len(widths) == len(self_attns)
+    assert len(widths) == len(dropout_rate)
+    mapping_d_ff = mapping_width * 3
+    d_ffs = []
+    for width in widths:
+        d_ffs.append(width * 3)
+
+    levels = []
+    for depth, width, d_ff, self_attn, dropout in zip(depths, widths, d_ffs, self_attns, dropout_rate):
+        if self_attn['type'] == 'global':
+            self_attn = itv2.GlobalAttentionSpec(self_attn.get('d_head', 64))
+        elif self_attn['type'] == 'neighborhood':
+            self_attn = itv2.NeighborhoodAttentionSpec(self_attn.get('d_head', 64), self_attn.get('kernel_size', 7))
+        elif self_attn['type'] == 'shifted-window':
+            self_attn = itv2.ShiftedWindowAttentionSpec(self_attn.get('d_head', 64), self_attn['window_size'])
+        elif self_attn['type'] == 'none':
+            self_attn = itv2.NoAttentionSpec()
+        else:
+            raise ValueError(f'unsupported self attention type {self_attn["type"]}')
+        levels.append(itv2.LevelSpec(depth, width, d_ff, self_attn, dropout))
+    mapping = itv2.MappingSpec(mapping_depth, mapping_width, mapping_d_ff, mapping_dropout_rate)
+    model = ImageTransformerDenoiserModelV2(
+        levels=levels,
+        mapping=mapping,
+        in_channels=in_channels,
+        out_channels=out_channels,
+        patch_size=patch_size,
+        num_classes=0,
+        mapping_cond_dim=0,
+    )
+
+    return model
+
+
+arch_name_to_object = {
+    'hdit': create_hdit_model,
+    'swinir': create_swinir_model,
+}

utils/create_degradation.py

@@ -0,0 +1,144 @@
+import math
+from functools import partial
+
+import cv2
+import numpy as np
+import torch
+from basicsr.data import degradations as degradations
+from basicsr.data.transforms import augment
+from basicsr.utils import img2tensor
+from torch.nn.functional import interpolate
+from torchvision.transforms import Compose
+from utils.basicsr_custom import (
+    random_mixed_kernels,
+    random_add_gaussian_noise,
+    random_add_jpg_compression,
+)
+
+
+def create_degradation(degradation):
+    if degradation == 'sr_bicubic_x8_gaussian_noise_005':
+        return Compose([
+            partial(down_scale, scale_factor=1.0 / 8.0, mode='bicubic'),
+            partial(add_gaussian_noise, std=0.05),
+            partial(interpolate, scale_factor=8.0, mode='nearest-exact'),
+            partial(torch.clip, min=0, max=1),
+            partial(torch.squeeze, dim=0),
+            lambda x: (x, None)
+
+        ])
+    elif degradation == 'gaussian_noise_035':
+        return Compose([
+            partial(add_gaussian_noise, std=0.35),
+            partial(torch.clip, min=0, max=1),
+            partial(torch.squeeze, dim=0),
+            lambda x: (x, None)
+
+        ])
+    elif degradation == 'colorization_gaussian_noise_025':
+        return Compose([
+            lambda x: torch.mean(x, dim=0, keepdim=True),
+            partial(add_gaussian_noise, std=0.25),
+            partial(torch.clip, min=0, max=1),
+            lambda x: (x, None)
+        ])
+    elif degradation == 'random_inpainting_gaussian_noise_01':
+        def inpainting_dps(x):
+            total = x.shape[1] ** 2
+            # random pixel sampling
+            l, h = [0.9, 0.9]
+            prob = np.random.uniform(l, h)
+            mask_vec = torch.ones([1, x.shape[1] * x.shape[1]])
+            samples = np.random.choice(x.shape[1] * x.shape[1], int(total * prob), replace=False)
+            mask_vec[:, samples] = 0
+            mask_b = mask_vec.view(1, x.shape[1], x.shape[1])
+            mask_b = mask_b.repeat(3, 1, 1)
+            mask = torch.ones_like(x, device=x.device)
+            mask[:, ...] = mask_b
+            return add_gaussian_noise(x * mask, 0.1).clip(0, 1), None
+
+        return inpainting_dps
+    elif degradation == 'difface':
+        def deg(x):
+            blur_kernel_size = 41
+            kernel_list = ['iso', 'aniso']
+            kernel_prob = [0.5, 0.5]
+            blur_sigma = [0.1, 15]
+            downsample_range = [0.8, 32]
+            noise_range = [0, 20]
+            jpeg_range = [30, 100]
+            gt_gray = True
+            gray_prob = 0.01
+            x = x.permute(1, 2, 0).numpy()[..., ::-1].astype(np.float32)
+            # random horizontal flip
+            img_gt = augment(x.copy(), hflip=True, rotation=False)
+            h, w, _ = img_gt.shape
+
+            # ------------------------ generate lq image ------------------------ #
+            # blur
+            kernel = degradations.random_mixed_kernels(
+                kernel_list,
+                kernel_prob,
+                blur_kernel_size,
+                blur_sigma,
+                blur_sigma, [-math.pi, math.pi],
+                noise_range=None)
+            img_lq = cv2.filter2D(img_gt, -1, kernel)
+            # downsample
+            scale = np.random.uniform(downsample_range[0], downsample_range[1])
+            img_lq = cv2.resize(img_lq, (int(w // scale), int(h // scale)), interpolation=cv2.INTER_LINEAR)
+            # noise
+            if noise_range is not None:
+                img_lq = random_add_gaussian_noise(img_lq, noise_range)
+            # jpeg compression
+            if jpeg_range is not None:
+                img_lq = random_add_jpg_compression(img_lq, jpeg_range)
+
+            # resize to original size
+            img_lq = cv2.resize(img_lq, (w, h), interpolation=cv2.INTER_LINEAR)
+
+            # random color jitter (only for lq)
+            # if self.color_jitter_prob is not None and (np.random.uniform() < self.color_jitter_prob):
+            #     img_lq = self.color_jitter(img_lq, self.color_jitter_shift)
+            # random to gray (only for lq)
+            if np.random.uniform() < gray_prob:
+                img_lq = cv2.cvtColor(img_lq, cv2.COLOR_BGR2GRAY)
+                img_lq = np.tile(img_lq[:, :, None], [1, 1, 3])
+                if gt_gray:  # whether convert GT to gray images
+                    img_gt = cv2.cvtColor(img_gt, cv2.COLOR_BGR2GRAY)
+                    img_gt = np.tile(img_gt[:, :, None], [1, 1, 3])  # repeat the color channels
+
+            # BGR to RGB, HWC to CHW, numpy to tensor
+            img_gt, img_lq = img2tensor([img_gt, img_lq], bgr2rgb=True, float32=True)
+
+            # random color jitter (pytorch version) (only for lq)
+            # if self.color_jitter_pt_prob is not None and (np.random.uniform() < self.color_jitter_pt_prob):
+            #     brightness = self.opt.get('brightness', (0.5, 1.5))
+            #     contrast = self.opt.get('contrast', (0.5, 1.5))
+            #     saturation = self.opt.get('saturation', (0, 1.5))
+            #     hue = self.opt.get('hue', (-0.1, 0.1))
+            #     img_lq = self.color_jitter_pt(img_lq, brightness, contrast, saturation, hue)
+
+            # round and clip
+            img_lq = torch.clamp((img_lq * 255.0).round(), 0, 255) / 255.
+
+            return img_lq, img_gt.clip(0, 1)
+
+        return deg
+    else:
+        raise NotImplementedError()
+
+
+def down_scale(x, scale_factor, mode):
+    with torch.no_grad():
+        return interpolate(x.unsqueeze(0),
+                           scale_factor=scale_factor,
+                           mode=mode,
+                           antialias=True,
+                           align_corners=False).clip(0, 1)
+
+
+def add_gaussian_noise(x, std):
+    with torch.no_grad():
+        x = x + torch.randn_like(x) * std
+        return x

utils/img_utils.py

@@ -0,0 +1,5 @@
+from torchvision.utils import make_grid
+
+
+def create_grid(img, normalize=False, num_images=5):
+    return make_grid(img[:num_images], padding=0, normalize=normalize, nrow=16)