space

.gitattributes

@@ -1,35 +1,35 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+__pycache__

LICENSE

@@ -0,0 +1,72 @@
+SUPIR Software License Agreement
+
+Copyright (c) 2024 by SupPixel Pty Ltd. All rights reserved.
+
+This work is jointly owned by SupPixel Pty Ltd, an Australian company (ACN 676 560 320). All rights reserved. The work is created by the author team of SUPIR at SupPixel Pty Ltd.
+
+This License Agreement ("Agreement") is made and entered into by and between SupPixel Pty Ltd (collectively, the "Licensor") and any person or entity who accesses, uses, or distributes the SUPIR software (the "Licensee").
+
+1. Definitions
+
+a. Licensed Software: Refers to the SUPIR software, including its source code, associated documentation, and any updates, upgrades, or modifications thereof. The Licensed Software includes both open-source components and proprietary components, as described herein.
+
+b. Neural Network Components: Includes the weights, biases, and architecture of the proprietary neural network(s) developed and trained by the Licensor. This excludes any pre-existing open-source neural networks used in conjunction with the Licensor's proprietary networks, for which the Licensor does not hold copyright.
+
+c. Commercial Use: Refers to any use of the Licensed Software or its components, in whole or in part, for the purpose of generating revenue, conducting business operations, or creating products or services for sale or profit. This includes, but is not limited to:
+
+- Deploying the software on a server to provide a service to users for a fee.
+- Using the software to process images that are intended for sale.
+- Integrating the software or its components into a commercial product that is sold to end-users.
+- Using the software to provide commercial consulting or contracting services, such as image analysis, enhancement, or modification for clients.
+- Using the software to process images that are subsequently used in training datasets for machine learning models or other types of data analysis. This includes datasets that are used internally, sold, or provided to third parties under a commercial agreement.
+- Licensing the software or its components to other businesses for commercial purposes.
+
+2. License Grant
+
+a. Open-Source Components: The Licensor grants the Licensee a non-exclusive, worldwide, royalty-free license to use, copy, modify, and distribute the source code of the open-source components of the Licensed Software, subject to the terms and conditions of their respective open-source licenses. The specific open-source licenses for each component are indicated in the source code files and documentation of the Licensed Software.
+
+b. Proprietary Neural Network Components: Notwithstanding the open-source license granted in Section 2(a) for open-source components, the Licensor retains all rights, title, and interest in and to the proprietary Neural Network Components developed and trained by the Licensor. The license granted herein does not include rights to use, copy, modify, distribute, or create derivative works of these proprietary Neural Network Components without obtaining express written permission from the Licensor.
+
+c. The source code of the Licensed Software is available as open-source on this GitHub repository. However, the open-source license does not grant any rights to the weights, biases, or architecture of the proprietary neural network(s) used in the Licensed Software, which remain the exclusive property of the Licensor.
+
+3. Commercial Use and Licensing
+
+a. The Licensor retains all ownership and commercial usage rights in the Licensed Software, including its neural network(s), regardless of the open-source availability of portions of the Licensed Software's source code. Under the open-source license, anyone may freely use, copy, modify, and distribute the source code of the Licensed Software for non-commercial purposes only. Similarly, the proprietary neural network weights, biases, and architecture may only be used for non-commercial purposes. Any commercial use, copying, modification, distribution, or creation of derivative works of the Licensed Software's source code or the proprietary neural network weights, biases, and architecture is strictly prohibited without express written permission from the Licensor.
+
+b. Any Commercial Use of the Licensed Software, including its source code and neural network(s), requires a separate commercial license agreement with the Licensor. The terms and conditions of such commercial licensing, including the scope of use, fees, and delivery of the relevant components, shall be negotiated and agreed upon between the Licensor and the Licensee in writing. For commercial licensing inquiries, please contact the Licensor at [email protected].
+
+4. Intellectual Property Rights
+
+a. The Licensor retains all intellectual property rights in the Licensed Software, including copyrights, patents, trademarks, trade secrets, and any other proprietary rights. The Licensee acknowledges that the Licensor is the sole owner of the Licensed Software and its components, including the proprietary Neural Network Components.
+
+b. The Licensee agrees not to challenge or contest the Licensor's ownership or intellectual property rights in the Licensed Software, nor to assist or encourage any third party in doing so.
+
+c. The Licensee shall promptly notify the Licensor of any actual or suspected infringement of the Licensor's intellectual property rights in the Licensed Software by third parties, and shall provide reasonable assistance to the Licensor in enforcing its rights against such infringement.
+
+5. Warranty Disclaimer
+
+The Licensed Software is provided "as is", without warranty of any kind, either express or implied, including, but not limited to, the implied warranties of merchantability, fitness for a particular purpose, and non-infringement. The Licensor disclaims all warranties, express or implied, regarding the accuracy, reliability, completeness, or usefulness of the Licensed Software or its components. The Licensor does not warrant that the Licensed Software will be error-free, uninterrupted, or free from defects or vulnerabilities.
+
+6. Limitation of Liability
+
+In no event shall the Licensor be liable for any claim, damages, or other liability, whether in an action of contract, tort, or otherwise, arising from, out of, or in connection with the Licensed Software or the use or other dealings in the Licensed Software. The Licensor shall not be liable for any direct, indirect, incidental, special, consequential, or exemplary damages, including but not limited to damages for loss of profits, goodwill, use, data, or other intangible losses, even if the Licensor has been advised of the possibility of such damages.
+
+7. Indemnification
+
+The Licensee agrees to indemnify, defend, and hold harmless the Licensor and its affiliates, officers, directors, employees, and agents from and against any and all claims, liabilities, damages, losses, costs, expenses, or fees (including reasonable attorneys' fees) arising from or relating to the Licensee's use of the Licensed Software or any violation of this Agreement by the Licensee.
+
+8. Governing Law and Jurisdiction
+
+This Agreement shall be governed by and construed in accordance with the laws of the State of New South Wales, Australia, without giving effect to any choice or conflict of law provision or rule. Any legal action or proceeding arising out of or relating to this Agreement shall be brought exclusively in the courts of New South Wales, Australia, and each party irrevocably submits to the jurisdiction of such courts.
+
+9. Entire Agreement
+
+This Agreement constitutes the entire agreement between the parties concerning the subject matter hereof and supersedes all prior or contemporaneous communications, understandings, and agreements, whether written or oral, between the parties regarding such subject matter.
+
+By accessing, using, or distributing the Licensed Software, the Licensee agrees to be bound by the terms and conditions of this Agreement. If the Licensee does not agree to the terms of this Agreement, the Licensee must not access, use, or distribute the Licensed Software.
+
+For inquiries or to obtain permission for Commercial Use, please contact:
+
+Dr. Jinjin Gu
+SupPixel Pty Ltd
+Email: [email protected]

README.md

@@ -1,13 +1,5 @@
----
-title: SUPIR Forge
-emoji: 📈
-colorFrom: red
-colorTo: red
-sdk: gradio
-sdk_version: 4.44.0
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SUPIR Backend
+
+The backend code for [forge-space-SUPIR](https://github.com/Haoming02/forge-space-SUPIR), that gets synced to [Haoming02/SUPIR-Forge](https://huggingface.co/spaces/Haoming02/SUPIR-Forge)
+
+> There is no point in downloading this, unless you are making a PR

SUPIR/models/SUPIR_model.py

@@ -0,0 +1,176 @@
+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
+
+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):
+ 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_min = cfg_scale
+ 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)
+ z_stage1 = self.encode_first_stage(x_stage1)
+
+ c, uc = self.prepare_condition(_z, p, p_p, n_p, N)
+
+ 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)
+
+ def prepare_condition(self, _z, p, p_p, n_p, N):
+ batch = {}
+ batch['original_size_as_tuple'] = torch.tensor([1024, 1024]).repeat(N, 1).to(_z.device)
+ batch['crop_coords_top_left'] = torch.tensor([0, 0]).repeat(N, 1).to(_z.device)
+ batch['target_size_as_tuple'] = torch.tensor([1024, 1024]).repeat(N, 1).to(_z.device)
+ batch['aesthetic_score'] = torch.tensor([9.0]).repeat(N, 1).to(_z.device)
+ batch['control'] = _z
+
+ batch_uc = copy.deepcopy(batch)
+ batch_uc['txt'] = [n_p for _ in p]
+
+ if not isinstance(p[0], list):
+ batch['txt'] = [''.join([_p, p_p]) for _p in p]
+ with torch.cuda.amp.autocast(dtype=self.ae_dtype):
+ c, uc = self.conditioner.get_unconditional_conditioning(batch, batch_uc)
+ else:
+ assert len(p) == 1, 'Support bs=1 only for local prompt conditioning.'
+ p_tiles = p[0]
+ c = []
+ for i, p_tile in enumerate(p_tiles):
+ batch['txt'] = [''.join([p_tile, p_p])]
+ with torch.cuda.amp.autocast(dtype=self.ae_dtype):
+ if i == 0:
+ _c, uc = self.conditioner.get_unconditional_conditioning(batch, batch_uc)
+ else:
+ _c, _ = self.conditioner.get_unconditional_conditioning(batch, None)
+ c.append(_c)
+ return c, uc

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,190 @@
+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:\n"{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:\n"{config_path}"')
+ return model
+
+
+def create_SUPIR_model(config_path, SDXL_CKPT, SUPIR_CKPT):
+ config = OmegaConf.load(config_path)
+ model = instantiate_from_config(config.model).cpu()
+ print(f'Loaded model config from:\n"{config_path}"')
+
+ sdxl_sd = load_state_dict(SDXL_CKPT)
+ model.load_state_dict(sdxl_sd, strict=False)
+ model.load_state_dict(load_state_dict(SUPIR_CKPT), strict=False)
+
+ return model, sdxl_sd
+
+
+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, fix_resize=None):
+ """
+ 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
+ if fix_resize is not None:
+ _upsacle = fix_resize / min(w, h)
+ w *= _upsacle
+ h *= _upsacle
+ w0, h0 = round(w), round(h)
+ 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
+ elif dtype_str == "fp8":
+ return torch.float8_e4m3fn
+ 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/face_restoration_helper.py

@@ -0,0 +1,514 @@
+import cv2
+import numpy as np
+import os
+import torch
+from torchvision.transforms.functional import normalize
+
+from facexlib.detection import init_detection_model
+from facexlib.parsing import init_parsing_model
+from facexlib.utils.misc import img2tensor, imwrite
+
+from .file import load_file_from_url
+
+
+def get_largest_face(det_faces, h, w):
+ def get_location(val, length):
+ if val < 0:
+ return 0
+ elif val > length:
+ return length
+ else:
+ return val
+
+ face_areas = []
+ for det_face in det_faces:
+ left = get_location(det_face[0], w)
+ right = get_location(det_face[2], w)
+ top = get_location(det_face[1], h)
+ bottom = get_location(det_face[3], h)
+ face_area = (right - left) * (bottom - top)
+ face_areas.append(face_area)
+ largest_idx = face_areas.index(max(face_areas))
+ return det_faces[largest_idx], largest_idx
+
+
+def get_center_face(det_faces, h=0, w=0, center=None):
+ if center is not None:
+ center = np.array(center)
+ else:
+ center = np.array([w / 2, h / 2])
+ center_dist = []
+ for det_face in det_faces:
+ face_center = np.array([(det_face[0] + det_face[2]) / 2, (det_face[1] + det_face[3]) / 2])
+ dist = np.linalg.norm(face_center - center)
+ center_dist.append(dist)
+ center_idx = center_dist.index(min(center_dist))
+ return det_faces[center_idx], center_idx
+
+
+class FaceRestoreHelper(object):
+ """Helper for the face restoration pipeline (base class)."""
+
+ def __init__(self,
+ upscale_factor,
+ face_size=512,
+ crop_ratio=(1, 1),
+ det_model='retinaface_resnet50',
+ save_ext='png',
+ template_3points=False,
+ pad_blur=False,
+ use_parse=False,
+ device=None):
+ self.template_3points = template_3points # improve robustness
+ self.upscale_factor = int(upscale_factor)
+ # the cropped face ratio based on the square face
+ self.crop_ratio = crop_ratio # (h, w)
+ assert (self.crop_ratio[0] >= 1 and self.crop_ratio[1] >= 1), 'crop ration only supports >=1'
+ self.face_size = (int(face_size * self.crop_ratio[1]), int(face_size * self.crop_ratio[0]))
+ self.det_model = det_model
+
+ if self.det_model == 'dlib':
+ # standard 5 landmarks for FFHQ faces with 1024 x 1024
+ self.face_template = np.array([[686.77227723, 488.62376238], [586.77227723, 493.59405941],
+ [337.91089109, 488.38613861], [437.95049505, 493.51485149],
+ [513.58415842, 678.5049505]])
+ self.face_template = self.face_template / (1024 // face_size)
+ elif self.template_3points:
+ self.face_template = np.array([[192, 240], [319, 240], [257, 371]])
+ else:
+ # standard 5 landmarks for FFHQ faces with 512 x 512
+ # facexlib
+ self.face_template = np.array([[192.98138, 239.94708], [318.90277, 240.1936], [256.63416, 314.01935],
+ [201.26117, 371.41043], [313.08905, 371.15118]])
+
+ # dlib: left_eye: 36:41 right_eye: 42:47 nose: 30,32,33,34 left mouth corner: 48 right mouth corner: 54
+ # self.face_template = np.array([[193.65928, 242.98541], [318.32558, 243.06108], [255.67984, 328.82894],
+ # [198.22603, 372.82502], [313.91018, 372.75659]])
+
+ self.face_template = self.face_template * (face_size / 512.0)
+ if self.crop_ratio[0] > 1:
+ self.face_template[:, 1] += face_size * (self.crop_ratio[0] - 1) / 2
+ if self.crop_ratio[1] > 1:
+ self.face_template[:, 0] += face_size * (self.crop_ratio[1] - 1) / 2
+ self.save_ext = save_ext
+ self.pad_blur = pad_blur
+ if self.pad_blur is True:
+ self.template_3points = False
+
+ self.all_landmarks_5 = []
+ self.det_faces = []
+ self.affine_matrices = []
+ self.inverse_affine_matrices = []
+ self.cropped_faces = []
+ self.restored_faces = []
+ self.pad_input_imgs = []
+
+ if device is None:
+ self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+ # self.device = get_device()
+ else:
+ self.device = device
+
+ # init face detection model
+ self.face_detector = init_detection_model(det_model, half=False, device=self.device)
+
+ # init face parsing model
+ self.use_parse = use_parse
+ self.face_parse = init_parsing_model(model_name='parsenet', device=self.device)
+
+ def set_upscale_factor(self, upscale_factor):
+ self.upscale_factor = upscale_factor
+
+ def read_image(self, img):
+ """img can be image path or cv2 loaded image."""
+ # self.input_img is Numpy array, (h, w, c), BGR, uint8, [0, 255]
+ if isinstance(img, str):
+ img = cv2.imread(img)
+
+ if np.max(img) > 256: # 16-bit image
+ img = img / 65535 * 255
+ if len(img.shape) == 2: # gray image
+ img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+ elif img.shape[2] == 4: # BGRA image with alpha channel
+ img = img[:, :, 0:3]
+
+ self.input_img = img
+ # self.is_gray = is_gray(img, threshold=10)
+ # if self.is_gray:
+ # print('Grayscale input: True')
+
+ if min(self.input_img.shape[:2]) < 512:
+ f = 512.0 / min(self.input_img.shape[:2])
+ self.input_img = cv2.resize(self.input_img, (0, 0), fx=f, fy=f, interpolation=cv2.INTER_LINEAR)
+
+ def init_dlib(self, detection_path, landmark5_path):
+ """Initialize the dlib detectors and predictors."""
+ try:
+ import dlib
+ except ImportError:
+ print('Please install dlib by running:' 'conda install -c conda-forge dlib')
+ detection_path = load_file_from_url(url=detection_path, model_dir='weights/dlib', progress=True, file_name=None)
+ landmark5_path = load_file_from_url(url=landmark5_path, model_dir='weights/dlib', progress=True, file_name=None)
+ face_detector = dlib.cnn_face_detection_model_v1(detection_path)
+ shape_predictor_5 = dlib.shape_predictor(landmark5_path)
+ return face_detector, shape_predictor_5
+
+ def get_face_landmarks_5_dlib(self,
+ only_keep_largest=False,
+ scale=1):
+ det_faces = self.face_detector(self.input_img, scale)
+
+ if len(det_faces) == 0:
+ print('No face detected. Try to increase upsample_num_times.')
+ return 0
+ else:
+ if only_keep_largest:
+ print('Detect several faces and only keep the largest.')
+ face_areas = []
+ for i in range(len(det_faces)):
+ face_area = (det_faces[i].rect.right() - det_faces[i].rect.left()) * (
+ det_faces[i].rect.bottom() - det_faces[i].rect.top())
+ face_areas.append(face_area)
+ largest_idx = face_areas.index(max(face_areas))
+ self.det_faces = [det_faces[largest_idx]]
+ else:
+ self.det_faces = det_faces
+
+ if len(self.det_faces) == 0:
+ return 0
+
+ for face in self.det_faces:
+ shape = self.shape_predictor_5(self.input_img, face.rect)
+ landmark = np.array([[part.x, part.y] for part in shape.parts()])
+ self.all_landmarks_5.append(landmark)
+
+ return len(self.all_landmarks_5)
+
+ def get_face_landmarks_5(self,
+ only_keep_largest=False,
+ only_center_face=False,
+ resize=None,
+ blur_ratio=0.01,
+ eye_dist_threshold=None):
+ if self.det_model == 'dlib':
+ return self.get_face_landmarks_5_dlib(only_keep_largest)
+
+ if resize is None:
+ scale = 1
+ input_img = self.input_img
+ else:
+ h, w = self.input_img.shape[0:2]
+ scale = resize / min(h, w)
+ scale = max(1, scale) # always scale up
+ h, w = int(h * scale), int(w * scale)
+ interp = cv2.INTER_AREA if scale < 1 else cv2.INTER_LINEAR
+ input_img = cv2.resize(self.input_img, (w, h), interpolation=interp)
+
+ with torch.no_grad():
+ bboxes = self.face_detector.detect_faces(input_img)
+
+ if bboxes is None or bboxes.shape[0] == 0:
+ return 0
+ else:
+ bboxes = bboxes / scale
+
+ for bbox in bboxes:
+ # remove faces with too small eye distance: side faces or too small faces
+ eye_dist = np.linalg.norm([bbox[6] - bbox[8], bbox[7] - bbox[9]])
+ if eye_dist_threshold is not None and (eye_dist < eye_dist_threshold):
+ continue
+
+ if self.template_3points:
+ landmark = np.array([[bbox[i], bbox[i + 1]] for i in range(5, 11, 2)])
+ else:
+ landmark = np.array([[bbox[i], bbox[i + 1]] for i in range(5, 15, 2)])
+ self.all_landmarks_5.append(landmark)
+ self.det_faces.append(bbox[0:5])
+
+ if len(self.det_faces) == 0:
+ return 0
+ if only_keep_largest:
+ h, w, _ = self.input_img.shape
+ self.det_faces, largest_idx = get_largest_face(self.det_faces, h, w)
+ self.all_landmarks_5 = [self.all_landmarks_5[largest_idx]]
+ elif only_center_face:
+ h, w, _ = self.input_img.shape
+ self.det_faces, center_idx = get_center_face(self.det_faces, h, w)
+ self.all_landmarks_5 = [self.all_landmarks_5[center_idx]]
+
+ # pad blurry images
+ if self.pad_blur:
+ self.pad_input_imgs = []
+ for landmarks in self.all_landmarks_5:
+ # get landmarks
+ eye_left = landmarks[0, :]
+ eye_right = landmarks[1, :]
+ eye_avg = (eye_left + eye_right) * 0.5
+ mouth_avg = (landmarks[3, :] + landmarks[4, :]) * 0.5
+ eye_to_eye = eye_right - eye_left
+ eye_to_mouth = mouth_avg - eye_avg
+
+ # Get the oriented crop rectangle
+ # x: half width of the oriented crop rectangle
+ x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
+ # - np.flipud(eye_to_mouth) * [-1, 1]: rotate 90 clockwise
+ # norm with the hypotenuse: get the direction
+ x /= np.hypot(*x) # get the hypotenuse of a right triangle
+ rect_scale = 1.5
+ x *= max(np.hypot(*eye_to_eye) * 2.0 * rect_scale, np.hypot(*eye_to_mouth) * 1.8 * rect_scale)
+ # y: half height of the oriented crop rectangle
+ y = np.flipud(x) * [-1, 1]
+
+ # c: center
+ c = eye_avg + eye_to_mouth * 0.1
+ # quad: (left_top, left_bottom, right_bottom, right_top)
+ quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
+ # qsize: side length of the square
+ qsize = np.hypot(*x) * 2
+ border = max(int(np.rint(qsize * 0.1)), 3)
+
+ # get pad
+ # pad: (width_left, height_top, width_right, height_bottom)
+ pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+ int(np.ceil(max(quad[:, 1]))))
+ pad = [
+ max(-pad[0] + border, 1),
+ max(-pad[1] + border, 1),
+ max(pad[2] - self.input_img.shape[0] + border, 1),
+ max(pad[3] - self.input_img.shape[1] + border, 1)
+ ]
+
+ if max(pad) > 1:
+ # pad image
+ pad_img = np.pad(self.input_img, ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
+ # modify landmark coords
+ landmarks[:, 0] += pad[0]
+ landmarks[:, 1] += pad[1]
+ # blur pad images
+ h, w, _ = pad_img.shape
+ y, x, _ = np.ogrid[:h, :w, :1]
+ mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0],
+ np.float32(w - 1 - x) / pad[2]),
+ 1.0 - np.minimum(np.float32(y) / pad[1],
+ np.float32(h - 1 - y) / pad[3]))
+ blur = int(qsize * blur_ratio)
+ if blur % 2 == 0:
+ blur += 1
+ blur_img = cv2.boxFilter(pad_img, 0, ksize=(blur, blur))
+ # blur_img = cv2.GaussianBlur(pad_img, (blur, blur), 0)
+
+ pad_img = pad_img.astype('float32')
+ pad_img += (blur_img - pad_img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
+ pad_img += (np.median(pad_img, axis=(0, 1)) - pad_img) * np.clip(mask, 0.0, 1.0)
+ pad_img = np.clip(pad_img, 0, 255) # float32, [0, 255]
+ self.pad_input_imgs.append(pad_img)
+ else:
+ self.pad_input_imgs.append(np.copy(self.input_img))
+
+ return len(self.all_landmarks_5)
+
+ def align_warp_face(self, save_cropped_path=None, border_mode='constant'):
+ """Align and warp faces with face template.
+ """
+ if self.pad_blur:
+ assert len(self.pad_input_imgs) == len(
+ self.all_landmarks_5), f'Mismatched samples: {len(self.pad_input_imgs)} and {len(self.all_landmarks_5)}'
+ for idx, landmark in enumerate(self.all_landmarks_5):
+ # use 5 landmarks to get affine matrix
+ # use cv2.LMEDS method for the equivalence to skimage transform
+ # ref: https://blog.csdn.net/yichxi/article/details/115827338
+ affine_matrix = cv2.estimateAffinePartial2D(landmark, self.face_template, method=cv2.LMEDS)[0]
+ self.affine_matrices.append(affine_matrix)
+ # warp and crop faces
+ if border_mode == 'constant':
+ border_mode = cv2.BORDER_CONSTANT
+ elif border_mode == 'reflect101':
+ border_mode = cv2.BORDER_REFLECT101
+ elif border_mode == 'reflect':
+ border_mode = cv2.BORDER_REFLECT
+ if self.pad_blur:
+ input_img = self.pad_input_imgs[idx]
+ else:
+ input_img = self.input_img
+ cropped_face = cv2.warpAffine(
+ input_img, affine_matrix, self.face_size, borderMode=border_mode, borderValue=(135, 133, 132)) # gray
+ self.cropped_faces.append(cropped_face)
+ # save the cropped face
+ if save_cropped_path is not None:
+ path = os.path.splitext(save_cropped_path)[0]
+ save_path = f'{path}_{idx:02d}.{self.save_ext}'
+ imwrite(cropped_face, save_path)
+
+ def get_inverse_affine(self, save_inverse_affine_path=None):
+ """Get inverse affine matrix."""
+ for idx, affine_matrix in enumerate(self.affine_matrices):
+ inverse_affine = cv2.invertAffineTransform(affine_matrix)
+ inverse_affine *= self.upscale_factor
+ self.inverse_affine_matrices.append(inverse_affine)
+ # save inverse affine matrices
+ if save_inverse_affine_path is not None:
+ path, _ = os.path.splitext(save_inverse_affine_path)
+ save_path = f'{path}_{idx:02d}.pth'
+ torch.save(inverse_affine, save_path)
+
+ def add_restored_face(self, restored_face, input_face=None):
+ # if self.is_gray:
+ # restored_face = bgr2gray(restored_face) # convert img into grayscale
+ # if input_face is not None:
+ # restored_face = adain_npy(restored_face, input_face) # transfer the color
+ self.restored_faces.append(restored_face)
+
+ def paste_faces_to_input_image(self, save_path=None, upsample_img=None, draw_box=False, face_upsampler=None):
+ h, w, _ = self.input_img.shape
+ h_up, w_up = int(h * self.upscale_factor), int(w * self.upscale_factor)
+
+ if upsample_img is None:
+ # simply resize the background
+ # upsample_img = cv2.resize(self.input_img, (w_up, h_up), interpolation=cv2.INTER_LANCZOS4)
+ upsample_img = cv2.resize(self.input_img, (w_up, h_up), interpolation=cv2.INTER_LINEAR)
+ else:
+ upsample_img = cv2.resize(upsample_img, (w_up, h_up), interpolation=cv2.INTER_LANCZOS4)
+
+ assert len(self.restored_faces) == len(
+ self.inverse_affine_matrices), ('length of restored_faces and affine_matrices are different.')
+
+ inv_mask_borders = []
+ for restored_face, inverse_affine in zip(self.restored_faces, self.inverse_affine_matrices):
+ if face_upsampler is not None:
+ restored_face = face_upsampler.enhance(restored_face, outscale=self.upscale_factor)[0]
+ inverse_affine /= self.upscale_factor
+ inverse_affine[:, 2] *= self.upscale_factor
+ face_size = (self.face_size[0] * self.upscale_factor, self.face_size[1] * self.upscale_factor)
+ else:
+ # Add an offset to inverse affine matrix, for more precise back alignment
+ if self.upscale_factor > 1:
+ extra_offset = 0.5 * self.upscale_factor
+ else:
+ extra_offset = 0
+ inverse_affine[:, 2] += extra_offset
+ face_size = self.face_size
+ inv_restored = cv2.warpAffine(restored_face, inverse_affine, (w_up, h_up))
+
+ # if draw_box or not self.use_parse: # use square parse maps
+ # mask = np.ones(face_size, dtype=np.float32)
+ # inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up))
+ # # remove the black borders
+ # inv_mask_erosion = cv2.erode(
+ # inv_mask, np.ones((int(2 * self.upscale_factor), int(2 * self.upscale_factor)), np.uint8))
+ # pasted_face = inv_mask_erosion[:, :, None] * inv_restored
+ # total_face_area = np.sum(inv_mask_erosion) # // 3
+ # # add border
+ # if draw_box:
+ # h, w = face_size
+ # mask_border = np.ones((h, w, 3), dtype=np.float32)
+ # border = int(1400/np.sqrt(total_face_area))
+ # mask_border[border:h-border, border:w-border,:] = 0
+ # inv_mask_border = cv2.warpAffine(mask_border, inverse_affine, (w_up, h_up))
+ # inv_mask_borders.append(inv_mask_border)
+ # if not self.use_parse:
+ # # compute the fusion edge based on the area of face
+ # w_edge = int(total_face_area**0.5) // 20
+ # erosion_radius = w_edge * 2
+ # inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
+ # blur_size = w_edge * 2
+ # inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
+ # if len(upsample_img.shape) == 2: # upsample_img is gray image
+ # upsample_img = upsample_img[:, :, None]
+ # inv_soft_mask = inv_soft_mask[:, :, None]
+
+ # always use square mask
+ mask = np.ones(face_size, dtype=np.float32)
+ inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up))
+ # remove the black borders
+ inv_mask_erosion = cv2.erode(
+ inv_mask, np.ones((int(2 * self.upscale_factor), int(2 * self.upscale_factor)), np.uint8))
+ pasted_face = inv_mask_erosion[:, :, None] * inv_restored
+ total_face_area = np.sum(inv_mask_erosion) # // 3
+ # add border
+ if draw_box:
+ h, w = face_size
+ mask_border = np.ones((h, w, 3), dtype=np.float32)
+ border = int(1400 / np.sqrt(total_face_area))
+ mask_border[border:h - border, border:w - border, :] = 0
+ inv_mask_border = cv2.warpAffine(mask_border, inverse_affine, (w_up, h_up))
+ inv_mask_borders.append(inv_mask_border)
+ # compute the fusion edge based on the area of face
+ w_edge = int(total_face_area ** 0.5) // 20
+ erosion_radius = w_edge * 2
+ inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
+ blur_size = w_edge * 2
+ inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
+ if len(upsample_img.shape) == 2: # upsample_img is gray image
+ upsample_img = upsample_img[:, :, None]
+ inv_soft_mask = inv_soft_mask[:, :, None]
+
+ # parse mask
+ if self.use_parse:
+ # inference
+ face_input = cv2.resize(restored_face, (512, 512), interpolation=cv2.INTER_LINEAR)
+ face_input = img2tensor(face_input.astype('float32') / 255., bgr2rgb=True, float32=True)
+ normalize(face_input, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+ face_input = torch.unsqueeze(face_input, 0).to(self.device)
+ with torch.no_grad():
+ out = self.face_parse(face_input)[0]
+ out = out.argmax(dim=1).squeeze().cpu().numpy()
+
+ parse_mask = np.zeros(out.shape)
+ MASK_COLORMAP = [0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 255, 0, 0, 0]
+ for idx, color in enumerate(MASK_COLORMAP):
+ parse_mask[out == idx] = color
+ # blur the mask
+ parse_mask = cv2.GaussianBlur(parse_mask, (101, 101), 11)
+ parse_mask = cv2.GaussianBlur(parse_mask, (101, 101), 11)
+ # remove the black borders
+ thres = 10
+ parse_mask[:thres, :] = 0
+ parse_mask[-thres:, :] = 0
+ parse_mask[:, :thres] = 0
+ parse_mask[:, -thres:] = 0
+ parse_mask = parse_mask / 255.
+
+ parse_mask = cv2.resize(parse_mask, face_size)
+ parse_mask = cv2.warpAffine(parse_mask, inverse_affine, (w_up, h_up), flags=3)
+ inv_soft_parse_mask = parse_mask[:, :, None]
+ # pasted_face = inv_restored
+ fuse_mask = (inv_soft_parse_mask < inv_soft_mask).astype('int')
+ inv_soft_mask = inv_soft_parse_mask * fuse_mask + inv_soft_mask * (1 - fuse_mask)
+
+ if len(upsample_img.shape) == 3 and upsample_img.shape[2] == 4: # alpha channel
+ alpha = upsample_img[:, :, 3:]
+ upsample_img = inv_soft_mask * pasted_face + (1 - inv_soft_mask) * upsample_img[:, :, 0:3]
+ upsample_img = np.concatenate((upsample_img, alpha), axis=2)
+ else:
+ upsample_img = inv_soft_mask * pasted_face + (1 - inv_soft_mask) * upsample_img
+
+ if np.max(upsample_img) > 256: # 16-bit image
+ upsample_img = upsample_img.astype(np.uint16)
+ else:
+ upsample_img = upsample_img.astype(np.uint8)
+
+ # draw bounding box
+ if draw_box:
+ # upsample_input_img = cv2.resize(input_img, (w_up, h_up))
+ img_color = np.ones([*upsample_img.shape], dtype=np.float32)
+ img_color[:, :, 0] = 0
+ img_color[:, :, 1] = 255
+ img_color[:, :, 2] = 0
+ for inv_mask_border in inv_mask_borders:
+ upsample_img = inv_mask_border * img_color + (1 - inv_mask_border) * upsample_img
+ # upsample_input_img = inv_mask_border * img_color + (1 - inv_mask_border) * upsample_input_img
+
+ if save_path is not None:
+ path = os.path.splitext(save_path)[0]
+ save_path = f'{path}.{self.save_ext}'
+ imwrite(upsample_img, save_path)
+ return upsample_img
+
+ def clean_all(self):
+ self.all_landmarks_5 = []
+ self.restored_faces = []
+ self.affine_matrices = []
+ self.cropped_faces = []
+ self.inverse_affine_matrices = []
+ self.det_faces = []
+ self.pad_input_imgs = []

SUPIR/utils/file.py

@@ -0,0 +1,79 @@
+import os
+from typing import List, Tuple
+
+from urllib.parse import urlparse
+from torch.hub import download_url_to_file, get_dir
+
+
+def load_file_list(file_list_path: str) -> List[str]:
+ files = []
+ # each line in file list contains a path of an image
+ with open(file_list_path, "r") as fin:
+ for line in fin:
+ path = line.strip()
+ if path:
+ files.append(path)
+ return files
+
+
+def list_image_files(
+ img_dir: str,
+ exts: Tuple[str]=(".jpg", ".png", ".jpeg"),
+ follow_links: bool=False,
+ log_progress: bool=False,
+ log_every_n_files: int=10000,
+ max_size: int=-1
+) -> List[str]:
+ files = []
+ for dir_path, _, file_names in os.walk(img_dir, followlinks=follow_links):
+ early_stop = False
+ for file_name in file_names:
+ if os.path.splitext(file_name)[1].lower() in exts:
+ if max_size >= 0 and len(files) >= max_size:
+ early_stop = True
+ break
+ files.append(os.path.join(dir_path, file_name))
+ if log_progress and len(files) % log_every_n_files == 0:
+ print(f"find {len(files)} images in {img_dir}")
+ if early_stop:
+ break
+ return files
+
+
+def get_file_name_parts(file_path: str) -> Tuple[str, str, str]:
+ parent_path, file_name = os.path.split(file_path)
+ stem, ext = os.path.splitext(file_name)
+ return parent_path, stem, ext
+
+
+# https://github.com/XPixelGroup/BasicSR/blob/master/basicsr/utils/download_util.py/
+def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
+ """Load file form http url, will download models if necessary.
+
+ Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
+
+ Args:
+ url (str): URL to be downloaded.
+ model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
+ Default: None.
+ progress (bool): Whether to show the download progress. Default: True.
+ file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
+
+ Returns:
+ str: The path to the downloaded file.
+ """
+ if model_dir is None: # use the pytorch hub_dir
+ hub_dir = get_dir()
+ model_dir = os.path.join(hub_dir, 'checkpoints')
+
+ os.makedirs(model_dir, exist_ok=True)
+
+ parts = urlparse(url)
+ filename = os.path.basename(parts.path)
+ if file_name is not None:
+ filename = file_name
+ cached_file = os.path.abspath(os.path.join(model_dir, filename))
+ if not os.path.exists(cached_file):
+ print(f'Downloading: "{url}" to {cached_file}\n')
+ download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
+ return cached_file

SUPIR/utils/tilevae.py

@@ -0,0 +1,971 @@
+# ------------------------------------------------------------------------
+#
+# 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
+
+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(
+ devices.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(
+ devices.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(devices.device)
+ devices.torch_gc()
+ gc.collect()
+
+ ret = fn(*args, **kwargs)
+
+ devices.torch_gc()
+ gc.collect()
+ if torch.cuda.is_available():
+ vram = torch.cuda.max_memory_allocated(devices.device) / 2**20
+ torch.cuda.reset_peak_memory_stats(devices.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=devices.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(devices.get_optimal_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)

app.py

@@ -0,0 +1,911 @@
+raise NotImplementedError(
+ "This Space is meant for SD-Webui-Forge\nhttps://github.com/lllyasviel/stable-diffusion-webui-forge"
+)
+
+import os
+import gradio as gr
+import argparse
+import numpy as np
+import torch
+import einops
+import copy
+import math
+import time
+import random
+import spaces
+import re
+import uuid
+
+from gradio_imageslider import ImageSlider
+from PIL import Image
+from SUPIR.util import HWC3, upscale_image, fix_resize, convert_dtype, create_SUPIR_model, load_QF_ckpt
+from huggingface_hub import hf_hub_download
+from pillow_heif import register_heif_opener
+
+register_heif_opener()
+
+max_64_bit_int = np.iinfo(np.int32).max
+
+hf_hub_download(repo_id="laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", filename="open_clip_pytorch_model.bin", local_dir="laion_CLIP-ViT-bigG-14-laion2B-39B-b160k")
+hf_hub_download(repo_id="camenduru/SUPIR", filename="sd_xl_base_1.0_0.9vae.safetensors", local_dir="yushan777_SUPIR")
+hf_hub_download(repo_id="camenduru/SUPIR", filename="SUPIR-v0F.ckpt", local_dir="yushan777_SUPIR")
+hf_hub_download(repo_id="camenduru/SUPIR", filename="SUPIR-v0Q.ckpt", local_dir="yushan777_SUPIR")
+hf_hub_download(repo_id="RunDiffusion/Juggernaut-XL-Lightning", filename="Juggernaut_RunDiffusionPhoto2_Lightning_4Steps.safetensors", local_dir="RunDiffusion_Juggernaut-XL-Lightning")
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--opt", type=str, default='options/SUPIR_v0.yaml')
+parser.add_argument("--ip", type=str, default='127.0.0.1')
+parser.add_argument("--port", type=int, default='6688')
+parser.add_argument("--no_llava", action='store_true', default=True)#False
+parser.add_argument("--use_image_slider", action='store_true', default=False)#False
+parser.add_argument("--log_history", action='store_true', default=False)
+parser.add_argument("--loading_half_params", action='store_true', default=False)#False
+parser.add_argument("--use_tile_vae", action='store_true', default=True)#False
+parser.add_argument("--encoder_tile_size", type=int, default=512)
+parser.add_argument("--decoder_tile_size", type=int, default=64)
+parser.add_argument("--load_8bit_llava", action='store_true', default=False)
+args = parser.parse_args()
+
+if torch.cuda.device_count() > 0:
+ SUPIR_device = 'cuda:0'
+
+ # Load SUPIR
+ model, default_setting = create_SUPIR_model(args.opt, SUPIR_sign='Q', load_default_setting=True)
+ if args.loading_half_params:
+ model = model.half()
+ if args.use_tile_vae:
+ model.init_tile_vae(encoder_tile_size=args.encoder_tile_size, decoder_tile_size=args.decoder_tile_size)
+ model = model.to(SUPIR_device)
+ model.first_stage_model.denoise_encoder_s1 = copy.deepcopy(model.first_stage_model.denoise_encoder)
+ model.current_model = 'v0-Q'
+ ckpt_Q, ckpt_F = load_QF_ckpt(args.opt)
+
+def check_upload(input_image):
+ if input_image is None:
+ raise gr.Error("Please provide an image to restore.")
+ return gr.update(visible = True)
+
+def update_seed(is_randomize_seed, seed):
+ if is_randomize_seed:
+ return random.randint(0, max_64_bit_int)
+ return seed
+
+def reset():
+ return [
+ None,
+ 0,
+ None,
+ None,
+ "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
+ "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
+ 1,
+ 1024,
+ 1,
+ 2,
+ 50,
+ -1.0,
+ 1.,
+ default_setting.s_cfg_Quality if torch.cuda.device_count() > 0 else 1.0,
+ True,
+ random.randint(0, max_64_bit_int),
+ 5,
+ 1.003,
+ "Wavelet",
+ "fp32",
+ "fp32",
+ 1.0,
+ True,
+ False,
+ default_setting.spt_linear_CFG_Quality if torch.cuda.device_count() > 0 else 1.0,
+ 0.,
+ "v0-Q",
+ "input",
+ 6
+ ]
+
+def check(input_image):
+ if input_image is None:
+ raise gr.Error("Please provide an image to restore.")
+
+@spaces.GPU(duration=420)
+def stage1_process(
+ input_image,
+ gamma_correction,
+ diff_dtype,
+ ae_dtype
+):
+ print('stage1_process ==>>')
+ if torch.cuda.device_count() == 0:
+ gr.Warning('Set this space to GPU config to make it work.')
+ return None, None
+ torch.cuda.set_device(SUPIR_device)
+ LQ = HWC3(np.array(Image.open(input_image)))
+ LQ = fix_resize(LQ, 512)
+ # stage1
+ LQ = np.array(LQ) / 255 * 2 - 1
+ LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
+
+ model.ae_dtype = convert_dtype(ae_dtype)
+ model.model.dtype = convert_dtype(diff_dtype)
+
+ LQ = model.batchify_denoise(LQ, is_stage1=True)
+ LQ = (LQ[0].permute(1, 2, 0) * 127.5 + 127.5).cpu().numpy().round().clip(0, 255).astype(np.uint8)
+ # gamma correction
+ LQ = LQ / 255.0
+ LQ = np.power(LQ, gamma_correction)
+ LQ *= 255.0
+ LQ = LQ.round().clip(0, 255).astype(np.uint8)
+ print('<<== stage1_process')
+ return LQ, gr.update(visible = True)
+
+def stage2_process(*args, **kwargs):
+ try:
+ return restore_in_Xmin(*args, **kwargs)
+ except Exception as e:
+ # NO_GPU_MESSAGE_INQUEUE
+ print("gradio.exceptions.Error 'No GPU is currently available for you after 60s'")
+ print('str(type(e)): ' + str(type(e))) # <class 'gradio.exceptions.Error'>
+ print('str(e): ' + str(e)) # You have exceeded your GPU quota...
+ try:
+ print('e.message: ' + e.message) # No GPU is currently available for you after 60s
+ except Exception as e2:
+ print('Failure')
+ if str(e).startswith("No GPU is currently available for you after 60s"):
+ print('Exception identified!!!')
+ #if str(type(e)) == "<class 'gradio.exceptions.Error'>":
+ #print('Exception of name ' + type(e).__name__)
+ raise e
+
+def restore_in_Xmin(
+ noisy_image,
+ rotation,
+ denoise_image,
+ prompt,
+ a_prompt,
+ n_prompt,
+ num_samples,
+ min_size,
+ downscale,
+ upscale,
+ edm_steps,
+ s_stage1,
+ s_stage2,
+ s_cfg,
+ randomize_seed,
+ seed,
+ s_churn,
+ s_noise,
+ color_fix_type,
+ diff_dtype,
+ ae_dtype,
+ gamma_correction,
+ linear_CFG,
+ linear_s_stage2,
+ spt_linear_CFG,
+ spt_linear_s_stage2,
+ model_select,
+ output_format,
+ allocation
+):
+ print("noisy_image:\n" + str(noisy_image))
+ print("denoise_image:\n" + str(denoise_image))
+ print("rotation: " + str(rotation))
+ print("prompt: " + str(prompt))
+ print("a_prompt: " + str(a_prompt))
+ print("n_prompt: " + str(n_prompt))
+ print("num_samples: " + str(num_samples))
+ print("min_size: " + str(min_size))
+ print("downscale: " + str(downscale))
+ print("upscale: " + str(upscale))
+ print("edm_steps: " + str(edm_steps))
+ print("s_stage1: " + str(s_stage1))
+ print("s_stage2: " + str(s_stage2))
+ print("s_cfg: " + str(s_cfg))
+ print("randomize_seed: " + str(randomize_seed))
+ print("seed: " + str(seed))
+ print("s_churn: " + str(s_churn))
+ print("s_noise: " + str(s_noise))
+ print("color_fix_type: " + str(color_fix_type))
+ print("diff_dtype: " + str(diff_dtype))
+ print("ae_dtype: " + str(ae_dtype))
+ print("gamma_correction: " + str(gamma_correction))
+ print("linear_CFG: " + str(linear_CFG))
+ print("linear_s_stage2: " + str(linear_s_stage2))
+ print("spt_linear_CFG: " + str(spt_linear_CFG))
+ print("spt_linear_s_stage2: " + str(spt_linear_s_stage2))
+ print("model_select: " + str(model_select))
+ print("GPU time allocation: " + str(allocation) + " min")
+ print("output_format: " + str(output_format))
+
+ input_format = re.sub(r"^.*\.([^\.]+)$", r"\1", noisy_image)
+
+ if input_format not in ['png', 'webp', 'jpg', 'jpeg', 'gif', 'bmp', 'heic']:
+ gr.Warning('Invalid image format. Please first convert into *.png, *.webp, *.jpg, *.jpeg, *.gif, *.bmp or *.heic.')
+ return None, None, None, None
+
+ if output_format == "input":
+ if noisy_image is None:
+ output_format = "png"
+ else:
+ output_format = input_format
+ print("final output_format: " + str(output_format))
+
+ if prompt is None:
+ prompt = ""
+
+ if a_prompt is None:
+ a_prompt = ""
+
+ if n_prompt is None:
+ n_prompt = ""
+
+ if prompt != "" and a_prompt != "":
+ a_prompt = prompt + ", " + a_prompt
+ else:
+ a_prompt = prompt + a_prompt
+ print("Final prompt: " + str(a_prompt))
+
+ denoise_image = np.array(Image.open(noisy_image if denoise_image is None else denoise_image))
+
+ if rotation == 90:
+ denoise_image = np.array(list(zip(*denoise_image[::-1])))
+ elif rotation == 180:
+ denoise_image = np.array(list(zip(*denoise_image[::-1])))
+ denoise_image = np.array(list(zip(*denoise_image[::-1])))
+ elif rotation == -90:
+ denoise_image = np.array(list(zip(*denoise_image))[::-1])
+
+ if 1 < downscale:
+ input_height, input_width, input_channel = denoise_image.shape
+ denoise_image = np.array(Image.fromarray(denoise_image).resize((input_width // downscale, input_height // downscale), Image.LANCZOS))
+
+ denoise_image = HWC3(denoise_image)
+
+ if torch.cuda.device_count() == 0:
+ gr.Warning('Set this space to GPU config to make it work.')
+ return [noisy_image, denoise_image], gr.update(label="Downloadable results in *." + output_format + " format", format = output_format, value = [denoise_image]), None, gr.update(visible=True)
+
+ if model_select != model.current_model:
+ print('load ' + model_select)
+ if model_select == 'v0-Q':
+ model.load_state_dict(ckpt_Q, strict=False)
+ elif model_select == 'v0-F':
+ model.load_state_dict(ckpt_F, strict=False)
+ model.current_model = model_select
+
+ model.ae_dtype = convert_dtype(ae_dtype)
+ model.model.dtype = convert_dtype(diff_dtype)
+
+ # Allocation
+ if allocation == 1:
+ return restore_in_1min(
+ noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
+ )
+ if allocation == 2:
+ return restore_in_2min(
+ noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
+ )
+ if allocation == 3:
+ return restore_in_3min(
+ noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
+ )
+ if allocation == 4:
+ return restore_in_4min(
+ noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
+ )
+ if allocation == 5:
+ return restore_in_5min(
+ noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
+ )
+ if allocation == 7:
+ return restore_in_7min(
+ noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
+ )
+ if allocation == 8:
+ return restore_in_8min(
+ noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
+ )
+ if allocation == 9:
+ return restore_in_9min(
+ noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
+ )
+ if allocation == 10:
+ return restore_in_10min(
+ noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
+ )
+ else:
+ return restore_in_6min(
+ noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
+ )
+
+@spaces.GPU(duration=59)
+def restore_in_1min(*args, **kwargs):
+ return restore_on_gpu(*args, **kwargs)
+
+@spaces.GPU(duration=119)
+def restore_in_2min(*args, **kwargs):
+ return restore_on_gpu(*args, **kwargs)
+
+@spaces.GPU(duration=179)
+def restore_in_3min(*args, **kwargs):
+ return restore_on_gpu(*args, **kwargs)
+
+@spaces.GPU(duration=239)
+def restore_in_4min(*args, **kwargs):
+ return restore_on_gpu(*args, **kwargs)
+
+@spaces.GPU(duration=299)
+def restore_in_5min(*args, **kwargs):
+ return restore_on_gpu(*args, **kwargs)
+
+@spaces.GPU(duration=359)
+def restore_in_6min(*args, **kwargs):
+ return restore_on_gpu(*args, **kwargs)
+
+@spaces.GPU(duration=419)
+def restore_in_7min(*args, **kwargs):
+ return restore_on_gpu(*args, **kwargs)
+
+@spaces.GPU(duration=479)
+def restore_in_8min(*args, **kwargs):
+ return restore_on_gpu(*args, **kwargs)
+
+@spaces.GPU(duration=539)
+def restore_in_9min(*args, **kwargs):
+ return restore_on_gpu(*args, **kwargs)
+
+@spaces.GPU(duration=599)
+def restore_in_10min(*args, **kwargs):
+ return restore_on_gpu(*args, **kwargs)
+
+def restore_on_gpu(
+ noisy_image,
+ input_image,
+ prompt,
+ a_prompt,
+ n_prompt,
+ num_samples,
+ min_size,
+ downscale,
+ upscale,
+ edm_steps,
+ s_stage1,
+ s_stage2,
+ s_cfg,
+ randomize_seed,
+ seed,
+ s_churn,
+ s_noise,
+ color_fix_type,
+ diff_dtype,
+ ae_dtype,
+ gamma_correction,
+ linear_CFG,
+ linear_s_stage2,
+ spt_linear_CFG,
+ spt_linear_s_stage2,
+ model_select,
+ output_format,
+ allocation
+):
+ start = time.time()
+ print('restore ==>>')
+
+ torch.cuda.set_device(SUPIR_device)
+
+ with torch.no_grad():
+ input_image = upscale_image(input_image, upscale, unit_resolution=32, min_size=min_size)
+ LQ = np.array(input_image) / 255.0
+ LQ = np.power(LQ, gamma_correction)
+ LQ *= 255.0
+ LQ = LQ.round().clip(0, 255).astype(np.uint8)
+ LQ = LQ / 255 * 2 - 1
+ LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
+ captions = ['']
+
+ samples = model.batchify_sample(LQ, captions, num_steps=edm_steps, restoration_scale=s_stage1, s_churn=s_churn,
+ s_noise=s_noise, cfg_scale=s_cfg, control_scale=s_stage2, seed=seed,
+ num_samples=num_samples, p_p=a_prompt, n_p=n_prompt, color_fix_type=color_fix_type,
+ use_linear_CFG=linear_CFG, use_linear_control_scale=linear_s_stage2,
+ cfg_scale_start=spt_linear_CFG, control_scale_start=spt_linear_s_stage2)
+
+ x_samples = (einops.rearrange(samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().round().clip(
+ 0, 255).astype(np.uint8)
+ results = [x_samples[i] for i in range(num_samples)]
+ torch.cuda.empty_cache()
+
+ # All the results have the same size
+ input_height, input_width, input_channel = np.array(input_image).shape
+ result_height, result_width, result_channel = np.array(results[0]).shape
+
+ print('<<== restore')
+ end = time.time()
+ secondes = int(end - start)
+ minutes = math.floor(secondes / 60)
+ secondes = secondes - (minutes * 60)
+ hours = math.floor(minutes / 60)
+ minutes = minutes - (hours * 60)
+ information = ("Start the process again if you want a different result. " if randomize_seed else "") + \
+ "If you don't get the image you wanted, add more details in the « Image description ». " + \
+ "Wait " + str(allocation) + " min before a new run to avoid quota penalty or use another computer. " + \
+ "The image" + (" has" if len(results) == 1 else "s have") + " been generated in " + \
+ ((str(hours) + " h, ") if hours != 0 else "") + \
+ ((str(minutes) + " min, ") if hours != 0 or minutes != 0 else "") + \
+ str(secondes) + " sec. " + \
+ "The new image resolution is " + str(result_width) + \
+ " pixels large and " + str(result_height) + \
+ " pixels high, so a resolution of " + f'{result_width * result_height:,}' + " pixels."
+ print(information)
+ try:
+ print("Initial resolution: " + f'{input_width * input_height:,}')
+ print("Final resolution: " + f'{result_width * result_height:,}')
+ print("edm_steps: " + str(edm_steps))
+ print("num_samples: " + str(num_samples))
+ print("downscale: " + str(downscale))
+ print("Estimated minutes: " + f'{(((result_width * result_height**(1/1.75)) * input_width * input_height * (edm_steps**(1/2)) * (num_samples**(1/2.5)))**(1/2.5)) / 25000:,}')
+ except Exception as e:
+ print('Exception of Estimation')
+
+ # Only one image can be shown in the slider
+ return [noisy_image] + [results[0]], gr.update(label="Downloadable results in *." + output_format + " format", format = output_format, value = results), gr.update(value = information, visible = True), gr.update(visible=True)
+
+def load_and_reset(param_setting):
+ print('load_and_reset ==>>')
+ if torch.cuda.device_count() == 0:
+ gr.Warning('Set this space to GPU config to make it work.')
+ return None, None, None, None, None, None, None, None, None, None, None, None, None, None
+ edm_steps = default_setting.edm_steps
+ s_stage2 = 1.0
+ s_stage1 = -1.0
+ s_churn = 5
+ s_noise = 1.003
+ a_prompt = 'Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - ' \
+ 'realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore ' \
+ 'detailing, hyper sharpness, perfect without deformations.'
+ n_prompt = 'painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, ' \
+ '3D render, unreal engine, blurring, dirty, messy, worst quality, low quality, frames, watermark, ' \
+ 'signature, jpeg artifacts, deformed, lowres, over-smooth'
+ color_fix_type = 'Wavelet'
+ spt_linear_s_stage2 = 0.0
+ linear_s_stage2 = False
+ linear_CFG = True
+ if param_setting == "Quality":
+ s_cfg = default_setting.s_cfg_Quality
+ spt_linear_CFG = default_setting.spt_linear_CFG_Quality
+ model_select = "v0-Q"
+ elif param_setting == "Fidelity":
+ s_cfg = default_setting.s_cfg_Fidelity
+ spt_linear_CFG = default_setting.spt_linear_CFG_Fidelity
+ model_select = "v0-F"
+ else:
+ raise NotImplementedError
+ gr.Info('The parameters are reset.')
+ print('<<== load_and_reset')
+ return edm_steps, s_cfg, s_stage2, s_stage1, s_churn, s_noise, a_prompt, n_prompt, color_fix_type, linear_CFG, \
+ linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select
+
+def log_information(result_gallery):
+ print('log_information')
+ if result_gallery is not None:
+ for i, result in enumerate(result_gallery):
+ print(result[0])
+
+def on_select_result(result_slider, result_gallery, evt: gr.SelectData):
+ print('on_select_result')
+ if result_gallery is not None:
+ for i, result in enumerate(result_gallery):
+ print(result[0])
+ return [result_slider[0], result_gallery[evt.index][0]]
+
+title_html = """
+ <h1><center>SUPIR</center></h1>
+ <big><center>Upscale your images up to x10 freely, without account, without watermark and download it</center></big>
+ <center><big><big>🤸<big><big><big><big><big><big>🤸</big></big></big></big></big></big></big></big></center>
+
+ <p>This is an online demo of SUPIR, a practicing model scaling for photo-realistic image restoration.
+ The content added by SUPIR is <b><u>imagination, not real-world information</u></b>.
+ SUPIR is for beauty and illustration only.
+ Most of the processes last few minutes.
+ If you want to upscale AI-generated images, be noticed that <i>PixArt Sigma</i> space can directly generate 5984x5984 images.
+ Due to Gradio issues, the generated image is slightly less satured than the original.
+ Please leave a <a href="https://huggingface.co/spaces/Fabrice-TIERCELIN/SUPIR/discussions/new">message in discussion</a> if you encounter issues.
+ You can also use <a href="https://huggingface.co/spaces/gokaygokay/AuraSR">AuraSR</a> to upscale x4.
+
+ <p><center><a href="https://arxiv.org/abs/2401.13627">Paper</a> &emsp; <a href="http://supir.xpixel.group/">Project Page</a> &emsp; <a href="https://huggingface.co/blog/MonsterMMORPG/supir-sota-image-upscale-better-than-magnific-ai">Local Install Guide</a></center></p>
+ <p><center><a style="display:inline-block" href='https://github.com/Fanghua-Yu/SUPIR'><img alt="GitHub Repo stars" src="https://img.shields.io/github/stars/Fanghua-Yu/SUPIR?style=social"></a></center></p>
+ """
+
+
+claim_md = """
+## **Piracy**
+The images are not stored but the logs are saved during a month.
+## **How to get SUPIR**
+You can get SUPIR on HuggingFace by [duplicating this space](https://huggingface.co/spaces/Fabrice-TIERCELIN/SUPIR?duplicate=true) and set GPU.
+You can also install SUPIR on your computer following [this tutorial](https://huggingface.co/blog/MonsterMMORPG/supir-sota-image-upscale-better-than-magnific-ai).
+You can install _Pinokio_ on your computer and then install _SUPIR_ into it. It should be quite easy if you have an Nvidia GPU.
+## **Terms of use**
+By using this service, users are required to agree to the following terms: The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research. Please submit a feedback to us if you get any inappropriate answer! We will collect those to keep improving our models. For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
+## **License**
+The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/Fanghua-Yu/SUPIR) of SUPIR.
+"""
+
+# Gradio interface
+with gr.Blocks() as interface:
+ if torch.cuda.device_count() == 0:
+ with gr.Row():
+ gr.HTML("""
+ <p style="background-color: red;"><big><big><big><b>⚠️To use SUPIR, <a href="https://huggingface.co/spaces/Fabrice-TIERCELIN/SUPIR?duplicate=true">duplicate this space</a> and set a GPU with 30 GB VRAM.</b>
+
+ You can't use SUPIR directly here because this space runs on a CPU, which is not enough for SUPIR. Please provide <a href="https://huggingface.co/spaces/Fabrice-TIERCELIN/SUPIR/discussions/new">feedback</a> if you have issues.
+ </big></big></big></p>
+ """)
+ gr.HTML(title_html)
+
+ input_image = gr.Image(label="Input (*.png, *.webp, *.jpeg, *.jpg, *.gif, *.bmp, *.heic)", show_label=True, type="filepath", height=600, elem_id="image-input")
+ rotation = gr.Radio([["No rotation", 0], ["⤵ Rotate +90°", 90], ["↩ Return 180°", 180], ["⤴ Rotate -90°", -90]], label="Orientation correction", info="Will apply the following rotation before restoring the image; the AI needs a good orientation to understand the content", value=0, interactive=True, visible=False)
+ with gr.Group():
+ prompt = gr.Textbox(label="Image description", info="Help the AI understand what the image represents; describe as much as possible, especially the details we can't see on the original image; you can write in any language", value="", placeholder="A 33 years old man, walking, in the street, Santiago, morning, Summer, photorealistic", lines=3)
+ prompt_hint = gr.HTML("You can use a <a href='"'https://huggingface.co/spaces/MaziyarPanahi/llava-llama-3-8b'"'>LlaVa space</a> to auto-generate the description of your image.")
+ upscale = gr.Radio([["x1", 1], ["x2", 2], ["x3", 3], ["x4", 4], ["x5", 5], ["x6", 6], ["x7", 7], ["x8", 8], ["x9", 9], ["x10", 10]], label="Upscale factor", info="Resolution x1 to x10", value=2, interactive=True)
+ output_format = gr.Radio([["As input", "input"], ["*.png", "png"], ["*.webp", "webp"], ["*.jpeg", "jpeg"], ["*.gif", "gif"], ["*.bmp", "bmp"]], label="Image format for result", info="File extention", value="input", interactive=True)
+ allocation = gr.Radio([["1 min", 1], ["2 min", 2], ["3 min", 3], ["4 min", 4], ["5 min", 5], ["6 min", 6], ["7 min (discouraged)", 7], ["8 min (discouraged)", 8], ["9 min (discouraged)", 9], ["10 min (discouraged)", 10]], label="GPU allocation time", info="lower=May abort run, higher=Quota penalty for next runs", value=5, interactive=True)
+
+ with gr.Accordion("Pre-denoising (optional)", open=False):
+ gamma_correction = gr.Slider(label="Gamma Correction", info = "lower=lighter, higher=darker", minimum=0.1, maximum=2.0, value=1.0, step=0.1)
+ denoise_button = gr.Button(value="Pre-denoise")
+ denoise_image = gr.Image(label="Denoised image", show_label=True, type="filepath", sources=[], interactive = False, height=600, elem_id="image-s1")
+ denoise_information = gr.HTML(value="If present, the denoised image will be used for the restoration instead of the input image.", visible=False)
+
+ with gr.Accordion("Advanced options", open=False):
+ a_prompt = gr.Textbox(label="Additional image description",
+ info="Completes the main image description",
+ value='Cinematic, High Contrast, highly detailed, taken using a Canon EOS R '
+ 'camera, hyper detailed photo - realistic maximum detail, 32k, Color '
+ 'Grading, ultra HD, extreme meticulous detailing, skin pore detailing, '
+ 'hyper sharpness, perfect without deformations.',
+ lines=3)
+ n_prompt = gr.Textbox(label="Negative image description",
+ info="Disambiguate by listing what the image does NOT represent",
+ value='painting, oil painting, illustration, drawing, art, sketch, anime, '
+ 'cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, '
+ 'worst quality, low quality, frames, watermark, signature, jpeg artifacts, '
+ 'deformed, lowres, over-smooth',
+ lines=3)
+ edm_steps = gr.Slider(label="Steps", info="lower=faster, higher=more details", minimum=1, maximum=200, value=default_setting.edm_steps if torch.cuda.device_count() > 0 else 1, step=1)
+ num_samples = gr.Slider(label="Num Samples", info="Number of generated results", minimum=1, maximum=4 if not args.use_image_slider else 1
+ , value=1, step=1)
+ min_size = gr.Slider(label="Minimum size", info="Minimum height, minimum width of the result", minimum=32, maximum=4096, value=1024, step=32)
+ downscale = gr.Radio([["/1", 1], ["/2", 2], ["/3", 3], ["/4", 4], ["/5", 5], ["/6", 6], ["/7", 7], ["/8", 8], ["/9", 9], ["/10", 10]], label="Pre-downscale factor", info="Reducing blurred image reduce the process time", value=1, interactive=True)
+ with gr.Row():
+ with gr.Column():
+ model_select = gr.Radio([["💃 Quality (v0-Q)", "v0-Q"], ["🎯 Fidelity (v0-F)", "v0-F"]], label="Model Selection", info="Pretrained model", value="v0-Q",
+ interactive=True)
+ with gr.Column():
+ color_fix_type = gr.Radio([["None", "None"], ["AdaIn (improve as a photo)", "AdaIn"], ["Wavelet (for JPEG artifacts)", "Wavelet"]], label="Color-Fix Type", info="AdaIn=Improve following a style, Wavelet=For JPEG artifacts", value="AdaIn",
+ interactive=True)
+ s_cfg = gr.Slider(label="Text Guidance Scale", info="lower=follow the image, higher=follow the prompt", minimum=1.0, maximum=15.0,
+ value=default_setting.s_cfg_Quality if torch.cuda.device_count() > 0 else 1.0, step=0.1)
+ s_stage2 = gr.Slider(label="Restoring Guidance Strength", minimum=0., maximum=1., value=1., step=0.05)
+ s_stage1 = gr.Slider(label="Pre-denoising Guidance Strength", minimum=-1.0, maximum=6.0, value=-1.0, step=1.0)
+ s_churn = gr.Slider(label="S-Churn", minimum=0, maximum=40, value=5, step=1)
+ s_noise = gr.Slider(label="S-Noise", minimum=1.0, maximum=1.1, value=1.003, step=0.001)
+ with gr.Row():
+ with gr.Column():
+ linear_CFG = gr.Checkbox(label="Linear CFG", value=True)
+ spt_linear_CFG = gr.Slider(label="CFG Start", minimum=1.0,
+ maximum=9.0, value=default_setting.spt_linear_CFG_Quality if torch.cuda.device_count() > 0 else 1.0, step=0.5)
+ with gr.Column():
+ linear_s_stage2 = gr.Checkbox(label="Linear Restoring Guidance", value=False)
+ spt_linear_s_stage2 = gr.Slider(label="Guidance Start", minimum=0.,
+ maximum=1., value=0., step=0.05)
+ with gr.Column():
+ diff_dtype = gr.Radio([["fp32 (precision)", "fp32"], ["fp16 (medium)", "fp16"], ["bf16 (speed)", "bf16"]], label="Diffusion Data Type", value="fp32",
+ interactive=True)
+ with gr.Column():
+ ae_dtype = gr.Radio([["fp32 (precision)", "fp32"], ["bf16 (speed)", "bf16"]], label="Auto-Encoder Data Type", value="fp32",
+ interactive=True)
+ randomize_seed = gr.Checkbox(label = "\U0001F3B2 Randomize seed", value = True, info = "If checked, result is always different")
+ seed = gr.Slider(label="Seed", minimum=0, maximum=max_64_bit_int, step=1, randomize=True)
+ with gr.Group():
+ param_setting = gr.Radio(["Quality", "Fidelity"], interactive=True, label="Presetting", value = "Quality")
+ restart_button = gr.Button(value="Apply presetting")
+
+ with gr.Column():
+ diffusion_button = gr.Button(value="🚀 Upscale/Restore", variant = "primary", elem_id = "process_button")
+ reset_btn = gr.Button(value="🧹 Reinit page", variant="stop", elem_id="reset_button", visible = False)
+
+ restore_information = gr.HTML(value = "Restart the process to get another result.", visible = False)
+ result_slider = ImageSlider(label = 'Comparator', show_label = False, interactive = False, elem_id = "slider1", show_download_button = False)
+ result_gallery = gr.Gallery(label = 'Downloadable results', show_label = True, interactive = False, elem_id = "gallery1")
+
+ gr.Examples(
+ examples = [
+ [
+ "./Examples/Example1.png",
+ 0,
+ None,
+ "Group of people, walking, happy, in the street, photorealistic, 8k, extremely detailled",
+ "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
+ "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
+ 2,
+ 1024,
+ 1,
+ 8,
+ 200,
+ -1,
+ 1,
+ 7.5,
+ False,
+ 42,
+ 5,
+ 1.003,
+ "AdaIn",
+ "fp16",
+ "bf16",
+ 1.0,
+ True,
+ 4,
+ False,
+ 0.,
+ "v0-Q",
+ "input",
+ 5
+ ],
+ [
+ "./Examples/Example2.jpeg",
+ 0,
+ None,
+ "La cabeza de un gato atigrado, en una casa, fotorrealista, 8k, extremadamente detallada",
+ "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
+ "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
+ 1,
+ 1024,
+ 1,
+ 1,
+ 200,
+ -1,
+ 1,
+ 7.5,
+ False,
+ 42,
+ 5,
+ 1.003,
+ "Wavelet",
+ "fp16",
+ "bf16",
+ 1.0,
+ True,
+ 4,
+ False,
+ 0.,
+ "v0-Q",
+ "input",
+ 4
+ ],
+ [
+ "./Examples/Example3.webp",
+ 0,
+ None,
+ "A red apple",
+ "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
+ "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
+ 1,
+ 1024,
+ 1,
+ 1,
+ 200,
+ -1,
+ 1,
+ 7.5,
+ False,
+ 42,
+ 5,
+ 1.003,
+ "Wavelet",
+ "fp16",
+ "bf16",
+ 1.0,
+ True,
+ 4,
+ False,
+ 0.,
+ "v0-Q",
+ "input",
+ 4
+ ],
+ [
+ "./Examples/Example3.webp",
+ 0,
+ None,
+ "A red marble",
+ "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
+ "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
+ 1,
+ 1024,
+ 1,
+ 1,
+ 200,
+ -1,
+ 1,
+ 7.5,
+ False,
+ 42,
+ 5,
+ 1.003,
+ "Wavelet",
+ "fp16",
+ "bf16",
+ 1.0,
+ True,
+ 4,
+ False,
+ 0.,
+ "v0-Q",
+ "input",
+ 4
+ ],
+ ],
+ run_on_click = True,
+ fn = stage2_process,
+ inputs = [
+ input_image,
+ rotation,
+ denoise_image,
+ prompt,
+ a_prompt,
+ n_prompt,
+ num_samples,
+ min_size,
+ downscale,
+ upscale,
+ edm_steps,
+ s_stage1,
+ s_stage2,
+ s_cfg,
+ randomize_seed,
+ seed,
+ s_churn,
+ s_noise,
+ color_fix_type,
+ diff_dtype,
+ ae_dtype,
+ gamma_correction,
+ linear_CFG,
+ linear_s_stage2,
+ spt_linear_CFG,
+ spt_linear_s_stage2,
+ model_select,
+ output_format,
+ allocation
+ ],
+ outputs = [
+ result_slider,
+ result_gallery,
+ restore_information,
+ reset_btn
+ ],
+ cache_examples = False,
+ )
+
+ with gr.Row():
+ gr.Markdown(claim_md)
+
+ input_image.upload(fn = check_upload, inputs = [
+ input_image
+ ], outputs = [
+ rotation
+ ], queue = False, show_progress = False)
+
+ denoise_button.click(fn = check, inputs = [
+ input_image
+ ], outputs = [], queue = False, show_progress = False).success(fn = stage1_process, inputs = [
+ input_image,
+ gamma_correction,
+ diff_dtype,
+ ae_dtype
+ ], outputs=[
+ denoise_image,
+ denoise_information
+ ])
+
+ diffusion_button.click(fn = update_seed, inputs = [
+ randomize_seed,
+ seed
+ ], outputs = [
+ seed
+ ], queue = False, show_progress = False).then(fn = check, inputs = [
+ input_image
+ ], outputs = [], queue = False, show_progress = False).success(fn=stage2_process, inputs = [
+ input_image,
+ rotation,
+ denoise_image,
+ prompt,
+ a_prompt,
+ n_prompt,
+ num_samples,
+ min_size,
+ downscale,
+ upscale,
+ edm_steps,
+ s_stage1,
+ s_stage2,
+ s_cfg,
+ randomize_seed,
+ seed,
+ s_churn,
+ s_noise,
+ color_fix_type,
+ diff_dtype,
+ ae_dtype,
+ gamma_correction,
+ linear_CFG,
+ linear_s_stage2,
+ spt_linear_CFG,
+ spt_linear_s_stage2,
+ model_select,
+ output_format,
+ allocation
+ ], outputs = [
+ result_slider,
+ result_gallery,
+ restore_information,
+ reset_btn
+ ]).success(fn = log_information, inputs = [
+ result_gallery
+ ], outputs = [], queue = False, show_progress = False)
+
+ result_gallery.change(on_select_result, [result_slider, result_gallery], result_slider)
+ result_gallery.select(on_select_result, [result_slider, result_gallery], result_slider)
+
+ restart_button.click(fn = load_and_reset, inputs = [
+ param_setting
+ ], outputs = [
+ edm_steps,
+ s_cfg,
+ s_stage2,
+ s_stage1,
+ s_churn,
+ s_noise,
+ a_prompt,
+ n_prompt,
+ color_fix_type,
+ linear_CFG,
+ linear_s_stage2,
+ spt_linear_CFG,
+ spt_linear_s_stage2,
+ model_select
+ ])
+
+ reset_btn.click(fn = reset, inputs = [], outputs = [
+ input_image,
+ rotation,
+ denoise_image,
+ prompt,
+ a_prompt,
+ n_prompt,
+ num_samples,
+ min_size,
+ downscale,
+ upscale,
+ edm_steps,
+ s_stage1,
+ s_stage2,
+ s_cfg,
+ randomize_seed,
+ seed,
+ s_churn,
+ s_noise,
+ color_fix_type,
+ diff_dtype,
+ ae_dtype,
+ gamma_correction,
+ linear_CFG,
+ linear_s_stage2,
+ spt_linear_CFG,
+ spt_linear_s_stage2,
+ model_select,
+ output_format,
+ allocation
+ ], queue = False, show_progress = False)
+
+ interface.queue(10).launch()

configs/clip_vit_config.json

@@ -0,0 +1,25 @@
+{
+ "_name_or_path": "openai/clip-vit-large-patch14",
+ "architectures": [
+ "CLIPTextModel"
+ ],
+ "attention_dropout": 0.0,
+ "bos_token_id": 0,
+ "dropout": 0.0,
+ "eos_token_id": 2,
+ "hidden_act": "quick_gelu",
+ "hidden_size": 768,
+ "initializer_factor": 1.0,
+ "initializer_range": 0.02,
+ "intermediate_size": 3072,
+ "layer_norm_eps": 1e-05,
+ "max_position_embeddings": 77,
+ "model_type": "clip_text_model",
+ "num_attention_heads": 12,
+ "num_hidden_layers": 12,
+ "pad_token_id": 1,
+ "projection_dim": 768,
+ "torch_dtype": "float32",
+ "transformers_version": "4.22.0.dev0",
+ "vocab_size": 49408
+}

configs/tokenizer/config.json

@@ -0,0 +1,171 @@
+{
+ "_name_or_path": "clip-vit-large-patch14/",
+ "architectures": [
+ "CLIPModel"
+ ],
+ "initializer_factor": 1.0,
+ "logit_scale_init_value": 2.6592,
+ "model_type": "clip",
+ "projection_dim": 768,
+ "text_config": {
+ "_name_or_path": "",
+ "add_cross_attention": false,
+ "architectures": null,
+ "attention_dropout": 0.0,
+ "bad_words_ids": null,
+ "bos_token_id": 0,
+ "chunk_size_feed_forward": 0,
+ "cross_attention_hidden_size": null,
+ "decoder_start_token_id": null,
+ "diversity_penalty": 0.0,
+ "do_sample": false,
+ "dropout": 0.0,
+ "early_stopping": false,
+ "encoder_no_repeat_ngram_size": 0,
+ "eos_token_id": 2,
+ "finetuning_task": null,
+ "forced_bos_token_id": null,
+ "forced_eos_token_id": null,
+ "hidden_act": "quick_gelu",
+ "hidden_size": 768,
+ "id2label": {
+ "0": "LABEL_0",
+ "1": "LABEL_1"
+ },
+ "initializer_factor": 1.0,
+ "initializer_range": 0.02,
+ "intermediate_size": 3072,
+ "is_decoder": false,
+ "is_encoder_decoder": false,
+ "label2id": {
+ "LABEL_0": 0,
+ "LABEL_1": 1
+ },
+ "layer_norm_eps": 1e-05,
+ "length_penalty": 1.0,
+ "max_length": 20,
+ "max_position_embeddings": 77,
+ "min_length": 0,
+ "model_type": "clip_text_model",
+ "no_repeat_ngram_size": 0,
+ "num_attention_heads": 12,
+ "num_beam_groups": 1,
+ "num_beams": 1,
+ "num_hidden_layers": 12,
+ "num_return_sequences": 1,
+ "output_attentions": false,
+ "output_hidden_states": false,
+ "output_scores": false,
+ "pad_token_id": 1,
+ "prefix": null,
+ "problem_type": null,
+ "projection_dim": 768,
+ "pruned_heads": {},
+ "remove_invalid_values": false,
+ "repetition_penalty": 1.0,
+ "return_dict": true,
+ "return_dict_in_generate": false,
+ "sep_token_id": null,
+ "task_specific_params": null,
+ "temperature": 1.0,
+ "tie_encoder_decoder": false,
+ "tie_word_embeddings": true,
+ "tokenizer_class": null,
+ "top_k": 50,
+ "top_p": 1.0,
+ "torch_dtype": null,
+ "torchscript": false,
+ "transformers_version": "4.16.0.dev0",
+ "use_bfloat16": false,
+ "vocab_size": 49408
+ },
+ "text_config_dict": {
+ "hidden_size": 768,
+ "intermediate_size": 3072,
+ "num_attention_heads": 12,
+ "num_hidden_layers": 12,
+ "projection_dim": 768
+ },
+ "torch_dtype": "float32",
+ "transformers_version": null,
+ "vision_config": {
+ "_name_or_path": "",
+ "add_cross_attention": false,
+ "architectures": null,
+ "attention_dropout": 0.0,
+ "bad_words_ids": null,
+ "bos_token_id": null,
+ "chunk_size_feed_forward": 0,
+ "cross_attention_hidden_size": null,
+ "decoder_start_token_id": null,
+ "diversity_penalty": 0.0,
+ "do_sample": false,
+ "dropout": 0.0,
+ "early_stopping": false,
+ "encoder_no_repeat_ngram_size": 0,
+ "eos_token_id": null,
+ "finetuning_task": null,
+ "forced_bos_token_id": null,
+ "forced_eos_token_id": null,
+ "hidden_act": "quick_gelu",
+ "hidden_size": 1024,
+ "id2label": {
+ "0": "LABEL_0",
+ "1": "LABEL_1"
+ },
+ "image_size": 224,
+ "initializer_factor": 1.0,
+ "initializer_range": 0.02,
+ "intermediate_size": 4096,
+ "is_decoder": false,
+ "is_encoder_decoder": false,
+ "label2id": {
+ "LABEL_0": 0,
+ "LABEL_1": 1
+ },
+ "layer_norm_eps": 1e-05,
+ "length_penalty": 1.0,
+ "max_length": 20,
+ "min_length": 0,
+ "model_type": "clip_vision_model",
+ "no_repeat_ngram_size": 0,
+ "num_attention_heads": 16,
+ "num_beam_groups": 1,
+ "num_beams": 1,
+ "num_hidden_layers": 24,
+ "num_return_sequences": 1,
+ "output_attentions": false,
+ "output_hidden_states": false,
+ "output_scores": false,
+ "pad_token_id": null,
+ "patch_size": 14,
+ "prefix": null,
+ "problem_type": null,
+ "projection_dim": 768,
+ "pruned_heads": {},
+ "remove_invalid_values": false,
+ "repetition_penalty": 1.0,
+ "return_dict": true,
+ "return_dict_in_generate": false,
+ "sep_token_id": null,
+ "task_specific_params": null,
+ "temperature": 1.0,
+ "tie_encoder_decoder": false,
+ "tie_word_embeddings": true,
+ "tokenizer_class": null,
+ "top_k": 50,
+ "top_p": 1.0,
+ "torch_dtype": null,
+ "torchscript": false,
+ "transformers_version": "4.16.0.dev0",
+ "use_bfloat16": false
+ },
+ "vision_config_dict": {
+ "hidden_size": 1024,
+ "intermediate_size": 4096,
+ "num_attention_heads": 16,
+ "num_hidden_layers": 24,
+ "patch_size": 14,
+ "projection_dim": 768
+ }
+}

configs/tokenizer/preprocessor_config.json

@@ -0,0 +1,19 @@
+{
+ "crop_size": 224,
+ "do_center_crop": true,
+ "do_normalize": true,
+ "do_resize": true,
+ "feature_extractor_type": "CLIPFeatureExtractor",
+ "image_mean": [
+ 0.48145466,
+ 0.4578275,
+ 0.40821073
+ ],
+ "image_std": [
+ 0.26862954,
+ 0.26130258,
+ 0.27577711
+ ],
+ "resample": 3,
+ "size": 224
+}

configs/tokenizer/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+ "bos_token": {
+ "content": "<|startoftext|>",
+ "single_word": false,
+ "lstrip": false,
+ "rstrip": false,
+ "normalized": true
+ },
+ "eos_token": {
+ "content": "<|endoftext|>",
+ "single_word": false,
+ "lstrip": false,
+ "rstrip": false,
+ "normalized": true
+ },
+ "unk_token": {
+ "content": "<|endoftext|>",
+ "single_word": false,
+ "lstrip": false,
+ "rstrip": false,
+ "normalized": true
+ },
+ "pad_token": "<|endoftext|>"
+}

configs/tokenizer/tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+ "unk_token": {
+ "content": "<|endoftext|>",
+ "single_word": false,
+ "lstrip": false,
+ "rstrip": false,
+ "normalized": true,
+ "__type": "AddedToken"
+ },
+ "bos_token": {
+ "content": "<|startoftext|>",
+ "single_word": false,
+ "lstrip": false,
+ "rstrip": false,
+ "normalized": true,
+ "__type": "AddedToken"
+ },
+ "eos_token": {
+ "content": "<|endoftext|>",
+ "single_word": false,
+ "lstrip": false,
+ "rstrip": false,
+ "normalized": true,
+ "__type": "AddedToken"
+ },
+ "pad_token": "<|endoftext|>",
+ "add_prefix_space": false,
+ "errors": "replace",
+ "do_lower_case": true,
+ "name_or_path": "openai/clip-vit-base-patch32",
+ "model_max_length": 77,
+ "special_tokens_map_file": "./special_tokens_map.json",
+ "tokenizer_class": "CLIPTokenizer"
+}

options/SUPIR_v0.yaml

@@ -0,0 +1,140 @@
+model:
+ target: SUPIR.models.SUPIR_model.SUPIRModel
+ params:
+ ae_dtype: bf16
+ diffusion_dtype: fp16
+ scale_factor: 0.13025
+ disable_first_stage_autocast: True
+ network_wrapper: sgm.modules.diffusionmodules.wrappers.ControlWrapper
+
+ denoiser_config:
+ target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiserWithControl
+ params:
+ num_idx: 1000
+ weighting_config:
+ target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+ scaling_config:
+ target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
+ discretization_config:
+ target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+
+ control_stage_config:
+ target: SUPIR.modules.SUPIR_v0.GLVControl
+ params:
+ adm_in_channels: 2816
+ num_classes: sequential
+ use_checkpoint: True
+ in_channels: 4
+ out_channels: 4
+ model_channels: 320
+ attention_resolutions: [4, 2]
+ num_res_blocks: 2
+ channel_mult: [1, 2, 4]
+ num_head_channels: 64
+ use_spatial_transformer: True
+ use_linear_in_transformer: True
+ transformer_depth: [1, 2, 10] # note: the first is unused (due to attn_res starting at 2) 32, 16, 8 --> 64, 32, 16
+ context_dim: 2048
+ spatial_transformer_attn_type: softmax-xformers
+ legacy: False
+ input_upscale: 1
+
+ network_config:
+ target: SUPIR.modules.SUPIR_v0.LightGLVUNet
+ params:
+ mode: XL-base
+ project_type: ZeroSFT
+ project_channel_scale: 2
+ adm_in_channels: 2816
+ num_classes: sequential
+ use_checkpoint: True
+ in_channels: 4
+ out_channels: 4
+ model_channels: 320
+ attention_resolutions: [4, 2]
+ num_res_blocks: 2
+ channel_mult: [1, 2, 4]
+ num_head_channels: 64
+ use_spatial_transformer: True
+ use_linear_in_transformer: True
+ transformer_depth: [1, 2, 10] # note: the first is unused (due to attn_res starting at 2) 32, 16, 8 --> 64, 32, 16
+ context_dim: 2048
+ spatial_transformer_attn_type: softmax-xformers
+ legacy: False
+
+ conditioner_config:
+ target: sgm.modules.GeneralConditionerWithControl
+ params:
+ emb_models:
+ # crossattn cond
+ - is_trainable: False
+ input_key: txt
+ target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
+ params:
+ layer: hidden
+ layer_idx: 11
+ # crossattn and vector cond
+ - is_trainable: False
+ input_key: txt
+ target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder2
+ params:
+ arch: ViT-bigG-14
+ version: laion2b_s39b_b160k
+ freeze: True
+ layer: penultimate
+ always_return_pooled: True
+ legacy: False
+ # vector cond
+ - is_trainable: False
+ input_key: original_size_as_tuple
+ target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+ params:
+ outdim: 256 # multiplied by two
+ # vector cond
+ - is_trainable: False
+ input_key: crop_coords_top_left
+ target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+ params:
+ outdim: 256 # multiplied by two
+ # vector cond
+ - is_trainable: False
+ input_key: target_size_as_tuple
+ target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+ params:
+ outdim: 256 # multiplied by two
+
+ first_stage_config:
+ target: sgm.models.autoencoder.AutoencoderKLInferenceWrapper
+ params:
+ ckpt_path: ~
+ embed_dim: 4
+ monitor: val/rec_loss
+ ddconfig:
+ attn_type: vanilla-xformers
+ double_z: true
+ z_channels: 4
+ resolution: 256
+ in_channels: 3
+ out_ch: 3
+ ch: 128
+ ch_mult: [1, 2, 4, 4]
+ num_res_blocks: 2
+ attn_resolutions: []
+ dropout: 0.0
+ lossconfig:
+ target: torch.nn.Identity
+
+ sampler_config:
+ target: sgm.modules.diffusionmodules.sampling.RestoreEDMSampler
+ params:
+ num_steps: 100
+ restore_cfg: 4.0
+ s_churn: 0
+ s_noise: 1.003
+ discretization_config:
+ target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+ guider_config:
+ target: sgm.modules.diffusionmodules.guiders.LinearCFG
+ params:
+ scale: 7.5
+ scale_min: 4.0

sgm/__init__.py

@@ -0,0 +1,4 @@
+from .models import AutoencodingEngine, DiffusionEngine
+from .util import get_configs_path, instantiate_from_config
+
+__version__ = "0.1.0"

sgm/lr_scheduler.py

@@ -0,0 +1,135 @@
+import numpy as np
+
+
+class LambdaWarmUpCosineScheduler:
+ """
+ note: use with a base_lr of 1.0
+ """
+
+ def __init__(
+ self,
+ warm_up_steps,
+ lr_min,
+ lr_max,
+ lr_start,
+ max_decay_steps,
+ verbosity_interval=0,
+ ):
+ self.lr_warm_up_steps = warm_up_steps
+ self.lr_start = lr_start
+ self.lr_min = lr_min
+ self.lr_max = lr_max
+ self.lr_max_decay_steps = max_decay_steps
+ self.last_lr = 0.0
+ self.verbosity_interval = verbosity_interval
+
+ def schedule(self, n, **kwargs):
+ if self.verbosity_interval > 0:
+ if n % self.verbosity_interval == 0:
+ print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
+ if n < self.lr_warm_up_steps:
+ lr = (
+ self.lr_max - self.lr_start
+ ) / self.lr_warm_up_steps * n + self.lr_start
+ self.last_lr = lr
+ return lr
+ else:
+ t = (n - self.lr_warm_up_steps) / (
+ self.lr_max_decay_steps - self.lr_warm_up_steps
+ )
+ t = min(t, 1.0)
+ lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
+ 1 + np.cos(t * np.pi)
+ )
+ self.last_lr = lr
+ return lr
+
+ def __call__(self, n, **kwargs):
+ return self.schedule(n, **kwargs)
+
+
+class LambdaWarmUpCosineScheduler2:
+ """
+ supports repeated iterations, configurable via lists
+ note: use with a base_lr of 1.0.
+ """
+
+ def __init__(
+ self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0
+ ):
+ assert (
+ len(warm_up_steps)
+ == len(f_min)
+ == len(f_max)
+ == len(f_start)
+ == len(cycle_lengths)
+ )
+ self.lr_warm_up_steps = warm_up_steps
+ self.f_start = f_start
+ self.f_min = f_min
+ self.f_max = f_max
+ self.cycle_lengths = cycle_lengths
+ self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
+ self.last_f = 0.0
+ self.verbosity_interval = verbosity_interval
+
+ def find_in_interval(self, n):
+ interval = 0
+ for cl in self.cum_cycles[1:]:
+ if n <= cl:
+ return interval
+ interval += 1
+
+ def schedule(self, n, **kwargs):
+ cycle = self.find_in_interval(n)
+ n = n - self.cum_cycles[cycle]
+ if self.verbosity_interval > 0:
+ if n % self.verbosity_interval == 0:
+ print(
+ f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+ f"current cycle {cycle}"
+ )
+ if n < self.lr_warm_up_steps[cycle]:
+ f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[
+ cycle
+ ] * n + self.f_start[cycle]
+ self.last_f = f
+ return f
+ else:
+ t = (n - self.lr_warm_up_steps[cycle]) / (
+ self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle]
+ )
+ t = min(t, 1.0)
+ f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (
+ 1 + np.cos(t * np.pi)
+ )
+ self.last_f = f
+ return f
+
+ def __call__(self, n, **kwargs):
+ return self.schedule(n, **kwargs)
+
+
+class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
+ def schedule(self, n, **kwargs):
+ cycle = self.find_in_interval(n)
+ n = n - self.cum_cycles[cycle]
+ if self.verbosity_interval > 0:
+ if n % self.verbosity_interval == 0:
+ print(
+ f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+ f"current cycle {cycle}"
+ )
+
+ if n < self.lr_warm_up_steps[cycle]:
+ f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[
+ cycle
+ ] * n + self.f_start[cycle]
+ self.last_f = f
+ return f
+ else:
+ f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (
+ self.cycle_lengths[cycle] - n
+ ) / (self.cycle_lengths[cycle])
+ self.last_f = f
+ return f

sgm/models/__init__.py

@@ -0,0 +1,2 @@
+from .autoencoder import AutoencodingEngine
+from .diffusion import DiffusionEngine

sgm/models/autoencoder.py

@@ -0,0 +1,335 @@
+import re
+from abc import abstractmethod
+from contextlib import contextmanager
+from typing import Any, Dict, Tuple, Union
+
+import pytorch_lightning as pl
+import torch
+from omegaconf import ListConfig
+from packaging import version
+from safetensors.torch import load_file as load_safetensors
+
+from ..modules.diffusionmodules.model import Decoder, Encoder
+from ..modules.distributions.distributions import DiagonalGaussianDistribution
+from ..modules.ema import LitEma
+from ..util import default, get_obj_from_str, instantiate_from_config
+
+
+class AbstractAutoencoder(pl.LightningModule):
+ """
+ This is the base class for all autoencoders, including image autoencoders, image autoencoders with discriminators,
+ unCLIP models, etc. Hence, it is fairly general, and specific features
+ (e.g. discriminator training, encoding, decoding) must be implemented in subclasses.
+ """
+
+ def __init__(
+ self,
+ ema_decay: Union[None, float] = None,
+ monitor: Union[None, str] = None,
+ input_key: str = "jpg",
+ ckpt_path: Union[None, str] = None,
+ ignore_keys: Union[Tuple, list, ListConfig] = (),
+ ):
+ super().__init__()
+ self.input_key = input_key
+ self.use_ema = ema_decay is not None
+ if monitor is not None:
+ self.monitor = monitor
+
+ if self.use_ema:
+ self.model_ema = LitEma(self, decay=ema_decay)
+ print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+ if ckpt_path is not None:
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+
+ if version.parse(torch.__version__) >= version.parse("2.0.0"):
+ self.automatic_optimization = False
+
+ def init_from_ckpt(
+ self, path: str, ignore_keys: Union[Tuple, list, ListConfig] = tuple()
+ ) -> None:
+ if path.endswith("ckpt"):
+ sd = torch.load(path, map_location="cpu")["state_dict"]
+ elif path.endswith("safetensors"):
+ sd = load_safetensors(path)
+ else:
+ raise NotImplementedError
+
+ keys = list(sd.keys())
+ for k in keys:
+ for ik in ignore_keys:
+ if re.match(ik, k):
+ print("Deleting key {} from state_dict.".format(k))
+ del sd[k]
+ missing, unexpected = self.load_state_dict(sd, strict=False)
+ print(
+ f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys"
+ )
+ if len(missing) > 0:
+ print(f"Missing Keys: {missing}")
+ if len(unexpected) > 0:
+ print(f"Unexpected Keys: {unexpected}")
+
+ @abstractmethod
+ def get_input(self, batch) -> Any:
+ raise NotImplementedError()
+
+ def on_train_batch_end(self, *args, **kwargs):
+ # for EMA computation
+ if self.use_ema:
+ self.model_ema(self)
+
+ @contextmanager
+ def ema_scope(self, context=None):
+ if self.use_ema:
+ self.model_ema.store(self.parameters())
+ self.model_ema.copy_to(self)
+ if context is not None:
+ print(f"{context}: Switched to EMA weights")
+ try:
+ yield None
+ finally:
+ if self.use_ema:
+ self.model_ema.restore(self.parameters())
+ if context is not None:
+ print(f"{context}: Restored training weights")
+
+ @abstractmethod
+ def encode(self, *args, **kwargs) -> torch.Tensor:
+ raise NotImplementedError("encode()-method of abstract base class called")
+
+ @abstractmethod
+ def decode(self, *args, **kwargs) -> torch.Tensor:
+ raise NotImplementedError("decode()-method of abstract base class called")
+
+ def instantiate_optimizer_from_config(self, params, lr, cfg):
+ print(f"loading >>> {cfg['target']} <<< optimizer from config")
+ return get_obj_from_str(cfg["target"])(
+ params, lr=lr, **cfg.get("params", dict())
+ )
+
+ def configure_optimizers(self) -> Any:
+ raise NotImplementedError()
+
+
+class AutoencodingEngine(AbstractAutoencoder):
+ """
+ Base class for all image autoencoders that we train, like VQGAN or AutoencoderKL
+ (we also restore them explicitly as special cases for legacy reasons).
+ Regularizations such as KL or VQ are moved to the regularizer class.
+ """
+
+ def __init__(
+ self,
+ *args,
+ encoder_config: Dict,
+ decoder_config: Dict,
+ loss_config: Dict,
+ regularizer_config: Dict,
+ optimizer_config: Union[Dict, None] = None,
+ lr_g_factor: float = 1.0,
+ **kwargs,
+ ):
+ super().__init__(*args, **kwargs)
+ # todo: add options to freeze encoder/decoder
+ self.encoder = instantiate_from_config(encoder_config)
+ self.decoder = instantiate_from_config(decoder_config)
+ self.loss = instantiate_from_config(loss_config)
+ self.regularization = instantiate_from_config(regularizer_config)
+ self.optimizer_config = default(
+ optimizer_config, {"target": "torch.optim.Adam"}
+ )
+ self.lr_g_factor = lr_g_factor
+
+ def get_input(self, batch: Dict) -> torch.Tensor:
+ # assuming unified data format, dataloader returns a dict.
+ # image tensors should be scaled to -1 ... 1 and in channels-first format (e.g., bchw instead if bhwc)
+ return batch[self.input_key]
+
+ def get_autoencoder_params(self) -> list:
+ params = (
+ list(self.encoder.parameters())
+ + list(self.decoder.parameters())
+ + list(self.regularization.get_trainable_parameters())
+ + list(self.loss.get_trainable_autoencoder_parameters())
+ )
+ return params
+
+ def get_discriminator_params(self) -> list:
+ params = list(self.loss.get_trainable_parameters()) # e.g., discriminator
+ return params
+
+ def get_last_layer(self):
+ return self.decoder.get_last_layer()
+
+ def encode(self, x: Any, return_reg_log: bool = False) -> Any:
+ z = self.encoder(x)
+ z, reg_log = self.regularization(z)
+ if return_reg_log:
+ return z, reg_log
+ return z
+
+ def decode(self, z: Any) -> torch.Tensor:
+ x = self.decoder(z)
+ return x
+
+ def forward(self, x: Any) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ z, reg_log = self.encode(x, return_reg_log=True)
+ dec = self.decode(z)
+ return z, dec, reg_log
+
+ def training_step(self, batch, batch_idx, optimizer_idx) -> Any:
+ x = self.get_input(batch)
+ z, xrec, regularization_log = self(x)
+
+ if optimizer_idx == 0:
+ # autoencode
+ aeloss, log_dict_ae = self.loss(
+ regularization_log,
+ x,
+ xrec,
+ optimizer_idx,
+ self.global_step,
+ last_layer=self.get_last_layer(),
+ split="train",
+ )
+
+ self.log_dict(
+ log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True
+ )
+ return aeloss
+
+ if optimizer_idx == 1:
+ # discriminator
+ discloss, log_dict_disc = self.loss(
+ regularization_log,
+ x,
+ xrec,
+ optimizer_idx,
+ self.global_step,
+ last_layer=self.get_last_layer(),
+ split="train",
+ )
+ self.log_dict(
+ log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True
+ )
+ return discloss
+
+ def validation_step(self, batch, batch_idx) -> Dict:
+ log_dict = self._validation_step(batch, batch_idx)
+ with self.ema_scope():
+ log_dict_ema = self._validation_step(batch, batch_idx, postfix="_ema")
+ log_dict.update(log_dict_ema)
+ return log_dict
+
+ def _validation_step(self, batch, batch_idx, postfix="") -> Dict:
+ x = self.get_input(batch)
+
+ z, xrec, regularization_log = self(x)
+ aeloss, log_dict_ae = self.loss(
+ regularization_log,
+ x,
+ xrec,
+ 0,
+ self.global_step,
+ last_layer=self.get_last_layer(),
+ split="val" + postfix,
+ )
+
+ discloss, log_dict_disc = self.loss(
+ regularization_log,
+ x,
+ xrec,
+ 1,
+ self.global_step,
+ last_layer=self.get_last_layer(),
+ split="val" + postfix,
+ )
+ self.log(f"val{postfix}/rec_loss", log_dict_ae[f"val{postfix}/rec_loss"])
+ log_dict_ae.update(log_dict_disc)
+ self.log_dict(log_dict_ae)
+ return log_dict_ae
+
+ def configure_optimizers(self) -> Any:
+ ae_params = self.get_autoencoder_params()
+ disc_params = self.get_discriminator_params()
+
+ opt_ae = self.instantiate_optimizer_from_config(
+ ae_params,
+ default(self.lr_g_factor, 1.0) * self.learning_rate,
+ self.optimizer_config,
+ )
+ opt_disc = self.instantiate_optimizer_from_config(
+ disc_params, self.learning_rate, self.optimizer_config
+ )
+
+ return [opt_ae, opt_disc], []
+
+ @torch.no_grad()
+ def log_images(self, batch: Dict, **kwargs) -> Dict:
+ log = dict()
+ x = self.get_input(batch)
+ _, xrec, _ = self(x)
+ log["inputs"] = x
+ log["reconstructions"] = xrec
+ with self.ema_scope():
+ _, xrec_ema, _ = self(x)
+ log["reconstructions_ema"] = xrec_ema
+ return log
+
+
+class AutoencoderKL(AutoencodingEngine):
+ def __init__(self, embed_dim: int, **kwargs):
+ ddconfig = kwargs.pop("ddconfig")
+ ckpt_path = kwargs.pop("ckpt_path", None)
+ ignore_keys = kwargs.pop("ignore_keys", ())
+ super().__init__(
+ encoder_config={"target": "torch.nn.Identity"},
+ decoder_config={"target": "torch.nn.Identity"},
+ regularizer_config={"target": "torch.nn.Identity"},
+ loss_config=kwargs.pop("lossconfig"),
+ **kwargs,
+ )
+ assert ddconfig["double_z"]
+ self.encoder = Encoder(**ddconfig)
+ self.decoder = Decoder(**ddconfig)
+ self.quant_conv = torch.nn.Conv2d(2 * ddconfig["z_channels"], 2 * embed_dim, 1)
+ self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+ self.embed_dim = embed_dim
+
+ if ckpt_path is not None:
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+
+ def encode(self, x):
+ assert (
+ not self.training
+ ), f"{self.__class__.__name__} only supports inference currently"
+ h = self.encoder(x)
+ moments = self.quant_conv(h)
+ posterior = DiagonalGaussianDistribution(moments)
+ return posterior
+
+ def decode(self, z, **decoder_kwargs):
+ z = self.post_quant_conv(z)
+ dec = self.decoder(z, **decoder_kwargs)
+ return dec
+
+
+class AutoencoderKLInferenceWrapper(AutoencoderKL):
+ def encode(self, x):
+ return super().encode(x).sample()
+
+
+class IdentityFirstStage(AbstractAutoencoder):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ def get_input(self, x: Any) -> Any:
+ return x
+
+ def encode(self, x: Any, *args, **kwargs) -> Any:
+ return x
+
+ def decode(self, x: Any, *args, **kwargs) -> Any:
+ return x

sgm/models/diffusion.py

@@ -0,0 +1,320 @@
+from contextlib import contextmanager
+from typing import Any, Dict, List, Tuple, Union
+
+import pytorch_lightning as pl
+import torch
+from omegaconf import ListConfig, OmegaConf
+from safetensors.torch import load_file as load_safetensors
+from torch.optim.lr_scheduler import LambdaLR
+
+from ..modules import UNCONDITIONAL_CONFIG
+from ..modules.diffusionmodules.wrappers import OPENAIUNETWRAPPER
+from ..modules.ema import LitEma
+from ..util import (
+ default,
+ disabled_train,
+ get_obj_from_str,
+ instantiate_from_config,
+ log_txt_as_img,
+)
+
+
+class DiffusionEngine(pl.LightningModule):
+ def __init__(
+ self,
+ network_config,
+ denoiser_config,
+ first_stage_config,
+ conditioner_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+ sampler_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+ optimizer_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+ scheduler_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+ loss_fn_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+ network_wrapper: Union[None, str] = None,
+ ckpt_path: Union[None, str] = None,
+ use_ema: bool = False,
+ ema_decay_rate: float = 0.9999,
+ scale_factor: float = 1.0,
+ disable_first_stage_autocast=False,
+ input_key: str = "jpg",
+ log_keys: Union[List, None] = None,
+ no_cond_log: bool = False,
+ compile_model: bool = False,
+ ):
+ super().__init__()
+ self.log_keys = log_keys
+ self.input_key = input_key
+ self.optimizer_config = default(
+ optimizer_config, {"target": "torch.optim.AdamW"}
+ )
+ model = instantiate_from_config(network_config)
+ self.model = get_obj_from_str(default(network_wrapper, OPENAIUNETWRAPPER))(
+ model, compile_model=compile_model
+ )
+
+ self.denoiser = instantiate_from_config(denoiser_config)
+ self.sampler = (
+ instantiate_from_config(sampler_config)
+ if sampler_config is not None
+ else None
+ )
+ self.conditioner = instantiate_from_config(
+ default(conditioner_config, UNCONDITIONAL_CONFIG)
+ )
+ self.scheduler_config = scheduler_config
+ self._init_first_stage(first_stage_config)
+
+ self.loss_fn = (
+ instantiate_from_config(loss_fn_config)
+ if loss_fn_config is not None
+ else None
+ )
+
+ self.use_ema = use_ema
+ if self.use_ema:
+ self.model_ema = LitEma(self.model, decay=ema_decay_rate)
+ print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+ self.scale_factor = scale_factor
+ self.disable_first_stage_autocast = disable_first_stage_autocast
+ self.no_cond_log = no_cond_log
+
+ if ckpt_path is not None:
+ self.init_from_ckpt(ckpt_path)
+
+ def init_from_ckpt(
+ self,
+ path: str,
+ ) -> None:
+ if path.endswith("ckpt"):
+ sd = torch.load(path, map_location="cpu")["state_dict"]
+ elif path.endswith("safetensors"):
+ sd = load_safetensors(path)
+ else:
+ raise NotImplementedError
+
+ missing, unexpected = self.load_state_dict(sd, strict=False)
+ print(
+ f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys"
+ )
+ if len(missing) > 0:
+ print(f"Missing Keys: {missing}")
+ if len(unexpected) > 0:
+ print(f"Unexpected Keys: {unexpected}")
+
+ def _init_first_stage(self, config):
+ model = instantiate_from_config(config).eval()
+ model.train = disabled_train
+ for param in model.parameters():
+ param.requires_grad = False
+ self.first_stage_model = model
+
+ def get_input(self, batch):
+ # assuming unified data format, dataloader returns a dict.
+ # image tensors should be scaled to -1 ... 1 and in bchw format
+ return batch[self.input_key]
+
+ @torch.no_grad()
+ def decode_first_stage(self, z):
+ z = 1.0 / self.scale_factor * z
+ with torch.autocast("cuda", enabled=not self.disable_first_stage_autocast):
+ out = self.first_stage_model.decode(z)
+ return out
+
+ @torch.no_grad()
+ def encode_first_stage(self, x):
+ with torch.autocast("cuda", enabled=not self.disable_first_stage_autocast):
+ z = self.first_stage_model.encode(x)
+ z = self.scale_factor * z
+ return z
+
+ def forward(self, x, batch):
+ loss = self.loss_fn(self.model, self.denoiser, self.conditioner, x, batch)
+ loss_mean = loss.mean()
+ loss_dict = {"loss": loss_mean}
+ return loss_mean, loss_dict
+
+ def shared_step(self, batch: Dict) -> Any:
+ x = self.get_input(batch)
+ x = self.encode_first_stage(x)
+ batch["global_step"] = self.global_step
+ loss, loss_dict = self(x, batch)
+ return loss, loss_dict
+
+ def training_step(self, batch, batch_idx):
+ loss, loss_dict = self.shared_step(batch)
+
+ self.log_dict(
+ loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=False
+ )
+
+ self.log(
+ "global_step",
+ self.global_step,
+ prog_bar=True,
+ logger=True,
+ on_step=True,
+ on_epoch=False,
+ )
+
+ # if self.scheduler_config is not None:
+ lr = self.optimizers().param_groups[0]["lr"]
+ self.log(
+ "lr_abs", lr, prog_bar=True, logger=True, on_step=True, on_epoch=False
+ )
+
+ return loss
+
+ def on_train_start(self, *args, **kwargs):
+ if self.sampler is None or self.loss_fn is None:
+ raise ValueError("Sampler and loss function need to be set for training.")
+
+ def on_train_batch_end(self, *args, **kwargs):
+ if self.use_ema:
+ self.model_ema(self.model)
+
+ @contextmanager
+ def ema_scope(self, context=None):
+ if self.use_ema:
+ self.model_ema.store(self.model.parameters())
+ self.model_ema.copy_to(self.model)
+ if context is not None:
+ print(f"{context}: Switched to EMA weights")
+ try:
+ yield None
+ finally:
+ if self.use_ema:
+ self.model_ema.restore(self.model.parameters())
+ if context is not None:
+ print(f"{context}: Restored training weights")
+
+ def instantiate_optimizer_from_config(self, params, lr, cfg):
+ return get_obj_from_str(cfg["target"])(
+ params, lr=lr, **cfg.get("params", dict())
+ )
+
+ def configure_optimizers(self):
+ lr = self.learning_rate
+ params = list(self.model.parameters())
+ for embedder in self.conditioner.embedders:
+ if embedder.is_trainable:
+ params = params + list(embedder.parameters())
+ opt = self.instantiate_optimizer_from_config(params, lr, self.optimizer_config)
+ if self.scheduler_config is not None:
+ scheduler = instantiate_from_config(self.scheduler_config)
+ print("Setting up LambdaLR scheduler...")
+ scheduler = [
+ {
+ "scheduler": LambdaLR(opt, lr_lambda=scheduler.schedule),
+ "interval": "step",
+ "frequency": 1,
+ }
+ ]
+ return [opt], scheduler
+ return opt
+
+ @torch.no_grad()
+ def sample(
+ self,
+ cond: Dict,
+ uc: Union[Dict, None] = None,
+ batch_size: int = 16,
+ shape: Union[None, Tuple, List] = None,
+ **kwargs,
+ ):
+ randn = torch.randn(batch_size, *shape).to(self.device)
+
+ denoiser = lambda input, sigma, c: self.denoiser(
+ self.model, input, sigma, c, **kwargs
+ )
+ samples = self.sampler(denoiser, randn, cond, uc=uc)
+ return samples
+
+ @torch.no_grad()
+ def log_conditionings(self, batch: Dict, n: int) -> Dict:
+ """
+ Defines heuristics to log different conditionings.
+ These can be lists of strings (text-to-image), tensors, ints, ...
+ """
+ image_h, image_w = batch[self.input_key].shape[2:]
+ log = dict()
+
+ for embedder in self.conditioner.embedders:
+ if (
+ (self.log_keys is None) or (embedder.input_key in self.log_keys)
+ ) and not self.no_cond_log:
+ x = batch[embedder.input_key][:n]
+ if isinstance(x, torch.Tensor):
+ if x.dim() == 1:
+ # class-conditional, convert integer to string
+ x = [str(x[i].item()) for i in range(x.shape[0])]
+ xc = log_txt_as_img((image_h, image_w), x, size=image_h // 4)
+ elif x.dim() == 2:
+ # size and crop cond and the like
+ x = [
+ "x".join([str(xx) for xx in x[i].tolist()])
+ for i in range(x.shape[0])
+ ]
+ xc = log_txt_as_img((image_h, image_w), x, size=image_h // 20)
+ else:
+ raise NotImplementedError()
+ elif isinstance(x, (List, ListConfig)):
+ if isinstance(x[0], str):
+ # strings
+ xc = log_txt_as_img((image_h, image_w), x, size=image_h // 20)
+ else:
+ raise NotImplementedError()
+ else:
+ raise NotImplementedError()
+ log[embedder.input_key] = xc
+ return log
+
+ @torch.no_grad()
+ def log_images(
+ self,
+ batch: Dict,
+ N: int = 8,
+ sample: bool = True,
+ ucg_keys: List[str] = None,
+ **kwargs,
+ ) -> Dict:
+ conditioner_input_keys = [e.input_key for e in self.conditioner.embedders]
+ if ucg_keys:
+ assert all(map(lambda x: x in conditioner_input_keys, ucg_keys)), (
+ "Each defined ucg key for sampling must be in the provided conditioner input keys,"
+ f"but we have {ucg_keys} vs. {conditioner_input_keys}"
+ )
+ else:
+ ucg_keys = conditioner_input_keys
+ log = dict()
+
+ x = self.get_input(batch)
+
+ c, uc = self.conditioner.get_unconditional_conditioning(
+ batch,
+ force_uc_zero_embeddings=ucg_keys
+ if len(self.conditioner.embedders) > 0
+ else [],
+ )
+
+ sampling_kwargs = {}
+
+ N = min(x.shape[0], N)
+ x = x.to(self.device)[:N]
+ log["inputs"] = x
+ z = self.encode_first_stage(x)
+ log["reconstructions"] = self.decode_first_stage(z)
+ log.update(self.log_conditionings(batch, N))
+
+ for k in c:
+ if isinstance(c[k], torch.Tensor):
+ c[k], uc[k] = map(lambda y: y[k][:N].to(self.device), (c, uc))
+
+ if sample:
+ with self.ema_scope("Plotting"):
+ samples = self.sample(
+ c, shape=z.shape[1:], uc=uc, batch_size=N, **sampling_kwargs
+ )
+ samples = self.decode_first_stage(samples)
+ log["samples"] = samples
+ return log

sgm/modules/__init__.py

@@ -0,0 +1,8 @@
+from .encoders.modules import GeneralConditioner
+from .encoders.modules import GeneralConditionerWithControl
+from .encoders.modules import PreparedConditioner
+
+UNCONDITIONAL_CONFIG = {
+ "target": "sgm.modules.GeneralConditioner",
+ "params": {"emb_models": []},
+}

sgm/modules/attention.py

@@ -0,0 +1,637 @@
+import math
+from inspect import isfunction
+from typing import Any, Optional
+
+import torch
+import torch.nn.functional as F
+# from einops._torch_specific import allow_ops_in_compiled_graph
+# allow_ops_in_compiled_graph()
+from einops import rearrange, repeat
+from packaging import version
+from torch import nn
+
+if version.parse(torch.__version__) >= version.parse("2.0.0"):
+ SDP_IS_AVAILABLE = True
+ from torch.backends.cuda import SDPBackend, sdp_kernel
+
+ BACKEND_MAP = {
+ SDPBackend.MATH: {
+ "enable_math": True,
+ "enable_flash": False,
+ "enable_mem_efficient": False,
+ },
+ SDPBackend.FLASH_ATTENTION: {
+ "enable_math": False,
+ "enable_flash": True,
+ "enable_mem_efficient": False,
+ },
+ SDPBackend.EFFICIENT_ATTENTION: {
+ "enable_math": False,
+ "enable_flash": False,
+ "enable_mem_efficient": True,
+ },
+ None: {"enable_math": True, "enable_flash": True, "enable_mem_efficient": True},
+ }
+else:
+ from contextlib import nullcontext
+
+ SDP_IS_AVAILABLE = False
+ sdp_kernel = nullcontext
+ BACKEND_MAP = {}
+ print(
+ f"No SDP backend available, likely because you are running in pytorch versions < 2.0. In fact, "
+ f"you are using PyTorch {torch.__version__}. You might want to consider upgrading."
+ )
+
+try:
+ import xformers
+ import xformers.ops
+
+ XFORMERS_IS_AVAILABLE = True
+except:
+ XFORMERS_IS_AVAILABLE = False
+ print("no module 'xformers'. Processing without...")
+
+from .diffusionmodules.util import checkpoint
+
+
+def exists(val):
+ return val is not None
+
+
+def uniq(arr):
+ return {el: True for el in arr}.keys()
+
+
+def default(val, d):
+ if exists(val):
+ return val
+ return d() if isfunction(d) else d
+
+
+def max_neg_value(t):
+ return -torch.finfo(t.dtype).max
+
+
+def init_(tensor):
+ dim = tensor.shape[-1]
+ std = 1 / math.sqrt(dim)
+ tensor.uniform_(-std, std)
+ return tensor
+
+
+# feedforward
+class GEGLU(nn.Module):
+ def __init__(self, dim_in, dim_out):
+ super().__init__()
+ self.proj = nn.Linear(dim_in, dim_out * 2)
+
+ def forward(self, x):
+ x, gate = self.proj(x).chunk(2, dim=-1)
+ return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
+ super().__init__()
+ inner_dim = int(dim * mult)
+ dim_out = default(dim_out, dim)
+ project_in = (
+ nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
+ if not glu
+ else GEGLU(dim, inner_dim)
+ )
+
+ self.net = nn.Sequential(
+ project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)
+ )
+
+ def forward(self, x):
+ return self.net(x)
+
+
+def zero_module(module):
+ """
+ Zero out the parameters of a module and return it.
+ """
+ for p in module.parameters():
+ p.detach().zero_()
+ return module
+
+
+def Normalize(in_channels):
+ return torch.nn.GroupNorm(
+ num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+ )
+
+
+class LinearAttention(nn.Module):
+ def __init__(self, dim, heads=4, dim_head=32):
+ super().__init__()
+ self.heads = heads
+ hidden_dim = dim_head * heads
+ self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+ self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+ def forward(self, x):
+ b, c, h, w = x.shape
+ qkv = self.to_qkv(x)
+ q, k, v = rearrange(
+ qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3
+ )
+ k = k.softmax(dim=-1)
+ context = torch.einsum("bhdn,bhen->bhde", k, v)
+ out = torch.einsum("bhde,bhdn->bhen", context, q)
+ out = rearrange(
+ out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w
+ )
+ return self.to_out(out)
+
+
+class SpatialSelfAttention(nn.Module):
+ def __init__(self, in_channels):
+ super().__init__()
+ self.in_channels = in_channels
+
+ self.norm = Normalize(in_channels)
+ self.q = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+ self.k = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+ self.v = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+ self.proj_out = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+
+ def forward(self, x):
+ h_ = x
+ h_ = self.norm(h_)
+ q = self.q(h_)
+ k = self.k(h_)
+ v = self.v(h_)
+
+ # compute attention
+ b, c, h, w = q.shape
+ q = rearrange(q, "b c h w -> b (h w) c")
+ k = rearrange(k, "b c h w -> b c (h w)")
+ w_ = torch.einsum("bij,bjk->bik", q, k)
+
+ w_ = w_ * (int(c) ** (-0.5))
+ w_ = torch.nn.functional.softmax(w_, dim=2)
+
+ # attend to values
+ v = rearrange(v, "b c h w -> b c (h w)")
+ w_ = rearrange(w_, "b i j -> b j i")
+ h_ = torch.einsum("bij,bjk->bik", v, w_)
+ h_ = rearrange(h_, "b c (h w) -> b c h w", h=h)
+ h_ = self.proj_out(h_)
+
+ return x + h_
+
+
+class CrossAttention(nn.Module):
+ def __init__(
+ self,
+ query_dim,
+ context_dim=None,
+ heads=8,
+ dim_head=64,
+ dropout=0.0,
+ backend=None,
+ ):
+ super().__init__()
+ inner_dim = dim_head * heads
+ context_dim = default(context_dim, query_dim)
+
+ self.scale = dim_head**-0.5
+ self.heads = heads
+
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+ self.to_out = nn.Sequential(
+ nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
+ )
+ self.backend = backend
+
+ def forward(
+ self,
+ x,
+ context=None,
+ mask=None,
+ additional_tokens=None,
+ n_times_crossframe_attn_in_self=0,
+ ):
+ h = self.heads
+
+ if additional_tokens is not None:
+ # get the number of masked tokens at the beginning of the output sequence
+ n_tokens_to_mask = additional_tokens.shape[1]
+ # add additional token
+ x = torch.cat([additional_tokens, x], dim=1)
+
+ q = self.to_q(x)
+ context = default(context, x)
+ k = self.to_k(context)
+ v = self.to_v(context)
+
+ if n_times_crossframe_attn_in_self:
+ # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439
+ assert x.shape[0] % n_times_crossframe_attn_in_self == 0
+ n_cp = x.shape[0] // n_times_crossframe_attn_in_self
+ k = repeat(
+ k[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp
+ )
+ v = repeat(
+ v[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp
+ )
+
+ q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))
+
+ ## old
+ """
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+ del q, k
+
+ if exists(mask):
+ mask = rearrange(mask, 'b ... -> b (...)')
+ max_neg_value = -torch.finfo(sim.dtype).max
+ mask = repeat(mask, 'b j -> (b h) () j', h=h)
+ sim.masked_fill_(~mask, max_neg_value)
+
+ # attention, what we cannot get enough of
+ sim = sim.softmax(dim=-1)
+
+ out = einsum('b i j, b j d -> b i d', sim, v)
+ """
+ ## new
+ with sdp_kernel(**BACKEND_MAP[self.backend]):
+ # print("dispatching into backend", self.backend, "q/k/v shape: ", q.shape, k.shape, v.shape)
+ out = F.scaled_dot_product_attention(
+ q, k, v, attn_mask=mask
+ ) # scale is dim_head ** -0.5 per default
+
+ del q, k, v
+ out = rearrange(out, "b h n d -> b n (h d)", h=h)
+
+ if additional_tokens is not None:
+ # remove additional token
+ out = out[:, n_tokens_to_mask:]
+ return self.to_out(out)
+
+
+class MemoryEfficientCrossAttention(nn.Module):
+ # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
+ def __init__(
+ self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0, **kwargs
+ ):
+ super().__init__()
+ # print(
+ # f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using "
+ # f"{heads} heads with a dimension of {dim_head}."
+ # )
+ inner_dim = dim_head * heads
+ context_dim = default(context_dim, query_dim)
+
+ self.heads = heads
+ self.dim_head = dim_head
+
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+ self.to_out = nn.Sequential(
+ nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
+ )
+ self.attention_op: Optional[Any] = None
+
+ def forward(
+ self,
+ x,
+ context=None,
+ mask=None,
+ additional_tokens=None,
+ n_times_crossframe_attn_in_self=0,
+ ):
+ if additional_tokens is not None:
+ # get the number of masked tokens at the beginning of the output sequence
+ n_tokens_to_mask = additional_tokens.shape[1]
+ # add additional token
+ x = torch.cat([additional_tokens, x], dim=1)
+ q = self.to_q(x)
+ context = default(context, x)
+ k = self.to_k(context)
+ v = self.to_v(context)
+
+ if n_times_crossframe_attn_in_self:
+ # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439
+ assert x.shape[0] % n_times_crossframe_attn_in_self == 0
+ # n_cp = x.shape[0]//n_times_crossframe_attn_in_self
+ k = repeat(
+ k[::n_times_crossframe_attn_in_self],
+ "b ... -> (b n) ...",
+ n=n_times_crossframe_attn_in_self,
+ )
+ v = repeat(
+ v[::n_times_crossframe_attn_in_self],
+ "b ... -> (b n) ...",
+ n=n_times_crossframe_attn_in_self,
+ )
+
+ b, _, _ = q.shape
+ q, k, v = map(
+ lambda t: t.unsqueeze(3)
+ .reshape(b, t.shape[1], self.heads, self.dim_head)
+ .permute(0, 2, 1, 3)
+ .reshape(b * self.heads, t.shape[1], self.dim_head)
+ .contiguous(),
+ (q, k, v),
+ )
+
+ # actually compute the attention, what we cannot get enough of
+ out = xformers.ops.memory_efficient_attention(
+ q, k, v, attn_bias=None, op=self.attention_op
+ )
+
+ # TODO: Use this directly in the attention operation, as a bias
+ if exists(mask):
+ raise NotImplementedError
+ out = (
+ out.unsqueeze(0)
+ .reshape(b, self.heads, out.shape[1], self.dim_head)
+ .permute(0, 2, 1, 3)
+ .reshape(b, out.shape[1], self.heads * self.dim_head)
+ )
+ if additional_tokens is not None:
+ # remove additional token
+ out = out[:, n_tokens_to_mask:]
+ return self.to_out(out)
+
+
+class BasicTransformerBlock(nn.Module):
+ ATTENTION_MODES = {
+ "softmax": CrossAttention, # vanilla attention
+ "softmax-xformers": MemoryEfficientCrossAttention, # ampere
+ }
+
+ def __init__(
+ self,
+ dim,
+ n_heads,
+ d_head,
+ dropout=0.0,
+ context_dim=None,
+ gated_ff=True,
+ checkpoint=False,
+ disable_self_attn=False,
+ attn_mode="softmax",
+ sdp_backend=None,
+ ):
+ super().__init__()
+ checkpoint = False
+
+ assert attn_mode in self.ATTENTION_MODES
+ if attn_mode != "softmax" and not XFORMERS_IS_AVAILABLE:
+ print(
+ f"Attention mode '{attn_mode}' is not available. Falling back to native attention. "
+ f"This is not a problem in Pytorch >= 2.0. FYI, you are running with PyTorch version {torch.__version__}"
+ )
+ attn_mode = "softmax"
+ elif attn_mode == "softmax" and not SDP_IS_AVAILABLE:
+ print(
+ "We do not support vanilla attention anymore, as it is too expensive. Sorry."
+ )
+ if not XFORMERS_IS_AVAILABLE:
+ assert (
+ False
+ ), "Please install xformers via e.g. 'pip install xformers==0.0.16'"
+ else:
+ print("Falling back to xformers efficient attention.")
+ attn_mode = "softmax-xformers"
+ attn_cls = self.ATTENTION_MODES[attn_mode]
+ if version.parse(torch.__version__) >= version.parse("2.0.0"):
+ assert sdp_backend is None or isinstance(sdp_backend, SDPBackend)
+ else:
+ assert sdp_backend is None
+ self.disable_self_attn = disable_self_attn
+ self.attn1 = attn_cls(
+ query_dim=dim,
+ heads=n_heads,
+ dim_head=d_head,
+ dropout=dropout,
+ context_dim=context_dim if self.disable_self_attn else None,
+ backend=sdp_backend,
+ ) # is a self-attention if not self.disable_self_attn
+ self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+ self.attn2 = attn_cls(
+ query_dim=dim,
+ context_dim=context_dim,
+ heads=n_heads,
+ dim_head=d_head,
+ dropout=dropout,
+ backend=sdp_backend,
+ ) # is self-attn if context is none
+ self.norm1 = nn.LayerNorm(dim)
+ self.norm2 = nn.LayerNorm(dim)
+ self.norm3 = nn.LayerNorm(dim)
+ self.checkpoint = checkpoint
+ if self.checkpoint:
+ print(f"{self.__class__.__name__} is using checkpointing")
+
+ def forward(
+ self, x, context=None, additional_tokens=None, n_times_crossframe_attn_in_self=0
+ ):
+ kwargs = {"x": x}
+
+ if context is not None:
+ kwargs.update({"context": context})
+
+ if additional_tokens is not None:
+ kwargs.update({"additional_tokens": additional_tokens})
+
+ if n_times_crossframe_attn_in_self:
+ kwargs.update(
+ {"n_times_crossframe_attn_in_self": n_times_crossframe_attn_in_self}
+ )
+
+ # return mixed_checkpoint(self._forward, kwargs, self.parameters(), self.checkpoint)
+ return checkpoint(
+ self._forward, (x, context), self.parameters(), self.checkpoint
+ )
+
+ def _forward(
+ self, x, context=None, additional_tokens=None, n_times_crossframe_attn_in_self=0
+ ):
+ x = (
+ self.attn1(
+ self.norm1(x),
+ context=context if self.disable_self_attn else None,
+ additional_tokens=additional_tokens,
+ n_times_crossframe_attn_in_self=n_times_crossframe_attn_in_self
+ if not self.disable_self_attn
+ else 0,
+ )
+ + x
+ )
+ x = (
+ self.attn2(
+ self.norm2(x), context=context, additional_tokens=additional_tokens
+ )
+ + x
+ )
+ x = self.ff(self.norm3(x)) + x
+ return x
+
+
+class BasicTransformerSingleLayerBlock(nn.Module):
+ ATTENTION_MODES = {
+ "softmax": CrossAttention, # vanilla attention
+ "softmax-xformers": MemoryEfficientCrossAttention # on the A100s not quite as fast as the above version
+ # (todo might depend on head_dim, check, falls back to semi-optimized kernels for dim!=[16,32,64,128])
+ }
+
+ def __init__(
+ self,
+ dim,
+ n_heads,
+ d_head,
+ dropout=0.0,
+ context_dim=None,
+ gated_ff=True,
+ checkpoint=True,
+ attn_mode="softmax",
+ ):
+ super().__init__()
+ assert attn_mode in self.ATTENTION_MODES
+ attn_cls = self.ATTENTION_MODES[attn_mode]
+ self.attn1 = attn_cls(
+ query_dim=dim,
+ heads=n_heads,
+ dim_head=d_head,
+ dropout=dropout,
+ context_dim=context_dim,
+ )
+ self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+ self.norm1 = nn.LayerNorm(dim)
+ self.norm2 = nn.LayerNorm(dim)
+ self.checkpoint = checkpoint
+
+ def forward(self, x, context=None):
+ return checkpoint(
+ self._forward, (x, context), self.parameters(), self.checkpoint
+ )
+
+ def _forward(self, x, context=None):
+ x = self.attn1(self.norm1(x), context=context) + x
+ x = self.ff(self.norm2(x)) + x
+ return x
+
+
+class SpatialTransformer(nn.Module):
+ """
+ Transformer block for image-like data.
+ First, project the input (aka embedding)
+ and reshape to b, t, d.
+ Then apply standard transformer action.
+ Finally, reshape to image
+ NEW: use_linear for more efficiency instead of the 1x1 convs
+ """
+
+ def __init__(
+ self,
+ in_channels,
+ n_heads,
+ d_head,
+ depth=1,
+ dropout=0.0,
+ context_dim=None,
+ disable_self_attn=False,
+ use_linear=False,
+ attn_type="softmax",
+ use_checkpoint=True,
+ # sdp_backend=SDPBackend.FLASH_ATTENTION
+ sdp_backend=None,
+ ):
+ super().__init__()
+ print(
+ f"constructing {self.__class__.__name__} of depth {depth} w/ {in_channels} channels and {n_heads} heads"
+ )
+ from omegaconf import ListConfig
+
+ if exists(context_dim) and not isinstance(context_dim, (list, ListConfig)):
+ context_dim = [context_dim]
+ if exists(context_dim) and isinstance(context_dim, list):
+ if depth != len(context_dim):
+ # print(
+ # f"WARNING: {self.__class__.__name__}: Found context dims {context_dim} of depth {len(context_dim)}, "
+ # f"which does not match the specified 'depth' of {depth}. Setting context_dim to {depth * [context_dim[0]]} now."
+ # )
+ # depth does not match context dims.
+ assert all(
+ map(lambda x: x == context_dim[0], context_dim)
+ ), "need homogenous context_dim to match depth automatically"
+ context_dim = depth * [context_dim[0]]
+ elif context_dim is None:
+ context_dim = [None] * depth
+ self.in_channels = in_channels
+ inner_dim = n_heads * d_head
+ self.norm = Normalize(in_channels)
+ if not use_linear:
+ self.proj_in = nn.Conv2d(
+ in_channels, inner_dim, kernel_size=1, stride=1, padding=0
+ )
+ else:
+ self.proj_in = nn.Linear(in_channels, inner_dim)
+
+ self.transformer_blocks = nn.ModuleList(
+ [
+ BasicTransformerBlock(
+ inner_dim,
+ n_heads,
+ d_head,
+ dropout=dropout,
+ context_dim=context_dim[d],
+ disable_self_attn=disable_self_attn,
+ attn_mode=attn_type,
+ checkpoint=use_checkpoint,
+ sdp_backend=sdp_backend,
+ )
+ for d in range(depth)
+ ]
+ )
+ if not use_linear:
+ self.proj_out = zero_module(
+ nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
+ )
+ else:
+ # self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
+ self.proj_out = zero_module(nn.Linear(inner_dim, in_channels))
+ self.use_linear = use_linear
+
+ def forward(self, x, context=None):
+ # note: if no context is given, cross-attention defaults to self-attention
+ if not isinstance(context, list):
+ context = [context]
+ b, c, h, w = x.shape
+ x_in = x
+ x = self.norm(x)
+ if not self.use_linear:
+ x = self.proj_in(x)
+ x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+ if self.use_linear:
+ x = self.proj_in(x)
+ for i, block in enumerate(self.transformer_blocks):
+ if i > 0 and len(context) == 1:
+ i = 0 # use same context for each block
+ x = block(x, context=context[i])
+ if self.use_linear:
+ x = self.proj_out(x)
+ x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+ if not self.use_linear:
+ x = self.proj_out(x)
+ return x + x_in

sgm/modules/autoencoding/losses/__init__.py

@@ -0,0 +1,246 @@
+from typing import Any, Union
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from ....util import default, instantiate_from_config
+from ..lpips.loss.lpips import LPIPS
+from ..lpips.model.model import NLayerDiscriminator, weights_init
+from ..lpips.vqperceptual import hinge_d_loss, vanilla_d_loss
+
+
+def adopt_weight(weight, global_step, threshold=0, value=0.0):
+ if global_step < threshold:
+ weight = value
+ return weight
+
+
+class LatentLPIPS(nn.Module):
+ def __init__(
+ self,
+ decoder_config,
+ perceptual_weight=1.0,
+ latent_weight=1.0,
+ scale_input_to_tgt_size=False,
+ scale_tgt_to_input_size=False,
+ perceptual_weight_on_inputs=0.0,
+ ):
+ super().__init__()
+ self.scale_input_to_tgt_size = scale_input_to_tgt_size
+ self.scale_tgt_to_input_size = scale_tgt_to_input_size
+ self.init_decoder(decoder_config)
+ self.perceptual_loss = LPIPS().eval()
+ self.perceptual_weight = perceptual_weight
+ self.latent_weight = latent_weight
+ self.perceptual_weight_on_inputs = perceptual_weight_on_inputs
+
+ def init_decoder(self, config):
+ self.decoder = instantiate_from_config(config)
+ if hasattr(self.decoder, "encoder"):
+ del self.decoder.encoder
+
+ def forward(self, latent_inputs, latent_predictions, image_inputs, split="train"):
+ log = dict()
+ loss = (latent_inputs - latent_predictions) ** 2
+ log[f"{split}/latent_l2_loss"] = loss.mean().detach()
+ image_reconstructions = None
+ if self.perceptual_weight > 0.0:
+ image_reconstructions = self.decoder.decode(latent_predictions)
+ image_targets = self.decoder.decode(latent_inputs)
+ perceptual_loss = self.perceptual_loss(
+ image_targets.contiguous(), image_reconstructions.contiguous()
+ )
+ loss = (
+ self.latent_weight * loss.mean()
+ + self.perceptual_weight * perceptual_loss.mean()
+ )
+ log[f"{split}/perceptual_loss"] = perceptual_loss.mean().detach()
+
+ if self.perceptual_weight_on_inputs > 0.0:
+ image_reconstructions = default(
+ image_reconstructions, self.decoder.decode(latent_predictions)
+ )
+ if self.scale_input_to_tgt_size:
+ image_inputs = torch.nn.functional.interpolate(
+ image_inputs,
+ image_reconstructions.shape[2:],
+ mode="bicubic",
+ antialias=True,
+ )
+ elif self.scale_tgt_to_input_size:
+ image_reconstructions = torch.nn.functional.interpolate(
+ image_reconstructions,
+ image_inputs.shape[2:],
+ mode="bicubic",
+ antialias=True,
+ )
+
+ perceptual_loss2 = self.perceptual_loss(
+ image_inputs.contiguous(), image_reconstructions.contiguous()
+ )
+ loss = loss + self.perceptual_weight_on_inputs * perceptual_loss2.mean()
+ log[f"{split}/perceptual_loss_on_inputs"] = perceptual_loss2.mean().detach()
+ return loss, log
+
+
+class GeneralLPIPSWithDiscriminator(nn.Module):
+ def __init__(
+ self,
+ disc_start: int,
+ logvar_init: float = 0.0,
+ pixelloss_weight=1.0,
+ disc_num_layers: int = 3,
+ disc_in_channels: int = 3,
+ disc_factor: float = 1.0,
+ disc_weight: float = 1.0,
+ perceptual_weight: float = 1.0,
+ disc_loss: str = "hinge",
+ scale_input_to_tgt_size: bool = False,
+ dims: int = 2,
+ learn_logvar: bool = False,
+ regularization_weights: Union[None, dict] = None,
+ ):
+ super().__init__()
+ self.dims = dims
+ if self.dims > 2:
+ print(
+ f"running with dims={dims}. This means that for perceptual loss calculation, "
+ f"the LPIPS loss will be applied to each frame independently. "
+ )
+ self.scale_input_to_tgt_size = scale_input_to_tgt_size
+ assert disc_loss in ["hinge", "vanilla"]
+ self.pixel_weight = pixelloss_weight
+ self.perceptual_loss = LPIPS().eval()
+ self.perceptual_weight = perceptual_weight
+ # output log variance
+ self.logvar = nn.Parameter(torch.ones(size=()) * logvar_init)
+ self.learn_logvar = learn_logvar
+
+ self.discriminator = NLayerDiscriminator(
+ input_nc=disc_in_channels, n_layers=disc_num_layers, use_actnorm=False
+ ).apply(weights_init)
+ self.discriminator_iter_start = disc_start
+ self.disc_loss = hinge_d_loss if disc_loss == "hinge" else vanilla_d_loss
+ self.disc_factor = disc_factor
+ self.discriminator_weight = disc_weight
+ self.regularization_weights = default(regularization_weights, {})
+
+ def get_trainable_parameters(self) -> Any:
+ return self.discriminator.parameters()
+
+ def get_trainable_autoencoder_parameters(self) -> Any:
+ if self.learn_logvar:
+ yield self.logvar
+ yield from ()
+
+ def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
+ if last_layer is not None:
+ nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
+ g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
+ else:
+ nll_grads = torch.autograd.grad(
+ nll_loss, self.last_layer[0], retain_graph=True
+ )[0]
+ g_grads = torch.autograd.grad(
+ g_loss, self.last_layer[0], retain_graph=True
+ )[0]
+
+ d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+ d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
+ d_weight = d_weight * self.discriminator_weight
+ return d_weight
+
+ def forward(
+ self,
+ regularization_log,
+ inputs,
+ reconstructions,
+ optimizer_idx,
+ global_step,
+ last_layer=None,
+ split="train",
+ weights=None,
+ ):
+ if self.scale_input_to_tgt_size:
+ inputs = torch.nn.functional.interpolate(
+ inputs, reconstructions.shape[2:], mode="bicubic", antialias=True
+ )
+
+ if self.dims > 2:
+ inputs, reconstructions = map(
+ lambda x: rearrange(x, "b c t h w -> (b t) c h w"),
+ (inputs, reconstructions),
+ )
+
+ rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous())
+ if self.perceptual_weight > 0:
+ p_loss = self.perceptual_loss(
+ inputs.contiguous(), reconstructions.contiguous()
+ )
+ rec_loss = rec_loss + self.perceptual_weight * p_loss
+
+ nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar
+ weighted_nll_loss = nll_loss
+ if weights is not None:
+ weighted_nll_loss = weights * nll_loss
+ weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]
+ nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
+
+ # now the GAN part
+ if optimizer_idx == 0:
+ # generator update
+ logits_fake = self.discriminator(reconstructions.contiguous())
+ g_loss = -torch.mean(logits_fake)
+
+ if self.disc_factor > 0.0:
+ try:
+ d_weight = self.calculate_adaptive_weight(
+ nll_loss, g_loss, last_layer=last_layer
+ )
+ except RuntimeError:
+ assert not self.training
+ d_weight = torch.tensor(0.0)
+ else:
+ d_weight = torch.tensor(0.0)
+
+ disc_factor = adopt_weight(
+ self.disc_factor, global_step, threshold=self.discriminator_iter_start
+ )
+ loss = weighted_nll_loss + d_weight * disc_factor * g_loss
+ log = dict()
+ for k in regularization_log:
+ if k in self.regularization_weights:
+ loss = loss + self.regularization_weights[k] * regularization_log[k]
+ log[f"{split}/{k}"] = regularization_log[k].detach().mean()
+
+ log.update(
+ {
+ "{}/total_loss".format(split): loss.clone().detach().mean(),
+ "{}/logvar".format(split): self.logvar.detach(),
+ "{}/nll_loss".format(split): nll_loss.detach().mean(),
+ "{}/rec_loss".format(split): rec_loss.detach().mean(),
+ "{}/d_weight".format(split): d_weight.detach(),
+ "{}/disc_factor".format(split): torch.tensor(disc_factor),
+ "{}/g_loss".format(split): g_loss.detach().mean(),
+ }
+ )
+
+ return loss, log
+
+ if optimizer_idx == 1:
+ # second pass for discriminator update
+ logits_real = self.discriminator(inputs.contiguous().detach())
+ logits_fake = self.discriminator(reconstructions.contiguous().detach())
+
+ disc_factor = adopt_weight(
+ self.disc_factor, global_step, threshold=self.discriminator_iter_start
+ )
+ d_loss = disc_factor * self.disc_loss(logits_real, logits_fake)
+
+ log = {
+ "{}/disc_loss".format(split): d_loss.clone().detach().mean(),
+ "{}/logits_real".format(split): logits_real.detach().mean(),
+ "{}/logits_fake".format(split): logits_fake.detach().mean(),
+ }
+ return d_loss, log

sgm/modules/autoencoding/lpips/loss/LICENSE

@@ -0,0 +1,23 @@
+Copyright (c) 2018, Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, Oliver Wang
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+ list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+ this list of conditions and the following disclaimer in the documentation
+ and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

sgm/modules/autoencoding/lpips/loss/lpips.py

@@ -0,0 +1,147 @@
+"""Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models"""
+
+from collections import namedtuple
+
+import torch
+import torch.nn as nn
+from torchvision import models
+
+from ..util import get_ckpt_path
+
+
+class LPIPS(nn.Module):
+ # Learned perceptual metric
+ def __init__(self, use_dropout=True):
+ super().__init__()
+ self.scaling_layer = ScalingLayer()
+ self.chns = [64, 128, 256, 512, 512] # vg16 features
+ self.net = vgg16(pretrained=True, requires_grad=False)
+ self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+ self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+ self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+ self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+ self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+ self.load_from_pretrained()
+ for param in self.parameters():
+ param.requires_grad = False
+
+ def load_from_pretrained(self, name="vgg_lpips"):
+ ckpt = get_ckpt_path(name, "sgm/modules/autoencoding/lpips/loss")
+ self.load_state_dict(
+ torch.load(ckpt, map_location=torch.device("cpu")), strict=False
+ )
+ print("loaded pretrained LPIPS loss from {}".format(ckpt))
+
+ @classmethod
+ def from_pretrained(cls, name="vgg_lpips"):
+ if name != "vgg_lpips":
+ raise NotImplementedError
+ model = cls()
+ ckpt = get_ckpt_path(name)
+ model.load_state_dict(
+ torch.load(ckpt, map_location=torch.device("cpu")), strict=False
+ )
+ return model
+
+ def forward(self, input, target):
+ in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
+ outs0, outs1 = self.net(in0_input), self.net(in1_input)
+ feats0, feats1, diffs = {}, {}, {}
+ lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+ for kk in range(len(self.chns)):
+ feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(
+ outs1[kk]
+ )
+ diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+
+ res = [
+ spatial_average(lins[kk].model(diffs[kk]), keepdim=True)
+ for kk in range(len(self.chns))
+ ]
+ val = res[0]
+ for l in range(1, len(self.chns)):
+ val += res[l]
+ return val
+
+
+class ScalingLayer(nn.Module):
+ def __init__(self):
+ super(ScalingLayer, self).__init__()
+ self.register_buffer(
+ "shift", torch.Tensor([-0.030, -0.088, -0.188])[None, :, None, None]
+ )
+ self.register_buffer(
+ "scale", torch.Tensor([0.458, 0.448, 0.450])[None, :, None, None]
+ )
+
+ def forward(self, inp):
+ return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+ """A single linear layer which does a 1x1 conv"""
+
+ def __init__(self, chn_in, chn_out=1, use_dropout=False):
+ super(NetLinLayer, self).__init__()
+ layers = (
+ [
+ nn.Dropout(),
+ ]
+ if (use_dropout)
+ else []
+ )
+ layers += [
+ nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False),
+ ]
+ self.model = nn.Sequential(*layers)
+
+
+class vgg16(torch.nn.Module):
+ def __init__(self, requires_grad=False, pretrained=True):
+ super(vgg16, self).__init__()
+ vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
+ self.slice1 = torch.nn.Sequential()
+ self.slice2 = torch.nn.Sequential()
+ self.slice3 = torch.nn.Sequential()
+ self.slice4 = torch.nn.Sequential()
+ self.slice5 = torch.nn.Sequential()
+ self.N_slices = 5
+ for x in range(4):
+ self.slice1.add_module(str(x), vgg_pretrained_features[x])
+ for x in range(4, 9):
+ self.slice2.add_module(str(x), vgg_pretrained_features[x])
+ for x in range(9, 16):
+ self.slice3.add_module(str(x), vgg_pretrained_features[x])
+ for x in range(16, 23):
+ self.slice4.add_module(str(x), vgg_pretrained_features[x])
+ for x in range(23, 30):
+ self.slice5.add_module(str(x), vgg_pretrained_features[x])
+ if not requires_grad:
+ for param in self.parameters():
+ param.requires_grad = False
+
+ def forward(self, X):
+ h = self.slice1(X)
+ h_relu1_2 = h
+ h = self.slice2(h)
+ h_relu2_2 = h
+ h = self.slice3(h)
+ h_relu3_3 = h
+ h = self.slice4(h)
+ h_relu4_3 = h
+ h = self.slice5(h)
+ h_relu5_3 = h
+ vgg_outputs = namedtuple(
+ "VggOutputs", ["relu1_2", "relu2_2", "relu3_3", "relu4_3", "relu5_3"]
+ )
+ out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+ return out
+
+
+def normalize_tensor(x, eps=1e-10):
+ norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True))
+ return x / (norm_factor + eps)
+
+
+def spatial_average(x, keepdim=True):
+ return x.mean([2, 3], keepdim=keepdim)

sgm/modules/autoencoding/lpips/model/LICENSE

@@ -0,0 +1,58 @@
+Copyright (c) 2017, Jun-Yan Zhu and Taesung Park
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+ list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+ this list of conditions and the following disclaimer in the documentation
+ and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+--------------------------- LICENSE FOR pix2pix --------------------------------
+BSD License
+
+For pix2pix software
+Copyright (c) 2016, Phillip Isola and Jun-Yan Zhu
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+ list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+ this list of conditions and the following disclaimer in the documentation
+ and/or other materials provided with the distribution.
+
+----------------------------- LICENSE FOR DCGAN --------------------------------
+BSD License
+
+For dcgan.torch software
+
+Copyright (c) 2015, Facebook, Inc. All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+Neither the name Facebook nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

sgm/modules/autoencoding/lpips/model/model.py

@@ -0,0 +1,88 @@
+import functools
+
+import torch.nn as nn
+
+from ..util import ActNorm
+
+
+def weights_init(m):
+ classname = m.__class__.__name__
+ if classname.find("Conv") != -1:
+ nn.init.normal_(m.weight.data, 0.0, 0.02)
+ elif classname.find("BatchNorm") != -1:
+ nn.init.normal_(m.weight.data, 1.0, 0.02)
+ nn.init.constant_(m.bias.data, 0)
+
+
+class NLayerDiscriminator(nn.Module):
+ """Defines a PatchGAN discriminator as in Pix2Pix
+ --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+ """
+
+ def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
+ """Construct a PatchGAN discriminator
+ Parameters:
+ input_nc (int) -- the number of channels in input images
+ ndf (int) -- the number of filters in the last conv layer
+ n_layers (int) -- the number of conv layers in the discriminator
+ norm_layer -- normalization layer
+ """
+ super(NLayerDiscriminator, self).__init__()
+ if not use_actnorm:
+ norm_layer = nn.BatchNorm2d
+ else:
+ norm_layer = ActNorm
+ if (
+ type(norm_layer) == functools.partial
+ ): # no need to use bias as BatchNorm2d has affine parameters
+ use_bias = norm_layer.func != nn.BatchNorm2d
+ else:
+ use_bias = norm_layer != nn.BatchNorm2d
+
+ kw = 4
+ padw = 1
+ sequence = [
+ nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),
+ nn.LeakyReLU(0.2, True),
+ ]
+ nf_mult = 1
+ nf_mult_prev = 1
+ for n in range(1, n_layers): # gradually increase the number of filters
+ nf_mult_prev = nf_mult
+ nf_mult = min(2**n, 8)
+ sequence += [
+ nn.Conv2d(
+ ndf * nf_mult_prev,
+ ndf * nf_mult,
+ kernel_size=kw,
+ stride=2,
+ padding=padw,
+ bias=use_bias,
+ ),
+ norm_layer(ndf * nf_mult),
+ nn.LeakyReLU(0.2, True),
+ ]
+
+ nf_mult_prev = nf_mult
+ nf_mult = min(2**n_layers, 8)
+ sequence += [
+ nn.Conv2d(
+ ndf * nf_mult_prev,
+ ndf * nf_mult,
+ kernel_size=kw,
+ stride=1,
+ padding=padw,
+ bias=use_bias,
+ ),
+ norm_layer(ndf * nf_mult),
+ nn.LeakyReLU(0.2, True),
+ ]
+
+ sequence += [
+ nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)
+ ] # output 1 channel prediction map
+ self.main = nn.Sequential(*sequence)
+
+ def forward(self, input):
+ """Standard forward."""
+ return self.main(input)

sgm/modules/autoencoding/lpips/util.py

@@ -0,0 +1,128 @@
+import hashlib
+import os
+
+import requests
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+
+URL_MAP = {"vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"}
+
+CKPT_MAP = {"vgg_lpips": "vgg.pth"}
+
+MD5_MAP = {"vgg_lpips": "d507d7349b931f0638a25a48a722f98a"}
+
+
+def download(url, local_path, chunk_size=1024):
+ os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+ with requests.get(url, stream=True) as r:
+ total_size = int(r.headers.get("content-length", 0))
+ with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+ with open(local_path, "wb") as f:
+ for data in r.iter_content(chunk_size=chunk_size):
+ if data:
+ f.write(data)
+ pbar.update(chunk_size)
+
+
+def md5_hash(path):
+ with open(path, "rb") as f:
+ content = f.read()
+ return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name, root, check=False):
+ assert name in URL_MAP
+ path = os.path.join(root, CKPT_MAP[name])
+ if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+ print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
+ download(URL_MAP[name], path)
+ md5 = md5_hash(path)
+ assert md5 == MD5_MAP[name], md5
+ return path
+
+
+class ActNorm(nn.Module):
+ def __init__(
+ self, num_features, logdet=False, affine=True, allow_reverse_init=False
+ ):
+ assert affine
+ super().__init__()
+ self.logdet = logdet
+ self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
+ self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
+ self.allow_reverse_init = allow_reverse_init
+
+ self.register_buffer("initialized", torch.tensor(0, dtype=torch.uint8))
+
+ def initialize(self, input):
+ with torch.no_grad():
+ flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
+ mean = (
+ flatten.mean(1)
+ .unsqueeze(1)
+ .unsqueeze(2)
+ .unsqueeze(3)
+ .permute(1, 0, 2, 3)
+ )
+ std = (
+ flatten.std(1)
+ .unsqueeze(1)
+ .unsqueeze(2)
+ .unsqueeze(3)
+ .permute(1, 0, 2, 3)
+ )
+
+ self.loc.data.copy_(-mean)
+ self.scale.data.copy_(1 / (std + 1e-6))
+
+ def forward(self, input, reverse=False):
+ if reverse:
+ return self.reverse(input)
+ if len(input.shape) == 2:
+ input = input[:, :, None, None]
+ squeeze = True
+ else:
+ squeeze = False
+
+ _, _, height, width = input.shape
+
+ if self.training and self.initialized.item() == 0:
+ self.initialize(input)
+ self.initialized.fill_(1)
+
+ h = self.scale * (input + self.loc)
+
+ if squeeze:
+ h = h.squeeze(-1).squeeze(-1)
+
+ if self.logdet:
+ log_abs = torch.log(torch.abs(self.scale))
+ logdet = height * width * torch.sum(log_abs)
+ logdet = logdet * torch.ones(input.shape[0]).to(input)
+ return h, logdet
+
+ return h
+
+ def reverse(self, output):
+ if self.training and self.initialized.item() == 0:
+ if not self.allow_reverse_init:
+ raise RuntimeError(
+ "Initializing ActNorm in reverse direction is "
+ "disabled by default. Use allow_reverse_init=True to enable."
+ )
+ else:
+ self.initialize(output)
+ self.initialized.fill_(1)
+
+ if len(output.shape) == 2:
+ output = output[:, :, None, None]
+ squeeze = True
+ else:
+ squeeze = False
+
+ h = output / self.scale - self.loc
+
+ if squeeze:
+ h = h.squeeze(-1).squeeze(-1)
+ return h

sgm/modules/autoencoding/lpips/vqperceptual.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn.functional as F
+
+
+def hinge_d_loss(logits_real, logits_fake):
+ loss_real = torch.mean(F.relu(1.0 - logits_real))
+ loss_fake = torch.mean(F.relu(1.0 + logits_fake))
+ d_loss = 0.5 * (loss_real + loss_fake)
+ return d_loss
+
+
+def vanilla_d_loss(logits_real, logits_fake):
+ d_loss = 0.5 * (
+ torch.mean(torch.nn.functional.softplus(-logits_real))
+ + torch.mean(torch.nn.functional.softplus(logits_fake))
+ )
+ return d_loss

sgm/modules/autoencoding/regularizers/__init__.py

@@ -0,0 +1,53 @@
+from abc import abstractmethod
+from typing import Any, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ....modules.distributions.distributions import DiagonalGaussianDistribution
+
+
+class AbstractRegularizer(nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]:
+ raise NotImplementedError()
+
+ @abstractmethod
+ def get_trainable_parameters(self) -> Any:
+ raise NotImplementedError()
+
+
+class DiagonalGaussianRegularizer(AbstractRegularizer):
+ def __init__(self, sample: bool = True):
+ super().__init__()
+ self.sample = sample
+
+ def get_trainable_parameters(self) -> Any:
+ yield from ()
+
+ def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]:
+ log = dict()
+ posterior = DiagonalGaussianDistribution(z)
+ if self.sample:
+ z = posterior.sample()
+ else:
+ z = posterior.mode()
+ kl_loss = posterior.kl()
+ kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
+ log["kl_loss"] = kl_loss
+ return z, log
+
+
+def measure_perplexity(predicted_indices, num_centroids):
+ # src: https://github.com/karpathy/deep-vector-quantization/blob/main/model.py
+ # eval cluster perplexity. when perplexity == num_embeddings then all clusters are used exactly equally
+ encodings = (
+ F.one_hot(predicted_indices, num_centroids).float().reshape(-1, num_centroids)
+ )
+ avg_probs = encodings.mean(0)
+ perplexity = (-(avg_probs * torch.log(avg_probs + 1e-10)).sum()).exp()
+ cluster_use = torch.sum(avg_probs > 0)
+ return perplexity, cluster_use

sgm/modules/diffusionmodules/__init__.py

@@ -0,0 +1,7 @@
+from .denoiser import Denoiser
+from .discretizer import Discretization
+from .loss import StandardDiffusionLoss
+from .model import Decoder, Encoder, Model
+from .openaimodel import UNetModel
+from .sampling import BaseDiffusionSampler
+from .wrappers import OpenAIWrapper

sgm/modules/diffusionmodules/denoiser.py

@@ -0,0 +1,73 @@
+import torch.nn as nn
+
+from ...util import append_dims, instantiate_from_config
+
+
+class Denoiser(nn.Module):
+ def __init__(self, weighting_config, scaling_config):
+ super().__init__()
+
+ self.weighting = instantiate_from_config(weighting_config)
+ self.scaling = instantiate_from_config(scaling_config)
+
+ def possibly_quantize_sigma(self, sigma):
+ return sigma
+
+ def possibly_quantize_c_noise(self, c_noise):
+ return c_noise
+
+ def w(self, sigma):
+ return self.weighting(sigma)
+
+ def __call__(self, network, input, sigma, cond):
+ sigma = self.possibly_quantize_sigma(sigma)
+ sigma_shape = sigma.shape
+ sigma = append_dims(sigma, input.ndim)
+ c_skip, c_out, c_in, c_noise = self.scaling(sigma)
+ c_noise = self.possibly_quantize_c_noise(c_noise.reshape(sigma_shape))
+ return network(input * c_in, c_noise, cond) * c_out + input * c_skip
+
+
+class DiscreteDenoiser(Denoiser):
+ def __init__(
+ self,
+ weighting_config,
+ scaling_config,
+ num_idx,
+ discretization_config,
+ do_append_zero=False,
+ quantize_c_noise=True,
+ flip=True,
+ ):
+ super().__init__(weighting_config, scaling_config)
+ sigmas = instantiate_from_config(discretization_config)(
+ num_idx, do_append_zero=do_append_zero, flip=flip
+ )
+ self.register_buffer("sigmas", sigmas)
+ self.quantize_c_noise = quantize_c_noise
+
+ def sigma_to_idx(self, sigma):
+ dists = sigma - self.sigmas[:, None]
+ return dists.abs().argmin(dim=0).view(sigma.shape)
+
+ def idx_to_sigma(self, idx):
+ return self.sigmas[idx]
+
+ def possibly_quantize_sigma(self, sigma):
+ return self.idx_to_sigma(self.sigma_to_idx(sigma))
+
+ def possibly_quantize_c_noise(self, c_noise):
+ if self.quantize_c_noise:
+ return self.sigma_to_idx(c_noise)
+ else:
+ return c_noise
+
+
+class DiscreteDenoiserWithControl(DiscreteDenoiser):
+ def __call__(self, network, input, sigma, cond, control_scale):
+ sigma = self.possibly_quantize_sigma(sigma)
+ sigma_shape = sigma.shape
+ sigma = append_dims(sigma, input.ndim)
+ c_skip, c_out, c_in, c_noise = self.scaling(sigma)
+ c_noise = self.possibly_quantize_c_noise(c_noise.reshape(sigma_shape))
+ return network(input * c_in, c_noise, cond, control_scale) * c_out + input * c_skip

sgm/modules/diffusionmodules/denoiser_scaling.py

@@ -0,0 +1,31 @@
+import torch
+
+
+class EDMScaling:
+ def __init__(self, sigma_data=0.5):
+ self.sigma_data = sigma_data
+
+ def __call__(self, sigma):
+ c_skip = self.sigma_data**2 / (sigma**2 + self.sigma_data**2)
+ c_out = sigma * self.sigma_data / (sigma**2 + self.sigma_data**2) ** 0.5
+ c_in = 1 / (sigma**2 + self.sigma_data**2) ** 0.5
+ c_noise = 0.25 * sigma.log()
+ return c_skip, c_out, c_in, c_noise
+
+
+class EpsScaling:
+ def __call__(self, sigma):
+ c_skip = torch.ones_like(sigma, device=sigma.device)
+ c_out = -sigma
+ c_in = 1 / (sigma**2 + 1.0) ** 0.5
+ c_noise = sigma.clone()
+ return c_skip, c_out, c_in, c_noise
+
+
+class VScaling:
+ def __call__(self, sigma):
+ c_skip = 1.0 / (sigma**2 + 1.0)
+ c_out = -sigma / (sigma**2 + 1.0) ** 0.5
+ c_in = 1.0 / (sigma**2 + 1.0) ** 0.5
+ c_noise = sigma.clone()
+ return c_skip, c_out, c_in, c_noise

sgm/modules/diffusionmodules/denoiser_weighting.py

@@ -0,0 +1,24 @@
+import torch
+
+class UnitWeighting:
+ def __call__(self, sigma):
+ return torch.ones_like(sigma, device=sigma.device)
+
+
+class EDMWeighting:
+ def __init__(self, sigma_data=0.5):
+ self.sigma_data = sigma_data
+
+ def __call__(self, sigma):
+ return (sigma**2 + self.sigma_data**2) / (sigma * self.sigma_data) ** 2
+
+
+class VWeighting(EDMWeighting):
+ def __init__(self):
+ super().__init__(sigma_data=1.0)
+
+
+class EpsWeighting:
+ def __call__(self, sigma):
+ return sigma**-2
+

sgm/modules/diffusionmodules/discretizer.py

@@ -0,0 +1,69 @@
+from abc import abstractmethod
+from functools import partial
+
+import numpy as np
+import torch
+
+from ...modules.diffusionmodules.util import make_beta_schedule
+from ...util import append_zero
+
+
+def generate_roughly_equally_spaced_steps(
+ num_substeps: int, max_step: int
+) -> np.ndarray:
+ return np.linspace(max_step - 1, 0, num_substeps, endpoint=False).astype(int)[::-1]
+
+
+class Discretization:
+ def __call__(self, n, do_append_zero=True, device="cpu", flip=False):
+ sigmas = self.get_sigmas(n, device=device)
+ sigmas = append_zero(sigmas) if do_append_zero else sigmas
+ return sigmas if not flip else torch.flip(sigmas, (0,))
+
+ @abstractmethod
+ def get_sigmas(self, n, device):
+ pass
+
+
+class EDMDiscretization(Discretization):
+ def __init__(self, sigma_min=0.02, sigma_max=80.0, rho=7.0):
+ self.sigma_min = sigma_min
+ self.sigma_max = sigma_max
+ self.rho = rho
+
+ def get_sigmas(self, n, device="cpu"):
+ ramp = torch.linspace(0, 1, n, device=device)
+ min_inv_rho = self.sigma_min ** (1 / self.rho)
+ max_inv_rho = self.sigma_max ** (1 / self.rho)
+ sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** self.rho
+ return sigmas
+
+
+class LegacyDDPMDiscretization(Discretization):
+ def __init__(
+ self,
+ linear_start=0.00085,
+ linear_end=0.0120,
+ num_timesteps=1000,
+ ):
+ super().__init__()
+ self.num_timesteps = num_timesteps
+ betas = make_beta_schedule(
+ "linear", num_timesteps, linear_start=linear_start, linear_end=linear_end
+ )
+ alphas = 1.0 - betas
+ self.alphas_cumprod = np.cumprod(alphas, axis=0)
+ self.to_torch = partial(torch.tensor, dtype=torch.float32)
+
+ def get_sigmas(self, n, device="cpu"):
+ if n < self.num_timesteps:
+ timesteps = generate_roughly_equally_spaced_steps(n, self.num_timesteps)
+ alphas_cumprod = self.alphas_cumprod[timesteps]
+ elif n == self.num_timesteps:
+ alphas_cumprod = self.alphas_cumprod
+ else:
+ raise ValueError
+
+ to_torch = partial(torch.tensor, dtype=torch.float32, device=device)
+ sigmas = to_torch((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
+ return torch.flip(sigmas, (0,))