Initial push

SUPIR/models/SUPIR_model.py

@@ -0,0 +1,183 @@
+import torch
+from ...sgm.models.diffusion import DiffusionEngine
+from ...sgm.util import instantiate_from_config
+import copy
+from ...sgm.modules.distributions.distributions import DiagonalGaussianDistribution
+import random
+from ...SUPIR.utils.colorfix import wavelet_reconstruction, adaptive_instance_normalization
+from pytorch_lightning import seed_everything
+from torch.nn.functional import interpolate
+from ...SUPIR.utils.tilevae import VAEHook
+import importlib
+import os
+
+class SUPIRModel(DiffusionEngine):
+    def __init__(self, control_stage_config, ae_dtype='fp32', diffusion_dtype='fp32', p_p='', n_p='', *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        control_model = instantiate_from_config(control_stage_config)
+        self.model.load_control_model(control_model)
+        self.first_stage_model.denoise_encoder = copy.deepcopy(self.first_stage_model.encoder)
+        self.sampler_config = kwargs['sampler_config']
+
+        assert (ae_dtype in ['fp32', 'fp16', 'bf16']) and (diffusion_dtype in ['fp32', 'fp16', 'bf16'])
+        if ae_dtype == 'fp32':
+            ae_dtype = torch.float32
+        elif ae_dtype == 'fp16':
+            raise RuntimeError('fp16 cause NaN in AE')
+        elif ae_dtype == 'bf16':
+            ae_dtype = torch.bfloat16
+
+        if diffusion_dtype == 'fp32':
+            diffusion_dtype = torch.float32
+        elif diffusion_dtype == 'fp16':
+            diffusion_dtype = torch.float16
+        elif diffusion_dtype == 'bf16':
+            diffusion_dtype = torch.bfloat16
+
+        self.ae_dtype = ae_dtype
+        self.model.dtype = diffusion_dtype
+
+        self.p_p = p_p
+        self.n_p = n_p
+
+    @torch.no_grad()
+    def encode_first_stage(self, x):
+        with torch.autocast("cuda", dtype=self.ae_dtype):
+            z = self.first_stage_model.encode(x)
+        z = self.scale_factor * z
+        return z
+
+    @torch.no_grad()
+    def encode_first_stage_with_denoise(self, x, use_sample=True, is_stage1=False):
+        with torch.autocast("cuda", dtype=self.ae_dtype):
+            if is_stage1:
+                h = self.first_stage_model.denoise_encoder_s1(x)
+            else:
+                h = self.first_stage_model.denoise_encoder(x)
+            moments = self.first_stage_model.quant_conv(h)
+            posterior = DiagonalGaussianDistribution(moments)
+            if use_sample:
+                z = posterior.sample()
+            else:
+                z = posterior.mode()
+        z = self.scale_factor * z
+        return z
+
+    @torch.no_grad()
+    def decode_first_stage(self, z):
+        z = 1.0 / self.scale_factor * z
+        with torch.autocast("cuda", dtype=self.ae_dtype):
+            out = self.first_stage_model.decode(z)
+        return out.float()
+
+    @torch.no_grad()
+    def batchify_denoise(self, x, is_stage1=False):
+        '''
+        [N, C, H, W], [-1, 1], RGB
+        '''
+        x = self.encode_first_stage_with_denoise(x, use_sample=False, is_stage1=is_stage1)
+        return self.decode_first_stage(x)
+
+    @torch.no_grad()
+    def batchify_sample(self, x, p, p_p='default', n_p='default', num_steps=100, restoration_scale=4.0, s_churn=0, s_noise=1.003, cfg_scale=4.0, seed=-1,
+                        num_samples=1, control_scale=1, color_fix_type='None', use_linear_CFG=False, use_linear_control_scale=False,
+                        cfg_scale_start=1.0, control_scale_start=0.0, **kwargs):
+        '''
+        [N, C], [-1, 1], RGB
+        '''
+        assert len(x) == len(p)
+        assert color_fix_type in ['Wavelet', 'AdaIn', 'None']
+
+        N = len(x)
+        if num_samples > 1:
+            assert N == 1
+            N = num_samples
+            x = x.repeat(N, 1, 1, 1)
+            p = p * N
+
+        if p_p == 'default':
+            p_p = self.p_p
+        if n_p == 'default':
+            n_p = self.n_p
+
+        self.sampler_config.params.num_steps = num_steps
+        if use_linear_CFG:
+            self.sampler_config.params.guider_config.params.scale_min = cfg_scale
+            self.sampler_config.params.guider_config.params.scale = cfg_scale_start
+        else:
+            self.sampler_config.params.guider_config.params.scale = cfg_scale
+        self.sampler_config.params.restore_cfg = restoration_scale
+        self.sampler_config.params.s_churn = s_churn
+        self.sampler_config.params.s_noise = s_noise
+        self.sampler = instantiate_from_config(self.sampler_config)
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        _z = self.encode_first_stage_with_denoise(x, use_sample=False)
+
+        x_stage1 = self.decode_first_stage(_z)
+        # x_stage1 = interpolate(x_stage1, scale_factor=scale_factor, mode='bilinear', antialias=True)
+        # _z = self.encode_first_stage_with_denoise(x_stage1)
+
+        z_stage1 = self.encode_first_stage(x_stage1)
+
+        batch = {}
+        batch['txt'] = [''.join([_p, p_p]) for _p in p]
+        batch['original_size_as_tuple'] = torch.tensor([1024, 1024]).repeat(N, 1).to(x.device)
+        batch['crop_coords_top_left'] = torch.tensor([0, 0]).repeat(N, 1).to(x.device)
+        batch['target_size_as_tuple'] = torch.tensor([1024, 1024]).repeat(N, 1).to(x.device)
+        batch['aesthetic_score'] = torch.tensor([9.0]).repeat(N, 1).to(x.device)
+        batch['control'] = _z
+
+        batch_uc = copy.deepcopy(batch)
+        batch_uc['txt'] = [n_p for _ in p]
+
+        with torch.cuda.amp.autocast(dtype=self.ae_dtype):
+            c, uc = self.conditioner.get_unconditional_conditioning(batch, batch_uc)
+
+        denoiser = lambda input, sigma, c, control_scale: self.denoiser(
+            self.model, input, sigma, c, control_scale, **kwargs
+        )
+
+        noised_z = torch.randn_like(_z).to(_z.device)
+
+        _samples = self.sampler(denoiser, noised_z, cond=c, uc=uc, x_center=z_stage1, control_scale=control_scale,
+                                use_linear_control_scale=use_linear_control_scale, control_scale_start=control_scale_start)
+        samples = self.decode_first_stage(_samples)
+        if color_fix_type == 'Wavelet':
+            samples = wavelet_reconstruction(samples, x_stage1)
+        elif color_fix_type == 'AdaIn':
+            samples = adaptive_instance_normalization(samples, x_stage1)
+        return samples
+
+    def init_tile_vae(self, encoder_tile_size=512, decoder_tile_size=64):
+        self.first_stage_model.denoise_encoder.original_forward = self.first_stage_model.denoise_encoder.forward
+        self.first_stage_model.encoder.original_forward = self.first_stage_model.encoder.forward
+        self.first_stage_model.decoder.original_forward = self.first_stage_model.decoder.forward
+        self.first_stage_model.denoise_encoder.forward = VAEHook(
+            self.first_stage_model.denoise_encoder, encoder_tile_size, is_decoder=False, fast_decoder=False,
+            fast_encoder=False, color_fix=False, to_gpu=True)
+        self.first_stage_model.encoder.forward = VAEHook(
+            self.first_stage_model.encoder, encoder_tile_size, is_decoder=False, fast_decoder=False,
+            fast_encoder=False, color_fix=False, to_gpu=True)
+        self.first_stage_model.decoder.forward = VAEHook(
+            self.first_stage_model.decoder, decoder_tile_size, is_decoder=True, fast_decoder=False,
+            fast_encoder=False, color_fix=False, to_gpu=True)
+
+
+if __name__ == '__main__':
+    from SUPIR.util import create_model, load_state_dict
+
+    model = create_model('../../options/dev/SUPIR_paper_version.yaml')
+
+    SDXL_CKPT = '/opt/data/private/AIGC_pretrain/SDXL_cache/sd_xl_base_1.0_0.9vae.safetensors'
+    SUPIR_CKPT = '/opt/data/private/AIGC_pretrain/SUPIR_cache/SUPIR-paper.ckpt'
+    model.load_state_dict(load_state_dict(SDXL_CKPT), strict=False)
+    model.load_state_dict(load_state_dict(SUPIR_CKPT), strict=False)
+    model = model.cuda()
+
+    x = torch.randn(1, 3, 512, 512).cuda()
+    p = ['a professional, detailed, high-quality photo']
+    samples = model.batchify_sample(x, p, num_steps=50, restoration_scale=4.0, s_churn=0, cfg_scale=4.0, seed=-1, num_samples=1)

SUPIR/modules/SUPIR_v0.py

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

SUPIR/modules/__init__.py

@@ -0,0 +1,11 @@
+SDXL_BASE_CHANNEL_DICT = {
+    'cond_output_channels': [320] * 4 + [640] * 3 + [1280] * 3,
+    'project_channels': [160] * 4 + [320] * 3 + [640] * 3,
+    'concat_channels': [320] * 2 + [640] * 3 + [1280] * 4 + [0]
+}
+
+SDXL_REFINE_CHANNEL_DICT = {
+    'cond_output_channels': [384] * 4 + [768] * 3 + [1536] * 6,
+    'project_channels': [192] * 4 + [384] * 3 + [768] * 6,
+    'concat_channels': [384] * 2 + [768] * 3 + [1536] * 7 + [0]
+}

SUPIR/util.py

@@ -0,0 +1,173 @@
+import os
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from torch.nn.functional import interpolate
+from omegaconf import OmegaConf
+from ..sgm.util import instantiate_from_config
+
+
+def get_state_dict(d):
+    return d.get('state_dict', d)
+
+
+def load_state_dict(ckpt_path, location='cpu'):
+    _, extension = os.path.splitext(ckpt_path)
+    if extension.lower() == ".safetensors":
+        import safetensors.torch
+        state_dict = safetensors.torch.load_file(ckpt_path, device=location)
+    else:
+        state_dict = get_state_dict(torch.load(ckpt_path, map_location=torch.device(location)))
+    state_dict = get_state_dict(state_dict)
+    print(f'Loaded state_dict from [{ckpt_path}]')
+    return state_dict
+
+
+def create_model(config_path):
+    config = OmegaConf.load(config_path)
+    model = instantiate_from_config(config.model).cpu()
+    print(f'Loaded model config from [{config_path}]')
+    return model
+
+
+def create_SUPIR_model(config_path, SUPIR_sign=None):
+    config = OmegaConf.load(config_path)
+    model = instantiate_from_config(config.model).cpu()
+    print(f'Loaded model config from [{config_path}]')
+    if config.SDXL_CKPT is not None:
+        model.load_state_dict(load_state_dict(config.SDXL_CKPT), strict=False)
+    if config.SUPIR_CKPT is not None:
+        model.load_state_dict(load_state_dict(config.SUPIR_CKPT), strict=False)
+    if SUPIR_sign is not None:
+        assert SUPIR_sign in ['F', 'Q']
+        if SUPIR_sign == 'F':
+            model.load_state_dict(load_state_dict(config.SUPIR_CKPT_F), strict=False)
+        elif SUPIR_sign == 'Q':
+            model.load_state_dict(load_state_dict(config.SUPIR_CKPT_Q), strict=False)
+    return model
+
+def load_QF_ckpt(config_path):
+    config = OmegaConf.load(config_path)
+    ckpt_F = torch.load(config.SUPIR_CKPT_F, map_location='cpu')
+    ckpt_Q = torch.load(config.SUPIR_CKPT_Q, map_location='cpu')
+    return ckpt_Q, ckpt_F
+
+
+def PIL2Tensor(img, upsacle=1, min_size=1024):
+    '''
+    PIL.Image -> Tensor[C, H, W], RGB, [-1, 1]
+    '''
+    # size
+    w, h = img.size
+    w *= upsacle
+    h *= upsacle
+    w0, h0 = round(w), round(h)
+    if min(w, h) < min_size:
+        _upsacle = min_size / min(w, h)
+        w *= _upsacle
+        h *= _upsacle
+    else:
+        _upsacle = 1
+    w = int(np.round(w / 64.0)) * 64
+    h = int(np.round(h / 64.0)) * 64
+    x = img.resize((w, h), Image.BICUBIC)
+    x = np.array(x).round().clip(0, 255).astype(np.uint8)
+    x = x / 255 * 2 - 1
+    x = torch.tensor(x, dtype=torch.float32).permute(2, 0, 1)
+    return x, h0, w0
+
+
+def Tensor2PIL(x, h0, w0):
+    '''
+    Tensor[C, H, W], RGB, [-1, 1] -> PIL.Image
+    '''
+    x = x.unsqueeze(0)
+    x = interpolate(x, size=(h0, w0), mode='bicubic')
+    x = (x.squeeze(0).permute(1, 2, 0) * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
+    return Image.fromarray(x)
+
+
+def HWC3(x):
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    assert x.ndim == 3
+    H, W, C = x.shape
+    assert C == 1 or C == 3 or C == 4
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    if C == 4:
+        color = x[:, :, 0:3].astype(np.float32)
+        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+        y = color * alpha + 255.0 * (1.0 - alpha)
+        y = y.clip(0, 255).astype(np.uint8)
+        return y
+
+
+def upscale_image(input_image, upscale, min_size=None, unit_resolution=64):
+    H, W, C = input_image.shape
+    H = float(H)
+    W = float(W)
+    H *= upscale
+    W *= upscale
+    if min_size is not None:
+        if min(H, W) < min_size:
+            _upsacle = min_size / min(W, H)
+            W *= _upsacle
+            H *= _upsacle
+    H = int(np.round(H / unit_resolution)) * unit_resolution
+    W = int(np.round(W / unit_resolution)) * unit_resolution
+    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if upscale > 1 else cv2.INTER_AREA)
+    img = img.round().clip(0, 255).astype(np.uint8)
+    return img
+
+
+def fix_resize(input_image, size=512, unit_resolution=64):
+    H, W, C = input_image.shape
+    H = float(H)
+    W = float(W)
+    upscale = size / min(H, W)
+    H *= upscale
+    W *= upscale
+    H = int(np.round(H / unit_resolution)) * unit_resolution
+    W = int(np.round(W / unit_resolution)) * unit_resolution
+    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if upscale > 1 else cv2.INTER_AREA)
+    img = img.round().clip(0, 255).astype(np.uint8)
+    return img
+
+
+
+def Numpy2Tensor(img):
+    '''
+    np.array[H, w, C] [0, 255] -> Tensor[C, H, W], RGB, [-1, 1]
+    '''
+    # size
+    img = np.array(img) / 255 * 2 - 1
+    img = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1)
+    return img
+
+
+def Tensor2Numpy(x, h0=None, w0=None):
+    '''
+    Tensor[C, H, W], RGB, [-1, 1] -> PIL.Image
+    '''
+    if h0 is not None and w0 is not None:
+        x = x.unsqueeze(0)
+        x = interpolate(x, size=(h0, w0), mode='bicubic')
+        x = x.squeeze(0)
+    x = (x.permute(1, 2, 0) * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
+    return x
+
+
+def convert_dtype(dtype_str):
+    if dtype_str == 'fp32':
+        return torch.float32
+    elif dtype_str == 'fp16':
+        return torch.float16
+    elif dtype_str == 'bf16':
+        return torch.bfloat16
+    else:
+        raise NotImplementedError

SUPIR/utils/colorfix.py

@@ -0,0 +1,120 @@
+'''
+# --------------------------------------------------------------------------------
+#   Color fixed script from Li Yi (https://github.com/pkuliyi2015/sd-webui-stablesr/blob/master/srmodule/colorfix.py)
+# --------------------------------------------------------------------------------
+'''
+
+import torch
+from PIL import Image
+from torch import Tensor
+from torch.nn import functional as F
+
+from torchvision.transforms import ToTensor, ToPILImage
+
+def adain_color_fix(target: Image, source: Image):
+    # Convert images to tensors
+    to_tensor = ToTensor()
+    target_tensor = to_tensor(target).unsqueeze(0)
+    source_tensor = to_tensor(source).unsqueeze(0)
+
+    # Apply adaptive instance normalization
+    result_tensor = adaptive_instance_normalization(target_tensor, source_tensor)
+
+    # Convert tensor back to image
+    to_image = ToPILImage()
+    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
+
+    return result_image
+
+def wavelet_color_fix(target: Image, source: Image):
+    # Convert images to tensors
+    to_tensor = ToTensor()
+    target_tensor = to_tensor(target).unsqueeze(0)
+    source_tensor = to_tensor(source).unsqueeze(0)
+
+    # Apply wavelet reconstruction
+    result_tensor = wavelet_reconstruction(target_tensor, source_tensor)
+
+    # Convert tensor back to image
+    to_image = ToPILImage()
+    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
+
+    return result_image
+
+def calc_mean_std(feat: Tensor, eps=1e-5):
+    """Calculate mean and std for adaptive_instance_normalization.
+    Args:
+        feat (Tensor): 4D tensor.
+        eps (float): A small value added to the variance to avoid
+            divide-by-zero. Default: 1e-5.
+    """
+    size = feat.size()
+    assert len(size) == 4, 'The input feature should be 4D tensor.'
+    b, c = size[:2]
+    feat_var = feat.reshape(b, c, -1).var(dim=2) + eps
+    feat_std = feat_var.sqrt().reshape(b, c, 1, 1)
+    feat_mean = feat.reshape(b, c, -1).mean(dim=2).reshape(b, c, 1, 1)
+    return feat_mean, feat_std
+
+def adaptive_instance_normalization(content_feat:Tensor, style_feat:Tensor):
+    """Adaptive instance normalization.
+    Adjust the reference features to have the similar color and illuminations
+    as those in the degradate features.
+    Args:
+        content_feat (Tensor): The reference feature.
+        style_feat (Tensor): The degradate features.
+    """
+    size = content_feat.size()
+    style_mean, style_std = calc_mean_std(style_feat)
+    content_mean, content_std = calc_mean_std(content_feat)
+    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
+    return normalized_feat * style_std.expand(size) + style_mean.expand(size)
+
+def wavelet_blur(image: Tensor, radius: int):
+    """
+    Apply wavelet blur to the input tensor.
+    """
+    # input shape: (1, 3, H, W)
+    # convolution kernel
+    kernel_vals = [
+        [0.0625, 0.125, 0.0625],
+        [0.125, 0.25, 0.125],
+        [0.0625, 0.125, 0.0625],
+    ]
+    kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device)
+    # add channel dimensions to the kernel to make it a 4D tensor
+    kernel = kernel[None, None]
+    # repeat the kernel across all input channels
+    kernel = kernel.repeat(3, 1, 1, 1)
+    image = F.pad(image, (radius, radius, radius, radius), mode='replicate')
+    # apply convolution
+    output = F.conv2d(image, kernel, groups=3, dilation=radius)
+    return output
+
+def wavelet_decomposition(image: Tensor, levels=5):
+    """
+    Apply wavelet decomposition to the input tensor.
+    This function only returns the low frequency & the high frequency.
+    """
+    high_freq = torch.zeros_like(image)
+    for i in range(levels):
+        radius = 2 ** i
+        low_freq = wavelet_blur(image, radius)
+        high_freq += (image - low_freq)
+        image = low_freq
+
+    return high_freq, low_freq
+
+def wavelet_reconstruction(content_feat:Tensor, style_feat:Tensor):
+    """
+    Apply wavelet decomposition, so that the content will have the same color as the style.
+    """
+    # calculate the wavelet decomposition of the content feature
+    content_high_freq, content_low_freq = wavelet_decomposition(content_feat)
+    del content_low_freq
+    # calculate the wavelet decomposition of the style feature
+    style_high_freq, style_low_freq = wavelet_decomposition(style_feat)
+    del style_high_freq
+    # reconstruct the content feature with the style's high frequency
+    return content_high_freq + style_low_freq
+

