From c325dec0d324637967d1e4fb6f9805e7730b6098 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sun, 7 Apr 2024 17:55:33 +0200 Subject: [PATCH] PixArt sigma Mentioned in #19 Also removes old example sampler since updating it would be too much overhead. Pytorch attention is still wonky. --- PixArt/conf.py | 44 +- PixArt/loader.py | 4 + PixArt/models/PixArt.py | 137 +-- PixArt/models/PixArtMS.py | 157 ++- PixArt/models/PixArt_blocks.py | 196 ++-- PixArt/nodes.py | 40 - PixArt/sampler.py | 72 -- PixArt/sampling/diffusion_utils.py | 88 -- PixArt/sampling/dpm_solver.py | 1356 ------------------------- PixArt/sampling/gaussian_diffusion.py | 908 ----------------- README.md | 9 + 11 files changed, 266 insertions(+), 2745 deletions(-) delete mode 100644 PixArt/sampler.py delete mode 100644 PixArt/sampling/diffusion_utils.py delete mode 100644 PixArt/sampling/dpm_solver.py delete mode 100644 PixArt/sampling/gaussian_diffusion.py diff --git a/PixArt/conf.py b/PixArt/conf.py index 5ba98b0..4b1c66f 100644 --- a/PixArt/conf.py +++ b/PixArt/conf.py @@ -1,40 +1,54 @@ """ List of all PixArt model types / settings """ + +sampling_settings = { + "beta_schedule" : "sqrt_linear", + "linear_start" : 0.0001, + "linear_end" : 0.02, + "timesteps" : 1000, +} + pixart_conf = { "PixArtMS_XL_2": { # models/PixArtMS "target": "PixArtMS", "unet_config": { "input_size" : 1024//8, - "lewei_scale" : 2, "depth" : 28, "num_heads" : 16, "patch_size" : 2, "hidden_size" : 1152, + "pe_interpolation": 2, }, - "sampling_settings": { - "beta_schedule" : "sqrt_linear", - "linear_start" : 0.0001, - "linear_end" : 0.02, - "timesteps" : 1000, + "sampling_settings" : sampling_settings, + }, + "PixArtMS_Sigma_XL_2": { + "target": "PixArtMSSigma", + "unet_config": { + "input_size" : 1024//8, + "token_num" : 300, + "depth" : 28, + "num_heads" : 16, + "patch_size" : 2, + "hidden_size" : 1152, + "micro_condition": False, + "pe_interpolation": 2, + "model_max_length": 300, }, + "sampling_settings" : sampling_settings, }, "PixArt_XL_2": { # models/PixArt "target": "PixArt", "unet_config": { "input_size" : 512//8, - "lewei_scale" : 1, + "token_num" : 120, "depth" : 28, "num_heads" : 16, "patch_size" : 2, "hidden_size" : 1152, + "pe_interpolation": 1, }, - "sampling_settings": { - "beta_schedule" : "sqrt_linear", - "linear_start" : 0.0001, - "linear_end" : 0.02, - "timesteps" : 1000, - }, + "sampling_settings" : sampling_settings, }, } @@ -77,3 +91,7 @@ pixart_res = { '3.62': [928, 256], '3.75': [960, 256], '3.88': [992, 256], '4.00': [1024,256] }, } +# These should be the same +pixart_res.update({ + "PixArtMS_Sigma_XL_2_512": pixart_res["PixArt_XL_2"], +}) diff --git a/PixArt/loader.py b/PixArt/loader.py index 67edee2..3e26711 100644 --- a/PixArt/loader.py +++ b/PixArt/loader.py @@ -82,6 +82,10 @@ def load_pixart(model_path, model_conf): elif model_conf.model_target == "PixArt": from .models.PixArt import PixArt model.diffusion_model = PixArt(**model_conf.unet_config) + elif model_conf.model_target == "PixArtMSSigma": + from .models.PixArtMS import PixArtMS + model.diffusion_model = PixArtMS(**model_conf.unet_config) + model.latent_format = comfy.latent_formats.SDXL() elif model_conf.model_target == "ControlPixArtMSHalf": from .models.PixArtMS import PixArtMS from .models.pixart_controlnet import ControlPixArtMSHalf diff --git a/PixArt/models/PixArt.py b/PixArt/models/PixArt.py index 13b6727..7cdcb70 100644 --- a/PixArt/models/PixArt.py +++ b/PixArt/models/PixArt.py @@ -18,31 +18,31 @@ from timm.models.vision_transformer import PatchEmbed, Mlp from .utils import auto_grad_checkpoint, to_2tuple -from .PixArt_blocks import t2i_modulate, CaptionEmbedder, WindowAttention, MultiHeadCrossAttention, T2IFinalLayer, TimestepEmbedder, LabelEmbedder, FinalLayer +from .PixArt_blocks import t2i_modulate, CaptionEmbedder, AttentionKVCompress, MultiHeadCrossAttention, T2IFinalLayer, TimestepEmbedder, LabelEmbedder, FinalLayer class PixArtBlock(nn.Module): """ A PixArt block with adaptive layer norm (adaLN-single) conditioning. """ - - def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, drop_path=0., window_size=0, input_size=None, use_rel_pos=False, **block_kwargs): + def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, drop_path=0, input_size=None, sampling=None, sr_ratio=1, qk_norm=False, **block_kwargs): super().__init__() - self.hidden_size = hidden_size self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.attn = WindowAttention(hidden_size, num_heads=num_heads, qkv_bias=True, - input_size=input_size if window_size == 0 else (window_size, window_size), - use_rel_pos=use_rel_pos, **block_kwargs) + self.attn = AttentionKVCompress( + hidden_size, num_heads=num_heads, qkv_bias=True, sampling=sampling, sr_ratio=sr_ratio, + qk_norm=qk_norm, **block_kwargs + ) self.cross_attn = MultiHeadCrossAttention(hidden_size, num_heads, **block_kwargs) self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) # to be compatible with lower version pytorch approx_gelu = lambda: nn.GELU(approximate="tanh") self.mlp = Mlp(in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu, drop=0) self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() - self.window_size = window_size self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size ** 0.5) + self.sampling = sampling + self.sr_ratio = sr_ratio - def forward(self, x, y, t, mask=None): + def forward(self, x, y, t, mask=None, **kwargs): B, N, C = x.shape shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None] + t.reshape(B, 6, -1)).chunk(6, dim=1) @@ -53,14 +53,11 @@ class PixArtBlock(nn.Module): return x -############################################################################# -# Core PixArt Model # -################################################################################# +### Core PixArt Model ### class PixArt(nn.Module): """ Diffusion model with a Transformer backbone. """ - def __init__( self, input_size=32, @@ -73,12 +70,12 @@ class PixArt(nn.Module): class_dropout_prob=0.1, pred_sigma=True, drop_path: float = 0., - window_size=0, - window_block_indexes=[], - use_rel_pos=False, caption_channels=4096, - lewei_scale=1.0, + pe_interpolation=1.0, config=None, + model_max_length=120, + qk_norm=False, + kv_compress_config=None, **kwargs, ): super().__init__() @@ -87,8 +84,8 @@ class PixArt(nn.Module): self.out_channels = in_channels * 2 if pred_sigma else in_channels self.patch_size = patch_size self.num_heads = num_heads - self.lewei_scale = lewei_scale, - self.dtype = torch.get_default_dtype() + self.pe_interpolation = pe_interpolation + self.depth = depth self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True) self.t_embedder = TimestepEmbedder(hidden_size) @@ -102,21 +99,32 @@ class PixArt(nn.Module): nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True) ) - self.y_embedder = CaptionEmbedder(in_channels=caption_channels, hidden_size=hidden_size, uncond_prob=class_dropout_prob, act_layer=approx_gelu) + self.y_embedder = CaptionEmbedder( + in_channels=caption_channels, hidden_size=hidden_size, uncond_prob=class_dropout_prob, + act_layer=approx_gelu, token_num=model_max_length + ) drop_path = [x.item() for x in torch.linspace(0, drop_path, depth)] # stochastic depth decay rule + self.kv_compress_config = kv_compress_config + if kv_compress_config is None: + self.kv_compress_config = { + 'sampling': None, + 'scale_factor': 1, + 'kv_compress_layer': [], + } self.blocks = nn.ModuleList([ - PixArtBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, drop_path=drop_path[i], - input_size=(input_size // patch_size, input_size // patch_size), - window_size=window_size if i in window_block_indexes else 0, - use_rel_pos=use_rel_pos if i in window_block_indexes else False) + PixArtBlock( + hidden_size, num_heads, mlp_ratio=mlp_ratio, drop_path=drop_path[i], + input_size=(input_size // patch_size, input_size // patch_size), + sampling=self.kv_compress_config['sampling'], + sr_ratio=int( + self.kv_compress_config['scale_factor'] + ) if i in self.kv_compress_config['kv_compress_layer'] else 1, + qk_norm=qk_norm, + ) for i in range(depth) ]) self.final_layer = T2IFinalLayer(hidden_size, patch_size, self.out_channels) - self.initialize_weights() - - print(f'Warning: lewei scale: {self.lewei_scale}, base size: {self.base_size}') - def forward_raw(self, x, t, y, mask=None, data_info=None): """ Original forward pass of PixArt. @@ -124,9 +132,13 @@ class PixArt(nn.Module): t: (N,) tensor of diffusion timesteps y: (N, 1, 120, C) tensor of class labels """ + x = x.to(self.dtype) + timestep = t.to(self.dtype) + y = y.to(self.dtype) + pos_embed = self.pos_embed.to(self.dtype) self.h, self.w = x.shape[-2]//self.patch_size, x.shape[-1]//self.patch_size - x = self.x_embedder(x) + self.pos_embed # (N, T, D), where T = H * W / patch_size ** 2 - t = self.t_embedder(t) # (N, D) + x = self.x_embedder(x) + pos_embed # (N, T, D), where T = H * W / patch_size ** 2 + t = self.t_embedder(timestep.to(x.dtype)) # (N, D) t0 = self.t_block(t) y = self.y_embedder(y, self.training) # (N, 1, L, D) if mask is not None: @@ -168,29 +180,6 @@ class PixArt(nn.Module): eps, rest = out[:, :self.in_channels], out[:, self.in_channels:] return eps - def forward_with_dpmsolver(self, x, t, y, mask=None, **kwargs): - """ - dpm solver donnot need variance prediction - """ - # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb - model_out = self.forward(x, t, y, mask) - return model_out.chunk(2, dim=1)[0] - - def forward_with_cfg(self, x, t, y, cfg_scale, **kwargs): - """ - Forward pass of PixArt, but also batches the unconditional forward pass for classifier-free guidance. - """ - # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb - half = x[: len(x) // 2] - combined = torch.cat([half, half], dim=0) - model_out = self.forward(combined, t, y, kwargs) - model_out = model_out['x'] if isinstance(model_out, dict) else model_out - eps, rest = model_out[:, :3], model_out[:, 3:] - cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0) - half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps) - eps = torch.cat([half_eps, half_eps], dim=0) - return torch.cat([eps, rest], dim=1) - def unpatchify(self, x): """ x: (N, T, patch_size**2 * C) @@ -206,44 +195,8 @@ class PixArt(nn.Module): imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p)) return imgs - def initialize_weights(self): - # Initialize transformer layers: - def _basic_init(module): - if isinstance(module, nn.Linear): - torch.nn.init.xavier_uniform_(module.weight) - if module.bias is not None: - nn.init.constant_(module.bias, 0) - self.apply(_basic_init) - - # Initialize (and freeze) pos_embed by sin-cos embedding: - pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5), lewei_scale=self.lewei_scale, base_size=self.base_size) - self.pos_embed.data.copy_(torch.from_numpy(pos_embed).unsqueeze(0).to(self.dtype)) - - # Initialize patch_embed like nn.Linear (instead of nn.Conv2d): - w = self.x_embedder.proj.weight.data - nn.init.xavier_uniform_(w.view([w.shape[0], -1])) - - # Initialize timestep embedding MLP: - nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) - nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) - nn.init.normal_(self.t_block[1].weight, std=0.02) - - # Initialize caption embedding MLP: - nn.init.normal_(self.y_embedder.y_proj.fc1.weight, std=0.02) - nn.init.normal_(self.y_embedder.y_proj.fc2.weight, std=0.02) - - # Zero-out adaLN modulation layers in PixArt blocks: - for block in self.blocks: - nn.init.constant_(block.cross_attn.proj.weight, 0) - nn.init.constant_(block.cross_attn.proj.bias, 0) - - # Zero-out output layers: - nn.init.constant_(self.final_layer.linear.weight, 0) - nn.init.constant_(self.final_layer.linear.bias, 0) - - -def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, lewei_scale=1.0, base_size=16): +def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, pe_interpolation=1.0, base_size=16): """ grid_size: int of the grid height and width return: @@ -251,8 +204,8 @@ def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens= """ if isinstance(grid_size, int): grid_size = to_2tuple(grid_size) - grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0]/base_size) / lewei_scale - grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1]/base_size) / lewei_scale + grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0]/base_size) / pe_interpolation + grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1]/base_size) / pe_interpolation grid = np.meshgrid(grid_w, grid_h) # here w goes first grid = np.stack(grid, axis=0) grid = grid.reshape([2, 1, grid_size[1], grid_size[0]]) diff --git a/PixArt/models/PixArtMS.py b/PixArt/models/PixArtMS.py index 66b3cc3..79b7614 100644 --- a/PixArt/models/PixArtMS.py +++ b/PixArt/models/PixArtMS.py @@ -15,12 +15,13 @@ from timm.models.layers import DropPath from timm.models.vision_transformer import Mlp from .utils import auto_grad_checkpoint, to_2tuple -from .PixArt_blocks import t2i_modulate, CaptionEmbedder, WindowAttention, MultiHeadCrossAttention, T2IFinalLayer, TimestepEmbedder, SizeEmbedder +from .PixArt_blocks import t2i_modulate, CaptionEmbedder, AttentionKVCompress, MultiHeadCrossAttention, T2IFinalLayer, TimestepEmbedder, SizeEmbedder from .PixArt import PixArt, get_2d_sincos_pos_embed class PatchEmbed(nn.Module): - """ 2D Image to Patch Embedding + """ + 2D Image to Patch Embedding """ def __init__( self, @@ -50,42 +51,39 @@ class PixArtMSBlock(nn.Module): """ A PixArt block with adaptive layer norm zero (adaLN-Zero) conditioning. """ - - def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, drop_path=0., window_size=0, input_size=None, use_rel_pos=False, **block_kwargs): + def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, drop_path=0., input_size=None, + sampling=None, sr_ratio=1, qk_norm=False, **block_kwargs): super().__init__() self.hidden_size = hidden_size self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.attn = WindowAttention(hidden_size, num_heads=num_heads, qkv_bias=True, - input_size=input_size if window_size == 0 else (window_size, window_size), - use_rel_pos=use_rel_pos, **block_kwargs) + self.attn = AttentionKVCompress( + hidden_size, num_heads=num_heads, qkv_bias=True, sampling=sampling, sr_ratio=sr_ratio, + qk_norm=qk_norm, **block_kwargs + ) self.cross_attn = MultiHeadCrossAttention(hidden_size, num_heads, **block_kwargs) self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) # to be compatible with lower version pytorch approx_gelu = lambda: nn.GELU(approximate="tanh") self.mlp = Mlp(in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu, drop=0) self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() - self.window_size = window_size self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size ** 0.5) - def forward(self, x, y, t, mask=None, **kwargs): + def forward(self, x, y, t, mask=None, HW=None, **kwargs): B, N, C = x.shape shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None] + t.reshape(B, 6, -1)).chunk(6, dim=1) - x = x + self.drop_path(gate_msa * self.attn(t2i_modulate(self.norm1(x), shift_msa, scale_msa))) + x = x + self.drop_path(gate_msa * self.attn(t2i_modulate(self.norm1(x), shift_msa, scale_msa), HW=HW)) x = x + self.cross_attn(x, y, mask) x = x + self.drop_path(gate_mlp * self.mlp(t2i_modulate(self.norm2(x), shift_mlp, scale_mlp))) return x -############################################################################# -# Core PixArt Model # -################################################################################# +### Core PixArt Model ### class PixArtMS(PixArt): """ Diffusion model with a Transformer backbone. """ - def __init__( self, input_size=32, @@ -99,12 +97,13 @@ class PixArtMS(PixArt): learn_sigma=True, pred_sigma=True, drop_path: float = 0., - window_size=0, - window_block_indexes=[], - use_rel_pos=False, caption_channels=4096, - lewei_scale=1., + pe_interpolation=1., config=None, + model_max_length=120, + micro_condition=True, + qk_norm=False, + kv_compress_config=None, **kwargs, ): super().__init__( @@ -119,11 +118,11 @@ class PixArtMS(PixArt): learn_sigma=learn_sigma, pred_sigma=pred_sigma, drop_path=drop_path, - window_size=window_size, - window_block_indexes=window_block_indexes, - use_rel_pos=use_rel_pos, - lewei_scale=lewei_scale, + pe_interpolation=pe_interpolation, config=config, + model_max_length=model_max_length, + qk_norm=qk_norm, + kv_compress_config=kv_compress_config, **kwargs, ) self.dtype = torch.get_default_dtype() @@ -134,20 +133,29 @@ class PixArtMS(PixArt): nn.Linear(hidden_size, 6 * hidden_size, bias=True) ) self.x_embedder = PatchEmbed(patch_size, in_channels, hidden_size, bias=True) - self.y_embedder = CaptionEmbedder(in_channels=caption_channels, hidden_size=hidden_size, uncond_prob=class_dropout_prob, act_layer=approx_gelu) - self.csize_embedder = SizeEmbedder(hidden_size//3) # c_size embed - self.ar_embedder = SizeEmbedder(hidden_size//3) # aspect ratio embed + self.y_embedder = CaptionEmbedder(in_channels=caption_channels, hidden_size=hidden_size, uncond_prob=class_dropout_prob, act_layer=approx_gelu, token_num=model_max_length) + self.micro_conditioning = micro_condition + if self.micro_conditioning: + self.csize_embedder = SizeEmbedder(hidden_size//3) # c_size embed + self.ar_embedder = SizeEmbedder(hidden_size//3) # aspect ratio embed drop_path = [x.item() for x in torch.linspace(0, drop_path, depth)] # stochastic depth decay rule + if kv_compress_config is None: + kv_compress_config = { + 'sampling': None, + 'scale_factor': 1, + 'kv_compress_layer': [], + } self.blocks = nn.ModuleList([ - PixArtMSBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, drop_path=drop_path[i], - input_size=(input_size // patch_size, input_size // patch_size), - window_size=window_size if i in window_block_indexes else 0, - use_rel_pos=use_rel_pos if i in window_block_indexes else False) + PixArtMSBlock( + hidden_size, num_heads, mlp_ratio=mlp_ratio, drop_path=drop_path[i], + input_size=(input_size // patch_size, input_size // patch_size), + sampling=kv_compress_config['sampling'], + sr_ratio=int(kv_compress_config['scale_factor']) if i in kv_compress_config['kv_compress_layer'] else 1, + qk_norm=qk_norm, + ) for i in range(depth) ]) self.final_layer = T2IFinalLayer(hidden_size, patch_size, self.out_channels) - self.training = False - self.initialize() def forward_raw(self, x, t, y, mask=None, data_info=None, **kwargs): """ @@ -157,16 +165,29 @@ class PixArtMS(PixArt): y: (N, 1, 120, C) tensor of class labels """ bs = x.shape[0] - c_size, ar = data_info['img_hw'], data_info['aspect_ratio'] + x = x.to(self.dtype) + timestep = t.to(self.dtype) + y = y.to(self.dtype) self.h, self.w = x.shape[-2]//self.patch_size, x.shape[-1]//self.patch_size - pos_embed = torch.from_numpy(get_2d_sincos_pos_embed(self.pos_embed.shape[-1], (self.h, self.w), lewei_scale=self.lewei_scale, base_size=self.base_size)).unsqueeze(0).to(x.device).to(self.dtype) + pos_embed = torch.from_numpy( + get_2d_sincos_pos_embed( + self.pos_embed.shape[-1], (self.h, self.w), pe_interpolation=self.pe_interpolation, + base_size=self.base_size + ) + ).unsqueeze(0).to(x.device).to(self.dtype) + x = self.x_embedder(x) + pos_embed # (N, T, D), where T = H * W / patch_size ** 2 - t = self.t_embedder(t) # (N, D) - csize = self.csize_embedder(c_size, bs) # (N, D) - ar = self.ar_embedder(ar, bs) # (N, D) - t = t + torch.cat([csize, ar], dim=1) + t = self.t_embedder(timestep) # (N, D) + + if self.micro_conditioning: + c_size, ar = data_info['img_hw'].to(self.dtype), data_info['aspect_ratio'].to(self.dtype) + csize = self.csize_embedder(c_size, bs) # (N, D) + ar = self.ar_embedder(ar, bs) # (N, D) + t = t + torch.cat([csize, ar], dim=1) + t0 = self.t_block(t) y = self.y_embedder(y, self.training) # (N, D) + if mask is not None: if mask.shape[0] != y.shape[0]: mask = mask.repeat(y.shape[0] // mask.shape[0], 1) @@ -177,9 +198,11 @@ class PixArtMS(PixArt): y_lens = [y.shape[2]] * y.shape[0] y = y.squeeze(1).view(1, -1, x.shape[-1]) for block in self.blocks: - x = auto_grad_checkpoint(block, x, y, t0, y_lens, **kwargs) # (N, T, D) #support grad checkpoint + x = auto_grad_checkpoint(block, x, y, t0, y_lens, (self.h, self.w), **kwargs) # (N, T, D) #support grad checkpoint + x = self.final_layer(x, t) # (N, T, patch_size ** 2 * out_channels) x = self.unpatchify(x) # (N, out_channels, H, W) + return x def forward(self, x, timesteps, context, img_hw=None, aspect_ratio=None, **kwargs): @@ -228,28 +251,6 @@ class PixArtMS(PixArt): eps, rest = out[:, :self.in_channels], out[:, self.in_channels:] return eps - def forward_with_dpmsolver(self, x, t, y, data_info, **kwargs): - """ - dpm solver donnot need variance prediction - """ - # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb - model_out = self.forward_raw(x, t, y, data_info=data_info, **kwargs) - return model_out.chunk(2, dim=1)[0] - - def forward_with_cfg(self, x, t, y, cfg_scale, data_info, **kwargs): - """ - Forward pass of PixArt, but also batches the unconditional forward pass for classifier-free guidance. - """ - # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb - half = x[: len(x) // 2] - combined = torch.cat([half, half], dim=0) - model_out = self.forward_raw(combined, t, y, data_info=data_info) - eps, rest = model_out[:, :3], model_out[:, 3:] - cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0) - half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps) - eps = torch.cat([half_eps, half_eps], dim=0) - return torch.cat([eps, rest], dim=1) - def unpatchify(self, x): """ x: (N, T, patch_size**2 * C) @@ -263,39 +264,3 @@ class PixArtMS(PixArt): x = torch.einsum('nhwpqc->nchpwq', x) imgs = x.reshape(shape=(x.shape[0], c, self.h * p, self.w * p)) return imgs - - def initialize(self): - # Initialize transformer layers: - def _basic_init(module): - if isinstance(module, nn.Linear): - torch.nn.init.xavier_uniform_(module.weight) - if module.bias is not None: - nn.init.constant_(module.bias, 0) - - self.apply(_basic_init) - - # Initialize patch_embed like nn.Linear (instead of nn.Conv2d): - w = self.x_embedder.proj.weight.data - nn.init.xavier_uniform_(w.view([w.shape[0], -1])) - - # Initialize timestep embedding MLP: - nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) - nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) - nn.init.normal_(self.t_block[1].weight, std=0.02) - nn.init.normal_(self.csize_embedder.mlp[0].weight, std=0.02) - nn.init.normal_(self.csize_embedder.mlp[2].weight, std=0.02) - nn.init.normal_(self.ar_embedder.mlp[0].weight, std=0.02) - nn.init.normal_(self.ar_embedder.mlp[2].weight, std=0.02) - - # Initialize caption embedding MLP: - nn.init.normal_(self.y_embedder.y_proj.fc1.weight, std=0.02) - nn.init.normal_(self.y_embedder.y_proj.fc2.weight, std=0.02) - - # Zero-out adaLN modulation layers in PixArt blocks: - for block in self.blocks: - nn.init.constant_(block.cross_attn.proj.weight, 0) - nn.init.constant_(block.cross_attn.proj.bias, 0) - - # Zero-out output layers: - nn.init.constant_(self.final_layer.linear.weight, 0) - nn.init.constant_(self.final_layer.linear.bias, 0) diff --git a/PixArt/models/PixArt_blocks.py b/PixArt/models/PixArt_blocks.py index 9e492d6..a7b604b 100644 --- a/PixArt/models/PixArt_blocks.py +++ b/PixArt/models/PixArt_blocks.py @@ -11,25 +11,29 @@ import math import torch import torch.nn as nn +import torch.nn.functional as F from timm.models.vision_transformer import Mlp, Attention as Attention_ -from einops import rearrange, repeat +from einops import rearrange -from .utils import add_decomposed_rel_pos from comfy import model_management - if model_management.xformers_enabled(): import xformers import xformers.ops +else: + print(""" +######################################## + PixArt: Not using xformers! + Expect images to be non-deterministic! + Batch sizes > 1 are most likely broken +######################################## +""") def modulate(x, shift, scale): return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) - def t2i_modulate(x, shift, scale): return x * (1 + scale) + shift -competent_attention_implementation = False - class MultiHeadCrossAttention(nn.Module): def __init__(self, d_model, num_heads, attn_drop=0., proj_drop=0., **block_kwargs): super(MultiHeadCrossAttention, self).__init__() @@ -49,32 +53,24 @@ class MultiHeadCrossAttention(nn.Module): # query/value: img tokens; key: condition; mask: if padding tokens B, N, C = x.shape + q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim) + kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim) + k, v = kv.unbind(2) + if model_management.xformers_enabled(): - q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim) - kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim) - k, v = kv.unbind(2) attn_bias = None if mask is not None: attn_bias = xformers.ops.fmha.BlockDiagonalMask.from_seqlens([N] * B, mask) - x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, attn_bias=attn_bias) - x = x.view(B, -1, C) - x = self.proj(x) - x = self.proj_drop(x) - return x + x = xformers.ops.memory_efficient_attention( + q, k, v, + p=self.attn_drop.p, + attn_bias=attn_bias + ) else: - global competent_attention_implementation - if not competent_attention_implementation: - print("""\nYou should REALLY consider installing/enabling xformers.\nAlternatively, open up ExtraModels/PixArt/models/PixArt_blocks.py and\n- Fix the attention map on line 77 if you know how to\n- Add scaled_dot_product_attention on line 150\n- Send a PR and remove this message on line 32/66-69\n""") - competent_attention_implementation = True - - q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim) - kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim) - k, v = kv.unbind(2) q, k, v = map(lambda t: t.permute(0, 2, 1, 3),(q, k, v),) - attn_mask = None if mask is not None and len(mask) > 1: - # This is probably wrong + # This is most definitely wrong, especially for B>1 attn_mask = torch.zeros( [1, q.shape[1], q.shape[2], v.shape[2]], dtype=q.dtype, @@ -82,26 +78,28 @@ class MultiHeadCrossAttention(nn.Module): ) attn_mask[:, :, (q.shape[2]//2):, mask[0]:] = True attn_mask[:, :, :(q.shape[2]//2), :mask[1]] = True - - x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=self.attn_drop.p) - x = x.permute(0, 2, 1, 3).contiguous() - x = x.view(B, -1, C) - x = self.proj(x) - x = self.proj_drop(x) - return x + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=attn_mask, + dropout_p=self.attn_drop.p + ).permute(0, 2, 1, 3).contiguous() + x = x.view(B, -1, C) + x = self.proj(x) + x = self.proj_drop(x) + return x -class WindowAttention(Attention_): - """Multi-head Attention block with relative position embeddings.""" +class AttentionKVCompress(Attention_): + """Multi-head Attention block with KV token compression and qk norm.""" def __init__( self, dim, num_heads=8, qkv_bias=True, - use_rel_pos=False, - rel_pos_zero_init=True, - input_size=None, + sampling='conv', + sr_ratio=1, + qk_norm=False, **block_kwargs, ): """ @@ -109,56 +107,94 @@ class WindowAttention(Attention_): dim (int): Number of input channels. num_heads (int): Number of attention heads. qkv_bias (bool: If True, add a learnable bias to query, key, value. - rel_pos (bool): If True, add relative positional embeddings to the attention map. - rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. - input_size (int or None): Input resolution for calculating the relative positional - parameter size. """ super().__init__(dim, num_heads=num_heads, qkv_bias=qkv_bias, **block_kwargs) - self.use_rel_pos = use_rel_pos - if self.use_rel_pos: - # initialize relative positional embeddings - self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, self.head_dim)) - self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, self.head_dim)) - - if not rel_pos_zero_init: - nn.init.trunc_normal_(self.rel_pos_h, std=0.02) - nn.init.trunc_normal_(self.rel_pos_w, std=0.02) - - def forward(self, x, mask=None): - B, N, C = x.shape # 2 4096 1152 - qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) - - if model_management.xformers_enabled(): - q, k, v = qkv.unbind(2) - - if getattr(self, 'fp32_attention', False): - q, k, v = q.float(), k.float(), v.float() - - attn_bias = None - if mask is not None: - attn_bias = torch.zeros([B * self.num_heads, q.shape[1], k.shape[1]], dtype=q.dtype, device=q.device) - attn_bias.masked_fill_(mask.squeeze(1).repeat(self.num_heads, 1, 1) == 0, float('-inf')) - # Switch between torch / xformers attention - x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, attn_bias=attn_bias) - x = x.view(B, N, C) - x = self.proj(x) - x = self.proj_drop(x) - return x + self.sampling=sampling # ['conv', 'ave', 'uniform', 'uniform_every'] + self.sr_ratio = sr_ratio + if sr_ratio > 1 and sampling == 'conv': + # Avg Conv Init. + self.sr = nn.Conv2d(dim, dim, groups=dim, kernel_size=sr_ratio, stride=sr_ratio) + self.sr.weight.data.fill_(1/sr_ratio**2) + self.sr.bias.data.zero_() + self.norm = nn.LayerNorm(dim) + if qk_norm: + self.q_norm = nn.LayerNorm(dim) + self.k_norm = nn.LayerNorm(dim) else: - q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() - q = q * self.scale - attn = q @ k.transpose(-2, -1) - attn = attn.softmax(dim=-1) - attn = self.attn_drop(attn) - x = attn @ v + def downsample_2d(self, tensor, H, W, scale_factor, sampling=None): + if sampling is None or scale_factor == 1: + return tensor + B, N, C = tensor.shape - x = x.transpose(1, 2).reshape(B, N, C) - x = self.proj(x) - x = self.proj_drop(x) - return x + if sampling == 'uniform_every': + return tensor[:, ::scale_factor], int(N // scale_factor) + + tensor = tensor.reshape(B, H, W, C).permute(0, 3, 1, 2) + new_H, new_W = int(H / scale_factor), int(W / scale_factor) + new_N = new_H * new_W + + if sampling == 'ave': + tensor = F.interpolate( + tensor, scale_factor=1 / scale_factor, mode='nearest' + ).permute(0, 2, 3, 1) + elif sampling == 'uniform': + tensor = tensor[:, :, ::scale_factor, ::scale_factor].permute(0, 2, 3, 1) + elif sampling == 'conv': + tensor = self.sr(tensor).reshape(B, C, -1).permute(0, 2, 1) + tensor = self.norm(tensor) + else: + raise ValueError + + return tensor.reshape(B, new_N, C).contiguous(), new_N + + def forward(self, x, mask=None, HW=None, block_id=None): + B, N, C = x.shape # 2 4096 1152 + new_N = N + if HW is None: + H = W = int(N ** 0.5) + else: + H, W = HW + qkv = self.qkv(x).reshape(B, N, 3, C) + + q, k, v = qkv.unbind(2) + dtype = q.dtype + q = self.q_norm(q) + k = self.k_norm(k) + + # KV compression + if self.sr_ratio > 1: + k, new_N = self.downsample_2d(k, H, W, self.sr_ratio, sampling=self.sampling) + v, new_N = self.downsample_2d(v, H, W, self.sr_ratio, sampling=self.sampling) + + q = q.reshape(B, N, self.num_heads, C // self.num_heads).to(dtype) + k = k.reshape(B, new_N, self.num_heads, C // self.num_heads).to(dtype) + v = v.reshape(B, new_N, self.num_heads, C // self.num_heads).to(dtype) + + attn_bias = None + if mask is not None: + attn_bias = torch.zeros([B * self.num_heads, q.shape[1], k.shape[1]], dtype=q.dtype, device=q.device) + attn_bias.masked_fill_(mask.squeeze(1).repeat(self.num_heads, 1, 1) == 0, float('-inf')) + # Switch between torch / xformers attention + if model_management.xformers_enabled(): + x = xformers.ops.memory_efficient_attention( + q, k, v, + p=self.attn_drop.p, + attn_bias=attn_bias + ) + else: + q, k, v = map(lambda t: t.transpose(1, 2),(q, k, v),) + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=attn_bias + ).transpose(1, 2).contiguous() + x = x.view(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x ################################################################################# diff --git a/PixArt/nodes.py b/PixArt/nodes.py index 08eeb17..502173d 100644 --- a/PixArt/nodes.py +++ b/PixArt/nodes.py @@ -7,7 +7,6 @@ from comfy import utils from .conf import pixart_conf, pixart_res from .lora import load_pixart_lora from .loader import load_pixart -from .sampler import sample_pixart class PixArtCheckpointLoader: @classmethod @@ -141,44 +140,6 @@ class PixArtControlNetCond: cond[c][1]["cn_hint"] = latent["samples"] * 0.18215 return (cond,) -class PixArtDPMSampler: - """ - The sampler from the reference code. - """ - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "model": ("MODEL", ), - "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), - "cfg": ("FLOAT", {"default": 4.5, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}), - "noise_schedule": (["linear","squaredcos_cap_v2"],{"default":"linear"}), - "noise_schedule_vp": (["linear","discrete"],{"default":"discrete"}), - "positive": ("CONDITIONING", ), - "negative": ("CONDITIONING", ), - "latent_image": ("LATENT", ), - } - } - RETURN_TYPES = ("LATENT",) - FUNCTION = "sample" - CATEGORY = "ExtraModels/PixArt" - TITLE = "PixArt DPM Sampler [Reference]" - - def sample(self, model, seed, steps, cfg, noise_schedule, noise_schedule_vp, positive, negative, latent_image): - samples = sample_pixart( - model = model, - seed = seed, - steps = steps, - cfg = cfg, - positive = positive, - negative = negative, - latent_image = latent_image["samples"], - noise_schedule = noise_schedule, - noise_schedule_vp = noise_schedule_vp, - ) - return ({"samples":samples},) - class PixArtT5TextEncode: """ Reference code, mostly to verify compatibility. @@ -236,7 +197,6 @@ NODE_CLASS_MAPPINGS = { "PixArtCheckpointLoader" : PixArtCheckpointLoader, "PixArtResolutionSelect" : PixArtResolutionSelect, "PixArtLoraLoader" : PixArtLoraLoader, - "PixArtDPMSampler" : PixArtDPMSampler, "PixArtT5TextEncode" : PixArtT5TextEncode, "PixArtResolutionCond" : PixArtResolutionCond, "PixArtControlNetCond" : PixArtControlNetCond, diff --git a/PixArt/sampler.py b/PixArt/sampler.py deleted file mode 100644 index 229a40f..0000000 --- a/PixArt/sampler.py +++ /dev/null @@ -1,72 +0,0 @@ -import torch -from .sampling import gaussian_diffusion as gd -from .sampling.dpm_solver import model_wrapper, DPM_Solver, NoiseScheduleVP - -from comfy.sample import prepare_sampling, prepare_noise -import comfy.utils -import latent_preview - -def sample_pixart(model, seed, steps, cfg, noise_schedule, noise_schedule_vp, positive, negative, latent_image): - """ - Mostly just a wrapper around the reference code. - """ - # prepare model - noise = prepare_noise(latent_image, seed) - real_model, _, _, _, models = prepare_sampling(model, noise.shape, positive, negative, noise_mask=None) - - # negative cond - cond = positive[0][0] - raw_uncond = negative[0][0] - - # Sampler seems to want the same dim for cond and uncond - # truncate uncond to the length of cond - # if shorter, pad uncond with y_null - null_y = real_model.diffusion_model.y_embedder.y_embedding[None].repeat(latent_image.shape[0], 1, 1) - uncond = null_y[:, :cond.shape[1], :] - uncond[:, :raw_uncond.shape[1], :] = raw_uncond[:, :cond.shape[1], :] - if raw_uncond.shape[1] > cond.shape[1]: - print("PixArt: Warning. Your negative prompt is too long.") - uncond[:, -1, :] = raw_uncond[:, -1, :] # add back EOS token - - # Move inputs - cond = cond.to(model.load_device).to(real_model.diffusion_model.dtype) - uncond = uncond.to(model.load_device).to(real_model.diffusion_model.dtype) - noise = noise.to(model.load_device).to(real_model.diffusion_model.dtype) - - # preview - pbar = comfy.utils.ProgressBar(steps) - previewer = latent_preview.get_previewer(model.load_device, model.model.latent_format) - - ## Noise schedule. - betas = torch.tensor(gd.get_named_beta_schedule(noise_schedule, 1000)) - noise_schedule = NoiseScheduleVP(schedule=noise_schedule_vp, betas=betas) - - ## Convert your discrete-time `model` to the continuous-time - ## noise prediction model. Here is an example for a diffusion model - ## `model` with the noise prediction type ("noise") . - model_fn = model_wrapper( - real_model.diffusion_model.forward, - noise_schedule, - model_type="noise", # 'noise', "x_start", "v", "score" - model_kwargs={}, - guidance_type="classifier-free", - condition=cond, - unconditional_condition=uncond, - guidance_scale=cfg, - ) - dpm_solver = DPM_Solver( - model_fn, - noise_schedule, - algorithm_type="dpmsolver++" - ) - samples = dpm_solver.sample( - noise, - steps=steps, - order=2, - skip_type="time_uniform", - method="multistep", - pbar=pbar, - previewer=previewer, - ) - - return samples.detach().cpu().float() * (1 / model.model.latent_format.scale_factor) diff --git a/PixArt/sampling/diffusion_utils.py b/PixArt/sampling/diffusion_utils.py deleted file mode 100644 index cedd4fa..0000000 --- a/PixArt/sampling/diffusion_utils.py +++ /dev/null @@ -1,88 +0,0 @@ -# Modified from OpenAI's diffusion repos -# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py -# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion -# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py - -import numpy as np -import torch as th - - -def normal_kl(mean1, logvar1, mean2, logvar2): - """ - Compute the KL divergence between two gaussians. - Shapes are automatically broadcasted, so batches can be compared to - scalars, among other use cases. - """ - tensor = None - for obj in (mean1, logvar1, mean2, logvar2): - if isinstance(obj, th.Tensor): - tensor = obj - break - assert tensor is not None, "at least one argument must be a Tensor" - - # Force variances to be Tensors. Broadcasting helps convert scalars to - # Tensors, but it does not work for th.exp(). - logvar1, logvar2 = [ - x if isinstance(x, th.Tensor) else th.tensor(x, device=tensor.device) - for x in (logvar1, logvar2) - ] - - return 0.5 * ( - -1.0 - + logvar2 - - logvar1 - + th.exp(logvar1 - logvar2) - + ((mean1 - mean2) ** 2) * th.exp(-logvar2) - ) - - -def approx_standard_normal_cdf(x): - """ - A fast approximation of the cumulative distribution function of the - standard normal. - """ - return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) - - -def continuous_gaussian_log_likelihood(x, *, means, log_scales): - """ - Compute the log-likelihood of a continuous Gaussian distribution. - :param x: the targets - :param means: the Gaussian mean Tensor. - :param log_scales: the Gaussian log stddev Tensor. - :return: a tensor like x of log probabilities (in nats). - """ - centered_x = x - means - inv_stdv = th.exp(-log_scales) - normalized_x = centered_x * inv_stdv - log_probs = th.distributions.Normal(th.zeros_like(x), th.ones_like(x)).log_prob(normalized_x) - return log_probs - - -def discretized_gaussian_log_likelihood(x, *, means, log_scales): - """ - Compute the log-likelihood of a Gaussian distribution discretizing to a - given image. - :param x: the target images. It is assumed that this was uint8 values, - rescaled to the range [-1, 1]. - :param means: the Gaussian mean Tensor. - :param log_scales: the Gaussian log stddev Tensor. - :return: a tensor like x of log probabilities (in nats). - """ - assert x.shape == means.shape == log_scales.shape - centered_x = x - means - inv_stdv = th.exp(-log_scales) - plus_in = inv_stdv * (centered_x + 1.0 / 255.0) - cdf_plus = approx_standard_normal_cdf(plus_in) - min_in = inv_stdv * (centered_x - 1.0 / 255.0) - cdf_min = approx_standard_normal_cdf(min_in) - log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) - log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) - cdf_delta = cdf_plus - cdf_min - log_probs = th.where( - x < -0.999, - log_cdf_plus, - th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), - ) - assert log_probs.shape == x.shape - return log_probs diff --git a/PixArt/sampling/dpm_solver.py b/PixArt/sampling/dpm_solver.py deleted file mode 100644 index 7538b1b..0000000 --- a/PixArt/sampling/dpm_solver.py +++ /dev/null @@ -1,1356 +0,0 @@ -import torch -from tqdm import tqdm - - -class NoiseScheduleVP: - def __init__( - self, - schedule='discrete', - betas=None, - alphas_cumprod=None, - continuous_beta_0=0.1, - continuous_beta_1=20., - dtype=torch.float32, - ): - """Create a wrapper class for the forward SDE (VP type). - - *** - Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t. - We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images. - *** - - The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ). - We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper). - Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have: - - log_alpha_t = self.marginal_log_mean_coeff(t) - sigma_t = self.marginal_std(t) - lambda_t = self.marginal_lambda(t) - - Moreover, as lambda(t) is an invertible function, we also support its inverse function: - - t = self.inverse_lambda(lambda_t) - - =============================================================== - - We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]). - - 1. For discrete-time DPMs: - - For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by: - t_i = (i + 1) / N - e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1. - We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3. - - Args: - betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details) - alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details) - - Note that we always have alphas_cumprod = cumprod(1 - betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`. - - **Important**: Please pay special attention for the args for `alphas_cumprod`: - The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that - q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ). - Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have - alpha_{t_n} = \sqrt{\hat{alpha_n}}, - and - log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}). - - - 2. For continuous-time DPMs: - - We support the linear VPSDE for the continuous time setting. The hyperparameters for the noise - schedule are the default settings in Yang Song's ScoreSDE: - - Args: - beta_min: A `float` number. The smallest beta for the linear schedule. - beta_max: A `float` number. The largest beta for the linear schedule. - T: A `float` number. The ending time of the forward process. - - =============================================================== - - Args: - schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs, - 'linear' for continuous-time DPMs. - Returns: - A wrapper object of the forward SDE (VP type). - - =============================================================== - - Example: - - # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1): - >>> ns = NoiseScheduleVP('discrete', betas=betas) - - # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1): - >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod) - - # For continuous-time DPMs (VPSDE), linear schedule: - >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.) - - """ - - if schedule not in ['discrete', 'linear']: - raise ValueError( - "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear'".format(schedule)) - - self.schedule = schedule - if schedule == 'discrete': - if betas is not None: - log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0) - else: - assert alphas_cumprod is not None - log_alphas = 0.5 * torch.log(alphas_cumprod) - self.T = 1. - self.log_alpha_array = self.numerical_clip_alpha(log_alphas).reshape((1, -1,)).to(dtype=dtype) - self.total_N = self.log_alpha_array.shape[1] - self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1)).to(dtype=dtype) - else: - self.T = 1. - self.total_N = 1000 - self.beta_0 = continuous_beta_0 - self.beta_1 = continuous_beta_1 - - def numerical_clip_alpha(self, log_alphas, clipped_lambda=-5.1): - """ - For some beta schedules such as cosine schedule, the log-SNR has numerical isssues. - We clip the log-SNR near t=T within -5.1 to ensure the stability. - Such a trick is very useful for diffusion models with the cosine schedule, such as i-DDPM, guided-diffusion and GLIDE. - """ - log_sigmas = 0.5 * torch.log(1. - torch.exp(2. * log_alphas)) - lambs = log_alphas - log_sigmas - idx = torch.searchsorted(torch.flip(lambs, [0]), clipped_lambda) - if idx > 0: - log_alphas = log_alphas[:-idx] - return log_alphas - - def marginal_log_mean_coeff(self, t): - """ - Compute log(alpha_t) of a given continuous-time label t in [0, T]. - """ - if self.schedule == 'discrete': - return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device), - self.log_alpha_array.to(t.device)).reshape((-1)) - elif self.schedule == 'linear': - return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0 - - def marginal_alpha(self, t): - """ - Compute alpha_t of a given continuous-time label t in [0, T]. - """ - return torch.exp(self.marginal_log_mean_coeff(t)) - - def marginal_std(self, t): - """ - Compute sigma_t of a given continuous-time label t in [0, T]. - """ - return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t))) - - def marginal_lambda(self, t): - """ - Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T]. - """ - log_mean_coeff = self.marginal_log_mean_coeff(t) - log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff)) - return log_mean_coeff - log_std - - def inverse_lambda(self, lamb): - """ - Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t. - """ - if self.schedule == 'linear': - tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb)) - Delta = self.beta_0 ** 2 + tmp - return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0) - elif self.schedule == 'discrete': - log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb) - t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]), - torch.flip(self.t_array.to(lamb.device), [1])) - return t.reshape((-1,)) - - -def model_wrapper( - model, - noise_schedule, - model_type="noise", - model_kwargs={}, - guidance_type="uncond", - condition=None, - unconditional_condition=None, - guidance_scale=1., - classifier_fn=None, - classifier_kwargs={}, -): - """Create a wrapper function for the noise prediction model. - - DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to - firstly wrap the model function to a noise prediction model that accepts the continuous time as the input. - - We support four types of the diffusion model by setting `model_type`: - - 1. "noise": noise prediction model. (Trained by predicting noise). - - 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0). - - 3. "v": velocity prediction model. (Trained by predicting the velocity). - The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2]. - - [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models." - arXiv preprint arXiv:2202.00512 (2022). - [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models." - arXiv preprint arXiv:2210.02303 (2022). - - 4. "score": marginal score function. (Trained by denoising score matching). - Note that the score function and the noise prediction model follows a simple relationship: - ``` - noise(x_t, t) = -sigma_t * score(x_t, t) - ``` - - We support three types of guided sampling by DPMs by setting `guidance_type`: - 1. "uncond": unconditional sampling by DPMs. - The input `model` has the following format: - `` - model(x, t_input, **model_kwargs) -> noise | x_start | v | score - `` - - 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier. - The input `model` has the following format: - `` - model(x, t_input, **model_kwargs) -> noise | x_start | v | score - `` - - The input `classifier_fn` has the following format: - `` - classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond) - `` - - [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis," - in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794. - - 3. "classifier-free": classifier-free guidance sampling by conditional DPMs. - The input `model` has the following format: - `` - model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score - `` - And if cond == `unconditional_condition`, the model output is the unconditional DPM output. - - [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance." - arXiv preprint arXiv:2207.12598 (2022). - - - The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999) - or continuous-time labels (i.e. epsilon to T). - - We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise: - `` - def model_fn(x, t_continuous) -> noise: - t_input = get_model_input_time(t_continuous) - return noise_pred(model, x, t_input, **model_kwargs) - `` - where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver. - - =============================================================== - - Args: - model: A diffusion model with the corresponding format described above. - noise_schedule: A noise schedule object, such as NoiseScheduleVP. - model_type: A `str`. The parameterization type of the diffusion model. - "noise" or "x_start" or "v" or "score". - model_kwargs: A `dict`. A dict for the other inputs of the model function. - guidance_type: A `str`. The type of the guidance for sampling. - "uncond" or "classifier" or "classifier-free". - condition: A pytorch tensor. The condition for the guided sampling. - Only used for "classifier" or "classifier-free" guidance type. - unconditional_condition: A pytorch tensor. The condition for the unconditional sampling. - Only used for "classifier-free" guidance type. - guidance_scale: A `float`. The scale for the guided sampling. - classifier_fn: A classifier function. Only used for the classifier guidance. - classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function. - Returns: - A noise prediction model that accepts the noised data and the continuous time as the inputs. - """ - - def get_model_input_time(t_continuous): - """ - Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time. - For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N]. - For continuous-time DPMs, we just use `t_continuous`. - """ - if noise_schedule.schedule == 'discrete': - return (t_continuous - 1. / noise_schedule.total_N) * 1000. - else: - return t_continuous - - def noise_pred_fn(x, t_continuous, cond=None): - t_input = get_model_input_time(t_continuous) - if cond is None: - output = model( - x = x, - timesteps = t_input, - context = None, - y = None, - **model_kwargs - ) - else: - output = model( - x = x, - timesteps = t_input, - context = cond, - y = None, - **model_kwargs - ) - if model_type == "noise": - return output - elif model_type == "x_start": - alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous) - return (x - expand_dims(alpha_t, x.dim()) * output) / expand_dims(sigma_t, x.dim()) - elif model_type == "v": - alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous) - return expand_dims(alpha_t, x.dim()) * output + expand_dims(sigma_t, x.dim()) * x - elif model_type == "score": - sigma_t = noise_schedule.marginal_std(t_continuous) - return -expand_dims(sigma_t, x.dim()) * output - - def cond_grad_fn(x, t_input): - """ - Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t). - """ - with torch.enable_grad(): - x_in = x.detach().requires_grad_(True) - log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs) - return torch.autograd.grad(log_prob.sum(), x_in)[0] - - def model_fn(x, t_continuous): - """ - The noise predicition model function that is used for DPM-Solver. - """ - if guidance_type == "uncond": - return noise_pred_fn(x, t_continuous) - elif guidance_type == "classifier": - assert classifier_fn is not None - t_input = get_model_input_time(t_continuous) - cond_grad = cond_grad_fn(x, t_input) - sigma_t = noise_schedule.marginal_std(t_continuous) - noise = noise_pred_fn(x, t_continuous) - return noise - guidance_scale * expand_dims(sigma_t, x.dim()) * cond_grad - elif guidance_type == "classifier-free": - if guidance_scale == 1. or unconditional_condition is None: - return noise_pred_fn(x, t_continuous, cond=condition) - else: - x_in = torch.cat([x] * 2) - t_in = torch.cat([t_continuous] * 2) - c_in = torch.cat([unconditional_condition, condition]) - noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2) - return noise_uncond + guidance_scale * (noise - noise_uncond) - - assert model_type in ["noise", "x_start", "v", "score"] - assert guidance_type in ["uncond", "classifier", "classifier-free"] - return model_fn - - -class DPM_Solver: - def __init__( - self, - model_fn, - noise_schedule, - algorithm_type="dpmsolver++", - correcting_x0_fn=None, - correcting_xt_fn=None, - thresholding_max_val=1., - dynamic_thresholding_ratio=0.995, - ): - """Construct a DPM-Solver. - - We support both DPM-Solver (`algorithm_type="dpmsolver"`) and DPM-Solver++ (`algorithm_type="dpmsolver++"`). - - We also support the "dynamic thresholding" method in Imagen[1]. For pixel-space diffusion models, you - can set both `algorithm_type="dpmsolver++"` and `correcting_x0_fn="dynamic_thresholding"` to use the - dynamic thresholding. The "dynamic thresholding" can greatly improve the sample quality for pixel-space - DPMs with large guidance scales. Note that the thresholding method is **unsuitable** for latent-space - DPMs (such as stable-diffusion). - - To support advanced algorithms in image-to-image applications, we also support corrector functions for - both x0 and xt. - - Args: - model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]): - `` - def model_fn(x, t_continuous): - return noise - `` - The shape of `x` is `(batch_size, **shape)`, and the shape of `t_continuous` is `(batch_size,)`. - noise_schedule: A noise schedule object, such as NoiseScheduleVP. - algorithm_type: A `str`. Either "dpmsolver" or "dpmsolver++". - correcting_x0_fn: A `str` or a function with the following format: - ``` - def correcting_x0_fn(x0, t): - x0_new = ... - return x0_new - ``` - This function is to correct the outputs of the data prediction model at each sampling step. e.g., - ``` - x0_pred = data_pred_model(xt, t) - if correcting_x0_fn is not None: - x0_pred = correcting_x0_fn(x0_pred, t) - xt_1 = update(x0_pred, xt, t) - ``` - If `correcting_x0_fn="dynamic_thresholding"`, we use the dynamic thresholding proposed in Imagen[1]. - correcting_xt_fn: A function with the following format: - ``` - def correcting_xt_fn(xt, t, step): - x_new = ... - return x_new - ``` - This function is to correct the intermediate samples xt at each sampling step. e.g., - ``` - xt = ... - xt = correcting_xt_fn(xt, t, step) - ``` - thresholding_max_val: A `float`. The max value for thresholding. - Valid only when use `dpmsolver++` and `correcting_x0_fn="dynamic_thresholding"`. - dynamic_thresholding_ratio: A `float`. The ratio for dynamic thresholding (see Imagen[1] for details). - Valid only when use `dpmsolver++` and `correcting_x0_fn="dynamic_thresholding"`. - - [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, - Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models - with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b. - """ - self.model = lambda x, t: model_fn(x, t.expand((x.shape[0]))) - self.noise_schedule = noise_schedule - assert algorithm_type in ["dpmsolver", "dpmsolver++"] - self.algorithm_type = algorithm_type - if correcting_x0_fn == "dynamic_thresholding": - self.correcting_x0_fn = self.dynamic_thresholding_fn - else: - self.correcting_x0_fn = correcting_x0_fn - self.correcting_xt_fn = correcting_xt_fn - self.dynamic_thresholding_ratio = dynamic_thresholding_ratio - self.thresholding_max_val = thresholding_max_val - - def dynamic_thresholding_fn(self, x0, t): - """ - The dynamic thresholding method. - """ - dims = x0.dim() - p = self.dynamic_thresholding_ratio - s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1) - s = expand_dims(torch.maximum(s, self.thresholding_max_val * torch.ones_like(s).to(s.device)), dims) - x0 = torch.clamp(x0, -s, s) / s - return x0 - - def noise_prediction_fn(self, x, t): - """ - Return the noise prediction model. - """ - return self.model(x, t) - - def data_prediction_fn(self, x, t): - """ - Return the data prediction model (with corrector). - """ - noise = self.noise_prediction_fn(x, t) - alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t) - x0 = (x - sigma_t * noise) / alpha_t - if self.correcting_x0_fn is not None: - x0 = self.correcting_x0_fn(x0, t) - return x0 - - def model_fn(self, x, t): - """ - Convert the model to the noise prediction model or the data prediction model. - """ - if self.algorithm_type == "dpmsolver++": - return self.data_prediction_fn(x, t) - else: - return self.noise_prediction_fn(x, t) - - def get_time_steps(self, skip_type, t_T, t_0, N, device): - """Compute the intermediate time steps for sampling. - - Args: - skip_type: A `str`. The type for the spacing of the time steps. We support three types: - - 'logSNR': uniform logSNR for the time steps. - - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.) - - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.) - t_T: A `float`. The starting time of the sampling (default is T). - t_0: A `float`. The ending time of the sampling (default is epsilon). - N: A `int`. The total number of the spacing of the time steps. - device: A torch device. - Returns: - A pytorch tensor of the time steps, with the shape (N + 1,). - """ - if skip_type == 'logSNR': - lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device)) - lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device)) - logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device) - return self.noise_schedule.inverse_lambda(logSNR_steps) - elif skip_type == 'time_uniform': - return torch.linspace(t_T, t_0, N + 1).to(device) - elif skip_type == 'time_quadratic': - t_order = 2 - t = torch.linspace(t_T ** (1. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device) - return t - else: - raise ValueError( - "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type)) - - def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device): - """ - Get the order of each step for sampling by the singlestep DPM-Solver. - - We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast". - Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is: - - If order == 1: - We take `steps` of DPM-Solver-1 (i.e. DDIM). - - If order == 2: - - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling. - - If steps % 2 == 0, we use K steps of DPM-Solver-2. - - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1. - - If order == 3: - - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling. - - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1. - - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1. - - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2. - - ============================================ - Args: - order: A `int`. The max order for the solver (2 or 3). - steps: A `int`. The total number of function evaluations (NFE). - skip_type: A `str`. The type for the spacing of the time steps. We support three types: - - 'logSNR': uniform logSNR for the time steps. - - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.) - - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.) - t_T: A `float`. The starting time of the sampling (default is T). - t_0: A `float`. The ending time of the sampling (default is epsilon). - device: A torch device. - Returns: - orders: A list of the solver order of each step. - """ - if order == 3: - K = steps // 3 + 1 - if steps % 3 == 0: - orders = [3, ] * (K - 2) + [2, 1] - elif steps % 3 == 1: - orders = [3, ] * (K - 1) + [1] - else: - orders = [3, ] * (K - 1) + [2] - elif order == 2: - if steps % 2 == 0: - K = steps // 2 - orders = [2, ] * K - else: - K = steps // 2 + 1 - orders = [2, ] * (K - 1) + [1] - elif order == 1: - K = 1 - orders = [1, ] * steps - else: - raise ValueError("'order' must be '1' or '2' or '3'.") - if skip_type == 'logSNR': - # To reproduce the results in DPM-Solver paper - timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device) - else: - timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[ - torch.cumsum(torch.tensor([0, ] + orders), 0).to(device)] - return timesteps_outer, orders - - def denoise_to_zero_fn(self, x, s): - """ - Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization. - """ - return self.data_prediction_fn(x, s) - - def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False): - """ - DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`. - - Args: - x: A pytorch tensor. The initial value at time `s`. - s: A pytorch tensor. The starting time, with the shape (1,). - t: A pytorch tensor. The ending time, with the shape (1,). - model_s: A pytorch tensor. The model function evaluated at time `s`. - If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. - return_intermediate: A `bool`. If true, also return the model value at time `s`. - Returns: - x_t: A pytorch tensor. The approximated solution at time `t`. - """ - ns = self.noise_schedule - dims = x.dim() - lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t) - h = lambda_t - lambda_s - log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t) - sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t) - alpha_t = torch.exp(log_alpha_t) - - if self.algorithm_type == "dpmsolver++": - phi_1 = torch.expm1(-h) - if model_s is None: - model_s = self.model_fn(x, s) - x_t = ( - sigma_t / sigma_s * x - - alpha_t * phi_1 * model_s - ) - if return_intermediate: - return x_t, {'model_s': model_s} - else: - return x_t - else: - phi_1 = torch.expm1(h) - if model_s is None: - model_s = self.model_fn(x, s) - x_t = ( - torch.exp(log_alpha_t - log_alpha_s) * x - - (sigma_t * phi_1) * model_s - ) - if return_intermediate: - return x_t, {'model_s': model_s} - else: - return x_t - - def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s=None, return_intermediate=False, - solver_type='dpmsolver'): - """ - Singlestep solver DPM-Solver-2 from time `s` to time `t`. - - Args: - x: A pytorch tensor. The initial value at time `s`. - s: A pytorch tensor. The starting time, with the shape (1,). - t: A pytorch tensor. The ending time, with the shape (1,). - r1: A `float`. The hyperparameter of the second-order solver. - model_s: A pytorch tensor. The model function evaluated at time `s`. - If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. - return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time). - solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. - The type slightly impacts the performance. We recommend to use 'dpmsolver' type. - Returns: - x_t: A pytorch tensor. The approximated solution at time `t`. - """ - if solver_type not in ['dpmsolver', 'taylor']: - raise ValueError("'solver_type' must be either 'dpmsolver' or 'taylor', got {}".format(solver_type)) - if r1 is None: - r1 = 0.5 - ns = self.noise_schedule - lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t) - h = lambda_t - lambda_s - lambda_s1 = lambda_s + r1 * h - s1 = ns.inverse_lambda(lambda_s1) - log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff( - s1), ns.marginal_log_mean_coeff(t) - sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t) - alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t) - - if self.algorithm_type == "dpmsolver++": - phi_11 = torch.expm1(-r1 * h) - phi_1 = torch.expm1(-h) - - if model_s is None: - model_s = self.model_fn(x, s) - x_s1 = ( - (sigma_s1 / sigma_s) * x - - (alpha_s1 * phi_11) * model_s - ) - model_s1 = self.model_fn(x_s1, s1) - if solver_type == 'dpmsolver': - x_t = ( - (sigma_t / sigma_s) * x - - (alpha_t * phi_1) * model_s - - (0.5 / r1) * (alpha_t * phi_1) * (model_s1 - model_s) - ) - elif solver_type == 'taylor': - x_t = ( - (sigma_t / sigma_s) * x - - (alpha_t * phi_1) * model_s - + (1. / r1) * (alpha_t * (phi_1 / h + 1.)) * (model_s1 - model_s) - ) - else: - phi_11 = torch.expm1(r1 * h) - phi_1 = torch.expm1(h) - - if model_s is None: - model_s = self.model_fn(x, s) - x_s1 = ( - torch.exp(log_alpha_s1 - log_alpha_s) * x - - (sigma_s1 * phi_11) * model_s - ) - model_s1 = self.model_fn(x_s1, s1) - if solver_type == 'dpmsolver': - x_t = ( - torch.exp(log_alpha_t - log_alpha_s) * x - - (sigma_t * phi_1) * model_s - - (0.5 / r1) * (sigma_t * phi_1) * (model_s1 - model_s) - ) - elif solver_type == 'taylor': - x_t = ( - torch.exp(log_alpha_t - log_alpha_s) * x - - (sigma_t * phi_1) * model_s - - (1. / r1) * (sigma_t * (phi_1 / h - 1.)) * (model_s1 - model_s) - ) - if return_intermediate: - return x_t, {'model_s': model_s, 'model_s1': model_s1} - else: - return x_t - - def singlestep_dpm_solver_third_update(self, x, s, t, r1=1. / 3., r2=2. / 3., model_s=None, model_s1=None, - return_intermediate=False, solver_type='dpmsolver'): - """ - Singlestep solver DPM-Solver-3 from time `s` to time `t`. - - Args: - x: A pytorch tensor. The initial value at time `s`. - s: A pytorch tensor. The starting time, with the shape (1,). - t: A pytorch tensor. The ending time, with the shape (1,). - r1: A `float`. The hyperparameter of the third-order solver. - r2: A `float`. The hyperparameter of the third-order solver. - model_s: A pytorch tensor. The model function evaluated at time `s`. - If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. - model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`). - If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it. - return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times). - solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. - The type slightly impacts the performance. We recommend to use 'dpmsolver' type. - Returns: - x_t: A pytorch tensor. The approximated solution at time `t`. - """ - if solver_type not in ['dpmsolver', 'taylor']: - raise ValueError("'solver_type' must be either 'dpmsolver' or 'taylor', got {}".format(solver_type)) - if r1 is None: - r1 = 1. / 3. - if r2 is None: - r2 = 2. / 3. - ns = self.noise_schedule - lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t) - h = lambda_t - lambda_s - lambda_s1 = lambda_s + r1 * h - lambda_s2 = lambda_s + r2 * h - s1 = ns.inverse_lambda(lambda_s1) - s2 = ns.inverse_lambda(lambda_s2) - log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff( - s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t) - sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std( - s2), ns.marginal_std(t) - alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t) - - if self.algorithm_type == "dpmsolver++": - phi_11 = torch.expm1(-r1 * h) - phi_12 = torch.expm1(-r2 * h) - phi_1 = torch.expm1(-h) - phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1. - phi_2 = phi_1 / h + 1. - phi_3 = phi_2 / h - 0.5 - - if model_s is None: - model_s = self.model_fn(x, s) - if model_s1 is None: - x_s1 = ( - (sigma_s1 / sigma_s) * x - - (alpha_s1 * phi_11) * model_s - ) - model_s1 = self.model_fn(x_s1, s1) - x_s2 = ( - (sigma_s2 / sigma_s) * x - - (alpha_s2 * phi_12) * model_s - + r2 / r1 * (alpha_s2 * phi_22) * (model_s1 - model_s) - ) - model_s2 = self.model_fn(x_s2, s2) - if solver_type == 'dpmsolver': - x_t = ( - (sigma_t / sigma_s) * x - - (alpha_t * phi_1) * model_s - + (1. / r2) * (alpha_t * phi_2) * (model_s2 - model_s) - ) - elif solver_type == 'taylor': - D1_0 = (1. / r1) * (model_s1 - model_s) - D1_1 = (1. / r2) * (model_s2 - model_s) - D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1) - D2 = 2. * (D1_1 - D1_0) / (r2 - r1) - x_t = ( - (sigma_t / sigma_s) * x - - (alpha_t * phi_1) * model_s - + (alpha_t * phi_2) * D1 - - (alpha_t * phi_3) * D2 - ) - else: - phi_11 = torch.expm1(r1 * h) - phi_12 = torch.expm1(r2 * h) - phi_1 = torch.expm1(h) - phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1. - phi_2 = phi_1 / h - 1. - phi_3 = phi_2 / h - 0.5 - - if model_s is None: - model_s = self.model_fn(x, s) - if model_s1 is None: - x_s1 = ( - (torch.exp(log_alpha_s1 - log_alpha_s)) * x - - (sigma_s1 * phi_11) * model_s - ) - model_s1 = self.model_fn(x_s1, s1) - x_s2 = ( - (torch.exp(log_alpha_s2 - log_alpha_s)) * x - - (sigma_s2 * phi_12) * model_s - - r2 / r1 * (sigma_s2 * phi_22) * (model_s1 - model_s) - ) - model_s2 = self.model_fn(x_s2, s2) - if solver_type == 'dpmsolver': - x_t = ( - (torch.exp(log_alpha_t - log_alpha_s)) * x - - (sigma_t * phi_1) * model_s - - (1. / r2) * (sigma_t * phi_2) * (model_s2 - model_s) - ) - elif solver_type == 'taylor': - D1_0 = (1. / r1) * (model_s1 - model_s) - D1_1 = (1. / r2) * (model_s2 - model_s) - D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1) - D2 = 2. * (D1_1 - D1_0) / (r2 - r1) - x_t = ( - (torch.exp(log_alpha_t - log_alpha_s)) * x - - (sigma_t * phi_1) * model_s - - (sigma_t * phi_2) * D1 - - (sigma_t * phi_3) * D2 - ) - - if return_intermediate: - return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2} - else: - return x_t - - def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpmsolver"): - """ - Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`. - - Args: - x: A pytorch tensor. The initial value at time `s`. - model_prev_list: A list of pytorch tensor. The previous computed model values. - t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,) - t: A pytorch tensor. The ending time, with the shape (1,). - solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. - The type slightly impacts the performance. We recommend to use 'dpmsolver' type. - Returns: - x_t: A pytorch tensor. The approximated solution at time `t`. - """ - if solver_type not in ['dpmsolver', 'taylor']: - raise ValueError("'solver_type' must be either 'dpmsolver' or 'taylor', got {}".format(solver_type)) - ns = self.noise_schedule - model_prev_1, model_prev_0 = model_prev_list[-2], model_prev_list[-1] - t_prev_1, t_prev_0 = t_prev_list[-2], t_prev_list[-1] - lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda( - t_prev_0), ns.marginal_lambda(t) - log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t) - sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t) - alpha_t = torch.exp(log_alpha_t) - - h_0 = lambda_prev_0 - lambda_prev_1 - h = lambda_t - lambda_prev_0 - r0 = h_0 / h - D1_0 = (1. / r0) * (model_prev_0 - model_prev_1) - if self.algorithm_type == "dpmsolver++": - phi_1 = torch.expm1(-h) - if solver_type == 'dpmsolver': - x_t = ( - (sigma_t / sigma_prev_0) * x - - (alpha_t * phi_1) * model_prev_0 - - 0.5 * (alpha_t * phi_1) * D1_0 - ) - elif solver_type == 'taylor': - x_t = ( - (sigma_t / sigma_prev_0) * x - - (alpha_t * phi_1) * model_prev_0 - + (alpha_t * (phi_1 / h + 1.)) * D1_0 - ) - else: - phi_1 = torch.expm1(h) - if solver_type == 'dpmsolver': - x_t = ( - (torch.exp(log_alpha_t - log_alpha_prev_0)) * x - - (sigma_t * phi_1) * model_prev_0 - - 0.5 * (sigma_t * phi_1) * D1_0 - ) - elif solver_type == 'taylor': - x_t = ( - (torch.exp(log_alpha_t - log_alpha_prev_0)) * x - - (sigma_t * phi_1) * model_prev_0 - - (sigma_t * (phi_1 / h - 1.)) * D1_0 - ) - return x_t - - def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpmsolver'): - """ - Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`. - - Args: - x: A pytorch tensor. The initial value at time `s`. - model_prev_list: A list of pytorch tensor. The previous computed model values. - t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,) - t: A pytorch tensor. The ending time, with the shape (1,). - solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. - The type slightly impacts the performance. We recommend to use 'dpmsolver' type. - Returns: - x_t: A pytorch tensor. The approximated solution at time `t`. - """ - ns = self.noise_schedule - model_prev_2, model_prev_1, model_prev_0 = model_prev_list - t_prev_2, t_prev_1, t_prev_0 = t_prev_list - lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda( - t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t) - log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t) - sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t) - alpha_t = torch.exp(log_alpha_t) - - h_1 = lambda_prev_1 - lambda_prev_2 - h_0 = lambda_prev_0 - lambda_prev_1 - h = lambda_t - lambda_prev_0 - r0, r1 = h_0 / h, h_1 / h - D1_0 = (1. / r0) * (model_prev_0 - model_prev_1) - D1_1 = (1. / r1) * (model_prev_1 - model_prev_2) - D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1) - D2 = (1. / (r0 + r1)) * (D1_0 - D1_1) - if self.algorithm_type == "dpmsolver++": - phi_1 = torch.expm1(-h) - phi_2 = phi_1 / h + 1. - phi_3 = phi_2 / h - 0.5 - x_t = ( - (sigma_t / sigma_prev_0) * x - - (alpha_t * phi_1) * model_prev_0 - + (alpha_t * phi_2) * D1 - - (alpha_t * phi_3) * D2 - ) - else: - phi_1 = torch.expm1(h) - phi_2 = phi_1 / h - 1. - phi_3 = phi_2 / h - 0.5 - x_t = ( - (torch.exp(log_alpha_t - log_alpha_prev_0)) * x - - (sigma_t * phi_1) * model_prev_0 - - (sigma_t * phi_2) * D1 - - (sigma_t * phi_3) * D2 - ) - return x_t - - def singlestep_dpm_solver_update(self, x, s, t, order, return_intermediate=False, solver_type='dpmsolver', r1=None, - r2=None): - """ - Singlestep DPM-Solver with the order `order` from time `s` to time `t`. - - Args: - x: A pytorch tensor. The initial value at time `s`. - s: A pytorch tensor. The starting time, with the shape (1,). - t: A pytorch tensor. The ending time, with the shape (1,). - order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3. - return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times). - solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. - The type slightly impacts the performance. We recommend to use 'dpmsolver' type. - r1: A `float`. The hyperparameter of the second-order or third-order solver. - r2: A `float`. The hyperparameter of the third-order solver. - Returns: - x_t: A pytorch tensor. The approximated solution at time `t`. - """ - if order == 1: - return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate) - elif order == 2: - return self.singlestep_dpm_solver_second_update(x, s, t, return_intermediate=return_intermediate, - solver_type=solver_type, r1=r1) - elif order == 3: - return self.singlestep_dpm_solver_third_update(x, s, t, return_intermediate=return_intermediate, - solver_type=solver_type, r1=r1, r2=r2) - else: - raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order)) - - def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpmsolver'): - """ - Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`. - - Args: - x: A pytorch tensor. The initial value at time `s`. - model_prev_list: A list of pytorch tensor. The previous computed model values. - t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,) - t: A pytorch tensor. The ending time, with the shape (1,). - order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3. - solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. - The type slightly impacts the performance. We recommend to use 'dpmsolver' type. - Returns: - x_t: A pytorch tensor. The approximated solution at time `t`. - """ - if order == 1: - return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1]) - elif order == 2: - return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type) - elif order == 3: - return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type) - else: - raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order)) - - def dpm_solver_adaptive(self, x, order, t_T, t_0, h_init=0.05, atol=0.0078, rtol=0.05, theta=0.9, t_err=1e-5, - solver_type='dpmsolver'): - """ - The adaptive step size solver based on singlestep DPM-Solver. - - Args: - x: A pytorch tensor. The initial value at time `t_T`. - order: A `int`. The (higher) order of the solver. We only support order == 2 or 3. - t_T: A `float`. The starting time of the sampling (default is T). - t_0: A `float`. The ending time of the sampling (default is epsilon). - h_init: A `float`. The initial step size (for logSNR). - atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1]. - rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05. - theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1]. - t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the - current time and `t_0` is less than `t_err`. The default setting is 1e-5. - solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. - The type slightly impacts the performance. We recommend to use 'dpmsolver' type. - Returns: - x_0: A pytorch tensor. The approximated solution at time `t_0`. - - [1] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, "Gotta go fast when generating data with score-based models," arXiv preprint arXiv:2105.14080, 2021. - """ - ns = self.noise_schedule - s = t_T * torch.ones((1,)).to(x) - lambda_s = ns.marginal_lambda(s) - lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x)) - h = h_init * torch.ones_like(s).to(x) - x_prev = x - nfe = 0 - if order == 2: - r1 = 0.5 - lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True) - higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1, - solver_type=solver_type, - **kwargs) - elif order == 3: - r1, r2 = 1. / 3., 2. / 3. - lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1, - return_intermediate=True, - solver_type=solver_type) - higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(x, s, t, r1=r1, r2=r2, - solver_type=solver_type, - **kwargs) - else: - raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order)) - while torch.abs((s - t_0)).mean() > t_err: - t = ns.inverse_lambda(lambda_s + h) - x_lower, lower_noise_kwargs = lower_update(x, s, t) - x_higher = higher_update(x, s, t, **lower_noise_kwargs) - delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev))) - norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True)) - E = norm_fn((x_higher - x_lower) / delta).max() - if torch.all(E <= 1.): - x = x_higher - s = t - x_prev = x_lower - lambda_s = ns.marginal_lambda(s) - h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s) - nfe += order - print('adaptive solver nfe', nfe) - return x - - def add_noise(self, x, t, noise=None): - """ - Compute the noised input xt = alpha_t * x + sigma_t * noise. - - Args: - x: A `torch.Tensor` with shape `(batch_size, *shape)`. - t: A `torch.Tensor` with shape `(t_size,)`. - Returns: - xt with shape `(t_size, batch_size, *shape)`. - """ - alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t) - if noise is None: - noise = torch.randn((t.shape[0], *x.shape), device=x.device) - x = x.reshape((-1, *x.shape)) - xt = expand_dims(alpha_t, x.dim()) * x + expand_dims(sigma_t, x.dim()) * noise - if t.shape[0] == 1: - return xt.squeeze(0) - else: - return xt - - def inverse(self, x, steps=20, t_start=None, t_end=None, order=2, skip_type='time_uniform', - method='multistep', lower_order_final=True, denoise_to_zero=False, solver_type='dpmsolver', - atol=0.0078, rtol=0.05, return_intermediate=False, - ): - """ - Inverse the sample `x` from time `t_start` to `t_end` by DPM-Solver. - For discrete-time DPMs, we use `t_start=1/N`, where `N` is the total time steps during training. - """ - t_0 = 1. / self.noise_schedule.total_N if t_start is None else t_start - t_T = self.noise_schedule.T if t_end is None else t_end - assert t_0 > 0 and t_T > 0, "Time range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array" - return self.sample(x, steps=steps, t_start=t_0, t_end=t_T, order=order, skip_type=skip_type, - method=method, lower_order_final=lower_order_final, denoise_to_zero=denoise_to_zero, - solver_type=solver_type, - atol=atol, rtol=rtol, return_intermediate=return_intermediate) - - def sample(self, x, steps=20, t_start=None, t_end=None, order=2, skip_type='time_uniform', - method='multistep', lower_order_final=True, denoise_to_zero=False, solver_type='dpmsolver', - atol=0.0078, rtol=0.05, return_intermediate=False, latent_scale_factor=1.0, pbar=None, previewer=None, - ): - """ - Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`. - - ===================================================== - - We support the following algorithms for both noise prediction model and data prediction model: - - 'singlestep': - Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver. - We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps). - The total number of function evaluations (NFE) == `steps`. - Given a fixed NFE == `steps`, the sampling procedure is: - - If `order` == 1: - - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM). - - If `order` == 2: - - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling. - - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2. - - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1. - - If `order` == 3: - - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling. - - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, and 1 step of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1. - - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1. - - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2. - - 'multistep': - Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`. - We initialize the first `order` values by lower order multistep solvers. - Given a fixed NFE == `steps`, the sampling procedure is: - Denote K = steps. - - If `order` == 1: - - We use K steps of DPM-Solver-1 (i.e. DDIM). - - If `order` == 2: - - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2. - - If `order` == 3: - - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, then (K - 2) step of multistep DPM-Solver-3. - - 'singlestep_fixed': - Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3). - We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE. - - 'adaptive': - Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper). - We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`. - You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs - (NFE) and the sample quality. - - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2. - - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3. - - ===================================================== - - Some advices for choosing the algorithm: - - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs: - Use singlestep DPM-Solver or DPM-Solver++ ("DPM-Solver-fast" in the paper) with `order = 3`. - e.g., DPM-Solver: - >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver") - >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3, - skip_type='time_uniform', method='singlestep') - e.g., DPM-Solver++: - >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++") - >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3, - skip_type='time_uniform', method='singlestep') - - For **guided sampling with large guidance scale** by DPMs: - Use multistep DPM-Solver with `algorithm_type="dpmsolver++"` and `order = 2`. - e.g. - >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++") - >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2, - skip_type='time_uniform', method='multistep') - - We support three types of `skip_type`: - - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images** - - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**. - - 'time_quadratic': quadratic time for the time steps. - - ===================================================== - Args: - x: A pytorch tensor. The initial value at time `t_start` - e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution. - steps: A `int`. The total number of function evaluations (NFE). - t_start: A `float`. The starting time of the sampling. - If `T` is None, we use self.noise_schedule.T (default is 1.0). - t_end: A `float`. The ending time of the sampling. - If `t_end` is None, we use 1. / self.noise_schedule.total_N. - e.g. if total_N == 1000, we have `t_end` == 1e-3. - For discrete-time DPMs: - - We recommend `t_end` == 1. / self.noise_schedule.total_N. - For continuous-time DPMs: - - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15. - order: A `int`. The order of DPM-Solver. - skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'. - method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'. - denoise_to_zero: A `bool`. Whether to denoise to time 0 at the final step. - Default is `False`. If `denoise_to_zero` is `True`, the total NFE is (`steps` + 1). - - This trick is firstly proposed by DDPM (https://arxiv.org/abs/2006.11239) and - score_sde (https://arxiv.org/abs/2011.13456). Such trick can improve the FID - for diffusion models sampling by diffusion SDEs for low-resolutional images - (such as CIFAR-10). However, we observed that such trick does not matter for - high-resolutional images. As it needs an additional NFE, we do not recommend - it for high-resolutional images. - lower_order_final: A `bool`. Whether to use lower order solvers at the final steps. - Only valid for `method=multistep` and `steps < 15`. We empirically find that - this trick is a key to stabilizing the sampling by DPM-Solver with very few steps - (especially for steps <= 10). So we recommend to set it to be `True`. - solver_type: A `str`. The taylor expansion type for the solver. `dpmsolver` or `taylor`. We recommend `dpmsolver`. - atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'. - rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'. - return_intermediate: A `bool`. Whether to save the xt at each step. - When set to `True`, method returns a tuple (x0, intermediates); when set to False, method returns only x0. - Returns: - x_end: A pytorch tensor. The approximated solution at time `t_end`. - - """ - t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end - t_T = self.noise_schedule.T if t_start is None else t_start - assert t_0 > 0 and t_T > 0, "Time range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array" - if return_intermediate: - assert method in ['multistep', 'singlestep', - 'singlestep_fixed'], "Cannot use adaptive solver when saving intermediate values" - if self.correcting_xt_fn is not None: - assert method in ['multistep', 'singlestep', - 'singlestep_fixed'], "Cannot use adaptive solver when correcting_xt_fn is not None" - device = x.device - intermediates = [] - with torch.no_grad(): - if method == 'adaptive': - x = self.dpm_solver_adaptive(x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol, - solver_type=solver_type) - elif method == 'multistep': - assert steps >= order - timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device) - assert timesteps.shape[0] - 1 == steps - # Init the initial values. - step = 0 - t = timesteps[step] - t_prev_list = [t] - model_prev_list = [self.model_fn(x, t)] - if self.correcting_xt_fn is not None: - x = self.correcting_xt_fn(x, t, step) - if return_intermediate: - intermediates.append(x) - # Init the first `order` values by lower order multistep DPM-Solver. - for step in range(1, order): - t = timesteps[step] - x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, t, step, - solver_type=solver_type) - if self.correcting_xt_fn is not None: - x = self.correcting_xt_fn(x, t, step) - if return_intermediate: - intermediates.append(x) - t_prev_list.append(t) - model_prev_list.append(self.model_fn(x, t)) - # Compute the remaining values by `order`-th order multistep DPM-Solver. - for step in tqdm(range(order, steps + 1)): - t = timesteps[step] - # We only use lower order for steps < 10 - if lower_order_final and steps < 10: - step_order = min(order, steps + 1 - step) - else: - step_order = order - x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, t, step_order, - solver_type=solver_type) - if self.correcting_xt_fn is not None: - x = self.correcting_xt_fn(x, t, step) - if return_intermediate: - intermediates.append(x) - for i in range(order - 1): - t_prev_list[i] = t_prev_list[i + 1] - model_prev_list[i] = model_prev_list[i + 1] - t_prev_list[-1] = t - # We do not need to evaluate the final model value. - if step < steps: - model_prev_list[-1] = self.model_fn(x, t) - # comfyui preview - if pbar: - preview_bytes = None - if previewer: - preview_bytes = previewer.decode_latent_to_preview_image("JPEG", x) - pbar.update_absolute(step, steps, preview_bytes) - - elif method in ['singlestep', 'singlestep_fixed']: - if method == 'singlestep': - timesteps_outer, orders = self.get_orders_and_timesteps_for_singlestep_solver(steps=steps, - order=order, - skip_type=skip_type, - t_T=t_T, t_0=t_0, - device=device) - elif method == 'singlestep_fixed': - K = steps // order - orders = [order, ] * K - timesteps_outer = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=K, device=device) - for step, order in enumerate(orders): - s, t = timesteps_outer[step], timesteps_outer[step + 1] - timesteps_inner = self.get_time_steps(skip_type=skip_type, t_T=s.item(), t_0=t.item(), N=order, - device=device) - lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner) - h = lambda_inner[-1] - lambda_inner[0] - r1 = None if order <= 1 else (lambda_inner[1] - lambda_inner[0]) / h - r2 = None if order <= 2 else (lambda_inner[2] - lambda_inner[0]) / h - x = self.singlestep_dpm_solver_update(x, s, t, order, solver_type=solver_type, r1=r1, r2=r2) - if self.correcting_xt_fn is not None: - x = self.correcting_xt_fn(x, t, step) - if return_intermediate: - intermediates.append(x) - else: - raise ValueError("Got wrong method {}".format(method)) - if denoise_to_zero: - t = torch.ones((1,)).to(device) * t_0 - x = self.denoise_to_zero_fn(x, t) - if self.correcting_xt_fn is not None: - x = self.correcting_xt_fn(x, t, step + 1) - if return_intermediate: - intermediates.append(x) - - if return_intermediate: - return x, intermediates - else: - return x - - -############################################################# -# other utility functions -############################################################# - -def interpolate_fn(x, xp, yp): - """ - A piecewise linear function y = f(x), using xp and yp as keypoints. - We implement f(x) in a differentiable way (i.e. applicable for autograd). - The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.) - - Args: - x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver). - xp: PyTorch tensor with shape [C, K], where K is the number of keypoints. - yp: PyTorch tensor with shape [C, K]. - Returns: - The function values f(x), with shape [N, C]. - """ - N, K = x.shape[0], xp.shape[1] - all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2) - sorted_all_x, x_indices = torch.sort(all_x, dim=2) - x_idx = torch.argmin(x_indices, dim=2) - cand_start_idx = x_idx - 1 - start_idx = torch.where( - torch.eq(x_idx, 0), - torch.tensor(1, device=x.device), - torch.where( - torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx, - ), - ) - end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1) - start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2) - end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2) - start_idx2 = torch.where( - torch.eq(x_idx, 0), - torch.tensor(0, device=x.device), - torch.where( - torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx, - ), - ) - y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1) - start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2) - end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2) - cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x) - return cand - - -def expand_dims(v, dims): - """ - Expand the tensor `v` to the dim `dims`. - - Args: - `v`: a PyTorch tensor with shape [N]. - `dim`: a `int`. - Returns: - a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`. - """ - return v[(...,) + (None,) * (dims - 1)] \ No newline at end of file diff --git a/PixArt/sampling/gaussian_diffusion.py b/PixArt/sampling/gaussian_diffusion.py deleted file mode 100644 index 5b4cf01..0000000 --- a/PixArt/sampling/gaussian_diffusion.py +++ /dev/null @@ -1,908 +0,0 @@ -# Modified from OpenAI's diffusion repos -# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py -# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion -# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py - - -import enum -import math - -import numpy as np -import torch as th -import torch.nn.functional as F - -from .diffusion_utils import discretized_gaussian_log_likelihood, normal_kl - - -def mean_flat(tensor): - """ - Take the mean over all non-batch dimensions. - """ - return tensor.mean(dim=list(range(1, len(tensor.shape)))) - - -class ModelMeanType(enum.Enum): - """ - Which type of output the model predicts. - """ - - PREVIOUS_X = enum.auto() # the model predicts x_{t-1} - START_X = enum.auto() # the model predicts x_0 - EPSILON = enum.auto() # the model predicts epsilon - - -class ModelVarType(enum.Enum): - """ - What is used as the model's output variance. - The LEARNED_RANGE option has been added to allow the model to predict - values between FIXED_SMALL and FIXED_LARGE, making its job easier. - """ - - LEARNED = enum.auto() - FIXED_SMALL = enum.auto() - FIXED_LARGE = enum.auto() - LEARNED_RANGE = enum.auto() - - -class LossType(enum.Enum): - MSE = enum.auto() # use raw MSE loss (and KL when learning variances) - RESCALED_MSE = ( - enum.auto() - ) # use raw MSE loss (with RESCALED_KL when learning variances) - KL = enum.auto() # use the variational lower-bound - RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB - - def is_vb(self): - return self == LossType.KL or self == LossType.RESCALED_KL - - -def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac): - betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) - warmup_time = int(num_diffusion_timesteps * warmup_frac) - betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64) - return betas - - -def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps): - """ - This is the deprecated API for creating beta schedules. - See get_named_beta_schedule() for the new library of schedules. - """ - if beta_schedule == "quad": - betas = ( - np.linspace( - beta_start ** 0.5, - beta_end ** 0.5, - num_diffusion_timesteps, - dtype=np.float64, - ) - ** 2 - ) - elif beta_schedule == "linear": - betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64) - elif beta_schedule == "warmup10": - betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1) - elif beta_schedule == "warmup50": - betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5) - elif beta_schedule == "const": - betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) - elif beta_schedule == "jsd": # 1/T, 1/(T-1), 1/(T-2), ..., 1 - betas = 1.0 / np.linspace( - num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64 - ) - else: - raise NotImplementedError(beta_schedule) - assert betas.shape == (num_diffusion_timesteps,) - return betas - - -def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): - """ - Get a pre-defined beta schedule for the given name. - The beta schedule library consists of beta schedules which remain similar - in the limit of num_diffusion_timesteps. - Beta schedules may be added, but should not be removed or changed once - they are committed to maintain backwards compatibility. - """ - if schedule_name == "linear": - # Linear schedule from Ho et al, extended to work for any number of - # diffusion steps. - scale = 1000 / num_diffusion_timesteps - return get_beta_schedule( - "linear", - beta_start=scale * 0.0001, - beta_end=scale * 0.02, - num_diffusion_timesteps=num_diffusion_timesteps, - ) - elif schedule_name == "squaredcos_cap_v2": - return betas_for_alpha_bar( - num_diffusion_timesteps, - lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, - ) - else: - raise NotImplementedError(f"unknown beta schedule: {schedule_name}") - - -def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, - which defines the cumulative product of (1-beta) over time from t = [0,1]. - :param num_diffusion_timesteps: the number of betas to produce. - :param alpha_bar: a lambda that takes an argument t from 0 to 1 and - produces the cumulative product of (1-beta) up to that - part of the diffusion process. - :param max_beta: the maximum beta to use; use values lower than 1 to - prevent singularities. - """ - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) - return np.array(betas) - - -class GaussianDiffusion: - """ - Utilities for training and sampling diffusion models. - Original ported from this codebase: - https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42 - :param betas: a 1-D numpy array of betas for each diffusion timestep, - starting at T and going to 1. - """ - - def __init__( - self, - *, - betas, - model_mean_type, - model_var_type, - loss_type, - snr=False - ): - - self.model_mean_type = model_mean_type - self.model_var_type = model_var_type - self.loss_type = loss_type - self.snr = snr - - # Use float64 for accuracy. - betas = np.array(betas, dtype=np.float64) - self.betas = betas - assert len(betas.shape) == 1, "betas must be 1-D" - assert (betas > 0).all() and (betas <= 1).all() - - self.num_timesteps = int(betas.shape[0]) - - alphas = 1.0 - betas - self.alphas_cumprod = np.cumprod(alphas, axis=0) - self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) - self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) - assert self.alphas_cumprod_prev.shape == (self.num_timesteps,) - - # calculations for diffusion q(x_t | x_{t-1}) and others - self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) - self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) - self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) - self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) - self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1) - - # calculations for posterior q(x_{t-1} | x_t, x_0) - self.posterior_variance = ( - betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) - ) - # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain - self.posterior_log_variance_clipped = np.log( - np.append(self.posterior_variance[1], self.posterior_variance[1:]) - ) if len(self.posterior_variance) > 1 else np.array([]) - - self.posterior_mean_coef1 = ( - betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) - ) - self.posterior_mean_coef2 = ( - (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod) - ) - - def q_mean_variance(self, x_start, t): - """ - Get the distribution q(x_t | x_0). - :param x_start: the [N x C x ...] tensor of noiseless inputs. - :param t: the number of diffusion steps (minus 1). Here, 0 means one step. - :return: A tuple (mean, variance, log_variance), all of x_start's shape. - """ - mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start - variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) - log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape) - return mean, variance, log_variance - - def q_sample(self, x_start, t, noise=None): - """ - Diffuse the data for a given number of diffusion steps. - In other words, sample from q(x_t | x_0). - :param x_start: the initial data batch. - :param t: the number of diffusion steps (minus 1). Here, 0 means one step. - :param noise: if specified, the split-out normal noise. - :return: A noisy version of x_start. - """ - if noise is None: - noise = th.randn_like(x_start) - assert noise.shape == x_start.shape - return ( - _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start - + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise - ) - - def q_posterior_mean_variance(self, x_start, x_t, t): - """ - Compute the mean and variance of the diffusion posterior: - q(x_{t-1} | x_t, x_0) - """ - assert x_start.shape == x_t.shape - posterior_mean = ( - _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start - + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t - ) - posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape) - posterior_log_variance_clipped = _extract_into_tensor( - self.posterior_log_variance_clipped, t, x_t.shape - ) - assert ( - posterior_mean.shape[0] - == posterior_variance.shape[0] - == posterior_log_variance_clipped.shape[0] - == x_start.shape[0] - ) - return posterior_mean, posterior_variance, posterior_log_variance_clipped - - def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None): - """ - Apply the model to get p(x_{t-1} | x_t), as well as a prediction of - the initial x, x_0. - :param model: the model, which takes a signal and a batch of timesteps - as input. - :param x: the [N x C x ...] tensor at time t. - :param t: a 1-D Tensor of timesteps. - :param clip_denoised: if True, clip the denoised signal into [-1, 1]. - :param denoised_fn: if not None, a function which applies to the - x_start prediction before it is used to sample. Applies before - clip_denoised. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :return: a dict with the following keys: - - 'mean': the model mean output. - - 'variance': the model variance output. - - 'log_variance': the log of 'variance'. - - 'pred_xstart': the prediction for x_0. - """ - if model_kwargs is None: - model_kwargs = {} - - B, C = x.shape[:2] - assert t.shape == (B,) - model_output = model(x, t, **model_kwargs) - if isinstance(model_output, tuple): - model_output, extra = model_output - else: - extra = None - - if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]: - assert model_output.shape == (B, C * 2, *x.shape[2:]) - model_output, model_var_values = th.split(model_output, C, dim=1) - min_log = _extract_into_tensor(self.posterior_log_variance_clipped, t, x.shape) - max_log = _extract_into_tensor(np.log(self.betas), t, x.shape) - # The model_var_values is [-1, 1] for [min_var, max_var]. - frac = (model_var_values + 1) / 2 - model_log_variance = frac * max_log + (1 - frac) * min_log - model_variance = th.exp(model_log_variance) - elif self.model_var_type in [ModelVarType.FIXED_LARGE, ModelVarType.FIXED_SMALL]: - model_variance, model_log_variance = { - # for fixedlarge, we set the initial (log-)variance like so - # to get a better decoder log likelihood. - ModelVarType.FIXED_LARGE: ( - np.append(self.posterior_variance[1], self.betas[1:]), - np.log(np.append(self.posterior_variance[1], self.betas[1:])), - ), - ModelVarType.FIXED_SMALL: ( - self.posterior_variance, - self.posterior_log_variance_clipped, - ), - }[self.model_var_type] - model_variance = _extract_into_tensor(model_variance, t, x.shape) - model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape) - else: - model_variance = th.zeros_like(model_output) - model_log_variance = th.zeros_like(model_output) - - def process_xstart(x): - if denoised_fn is not None: - x = denoised_fn(x) - if clip_denoised: - return x.clamp(-1, 1) - return x - - if self.model_mean_type == ModelMeanType.START_X: - pred_xstart = process_xstart(model_output) - else: - pred_xstart = process_xstart( - self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output) - ) - model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t) - - assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape - return { - "mean": model_mean, - "variance": model_variance, - "log_variance": model_log_variance, - "pred_xstart": pred_xstart, - "extra": extra, - } - - def _predict_xstart_from_eps(self, x_t, t, eps): - assert x_t.shape == eps.shape - return ( - _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps - ) - - def _predict_eps_from_xstart(self, x_t, t, pred_xstart): - return ( - _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart - ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) - - def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): - """ - Compute the mean for the previous step, given a function cond_fn that - computes the gradient of a conditional log probability with respect to - x. In particular, cond_fn computes grad(log(p(y|x))), and we want to - condition on y. - This uses the conditioning strategy from Sohl-Dickstein et al. (2015). - """ - gradient = cond_fn(x, t, **model_kwargs) - new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() - return new_mean - - def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): - """ - Compute what the p_mean_variance output would have been, should the - model's score function be conditioned by cond_fn. - See condition_mean() for details on cond_fn. - Unlike condition_mean(), this instead uses the conditioning strategy - from Song et al (2020). - """ - alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) - - eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"]) - eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs) - - out = p_mean_var.copy() - out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps) - out["mean"], _, _ = self.q_posterior_mean_variance(x_start=out["pred_xstart"], x_t=x, t=t) - return out - - def p_sample( - self, - model, - x, - t, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - ): - """ - Sample x_{t-1} from the model at the given timestep. - :param model: the model to sample from. - :param x: the current tensor at x_{t-1}. - :param t: the value of t, starting at 0 for the first diffusion step. - :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. - :param denoised_fn: if not None, a function which applies to the - x_start prediction before it is used to sample. - :param cond_fn: if not None, this is a gradient function that acts - similarly to the model. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :return: a dict containing the following keys: - - 'sample': a random sample from the model. - - 'pred_xstart': a prediction of x_0. - """ - out = self.p_mean_variance( - model, - x, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - model_kwargs=model_kwargs, - ) - noise = th.randn_like(x) - nonzero_mask = ( - (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) - ) # no noise when t == 0 - if cond_fn is not None: - out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs) - sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise - return {"sample": sample, "pred_xstart": out["pred_xstart"]} - - def p_sample_loop( - self, - model, - shape, - noise=None, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - device=None, - progress=False, - ): - """ - Generate samples from the model. - :param model: the model module. - :param shape: the shape of the samples, (N, C, H, W). - :param noise: if specified, the noise from the encoder to sample. - Should be of the same shape as `shape`. - :param clip_denoised: if True, clip x_start predictions to [-1, 1]. - :param denoised_fn: if not None, a function which applies to the - x_start prediction before it is used to sample. - :param cond_fn: if not None, this is a gradient function that acts - similarly to the model. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :param device: if specified, the device to create the samples on. - If not specified, use a model parameter's device. - :param progress: if True, show a tqdm progress bar. - :return: a non-differentiable batch of samples. - """ - final = None - for sample in self.p_sample_loop_progressive( - model, - shape, - noise=noise, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - cond_fn=cond_fn, - model_kwargs=model_kwargs, - device=device, - progress=progress, - ): - final = sample - return final["sample"] - - def p_sample_loop_progressive( - self, - model, - shape, - noise=None, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - device=None, - progress=False, - ): - """ - Generate samples from the model and yield intermediate samples from - each timestep of diffusion. - Arguments are the same as p_sample_loop(). - Returns a generator over dicts, where each dict is the return value of - p_sample(). - """ - if device is None: - device = next(model.parameters()).device - assert isinstance(shape, (tuple, list)) - if noise is not None: - img = noise - else: - img = th.randn(*shape, device=device) - indices = list(range(self.num_timesteps))[::-1] - - if progress: - # Lazy import so that we don't depend on tqdm. - from tqdm.auto import tqdm - - indices = tqdm(indices) - - for i in indices: - t = th.tensor([i] * shape[0], device=device) - with th.no_grad(): - out = self.p_sample( - model, - img, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - cond_fn=cond_fn, - model_kwargs=model_kwargs, - ) - yield out - img = out["sample"] - - def ddim_sample( - self, - model, - x, - t, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - eta=0.0, - ): - """ - Sample x_{t-1} from the model using DDIM. - Same usage as p_sample(). - """ - out = self.p_mean_variance( - model, - x, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - model_kwargs=model_kwargs, - ) - if cond_fn is not None: - out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) - - # Usually our model outputs epsilon, but we re-derive it - # in case we used x_start or x_prev prediction. - eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) - - alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) - alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) - sigma = ( - eta - * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) - * th.sqrt(1 - alpha_bar / alpha_bar_prev) - ) - # Equation 12. - noise = th.randn_like(x) - mean_pred = ( - out["pred_xstart"] * th.sqrt(alpha_bar_prev) - + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps - ) - nonzero_mask = ( - (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) - ) # no noise when t == 0 - sample = mean_pred + nonzero_mask * sigma * noise - return {"sample": sample, "pred_xstart": out["pred_xstart"]} - - def ddim_reverse_sample( - self, - model, - x, - t, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - eta=0.0, - ): - """ - Sample x_{t+1} from the model using DDIM reverse ODE. - """ - assert eta == 0.0, "Reverse ODE only for deterministic path" - out = self.p_mean_variance( - model, - x, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - model_kwargs=model_kwargs, - ) - if cond_fn is not None: - out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) - # Usually our model outputs epsilon, but we re-derive it - # in case we used x_start or x_prev prediction. - eps = ( - _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x - - out["pred_xstart"] - ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape) - alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape) - - # Equation 12. reversed - mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps - - return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]} - - def ddim_sample_loop( - self, - model, - shape, - noise=None, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - device=None, - progress=False, - eta=0.0, - ): - """ - Generate samples from the model using DDIM. - Same usage as p_sample_loop(). - """ - final = None - for sample in self.ddim_sample_loop_progressive( - model, - shape, - noise=noise, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - cond_fn=cond_fn, - model_kwargs=model_kwargs, - device=device, - progress=progress, - eta=eta, - ): - final = sample - return final["sample"] - - def ddim_sample_loop_progressive( - self, - model, - shape, - noise=None, - clip_denoised=True, - denoised_fn=None, - cond_fn=None, - model_kwargs=None, - device=None, - progress=False, - eta=0.0, - ): - """ - Use DDIM to sample from the model and yield intermediate samples from - each timestep of DDIM. - Same usage as p_sample_loop_progressive(). - """ - if device is None: - device = next(model.parameters()).device - assert isinstance(shape, (tuple, list)) - if noise is not None: - img = noise - else: - img = th.randn(*shape, device=device) - indices = list(range(self.num_timesteps))[::-1] - - if progress: - # Lazy import so that we don't depend on tqdm. - from tqdm.auto import tqdm - - indices = tqdm(indices) - - for i in indices: - t = th.tensor([i] * shape[0], device=device) - with th.no_grad(): - out = self.ddim_sample( - model, - img, - t, - clip_denoised=clip_denoised, - denoised_fn=denoised_fn, - cond_fn=cond_fn, - model_kwargs=model_kwargs, - eta=eta, - ) - yield out - img = out["sample"] - - def _vb_terms_bpd( - self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None - ): - """ - Get a term for the variational lower-bound. - The resulting units are bits (rather than nats, as one might expect). - This allows for comparison to other papers. - :return: a dict with the following keys: - - 'output': a shape [N] tensor of NLLs or KLs. - - 'pred_xstart': the x_0 predictions. - """ - true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance( - x_start=x_start, x_t=x_t, t=t - ) - out = self.p_mean_variance( - model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs - ) - kl = normal_kl( - true_mean, true_log_variance_clipped, out["mean"], out["log_variance"] - ) - kl = mean_flat(kl) / np.log(2.0) - - decoder_nll = -discretized_gaussian_log_likelihood( - x_start, means=out["mean"], log_scales=0.5 * out["log_variance"] - ) - assert decoder_nll.shape == x_start.shape - decoder_nll = mean_flat(decoder_nll) / np.log(2.0) - - # At the first timestep return the decoder NLL, - # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t)) - output = th.where((t == 0), decoder_nll, kl) - return {"output": output, "pred_xstart": out["pred_xstart"]} - - def training_losses(self, model, x_start, t, model_kwargs=None, noise=None): - """ - Compute training losses for a single timestep. - :param model: the model to evaluate loss on. - :param x_start: the [N x C x ...] tensor of inputs. - :param t: a batch of timestep indices. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :param noise: if specified, the specific Gaussian noise to try to remove. - :return: a dict with the key "loss" containing a tensor of shape [N]. - Some mean or variance settings may also have other keys. - """ - if model_kwargs is None: - model_kwargs = {} - if noise is None: - noise = th.randn_like(x_start) - x_t = self.q_sample(x_start, t, noise=noise) - - terms = {} - - if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL: - terms["loss"] = self._vb_terms_bpd( - model=model, - x_start=x_start, - x_t=x_t, - t=t, - clip_denoised=False, - model_kwargs=model_kwargs, - )["output"] - if self.loss_type == LossType.RESCALED_KL: - terms["loss"] *= self.num_timesteps - elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE: - model_output = model(x_t, t, **model_kwargs) - if isinstance(model_output, dict) and model_output.get('x', None) is not None: - output = model_output['x'] - else: - output = model_output - - if self.model_var_type in [ - ModelVarType.LEARNED, - ModelVarType.LEARNED_RANGE, - ]: - B, C = x_t.shape[:2] - assert output.shape == (B, C * 2, *x_t.shape[2:]) - output, model_var_values = th.split(output, C, dim=1) - # Learn the variance using the variational bound, but don't let it affect our mean prediction. - frozen_out = th.cat([output.detach(), model_var_values], dim=1) - terms["vb"] = self._vb_terms_bpd( - model=lambda *args, r=frozen_out: r, - x_start=x_start, - x_t=x_t, - t=t, - clip_denoised=False, - )["output"] - if self.loss_type == LossType.RESCALED_MSE: - # Divide by 1000 for equivalence with initial implementation. - # Without a factor of 1/1000, the VB term hurts the MSE term. - terms["vb"] *= self.num_timesteps / 1000.0 - - target = { - ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance( - x_start=x_start, x_t=x_t, t=t - )[0], - ModelMeanType.START_X: x_start, - ModelMeanType.EPSILON: noise, - }[self.model_mean_type] - assert output.shape == target.shape == x_start.shape - if self.snr: - if self.model_mean_type == ModelMeanType.START_X: - pred_noise = self._predict_eps_from_xstart(x_t=x_t, t=t, pred_xstart=output) - pred_startx = output - elif self.model_mean_type == ModelMeanType.EPSILON: - pred_noise = output - pred_startx = self._predict_xstart_from_eps(x_t=x_t, t=t, eps=output) - # terms["mse_eps"] = mean_flat((noise - pred_noise) ** 2) - # terms["mse_x0"] = mean_flat((x_start - pred_startx) ** 2) - - t = t[:, None, None, None].expand(pred_startx.shape) # [128, 4, 32, 32] - # best - target = th.where(t > 249, noise, x_start) - output = th.where(t > 249, pred_noise, pred_startx) - loss = (target - output) ** 2 - if model_kwargs.get('mask_ratio', False) and model_kwargs['mask_ratio'] > 0: - assert 'mask' in model_output - loss = F.avg_pool2d(loss.mean(dim=1), model.model.module.patch_size).flatten(1) - mask = model_output['mask'] - unmask = 1 - mask - terms['mse'] = mean_flat(loss * unmask) * unmask.shape[1]/unmask.sum(1) - if model_kwargs['mask_loss_coef'] > 0: - terms['mae'] = model_kwargs['mask_loss_coef'] * mean_flat(loss * mask) * mask.shape[1]/mask.sum(1) - else: - terms["mse"] = mean_flat(loss) - if "vb" in terms: - terms["loss"] = terms["mse"] + terms["vb"] - else: - terms["loss"] = terms["mse"] - if "mae" in terms: - terms["loss"] = terms["loss"] + terms["mae"] - else: - raise NotImplementedError(self.loss_type) - - return terms - - def _prior_bpd(self, x_start): - """ - Get the prior KL term for the variational lower-bound, measured in - bits-per-dim. - This term can't be optimized, as it only depends on the encoder. - :param x_start: the [N x C x ...] tensor of inputs. - :return: a batch of [N] KL values (in bits), one per batch element. - """ - batch_size = x_start.shape[0] - t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) - qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) - kl_prior = normal_kl( - mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0 - ) - return mean_flat(kl_prior) / np.log(2.0) - - def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None): - """ - Compute the entire variational lower-bound, measured in bits-per-dim, - as well as other related quantities. - :param model: the model to evaluate loss on. - :param x_start: the [N x C x ...] tensor of inputs. - :param clip_denoised: if True, clip denoised samples. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :return: a dict containing the following keys: - - total_bpd: the total variational lower-bound, per batch element. - - prior_bpd: the prior term in the lower-bound. - - vb: an [N x T] tensor of terms in the lower-bound. - - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep. - - mse: an [N x T] tensor of epsilon MSEs for each timestep. - """ - device = x_start.device - batch_size = x_start.shape[0] - - vb = [] - xstart_mse = [] - mse = [] - for t in list(range(self.num_timesteps))[::-1]: - t_batch = th.tensor([t] * batch_size, device=device) - noise = th.randn_like(x_start) - x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise) - # Calculate VLB term at the current timestep - with th.no_grad(): - out = self._vb_terms_bpd( - model, - x_start=x_start, - x_t=x_t, - t=t_batch, - clip_denoised=clip_denoised, - model_kwargs=model_kwargs, - ) - vb.append(out["output"]) - xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2)) - eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"]) - mse.append(mean_flat((eps - noise) ** 2)) - - vb = th.stack(vb, dim=1) - xstart_mse = th.stack(xstart_mse, dim=1) - mse = th.stack(mse, dim=1) - - prior_bpd = self._prior_bpd(x_start) - total_bpd = vb.sum(dim=1) + prior_bpd - return { - "total_bpd": total_bpd, - "prior_bpd": prior_bpd, - "vb": vb, - "xstart_mse": xstart_mse, - "mse": mse, - } - - -def _extract_into_tensor(arr, timesteps, broadcast_shape): - """ - Extract values from a 1-D numpy array for a batch of indices. - :param arr: the 1-D numpy array. - :param timesteps: a tensor of indices into the array to extract. - :param broadcast_shape: a larger shape of K dimensions with the batch - dimension equal to the length of timesteps. - :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims. - """ - res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float() - while len(res.shape) < len(broadcast_shape): - res = res[..., None] - return res + th.zeros(broadcast_shape, device=timesteps.device) diff --git a/README.md b/README.md index 273ecd3..84a9703 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,15 @@ Limitations: ![PixArtT12](https://github.com/city96/ComfyUI_ExtraModels/assets/125218114/eb1a02f9-6114-47eb-a066-261c39c55615) +### PixArt Sigma + +Currently only the early 256 and 512 checkpoints are released, out of which only the 512 is supported in this custom node pack. + +You can find the [512 checkpoint here](https://huggingface.co/PixArt-alpha/PixArt-Sigma/blob/main/PixArt-Sigma-XL-2-512-MS.pth). Place it in your models folder and select the appropriate type in the model loader / resolution selection node. + +> [!IMPORTANT] +> Make sure to select an SDXL VAE for PixArt Sigma! + ### PixArt LCM The LCM model also works if you're on the latest version. To use it: