himanshu-skid19/Unconditional_Image_Generation_Using_Diffusion_Models

1import streamlit as st2from PIL import Image, ImageOps3import torch4from matplotlib.image import imread5import numpy as np6import tensorflow as tf7import math8import torch.nn.functional as F9from tqdm.auto import tqdm10from torchvision import transforms11import matplotlib.pyplot as plt12 13from torch import nn14img_size = 6415BATCH_SIZE = 6416device = torch.device("cuda" if torch.cuda.is_available() else "cpu")17 18 19class Block(nn.Module):20    def __init__(self, in_ch, out_ch, time_emb_dim, up=False):21        super().__init__()22        self.time_mlp =  nn.Linear(time_emb_dim, out_ch)23        if up:24            self.conv1 = nn.Conv2d(2*in_ch, out_ch, 3, padding=1)25            self.transform = nn.ConvTranspose2d(out_ch, out_ch, 4, 2, 1)26            self.Upsample = nn.Upsample(scale_factor = 2, mode ='bilinear')27 28        else:29            self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)30            self.transform = nn.Conv2d(out_ch, out_ch, 4, 2, 1)31            self.maxpool = nn.MaxPool2d(4, 2, 1)32        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)33        self.bnorm1 = nn.BatchNorm2d(out_ch)34        self.bnorm2 = nn.BatchNorm2d(out_ch)35        self.silu  = nn.SiLU()36        self.relu = nn.ReLU()37 38    def forward(self, x, t, ):39        # First Conv40        h = (self.silu(self.bnorm1(self.conv1(x))))41        # Time embedding42        time_emb = self.relu(self.time_mlp(t))43        # Extend last 2 dimensions44        time_emb = time_emb[(..., ) + (None, ) * 2]45        # Add time channel46        h = h + time_emb47        # Second Conv48        h =  (self.silu(self.bnorm2(self.conv2(h))))49        # Down or Upsample50        return self.transform(h)51 52 53class SinusoidalPositionEmbeddings(nn.Module):54    def __init__(self, dim):55        super().__init__()56        self.dim = dim57 58    def forward(self, time):59        device = time.device60        half_dim = self.dim // 261        embeddings = math.log(10000) / (half_dim - 1)62        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)63        embeddings = time[:, None] * embeddings[None, :]64        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)65        # TODO: Double check the ordering here66        return embeddings67 68 69class SimpleUnet(nn.Module):70    """71    A simplified variant of the Unet architecture.72    """73    def __init__(self):74        super().__init__()75        image_channels = 376        down_channels = (32, 64, 128, 256, 512)77        up_channels = (512, 256, 128, 64, 32)78        out_dim = 379        time_emb_dim = 3280 81        # Time embedding82        self.time_mlp = nn.Sequential(83                SinusoidalPositionEmbeddings(time_emb_dim),84                nn.Linear(time_emb_dim, time_emb_dim),85                nn.ReLU()86            )87 88        # Initial projection89        self.conv0 = nn.Conv2d(image_channels, down_channels[0], 3, padding=1)90 91        # Downsample92        self.downs = nn.ModuleList([Block(down_channels[i], down_channels[i+1], \93                                    time_emb_dim) \94                    for i in range(len(down_channels)-1)])95        # Upsample96        self.ups = nn.ModuleList([Block(up_channels[i], up_channels[i+1], \97                                        time_emb_dim, up=True) \98                    for i in range(len(up_channels)-1)])99 100        # Edit: Corrected a bug found by Jakub C (see YouTube comment)101        self.output = nn.Conv2d(up_channels[-1], out_dim, 1)102 103    def forward(self, x, timestep):104        # Embedd time105        t = self.time_mlp(timestep)106        # Initial conv107        x = self.conv0(x)108        # Unet109        residual_inputs = []110        for down in self.downs:111            x = down(x, t)112            residual_inputs.append(x)113        for up in self.ups:114            residual_x = residual_inputs.pop()115            # Add residual x as additional channels116            x = torch.cat((x, residual_x), dim=1)117            x = up(x, t)118        return self.output(x)119 120model = SimpleUnet()121 122 123def linear_beta_schedule(timesteps):124    beta_start = 0.0001125    beta_end = 0.02126    return torch.linspace(beta_start, beta_end, timesteps)127 128timesteps= 300129betas = linear_beta_schedule(timesteps=timesteps)130 131alphas = 1. - betas132alphas_cumprod = torch.cumprod(alphas, axis=0)133alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)134sqrt_recip_alphas = torch.sqrt(1.0 / alphas)135sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)136sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)137posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)138 139 140def extract(a, t, x_shape):141    batch_size = t.shape[0]142    out = a.gather(-1, t.cpu())143    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)144 145@torch.no_grad()146def p_sample(model, x, t, t_index):147    betas_t = extract(betas, t, x.shape)148    sqrt_one_minus_alphas_cumprod_t = extract(149        sqrt_one_minus_alphas_cumprod, t, x.shape150    )151    sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)152 153    # Equation 11 in the paper154    # Use our model (noise predictor) to predict the mean155    model_mean = sqrt_recip_alphas_t * (156        x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t157    )158 159    if t_index == 0:160        return model_mean161    else:162        posterior_variance_t = extract(posterior_variance, t, x.shape)163        noise = torch.randn_like(x)164        # Algorithm 2 line 4:165        return model_mean + torch.sqrt(posterior_variance_t) * noise166 167# Algorithm 2 but save all images:168@torch.no_grad()169def p_sample_loop(model, shape):170    device = next(model.parameters()).device171 172    b = shape[0]173    # start from pure noise (for each example in the batch)174    img = torch.randn(shape, device=device)175    imgs = []176 177    for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=1):178        img = p_sample(model, img, torch.full((b,), i, device=device, dtype=torch.long), 3)179        imgs.append(img.cpu().numpy())180    return imgs181 182@torch.no_grad()183def sample(model, image_size, batch_size=16, channels=3):184    return p_sample_loop(model, shape=(batch_size, channels, image_size, image_size))185 186 187 188model = SimpleUnet()189 190st.title("Generatig images using a diffusion model")191model.load_state_dict(torch.load("new_linear_model_1090.pt", map_location=torch.device('cpu')))192 193 194if(st.button("Click to generate image")):195    samples = sample(model, image_size=img_size, batch_size=64, channels=3)196    for i in range(1):197        reverse_transforms = transforms.Compose([198            transforms.Lambda(lambda t: (t + 1) / 2),199            transforms.Lambda(lambda t: t.permute(1, 2, 0)), # CHW to HWC200            transforms.Lambda(lambda t: t * 255.),201            transforms.Lambda(lambda t: t.numpy().astype(np.uint8)),202            transforms.ToPILImage(),203        ])204        img  = reverse_transforms(torch.Tensor((samples[-1][i].reshape(3, img_size, img_size))))205        206        st.image(plt.imshow(img))