SUPIR/utils/devices.py

@@ -0,0 +1,138 @@
+import sys
+import contextlib
+from functools import lru_cache
+
+import torch
+#from modules import errors
+
+if sys.platform == "darwin":
+    from modules import mac_specific
+
+
+def has_mps() -> bool:
+    if sys.platform != "darwin":
+        return False
+    else:
+        return mac_specific.has_mps
+
+
+def get_cuda_device_string():
+    return "cuda"
+
+
+def get_optimal_device_name():
+    if torch.cuda.is_available():
+        return get_cuda_device_string()
+
+    if has_mps():
+        return "mps"
+
+    return "cpu"
+
+
+def get_optimal_device():
+    return torch.device(get_optimal_device_name())
+
+
+def get_device_for(task):
+    return get_optimal_device()
+
+
+def torch_gc():
+
+    if torch.cuda.is_available():
+        with torch.cuda.device(get_cuda_device_string()):
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+    if has_mps():
+        mac_specific.torch_mps_gc()
+
+
+def enable_tf32():
+    if torch.cuda.is_available():
+
+        # enabling benchmark option seems to enable a range of cards to do fp16 when they otherwise can't
+        # see https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/4407
+        if any(torch.cuda.get_device_capability(devid) == (7, 5) for devid in range(0, torch.cuda.device_count())):
+            torch.backends.cudnn.benchmark = True
+
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+
+enable_tf32()
+#errors.run(enable_tf32, "Enabling TF32")
+
+cpu = torch.device("cpu")
+device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = torch.device("cuda")
+dtype = torch.float16
+dtype_vae = torch.float16
+dtype_unet = torch.float16
+unet_needs_upcast = False
+
+
+def cond_cast_unet(input):
+    return input.to(dtype_unet) if unet_needs_upcast else input
+
+
+def cond_cast_float(input):
+    return input.float() if unet_needs_upcast else input
+
+
+def randn(seed, shape):
+    torch.manual_seed(seed)
+    return torch.randn(shape, device=device)
+
+
+def randn_without_seed(shape):
+    return torch.randn(shape, device=device)
+
+
+def autocast(disable=False):
+    if disable:
+        return contextlib.nullcontext()
+
+    return torch.autocast("cuda")
+
+
+def without_autocast(disable=False):
+    return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
+
+
+class NansException(Exception):
+    pass
+
+
+def test_for_nans(x, where):
+    if not torch.all(torch.isnan(x)).item():
+        return
+
+    if where == "unet":
+        message = "A tensor with all NaNs was produced in Unet."
+
+    elif where == "vae":
+        message = "A tensor with all NaNs was produced in VAE."
+
+    else:
+        message = "A tensor with all NaNs was produced."
+
+    message += " Use --disable-nan-check commandline argument to disable this check."
+
+    raise NansException(message)
+
+
+@lru_cache
+def first_time_calculation():
+    """
+    just do any calculation with pytorch layers - the first time this is done it allocaltes about 700MB of memory and
+    spends about 2.7 seconds doing that, at least wih NVidia.
+    """
+
+    x = torch.zeros((1, 1)).to(device, dtype)
+    linear = torch.nn.Linear(1, 1).to(device, dtype)
+    linear(x)
+
+    x = torch.zeros((1, 1, 3, 3)).to(device, dtype)
+    conv2d = torch.nn.Conv2d(1, 1, (3, 3)).to(device, dtype)
+    conv2d(x)

