from typing import Any import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange from collections import defaultdict import torch as th import numpy as np import math from tqdm import tqdm from PIL import Image class GaussianDiffusion: def __init__(self, model, noise_steps, beta_0, beta_T, image_size, channels=3, schedule="linear"): """ suggested betas for: * linear schedule: 1e-4, 0.02 model: the model to be trained (nn.Module) noise_steps: the number of steps to apply noise (int) beta_0: the initial value of beta (float) beta_T: the final value of beta (float) image_size: the size of the image (int, int) """ self.device = 'cpu' self.channels = channels self.model = model self.noise_steps = noise_steps self.beta_0 = beta_0 self.beta_T = beta_T self.image_size = image_size self.betas = self.beta_schedule(schedule=schedule) self.alphas = 1.0 - self.betas # cumulative product of alphas, so we can optimize forward process calculation self.alpha_hat = torch.cumprod(self.alphas, dim=0) def beta_schedule(self, schedule="cosine"): if schedule == "linear": return torch.linspace(self.beta_0, self.beta_T, self.noise_steps).to(self.device) elif schedule == "cosine": return self.betas_for_cosine(self.noise_steps) elif schedule == "sigmoid": return self.betas_for_sigmoid(self.noise_steps) @staticmethod def sigmoid(x): return 1 / (1 + np.exp(-x)) def betas_for_sigmoid(self, num_diffusion_timesteps, start=-3,end=3, tau=1.0, clip_min = 1e-9): betas = [] v_start = self.sigmoid(start/tau) v_end = self.sigmoid(end/tau) for t in range(num_diffusion_timesteps): t_float = float(t/num_diffusion_timesteps) output0 = self.sigmoid((t_float* (end-start)+start)/tau) output = (v_end-output0) / (v_end-v_start) betas.append(np.clip(output*.2, clip_min,.2)) return torch.flip(torch.tensor(betas).to(self.device),dims=[0]).float() def betas_for_cosine(self,num_steps,start=0,end=1,tau=1,clip_min=1e-9): v_start = math.cos(start*math.pi / 2) ** (2 * tau) betas = [] v_end = math.cos(end* math.pi/2) ** 2*tau for t in range(num_steps): t_float = float(t)/num_steps output = math.cos((t_float* (end-start)+start)*math.pi/2)**(2*tau) output = (v_end - output) / (v_end-v_start) betas.append(np.clip(output*.2,clip_min,.2)) return torch.flip(torch.tensor(betas).to(self.device),dims=[0]).float() def sample_time_steps(self, batch_size=1): return torch.randint(0, self.noise_steps, (batch_size,)).to(self.device) def to(self,device): self.device = device self.betas = self.betas.to(device) self.alphas = self.alphas.to(device) self.alpha_hat = self.alpha_hat.to(device) def q(self, x, t): """ Forward process """ pass def p(self, x, t): """ Backward process """ pass def apply_noise(self, x, t): # force x to be (batch_size, image_width, image_height, channels) if len(x.shape) == 3: x = x.unsqueeze(0) if type(t) == int: t = torch.tensor([t]) #print(f'Shape -> {x.shape}, len -> {len(x.shape)}') sqrt_alpha_hat = torch.sqrt(torch.tensor([self.alpha_hat[t_] for t_ in t]).to(self.device)) sqrt_one_minus_alpha_hat = torch.sqrt(torch.tensor([1.0 - self.alpha_hat[t_] for t_ in t]).to(self.device)) # standard normal distribution epsilon = torch.randn_like(x).to(self.device) # Eq 2. in DDPM paper #noisy_image = sqrt_alpha_hat * x + sqrt_one_minus_alpha_hat * epsilon """print(f''' Shape of x {x.shape} Shape of sqrt {sqrt_one_minus_alpha_hat.shape}''')""" try: #print(x.shape) #noisy_image = torch.einsum("b,bwhc->bwhc", sqrt_alpha_hat, x.to(self.device)) + torch.einsum("b,bwhc->bwhc", sqrt_one_minus_alpha_hat, epsilon) noisy_image = torch.einsum("b,bcwh->bcwh", sqrt_alpha_hat, x.to(self.device)) + torch.einsum("b,bcwh->bcwh", sqrt_one_minus_alpha_hat, epsilon) except: print(f'Failed image: shape {x.shape}') #print(f'Noisy image -> {noisy_image.shape}') # returning noisy iamge and the noise which was added to the image #return noisy_image, epsilon #return torch.clip(noisy_image, -1.0, 1.0), epsilon return noisy_image, epsilon @staticmethod def normalize_image(x): # normalize image to [-1, 1] return x / 255.0 * 2.0 - 1.0 @staticmethod def denormalize_image(x): # denormalize image to [0, 255] return (x + 1.0) / 2.0 * 255.0 def sample_step(self, x, t, cond): batch_size = x.shape[0] device = x.device z = torch.randn_like(x) if t >= 1 else torch.zeros_like(x) z = z.to(device) alpha = self.alphas[t] one_over_sqrt_alpha = 1.0 / torch.sqrt(alpha) one_minus_alpha = 1.0 - alpha sqrt_one_minus_alpha_hat = torch.sqrt(1.0 - self.alpha_hat[t]) beta_hat = (1 - self.alpha_hat[t-1]) / (1 - self.alpha_hat[t]) * self.betas[t] beta = self.betas[t] # should we reshape the params to (batch_size, 1, 1, 1) ? # we can either use beta_hat or beta_t # std = torch.sqrt(beta_hat) std = torch.sqrt(beta) # mean + variance * z if cond is not None: predicted_noise = self.model(x, torch.tensor([t]).repeat(batch_size).to(device), cond) else: predicted_noise = self.model(x, torch.tensor([t]).repeat(batch_size).to(device)) mean = one_over_sqrt_alpha * (x - one_minus_alpha / sqrt_one_minus_alpha_hat * predicted_noise) x_t_minus_1 = mean + std * z return x_t_minus_1 def sample(self, num_samples, show_progress=True): """ Sample from the model """ cond = None if self.model.is_conditional: # cond is arange() assert num_samples <= self.model.num_classes, "num_samples must be less than or equal to the number of classes" cond = torch.arange(self.model.num_classes)[:num_samples].to(self.device) cond = rearrange(cond, 'i -> i ()') self.model.eval() image_versions = [] with torch.no_grad(): x = torch.randn(num_samples, self.channels, *self.image_size).to(self.device) it = reversed(range(1, self.noise_steps)) if show_progress: it = tqdm(it) for t in it: image_versions.append(self.denormalize_image(torch.clip(x, -1, 1)).clone().squeeze(0)) x = self.sample_step(x, t, cond) self.model.train() x = torch.clip(x, -1.0, 1.0) return self.denormalize_image(x), image_versions def validate(self, dataloader): """ Calculate the loss on the validation set """ self.model.eval() acc_loss = 0 with torch.no_grad(): for (image, cond) in dataloader: t = self.sample_time_steps(batch_size=image.shape[0]) noisy_image, added_noise = self.apply_noise(image, t) noisy_image = noisy_image.to(self.device) added_noise = added_noise.to(self.device) cond = cond.to(self.device) predicted_noise = self.model(noisy_image, t, cond) loss = nn.MSELoss()(predicted_noise, added_noise) acc_loss += loss.item() self.model.train() return acc_loss / len(dataloader) class DiffusionImageAPI: def __init__(self, diffusion_model): self.diffusion_model = diffusion_model def get_noisy_image(self, image, t): x = torch.tensor(np.array(image)) x = self.diffusion_model.normalize_image(x) y, _ = self.diffusion_model.apply_noise(x, t) y = self.diffusion_model.denormalize_image(y) #print(f"Shape of Image: {y.shape}") return Image.fromarray(y.squeeze(0).numpy().astype(np.uint8)) def get_noisy_images(self, image, time_steps): """ image: the image to be processed PIL.Image time_steps: the number of time steps to apply noise (int) """ return [self.get_noisy_image(image, int(t)) for t in time_steps] def tensor_to_image(self, tensor): return Image.fromarray(tensor.cpu().numpy().astype(np.uint8)) str_to_act = defaultdict(lambda: nn.SiLU()) str_to_act.update({ "relu": nn.ReLU(), "silu": nn.SiLU(), "gelu": nn.GELU(), }) class SinusoidalPositionalEncoding(nn.Module): def __init__(self, dim): super().__init__() self.dim = dim def forward(self, t): device = t.device t = t.unsqueeze(-1) inv_freq = 1.0 / (10000 ** (torch.arange(0, self.dim, 2, device=device).float() / self.dim)) sin_enc = torch.sin(t.repeat(1, self.dim // 2) * inv_freq) cos_enc = torch.cos(t.repeat(1, self.dim // 2) * inv_freq) pos_enc = torch.cat([sin_enc, cos_enc], dim=-1) return pos_enc class TimeEmbedding(nn.Module): def __init__(self, model_dim: int, emb_dim: int, act="silu"): super().__init__() self.lin = nn.Linear(model_dim, emb_dim) self.act = str_to_act[act] self.lin2 = nn.Linear(emb_dim, emb_dim) def forward(self, x): x = self.lin(x) x = self.act(x) x = self.lin2(x) return x class ConvBlock(nn.Module): def __init__(self, in_channels, out_channels, act="silu", dropout=None, zero=False): super().__init__() self.norm = nn.GroupNorm( num_groups=32, num_channels=in_channels, ) self.act = str_to_act[act] if dropout is not None: self.dropout = nn.Dropout(dropout) self.conv = nn.Conv2d( in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1, ) if zero: self.conv.weight.data.zero_() def forward(self, x): x = self.norm(x) x = self.act(x) if hasattr(self, "dropout"): x = self.dropout(x) x = self.conv(x) return x class EmbeddingBlock(nn.Module): def __init__(self, channels: int, emb_dim: int, act="silu"): super().__init__() self.act = str_to_act[act] self.lin = nn.Linear(emb_dim, channels) def forward(self, x): x = self.act(x) x = self.lin(x) return x class ResBlock(nn.Module): def __init__(self, channels: int, emb_dim: int, dropout: float = 0, out_channels=None): """A resblock with a time embedding and an optional change in channel count """ if out_channels is None: out_channels = channels super().__init__() self.conv1 = ConvBlock(channels, out_channels) self.emb = EmbeddingBlock(out_channels, emb_dim) self.conv2 = ConvBlock(out_channels, out_channels, dropout=dropout, zero=True) if channels != out_channels: self.skip_connection = nn.Conv2d(channels, out_channels, kernel_size=1) else: self.skip_connection = nn.Identity() def forward(self, x, t): original = x x = self.conv1(x) t = self.emb(t) # t: (batch_size, time_embedding_dim) = (batch_size, out_channels) # x: (batch_size, out_channels, height, width) # we repeat the time embedding to match the shape of x t = t.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, x.shape[2], x.shape[3]) x = x + t x = self.conv2(x) x = x + self.skip_connection(original) return x class SelfAttentionBlock(nn.Module): def __init__(self, channels, num_heads=1): super().__init__() self.channels = channels self.num_heads = num_heads self.norm = nn.GroupNorm(32, channels) self.attention = nn.MultiheadAttention( embed_dim=channels, num_heads=num_heads, dropout=0, batch_first=True, bias=True, ) def forward(self, x): h, w = x.shape[-2:] original = x x = self.norm(x) x = rearrange(x, "b c h w -> b (h w) c") x = self.attention(x, x, x)[0] x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w) return x + original class Downsample(nn.Module): def __init__(self, channels): super().__init__() # ddpm uses maxpool # self.down = nn.MaxPool2d # iddpm uses strided conv self.down = nn.Conv2d( in_channels=channels, out_channels=channels, kernel_size=3, stride=2, padding=1, ) def forward(self, x): return self.down(x) class DownBlock(nn.Module): """According to U-Net paper 'The contracting path follows the typical architecture of a convolutional network. It consists of the repeated application of two 3x3 convolutions (unpadded convolutions), each followed by a rectified linear unit (ReLU) and a 2x2 max pooling operation with stride 2 for downsampling. At each downsampling step we double the number of feature channels.' """ def __init__(self, in_channels, out_channels, time_embedding_dim, use_attn=False, dropout=0, downsample=True, width=1): """in_channels will typically be half of out_channels""" super().__init__() self.width = width self.use_attn = use_attn self.do_downsample = downsample self.blocks = nn.ModuleList() for _ in range(width): self.blocks.append(ResBlock( channels=in_channels, out_channels=out_channels, emb_dim=time_embedding_dim, dropout=dropout, )) if self.use_attn: self.blocks.append(SelfAttentionBlock( channels=out_channels, )) in_channels = out_channels if self.do_downsample: self.downsample = Downsample(out_channels) def forward(self, x, t): for block in self.blocks: if isinstance(block, ResBlock): x = block(x, t) elif isinstance(block, SelfAttentionBlock): x = block(x) residual = x if self.do_downsample: x = self.downsample(x) return x, residual class Upsample(nn.Module): def __init__(self, channels): super().__init__() self.upsample = nn.Upsample(scale_factor=2) self.conv = nn.Conv2d( in_channels=channels, out_channels=channels, kernel_size=3, padding=1, ) def forward(self, x): x = self.upsample(x) x = self.conv(x) return x class UpBlock(nn.Module): """According to U-Net paper Every step in the expansive path consists of an upsampling of the feature map followed by a 2x2 convolution (“up-convolution”) that halves the number of feature channels, a concatenation with the correspondingly cropped feature map from the contracting path, and two 3x3 convolutions, each followed by a ReLU. """ def __init__(self, in_channels, out_channels, time_embedding_dim, use_attn=False, dropout=0, upsample=True, width=1): """in_channels will typically be double of out_channels """ super().__init__() self.use_attn = use_attn self.do_upsample = upsample self.blocks = nn.ModuleList() for _ in range(width): self.blocks.append(ResBlock( channels=in_channels, out_channels=out_channels, emb_dim=time_embedding_dim, dropout=dropout, )) if self.use_attn: self.blocks.append(SelfAttentionBlock( channels=out_channels, )) in_channels = out_channels if self.do_upsample: self.upsample = Upsample(out_channels) def forward(self, x, t): for block in self.blocks: if isinstance(block, ResBlock): x = block(x, t) elif isinstance(block, SelfAttentionBlock): x = block(x) if self.do_upsample: x = self.upsample(x) return x class Bottleneck(nn.Module): def __init__(self, channels, dropout, time_embedding_dim): super().__init__() in_channels = channels out_channels = channels self.resblock_1 = ResBlock( channels=in_channels, out_channels=out_channels, dropout=dropout, emb_dim=time_embedding_dim ) self.attention_block = SelfAttentionBlock( channels=out_channels, ) self.resblock_2 = ResBlock( channels=out_channels, out_channels=out_channels, dropout=dropout, emb_dim=time_embedding_dim ) def forward(self, x, t): x = self.resblock_1(x, t) x = self.attention_block(x) x = self.resblock_2(x, t) return x class Unet(nn.Module): def __init__( self, image_channels=3, res_block_width=2, starting_channels=128, dropout=0, channel_mults=(1, 2, 2, 4, 4), attention_layers=(False, False, False, True, False) ): super().__init__() self.is_conditional = False self.image_channels = image_channels self.starting_channels = starting_channels time_embedding_dim = 4 * starting_channels self.time_encoding = SinusoidalPositionalEncoding(dim=starting_channels) self.time_embedding = TimeEmbedding(model_dim=starting_channels, emb_dim=time_embedding_dim) self.input = nn.Conv2d(3, starting_channels, kernel_size=3, padding=1) current_channel_count = starting_channels input_channel_counts = [] self.contracting_path = nn.ModuleList([]) for i, channel_multiplier in enumerate(channel_mults): is_last_layer = i == len(channel_mults) - 1 next_channel_count = channel_multiplier * starting_channels self.contracting_path.append(DownBlock( in_channels=current_channel_count, out_channels=next_channel_count, time_embedding_dim=time_embedding_dim, use_attn=attention_layers[i], dropout=dropout, downsample=not is_last_layer, width=res_block_width, )) current_channel_count = next_channel_count input_channel_counts.append(current_channel_count) self.bottleneck = Bottleneck(channels=current_channel_count, time_embedding_dim=time_embedding_dim, dropout=dropout) self.expansive_path = nn.ModuleList([]) for i, channel_multiplier in enumerate(reversed(channel_mults)): next_channel_count = channel_multiplier * starting_channels self.expansive_path.append(UpBlock( in_channels=current_channel_count + input_channel_counts.pop(), out_channels=next_channel_count, time_embedding_dim=time_embedding_dim, use_attn=list(reversed(attention_layers))[i], dropout=dropout, upsample=i != len(channel_mults) - 1, width=res_block_width, )) current_channel_count = next_channel_count last_conv = nn.Conv2d( in_channels=starting_channels, out_channels=image_channels, kernel_size=3, padding=1, ) last_conv.weight.data.zero_() self.head = nn.Sequential( nn.GroupNorm(32, starting_channels), nn.SiLU(), last_conv, ) def forward(self, x, t): t = self.time_encoding(t) return self._forward(x, t) def _forward(self, x, t): t = self.time_embedding(t) x = self.input(x) residuals = [] for contracting_block in self.contracting_path: x, residual = contracting_block(x, t) residuals.append(residual) x = self.bottleneck(x, t) for expansive_block in self.expansive_path: # Add the residual residual = residuals.pop() x = torch.cat([x, residual], dim=1) x = expansive_block(x, t) x = self.head(x) return x class ConditionalUnet(nn.Module): def __init__(self, unet, num_classes): super().__init__() self.is_conditional = True self.unet = unet self.num_classes = num_classes self.class_embedding = nn.Embedding(num_classes, unet.starting_channels) def forward(self, x, t, cond=None): # cond: (batch_size, n), where n is the number of classes that we are conditioning on t = self.unet.time_encoding(t) if cond is not None: cond = self.class_embedding(cond) # sum across the classes so we get a single vector representing the set of classes cond = cond.sum(dim=1) t += cond return self.unet._forward(x, t)