SUPIR/utils/tilevae.py

@@ -0,0 +1,974 @@
+# ------------------------------------------------------------------------
+#
+#   Ultimate VAE Tile Optimization
+#
+#   Introducing a revolutionary new optimization designed to make
+#   the VAE work with giant images on limited VRAM!
+#   Say goodbye to the frustration of OOM and hello to seamless output!
+#
+# ------------------------------------------------------------------------
+#
+#   This script is a wild hack that splits the image into tiles,
+#   encodes each tile separately, and merges the result back together.
+#
+#   Advantages:
+#   - The VAE can now work with giant images on limited VRAM
+#       (~10 GB for 8K images!)
+#   - The merged output is completely seamless without any post-processing.
+#
+#   Drawbacks:
+#   - Giant RAM needed. To store the intermediate results for a 4096x4096
+#       images, you need 32 GB RAM it consumes ~20GB); for 8192x8192
+#       you need 128 GB RAM machine (it consumes ~100 GB)
+#   - NaNs always appear in for 8k images when you use fp16 (half) VAE
+#       You must use --no-half-vae to disable half VAE for that giant image.
+#   - Slow speed. With default tile size, it takes around 50/200 seconds
+#       to encode/decode a 4096x4096 image; and 200/900 seconds to encode/decode
+#       a 8192x8192 image. (The speed is limited by both the GPU and the CPU.)
+#   - The gradient calculation is not compatible with this hack. It
+#       will break any backward() or torch.autograd.grad() that passes VAE.
+#       (But you can still use the VAE to generate training data.)
+#
+#   How it works:
+#   1) The image is split into tiles.
+#       - To ensure perfect results, each tile is padded with 32 pixels
+#           on each side.
+#       - Then the conv2d/silu/upsample/downsample can produce identical
+#           results to the original image without splitting.
+#   2) The original forward is decomposed into a task queue and a task worker.
+#       - The task queue is a list of functions that will be executed in order.
+#       - The task worker is a loop that executes the tasks in the queue.
+#   3) The task queue is executed for each tile.
+#       - Current tile is sent to GPU.
+#       - local operations are directly executed.
+#       - Group norm calculation is temporarily suspended until the mean
+#           and var of all tiles are calculated.
+#       - The residual is pre-calculated and stored and addded back later.
+#       - When need to go to the next tile, the current tile is send to cpu.
+#   4) After all tiles are processed, tiles are merged on cpu and return.
+#
+#   Enjoy!
+#
+#   @author: LI YI @ Nanyang Technological University - Singapore
+#   @date: 2023-03-02
+#   @license: MIT License
+#
+#   Please give me a star if you like this project!
+#
+# -------------------------------------------------------------------------
+
+import gc
+from time import time
+import math
+from tqdm import tqdm
+
+import torch
+import torch.version
+import torch.nn.functional as F
+from einops import rearrange
+from diffusers.utils.import_utils import is_xformers_available
+
+#import SUPIR.utils.devices as devices
+
+import comfy.model_management
+device = comfy.model_management.get_torch_device()
+
+try:
+    import xformers
+    import xformers.ops
+except ImportError:
+    pass
+
+sd_flag = True
+
+def get_recommend_encoder_tile_size():
+    if torch.cuda.is_available():
+        total_memory = torch.cuda.get_device_properties(
+            device).total_memory // 2**20
+        if total_memory > 16*1000:
+            ENCODER_TILE_SIZE = 3072
+        elif total_memory > 12*1000:
+            ENCODER_TILE_SIZE = 2048
+        elif total_memory > 8*1000:
+            ENCODER_TILE_SIZE = 1536
+        else:
+            ENCODER_TILE_SIZE = 960
+    else:
+        ENCODER_TILE_SIZE = 512
+    return ENCODER_TILE_SIZE
+
+
+def get_recommend_decoder_tile_size():
+    if torch.cuda.is_available():
+        total_memory = torch.cuda.get_device_properties(
+            device).total_memory // 2**20
+        if total_memory > 30*1000:
+            DECODER_TILE_SIZE = 256
+        elif total_memory > 16*1000:
+            DECODER_TILE_SIZE = 192
+        elif total_memory > 12*1000:
+            DECODER_TILE_SIZE = 128
+        elif total_memory > 8*1000:
+            DECODER_TILE_SIZE = 96
+        else:
+            DECODER_TILE_SIZE = 64
+    else:
+        DECODER_TILE_SIZE = 64
+    return DECODER_TILE_SIZE
+
+
+if 'global const':
+    DEFAULT_ENABLED = False
+    DEFAULT_MOVE_TO_GPU = False
+    DEFAULT_FAST_ENCODER = True
+    DEFAULT_FAST_DECODER = True
+    DEFAULT_COLOR_FIX = 0
+    DEFAULT_ENCODER_TILE_SIZE = get_recommend_encoder_tile_size()
+    DEFAULT_DECODER_TILE_SIZE = get_recommend_decoder_tile_size()
+
+
+# inplace version of silu
+def inplace_nonlinearity(x):
+    # Test: fix for Nans
+    return F.silu(x, inplace=True)
+
+# extracted from ldm.modules.diffusionmodules.model
+
+# from diffusers lib
+def attn_forward_new(self, h_):
+    batch_size, channel, height, width = h_.shape
+    hidden_states = h_.view(batch_size, channel, height * width).transpose(1, 2)
+
+    attention_mask = None
+    encoder_hidden_states = None
+    batch_size, sequence_length, _ = hidden_states.shape
+    attention_mask = self.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+    query = self.to_q(hidden_states)
+
+    if encoder_hidden_states is None:
+        encoder_hidden_states = hidden_states
+    elif self.norm_cross:
+        encoder_hidden_states = self.norm_encoder_hidden_states(encoder_hidden_states)
+
+    key = self.to_k(encoder_hidden_states)
+    value = self.to_v(encoder_hidden_states)
+
+    query = self.head_to_batch_dim(query)
+    key = self.head_to_batch_dim(key)
+    value = self.head_to_batch_dim(value)
+
+    attention_probs = self.get_attention_scores(query, key, attention_mask)
+    hidden_states = torch.bmm(attention_probs, value)
+    hidden_states = self.batch_to_head_dim(hidden_states)
+
+    # linear proj
+    hidden_states = self.to_out[0](hidden_states)
+    # dropout
+    hidden_states = self.to_out[1](hidden_states)
+
+    hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+    return hidden_states
+
+def attn_forward_new_pt2_0(self, hidden_states,):
+    scale = 1
+    attention_mask = None
+    encoder_hidden_states = None
+
+    input_ndim = hidden_states.ndim
+
+    if input_ndim == 4:
+        batch_size, channel, height, width = hidden_states.shape
+        hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+    batch_size, sequence_length, _ = (
+        hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+    )
+
+    if attention_mask is not None:
+        attention_mask = self.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+        # scaled_dot_product_attention expects attention_mask shape to be
+        # (batch, heads, source_length, target_length)
+        attention_mask = attention_mask.view(batch_size, self.heads, -1, attention_mask.shape[-1])
+
+    if self.group_norm is not None:
+        hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+    query = self.to_q(hidden_states, scale=scale)
+
+    if encoder_hidden_states is None:
+        encoder_hidden_states = hidden_states
+    elif self.norm_cross:
+        encoder_hidden_states = self.norm_encoder_hidden_states(encoder_hidden_states)
+
+    key = self.to_k(encoder_hidden_states, scale=scale)
+    value = self.to_v(encoder_hidden_states, scale=scale)
+
+    inner_dim = key.shape[-1]
+    head_dim = inner_dim // self.heads
+
+    query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+
+    key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+    value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+
+    # the output of sdp = (batch, num_heads, seq_len, head_dim)
+    # TODO: add support for attn.scale when we move to Torch 2.1
+    hidden_states = F.scaled_dot_product_attention(
+        query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+    )
+
+    hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
+    hidden_states = hidden_states.to(query.dtype)
+
+    # linear proj
+    hidden_states = self.to_out[0](hidden_states, scale=scale)
+    # dropout
+    hidden_states = self.to_out[1](hidden_states)
+
+    if input_ndim == 4:
+        hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+    return hidden_states
+
+def attn_forward_new_xformers(self, hidden_states):
+    scale = 1
+    attention_op = None
+    attention_mask = None
+    encoder_hidden_states = None
+
+    input_ndim = hidden_states.ndim
+
+    if input_ndim == 4:
+        batch_size, channel, height, width = hidden_states.shape
+        hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+    batch_size, key_tokens, _ = (
+        hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+    )
+
+    attention_mask = self.prepare_attention_mask(attention_mask, key_tokens, batch_size)
+    if attention_mask is not None:
+        # expand our mask's singleton query_tokens dimension:
+        #   [batch*heads,            1, key_tokens] ->
+        #   [batch*heads, query_tokens, key_tokens]
+        # so that it can be added as a bias onto the attention scores that xformers computes:
+        #   [batch*heads, query_tokens, key_tokens]
+        # we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
+        _, query_tokens, _ = hidden_states.shape
+        attention_mask = attention_mask.expand(-1, query_tokens, -1)
+
+    if self.group_norm is not None:
+        hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+    query = self.to_q(hidden_states, scale=scale)
+
+    if encoder_hidden_states is None:
+        encoder_hidden_states = hidden_states
+    elif self.norm_cross:
+        encoder_hidden_states = self.norm_encoder_hidden_states(encoder_hidden_states)
+
+    key = self.to_k(encoder_hidden_states, scale=scale)
+    value = self.to_v(encoder_hidden_states, scale=scale)
+
+    query = self.head_to_batch_dim(query).contiguous()
+    key = self.head_to_batch_dim(key).contiguous()
+    value = self.head_to_batch_dim(value).contiguous()
+
+    hidden_states = xformers.ops.memory_efficient_attention(
+        query, key, value, attn_bias=attention_mask, op=attention_op#, scale=scale
+    )
+    hidden_states = hidden_states.to(query.dtype)
+    hidden_states = self.batch_to_head_dim(hidden_states)
+
+    # linear proj
+    hidden_states = self.to_out[0](hidden_states, scale=scale)
+    # dropout
+    hidden_states = self.to_out[1](hidden_states)
+
+    if input_ndim == 4:
+        hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+    return hidden_states
+
+def attn_forward(self, h_):
+    q = self.q(h_)
+    k = self.k(h_)
+    v = self.v(h_)
+
+    # compute attention
+    b, c, h, w = q.shape
+    q = q.reshape(b, c, h*w)
+    q = q.permute(0, 2, 1)   # b,hw,c
+    k = k.reshape(b, c, h*w)  # b,c,hw
+    w_ = torch.bmm(q, k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+    w_ = w_ * (int(c)**(-0.5))
+    w_ = torch.nn.functional.softmax(w_, dim=2)
+
+    # attend to values
+    v = v.reshape(b, c, h*w)
+    w_ = w_.permute(0, 2, 1)   # b,hw,hw (first hw of k, second of q)
+    # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+    h_ = torch.bmm(v, w_)
+    h_ = h_.reshape(b, c, h, w)
+
+    h_ = self.proj_out(h_)
+
+    return h_
+
+
+def xformer_attn_forward(self, h_):
+    q = self.q(h_)
+    k = self.k(h_)
+    v = self.v(h_)
+
+    # compute attention
+    B, C, H, W = q.shape
+    q, k, v = map(lambda x: rearrange(x, 'b c h w -> b (h w) c'), (q, k, v))
+
+    q, k, v = map(
+        lambda t: t.unsqueeze(3)
+        .reshape(B, t.shape[1], 1, C)
+        .permute(0, 2, 1, 3)
+        .reshape(B * 1, t.shape[1], C)
+        .contiguous(),
+        (q, k, v),
+    )
+    out = xformers.ops.memory_efficient_attention(
+        q, k, v, attn_bias=None, op=self.attention_op)
+
+    out = (
+        out.unsqueeze(0)
+        .reshape(B, 1, out.shape[1], C)
+        .permute(0, 2, 1, 3)
+        .reshape(B, out.shape[1], C)
+    )
+    out = rearrange(out, 'b (h w) c -> b c h w', b=B, h=H, w=W, c=C)
+    out = self.proj_out(out)
+    return out
+
+
+def attn2task(task_queue, net):
+    if False: #isinstance(net, AttnBlock):
+        task_queue.append(('store_res', lambda x: x))
+        task_queue.append(('pre_norm', net.norm))
+        task_queue.append(('attn', lambda x, net=net: attn_forward(net, x)))
+        task_queue.append(['add_res', None])
+    elif False: #isinstance(net, MemoryEfficientAttnBlock):
+        task_queue.append(('store_res', lambda x: x))
+        task_queue.append(('pre_norm', net.norm))
+        task_queue.append(
+            ('attn', lambda x, net=net: xformer_attn_forward(net, x)))
+        task_queue.append(['add_res', None])
+    else:
+        task_queue.append(('store_res', lambda x: x))
+        task_queue.append(('pre_norm', net.norm))
+        if is_xformers_available:
+            # task_queue.append(('attn', lambda x, net=net: attn_forward_new_xformers(net, x)))
+            task_queue.append(
+                ('attn', lambda x, net=net: xformer_attn_forward(net, x)))
+        elif hasattr(F, "scaled_dot_product_attention"):
+            task_queue.append(('attn', lambda x, net=net: attn_forward_new_pt2_0(net, x)))
+        else:
+            task_queue.append(('attn', lambda x, net=net: attn_forward_new(net, x)))
+        task_queue.append(['add_res', None])
+
+def resblock2task(queue, block):
+    """
+    Turn a ResNetBlock into a sequence of tasks and append to the task queue
+
+    @param queue: the target task queue
+    @param block: ResNetBlock
+
+    """
+    if block.in_channels != block.out_channels:
+        if sd_flag:
+            if block.use_conv_shortcut:
+                queue.append(('store_res', block.conv_shortcut))
+            else:
+                queue.append(('store_res', block.nin_shortcut))
+        else:
+            if block.use_in_shortcut:
+                queue.append(('store_res', block.conv_shortcut))
+            else:
+                queue.append(('store_res', block.nin_shortcut))
+
+    else:
+        queue.append(('store_res', lambda x: x))
+    queue.append(('pre_norm', block.norm1))
+    queue.append(('silu', inplace_nonlinearity))
+    queue.append(('conv1', block.conv1))
+    queue.append(('pre_norm', block.norm2))
+    queue.append(('silu', inplace_nonlinearity))
+    queue.append(('conv2', block.conv2))
+    queue.append(['add_res', None])
+
+
+def build_sampling(task_queue, net, is_decoder):
+    """
+    Build the sampling part of a task queue
+    @param task_queue: the target task queue
+    @param net: the network
+    @param is_decoder: currently building decoder or encoder
+    """
+    if is_decoder:
+        if sd_flag:
+            resblock2task(task_queue, net.mid.block_1)
+            attn2task(task_queue, net.mid.attn_1)
+            print(task_queue)
+            resblock2task(task_queue, net.mid.block_2)
+            resolution_iter = reversed(range(net.num_resolutions))
+            block_ids = net.num_res_blocks + 1
+            condition = 0
+            module = net.up
+            func_name = 'upsample'
+        else:
+            resblock2task(task_queue, net.mid_block.resnets[0])
+            attn2task(task_queue, net.mid_block.attentions[0])
+            resblock2task(task_queue, net.mid_block.resnets[1])
+            resolution_iter = (range(len(net.up_blocks)))  # net.num_resolutions = 3
+            block_ids = 2 + 1
+            condition = len(net.up_blocks) - 1
+            module = net.up_blocks
+            func_name = 'upsamplers'
+    else:
+        if sd_flag:
+            resolution_iter = range(net.num_resolutions)
+            block_ids = net.num_res_blocks
+            condition = net.num_resolutions - 1
+            module = net.down
+            func_name = 'downsample'
+        else:
+            resolution_iter = range(len(net.down_blocks))
+            block_ids = 2
+            condition = len(net.down_blocks) - 1
+            module = net.down_blocks
+            func_name = 'downsamplers'
+
+    for i_level in resolution_iter:
+        for i_block in range(block_ids):
+            if sd_flag:
+                resblock2task(task_queue, module[i_level].block[i_block])
+            else:
+                resblock2task(task_queue, module[i_level].resnets[i_block])
+        if i_level != condition:
+            if sd_flag:
+                task_queue.append((func_name, getattr(module[i_level], func_name)))
+            else:
+                if is_decoder:
+                    task_queue.append((func_name, module[i_level].upsamplers[0]))
+                else:
+                    task_queue.append((func_name, module[i_level].downsamplers[0]))
+
+    if not is_decoder:
+        if sd_flag:
+            resblock2task(task_queue, net.mid.block_1)
+            attn2task(task_queue, net.mid.attn_1)
+            resblock2task(task_queue, net.mid.block_2)
+        else:
+            resblock2task(task_queue, net.mid_block.resnets[0])
+            attn2task(task_queue, net.mid_block.attentions[0])
+            resblock2task(task_queue, net.mid_block.resnets[1])
+
+
+def build_task_queue(net, is_decoder):
+    """
+    Build a single task queue for the encoder or decoder
+    @param net: the VAE decoder or encoder network
+    @param is_decoder: currently building decoder or encoder
+    @return: the task queue
+    """
+    task_queue = []
+    task_queue.append(('conv_in', net.conv_in))
+
+    # construct the sampling part of the task queue
+    # because encoder and decoder share the same architecture, we extract the sampling part
+    build_sampling(task_queue, net, is_decoder)
+    if is_decoder and not sd_flag:
+        net.give_pre_end = False
+        net.tanh_out = False
+
+    if not is_decoder or not net.give_pre_end:
+        if sd_flag:
+            task_queue.append(('pre_norm', net.norm_out))
+        else:
+            task_queue.append(('pre_norm', net.conv_norm_out))
+        task_queue.append(('silu', inplace_nonlinearity))
+        task_queue.append(('conv_out', net.conv_out))
+        if is_decoder and net.tanh_out:
+            task_queue.append(('tanh', torch.tanh))
+
+    return task_queue
+
+
+def clone_task_queue(task_queue):
+    """
+    Clone a task queue
+    @param task_queue: the task queue to be cloned
+    @return: the cloned task queue
+    """
+    return [[item for item in task] for task in task_queue]
+
+
+def get_var_mean(input, num_groups, eps=1e-6):
+    """
+    Get mean and var for group norm
+    """
+    b, c = input.size(0), input.size(1)
+    channel_in_group = int(c/num_groups)
+    input_reshaped = input.contiguous().view(
+        1, int(b * num_groups), channel_in_group, *input.size()[2:])
+    var, mean = torch.var_mean(
+        input_reshaped, dim=[0, 2, 3, 4], unbiased=False)
+    return var, mean
+
+
+def custom_group_norm(input, num_groups, mean, var, weight=None, bias=None, eps=1e-6):
+    """
+    Custom group norm with fixed mean and var
+
+    @param input: input tensor
+    @param num_groups: number of groups. by default, num_groups = 32
+    @param mean: mean, must be pre-calculated by get_var_mean
+    @param var: var, must be pre-calculated by get_var_mean
+    @param weight: weight, should be fetched from the original group norm
+    @param bias: bias, should be fetched from the original group norm
+    @param eps: epsilon, by default, eps = 1e-6 to match the original group norm
+
+    @return: normalized tensor
+    """
+    b, c = input.size(0), input.size(1)
+    channel_in_group = int(c/num_groups)
+    input_reshaped = input.contiguous().view(
+        1, int(b * num_groups), channel_in_group, *input.size()[2:])
+
+    out = F.batch_norm(input_reshaped, mean, var, weight=None, bias=None,
+                       training=False, momentum=0, eps=eps)
+
+    out = out.view(b, c, *input.size()[2:])
+
+    # post affine transform
+    if weight is not None:
+        out *= weight.view(1, -1, 1, 1)
+    if bias is not None:
+        out += bias.view(1, -1, 1, 1)
+    return out
+
+
+def crop_valid_region(x, input_bbox, target_bbox, is_decoder):
+    """
+    Crop the valid region from the tile
+    @param x: input tile
+    @param input_bbox: original input bounding box
+    @param target_bbox: output bounding box
+    @param scale: scale factor
+    @return: cropped tile
+    """
+    padded_bbox = [i * 8 if is_decoder else i//8 for i in input_bbox]
+    margin = [target_bbox[i] - padded_bbox[i] for i in range(4)]
+    return x[:, :, margin[2]:x.size(2)+margin[3], margin[0]:x.size(3)+margin[1]]
+
+# ↓↓↓ https://github.com/Kahsolt/stable-diffusion-webui-vae-tile-infer ↓↓↓
+
+
+def perfcount(fn):
+    def wrapper(*args, **kwargs):
+        ts = time()
+
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats(device)
+        comfy.model_management.soft_empty_cache()
+        gc.collect()
+
+        ret = fn(*args, **kwargs)
+
+        comfy.model_management.soft_empty_cache()
+        gc.collect()
+        if torch.cuda.is_available():
+            vram = torch.cuda.max_memory_allocated(device) / 2**20
+            torch.cuda.reset_peak_memory_stats(device)
+            print(
+                f'[Tiled VAE]: Done in {time() - ts:.3f}s, max VRAM alloc {vram:.3f} MB')
+        else:
+            print(f'[Tiled VAE]: Done in {time() - ts:.3f}s')
+
+        return ret
+    return wrapper
+
+# copy end :)
+
+
+class GroupNormParam:
+    def __init__(self):
+        self.var_list = []
+        self.mean_list = []
+        self.pixel_list = []
+        self.weight = None
+        self.bias = None
+
+    def add_tile(self, tile, layer):
+        var, mean = get_var_mean(tile, 32)
+        # For giant images, the variance can be larger than max float16
+        # In this case we create a copy to float32
+        if var.dtype == torch.float16 and var.isinf().any():
+            fp32_tile = tile.float()
+            var, mean = get_var_mean(fp32_tile, 32)
+        # ============= DEBUG: test for infinite =============
+        # if torch.isinf(var).any():
+        #    print('var: ', var)
+        # ====================================================
+        self.var_list.append(var)
+        self.mean_list.append(mean)
+        self.pixel_list.append(
+            tile.shape[2]*tile.shape[3])
+        if hasattr(layer, 'weight'):
+            self.weight = layer.weight
+            self.bias = layer.bias
+        else:
+            self.weight = None
+            self.bias = None
+
+    def summary(self):
+        """
+        summarize the mean and var and return a function
+        that apply group norm on each tile
+        """
+        if len(self.var_list) == 0:
+            return None
+        var = torch.vstack(self.var_list)
+        mean = torch.vstack(self.mean_list)
+        max_value = max(self.pixel_list)
+        pixels = torch.tensor(
+            self.pixel_list, dtype=torch.float32, device=device) / max_value
+        sum_pixels = torch.sum(pixels)
+        pixels = pixels.unsqueeze(
+            1) / sum_pixels
+        var = torch.sum(
+            var * pixels, dim=0)
+        mean = torch.sum(
+            mean * pixels, dim=0)
+        return lambda x:  custom_group_norm(x, 32, mean, var, self.weight, self.bias)
+
+    @staticmethod
+    def from_tile(tile, norm):
+        """
+        create a function from a single tile without summary
+        """
+        var, mean = get_var_mean(tile, 32)
+        if var.dtype == torch.float16 and var.isinf().any():
+            fp32_tile = tile.float()
+            var, mean = get_var_mean(fp32_tile, 32)
+            # if it is a macbook, we need to convert back to float16
+            if var.device.type == 'mps':
+                # clamp to avoid overflow
+                var = torch.clamp(var, 0, 60000)
+                var = var.half()
+                mean = mean.half()
+        if hasattr(norm, 'weight'):
+            weight = norm.weight
+            bias = norm.bias
+        else:
+            weight = None
+            bias = None
+
+        def group_norm_func(x, mean=mean, var=var, weight=weight, bias=bias):
+            return custom_group_norm(x, 32, mean, var, weight, bias, 1e-6)
+        return group_norm_func
+
+
+class VAEHook:
+    def __init__(self, net, tile_size, is_decoder, fast_decoder, fast_encoder, color_fix, to_gpu=False):
+        self.net = net                  # encoder | decoder
+        self.tile_size = tile_size
+        self.is_decoder = is_decoder
+        self.fast_mode = (fast_encoder and not is_decoder) or (
+            fast_decoder and is_decoder)
+        self.color_fix = color_fix and not is_decoder
+        self.to_gpu = to_gpu
+        self.pad = 11 if is_decoder else 32
+
+    def __call__(self, x):
+        B, C, H, W = x.shape
+        original_device = next(self.net.parameters()).device
+        try:
+            if self.to_gpu:
+                self.net.to(device)
+            if max(H, W) <= self.pad * 2 + self.tile_size:
+                print("[Tiled VAE]: the input size is tiny and unnecessary to tile.")
+                return self.net.original_forward(x)
+            else:
+                return self.vae_tile_forward(x)
+        finally:
+            self.net.to(original_device)
+
+    def get_best_tile_size(self, lowerbound, upperbound):
+        """
+        Get the best tile size for GPU memory
+        """
+        divider = 32
+        while divider >= 2:
+            remainer = lowerbound % divider
+            if remainer == 0:
+                return lowerbound
+            candidate = lowerbound - remainer + divider
+            if candidate <= upperbound:
+                return candidate
+            divider //= 2
+        return lowerbound
+
+    def split_tiles(self, h, w):
+        """
+        Tool function to split the image into tiles
+        @param h: height of the image
+        @param w: width of the image
+        @return: tile_input_bboxes, tile_output_bboxes
+        """
+        tile_input_bboxes, tile_output_bboxes = [], []
+        tile_size = self.tile_size
+        pad = self.pad
+        num_height_tiles = math.ceil((h - 2 * pad) / tile_size)
+        num_width_tiles = math.ceil((w - 2 * pad) / tile_size)
+        # If any of the numbers are 0, we let it be 1
+        # This is to deal with long and thin images
+        num_height_tiles = max(num_height_tiles, 1)
+        num_width_tiles = max(num_width_tiles, 1)
+
+        # Suggestions from https://github.com/Kahsolt: auto shrink the tile size
+        real_tile_height = math.ceil((h - 2 * pad) / num_height_tiles)
+        real_tile_width = math.ceil((w - 2 * pad) / num_width_tiles)
+        real_tile_height = self.get_best_tile_size(real_tile_height, tile_size)
+        real_tile_width = self.get_best_tile_size(real_tile_width, tile_size)
+
+        print(f'[Tiled VAE]: split to {num_height_tiles}x{num_width_tiles} = {num_height_tiles*num_width_tiles} tiles. ' +
+              f'Optimal tile size {real_tile_width}x{real_tile_height}, original tile size {tile_size}x{tile_size}')
+
+        for i in range(num_height_tiles):
+            for j in range(num_width_tiles):
+                # bbox: [x1, x2, y1, y2]
+                # the padding is is unnessary for image borders. So we directly start from (32, 32)
+                input_bbox = [
+                    pad + j * real_tile_width,
+                    min(pad + (j + 1) * real_tile_width, w),
+                    pad + i * real_tile_height,
+                    min(pad + (i + 1) * real_tile_height, h),
+                ]
+
+                # if the output bbox is close to the image boundary, we extend it to the image boundary
+                output_bbox = [
+                    input_bbox[0] if input_bbox[0] > pad else 0,
+                    input_bbox[1] if input_bbox[1] < w - pad else w,
+                    input_bbox[2] if input_bbox[2] > pad else 0,
+                    input_bbox[3] if input_bbox[3] < h - pad else h,
+                ]
+
+                # scale to get the final output bbox
+                output_bbox = [x * 8 if self.is_decoder else x // 8 for x in output_bbox]
+                tile_output_bboxes.append(output_bbox)
+
+                # indistinguishable expand the input bbox by pad pixels
+                tile_input_bboxes.append([
+                    max(0, input_bbox[0] - pad),
+                    min(w, input_bbox[1] + pad),
+                    max(0, input_bbox[2] - pad),
+                    min(h, input_bbox[3] + pad),
+                ])
+
+        return tile_input_bboxes, tile_output_bboxes
+
+    @torch.no_grad()
+    def estimate_group_norm(self, z, task_queue, color_fix):
+        device = z.device
+        tile = z
+        last_id = len(task_queue) - 1
+        while last_id >= 0 and task_queue[last_id][0] != 'pre_norm':
+            last_id -= 1
+        if last_id <= 0 or task_queue[last_id][0] != 'pre_norm':
+            raise ValueError('No group norm found in the task queue')
+        # estimate until the last group norm
+        for i in range(last_id + 1):
+            task = task_queue[i]
+            if task[0] == 'pre_norm':
+                group_norm_func = GroupNormParam.from_tile(tile, task[1])
+                task_queue[i] = ('apply_norm', group_norm_func)
+                if i == last_id:
+                    return True
+                tile = group_norm_func(tile)
+            elif task[0] == 'store_res':
+                task_id = i + 1
+                while task_id < last_id and task_queue[task_id][0] != 'add_res':
+                    task_id += 1
+                if task_id >= last_id:
+                    continue
+                task_queue[task_id][1] = task[1](tile)
+            elif task[0] == 'add_res':
+                tile += task[1].to(device)
+                task[1] = None
+            elif color_fix and task[0] == 'downsample':
+                for j in range(i, last_id + 1):
+                    if task_queue[j][0] == 'store_res':
+                        task_queue[j] = ('store_res_cpu', task_queue[j][1])
+                return True
+            else:
+                tile = task[1](tile)
+            try:
+                devices.test_for_nans(tile, "vae")
+            except:
+                print(f'Nan detected in fast mode estimation. Fast mode disabled.')
+                return False
+
+        raise IndexError('Should not reach here')
+
+    @perfcount
+    @torch.no_grad()
+    def vae_tile_forward(self, z):
+        """
+        Decode a latent vector z into an image in a tiled manner.
+        @param z: latent vector
+        @return: image
+        """
+        device = next(self.net.parameters()).device
+        dtype = z.dtype
+        net = self.net
+        tile_size = self.tile_size
+        is_decoder = self.is_decoder
+
+        z = z.detach() # detach the input to avoid backprop
+
+        N, height, width = z.shape[0], z.shape[2], z.shape[3]
+        net.last_z_shape = z.shape
+
+        # Split the input into tiles and build a task queue for each tile
+        print(f'[Tiled VAE]: input_size: {z.shape}, tile_size: {tile_size}, padding: {self.pad}')
+
+        in_bboxes, out_bboxes = self.split_tiles(height, width)
+
+        # Prepare tiles by split the input latents
+        tiles = []
+        for input_bbox in in_bboxes:
+            tile = z[:, :, input_bbox[2]:input_bbox[3], input_bbox[0]:input_bbox[1]].cpu()
+            tiles.append(tile)
+
+        num_tiles = len(tiles)
+        num_completed = 0
+
+        # Build task queues
+        single_task_queue = build_task_queue(net, is_decoder)
+        #print(single_task_queue)
+        if self.fast_mode:
+            # Fast mode: downsample the input image to the tile size,
+            # then estimate the group norm parameters on the downsampled image
+            scale_factor = tile_size / max(height, width)
+            z = z.to(device)
+            downsampled_z = F.interpolate(z, scale_factor=scale_factor, mode='nearest-exact')
+            # use nearest-exact to keep statictics as close as possible
+            print(f'[Tiled VAE]: Fast mode enabled, estimating group norm parameters on {downsampled_z.shape[3]} x {downsampled_z.shape[2]} image')
+
+            # ======= Special thanks to @Kahsolt for distribution shift issue ======= #
+            # The downsampling will heavily distort its mean and std, so we need to recover it.
+            std_old, mean_old = torch.std_mean(z, dim=[0, 2, 3], keepdim=True)
+            std_new, mean_new = torch.std_mean(downsampled_z, dim=[0, 2, 3], keepdim=True)
+            downsampled_z = (downsampled_z - mean_new) / std_new * std_old + mean_old
+            del std_old, mean_old, std_new, mean_new
+            # occasionally the std_new is too small or too large, which exceeds the range of float16
+            # so we need to clamp it to max z's range.
+            downsampled_z = torch.clamp_(downsampled_z, min=z.min(), max=z.max())
+            estimate_task_queue = clone_task_queue(single_task_queue)
+            if self.estimate_group_norm(downsampled_z, estimate_task_queue, color_fix=self.color_fix):
+                single_task_queue = estimate_task_queue
+            del downsampled_z
+
+        task_queues = [clone_task_queue(single_task_queue) for _ in range(num_tiles)]
+
+        # Dummy result
+        result = None
+        result_approx = None
+        #try:
+        #    with devices.autocast():
+        #        result_approx = torch.cat([F.interpolate(cheap_approximation(x).unsqueeze(0), scale_factor=opt_f, mode='nearest-exact') for x in z], dim=0).cpu()
+        #except: pass
+        # Free memory of input latent tensor
+        del z
+
+        # Task queue execution
+        pbar = tqdm(total=num_tiles * len(task_queues[0]), desc=f"[Tiled VAE]: Executing {'Decoder' if is_decoder else 'Encoder'} Task Queue: ")
+
+        # execute the task back and forth when switch tiles so that we always
+        # keep one tile on the GPU to reduce unnecessary data transfer
+        forward = True
+        interrupted = False
+        #state.interrupted = interrupted
+        while True:
+            #if state.interrupted: interrupted = True ; break
+
+            group_norm_param = GroupNormParam()
+            for i in range(num_tiles) if forward else reversed(range(num_tiles)):
+                #if state.interrupted: interrupted = True ; break
+
+                tile = tiles[i].to(device)
+                input_bbox = in_bboxes[i]
+                task_queue = task_queues[i]
+
+                interrupted = False
+                while len(task_queue) > 0:
+                    #if state.interrupted: interrupted = True ; break
+
+                    # DEBUG: current task
+                    # print('Running task: ', task_queue[0][0], ' on tile ', i, '/', num_tiles, ' with shape ', tile.shape)
+                    task = task_queue.pop(0)
+                    if task[0] == 'pre_norm':
+                        group_norm_param.add_tile(tile, task[1])
+                        break
+                    elif task[0] == 'store_res' or task[0] == 'store_res_cpu':
+                        task_id = 0
+                        res = task[1](tile)
+                        if not self.fast_mode or task[0] == 'store_res_cpu':
+                            res = res.cpu()
+                        while task_queue[task_id][0] != 'add_res':
+                            task_id += 1
+                        task_queue[task_id][1] = res
+                    elif task[0] == 'add_res':
+                        tile += task[1].to(device)
+                        task[1] = None
+                    else:
+                        tile = task[1](tile)
+                        #print(tiles[i].shape, tile.shape, task)
+                    pbar.update(1)
+
+                if interrupted: break
+
+                # check for NaNs in the tile.
+                # If there are NaNs, we abort the process to save user's time
+                #devices.test_for_nans(tile, "vae")
+
+                #print(tiles[i].shape, tile.shape, i, num_tiles)
+                if len(task_queue) == 0:
+                    tiles[i] = None
+                    num_completed += 1
+                    if result is None:      # NOTE: dim C varies from different cases, can only be inited dynamically
+                        result = torch.zeros((N, tile.shape[1], height * 8 if is_decoder else height // 8, width * 8 if is_decoder else width // 8), device=device, requires_grad=False)
+                    result[:, :, out_bboxes[i][2]:out_bboxes[i][3], out_bboxes[i][0]:out_bboxes[i][1]] = crop_valid_region(tile, in_bboxes[i], out_bboxes[i], is_decoder)
+                    del tile
+                elif i == num_tiles - 1 and forward:
+                    forward = False
+                    tiles[i] = tile
+                elif i == 0 and not forward:
+                    forward = True
+                    tiles[i] = tile
+                else:
+                    tiles[i] = tile.cpu()
+                    del tile
+
+            if interrupted: break
+            if num_completed == num_tiles: break
+
+            # insert the group norm task to the head of each task queue
+            group_norm_func = group_norm_param.summary()
+            if group_norm_func is not None:
+                for i in range(num_tiles):
+                    task_queue = task_queues[i]
+                    task_queue.insert(0, ('apply_norm', group_norm_func))
+
+        # Done!
+        pbar.close()
+        return result.to(dtype) if result is not None else result_approx.to(device)