{"metadata":{"kernelspec":{"display_name":"Diffution_model","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.16"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":3649,"databundleVersionId":46718,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"ab63a906","cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nfrom torchvision.transforms import v2\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nimport math\n","metadata":{},"outputs":[],"execution_count":null},{"id":"3b8b71cb","cell_type":"code","source":"DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nprint(f'Using device {DEVICE}')","metadata":{},"outputs":[],"execution_count":null},{"id":"0932fff6","cell_type":"code","source":"import torch\n\nprint(f\"PyTorch Version: {torch.__version__}\")\nprint(\"---\")\nprint(f\"Is CUDA available? {torch.cuda.is_available()}\")\nif torch.cuda.is_available():\n    print(f\"PyTorch's Compiled CUDA Version: {torch.version.cuda}\")\n    print(f\"Number of GPUs: {torch.cuda.device_count()}\")\n    print(f\"Current GPU Index: {torch.cuda.current_device()}\")\n    print(f\"Current GPU Name: {torch.cuda.get_device_name(torch.cuda.current_device())}\")","metadata":{},"outputs":[],"execution_count":null},{"id":"131b2b60","cell_type":"code","source":"img_size = 32\nbatch_size = 128\n\ntransforms = v2.Compose([\n    v2.Resize((img_size, img_size)), # resize to 32*32\n    v2.ToTensor(),\n    v2.Lambda(lambda t:(t*2)-1) # change all the grid values into the range [-1.0, 1.0]\n    ])\n\ndataset = torchvision.datasets.CIFAR10(root='./data', download=True, transform=transforms)\ndataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)\n\n\n","metadata":{},"outputs":[],"execution_count":null},{"id":"1a731837","cell_type":"code","source":"def show_img(images, title=\"\"):\n    images = (images+1) /2\n    images = images.clamp(0, 1)\n\n    grid_img = torchvision.utils.make_grid(images, nrow = 8)\n    plt.imshow(grid_img.permute(1, 2, 0).cpu().numpy())\n    plt.title(title)\n    plt.axis('off')\n    plt.show()\n\n\nimage_batch , _= next(iter(dataloader))\nshow_img(image_batch[:16])","metadata":{},"outputs":[],"execution_count":null},{"id":"3c733412","cell_type":"code","source":"timestamps = 300\n\ndef linear_beta_schedule(timestamps, start=0.0001, end=0.02):\n    return torch.linspace(start, end, timestamps)\n\nbetas = linear_beta_schedule(timestamps=timestamps) # this is variance schedule, defines amt of noise that is added till the timestamps\n\n# alp = 1. - betas # amt of image content remain\n# alphae_cumprod =  torch.cumprod(alp, axis=0) # total amt of image content/signal left after t-steps\n\n# sqrt_alp_cumprod = torch.sqrt(alphae_cumprod)\n# sqrt_one_minus_alp_cumprod = torch.sqrt(1. - alphae_cumprod)\n\n\nalphas = 1. - betas\nalphas_cumprod = torch.cumprod(alphas, axis=0)\nalphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)\nsqrt_recip_alphas = torch.sqrt(1.0 / alphas)\nsqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)\nsqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)\nposterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)\n\ndef get_i_from_list(vals, t, x_shape):\n    batch_size = t.shape[0]\n    out = vals.gather(-1, t.cpu())\n    return out.reshape(batch_size, *((1,)*(len(x_shape)-1))).to(t.device)\n\ndef forward_diff_sample(x_0, t, device=DEVICE):\n    noise = torch.randn_like(x_0)  # random noise with normal distribution\n\n    # sqrt(alpha_cumprod) and sqrt(1 - alpha_cumprod) for the given timestep t\n    sqrt_alphas_cumprod_t = get_i_from_list(sqrt_alphas_cumprod, t, x_0.shape)\n    sqrt_one_minus_alphas_cumprod_t = get_i_from_list(sqrt_one_minus_alphas_cumprod, t, x_0.shape)\n    \n    # Calculate the noisy image\n    # mean + variance * noise\n    noisy_image = (sqrt_alphas_cumprod_t.to(device)*x_0.to(device)) + (sqrt_one_minus_alphas_cumprod_t.to(device)*noise.to(device))\n\n    return noisy_image, noise.to(device)\n\n","metadata":{},"outputs":[],"execution_count":null},{"id":"4a2b2fbf","cell_type":"code","source":"show_img(image_batch[0])","metadata":{},"outputs":[],"execution_count":null},{"id":"702b1081","cell_type":"code","source":"single_image = image_batch[0].unsqueeze(0)\n\n# Define timesteps to visualize\ntimesteps_to_show = [0, 50, 100, 150, 199, 299]\nnoisy_images = []\n\n# Generate noisy versions of the image at different timesteps\nfor t_val in timesteps_to_show:\n    t = torch.tensor([t_val], dtype=torch.long)\n    noisy_img, _ = forward_diff_sample(single_image, t)\n    noisy_images.append(noisy_img)\n\n# Concatenate for visualization\nnoisy_images_tensor = torch.cat(noisy_images, dim=0)\n\n# Show the results\nshow_img(noisy_images_tensor)\n","metadata":{},"outputs":[],"execution_count":null},{"id":"09ce2372","cell_type":"code","source":"show_img(image_batch[0]), show_img(noisy_img)","metadata":{},"outputs":[],"execution_count":null},{"id":"ec00883d","cell_type":"markdown","source":"### Why a U-Net?\n- Symmetrical Structure: It has a downsampling path (encoder) to capture contextual information and an upsampling path (decoder) to reconstruct the image.\n\n- Skip Connections: It connects layers from the downsampling path directly to corresponding layers in the upsampling path. This is crucial as it allows the model to reuse high-resolution feature maps from the early layers, helping it reconstruct fine details in the output.","metadata":{}},{"id":"84d2253d","cell_type":"code","source":"class SPE(nn.Module): # sinusoidal PE that is used in OG Trnasformer paper\n    '''\n    Instead of just telling the artist \"the noise level is 150\", the assistant gives them a rich description:\n    \"This noise level is medium-high, which means you should focus on finding large shapes rather than fine textures, \n    and the color palette might be washed out.\"\n    '''\n    def __init__(self, dim):\n        super().__init__()\n        self.dim = dim\n\n    def forward(self, time): # inputs time(number) and converts it into embedding that contains rich features(time embedding)\n        device = time.device\n\n        half_dim = self.dim // 2\n\n        embedding = math.log(10000)/(half_dim-1)\n        embedding = torch.exp(torch.arange(half_dim, device=DEVICE)*-embedding)\n        embedding = time[:, None] * embedding[None, :]\n        embedding = torch.cat((embedding.sin(), embedding.cos()), dim=-1)\n        return embedding\n    \nclass Block(nn.Module): # bulding block of U-net\n    def __init__(self, in_ch, out_ch, time_emb_dim, up=False):\n        super().__init__()\n        self.time_mlp = nn.Linear(time_emb_dim, out_ch)\n\n        if up:\n            self.conv1 = nn.Conv2d(2*in_ch, out_ch, 3, padding=1)\n            self.transform = nn.ConvTranspose2d(out_ch, out_ch, 4, 2, 1)\n        else:\n            self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)\n            self.transform = nn.Conv2d(out_ch, out_ch, 4, 2, 1)\n\n        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)\n        self.norm1=nn.GroupNorm(8, out_ch)\n        self.norm2=nn.GroupNorm(8, out_ch)\n        self.silu = nn.SiLU()\n\n    def forward(self, x, t): # x - current state, t - noise level\n        h = self.norm1(self.silu(self.conv1(x))) # convolve to activation fxn to normalisation\n\n        time_emb = self.silu(self.time_mlp(t))\n        time_emb = time_emb.unsqueeze(-1).unsqueeze(-1) # Reshape to (B, C, 1, 1)\n        \n        # Add time embedding\n        h = h + time_emb # understanding of image + understand of noise level, they way it will paint will vary\n        \n        # Second convolution\n        h = self.norm2(self.silu(self.conv2(h)))\n        \n        # Downsample or Upsample, depends on weather to zoom in and work on details....or zoom out and see the bigger picture\n        return self.transform(h)\n\n\nclass UnetArch(nn.Module): \n    def __init__(self, image_channels=3, time_emb_dim=32):\n        super().__init__()\n        \n        down_channels = (64, 128, 256, 512, 1024)\n        up_channels = (1024, 512, 256, 128, 64)\n        \n        # Time embedding\n        self.time_mlp = nn.Sequential(\n            SPE(time_emb_dim),\n            nn.Linear(time_emb_dim, time_emb_dim),\n            nn.ReLU()\n        )\n        \n        # Initial projection\n        self.conv0 = nn.Conv2d(image_channels, down_channels[0], 3, padding=1)\n\n        # Downsampling path\n        self.downs = nn.ModuleList([Block(down_channels[i], down_channels[i+1], time_emb_dim) for i in range(len(down_channels)-1)]) \n        \n        # Upsampling path\n        self.ups = nn.ModuleList([Block(up_channels[i], up_channels[i+1], time_emb_dim, up=True) for i in range(len(up_channels)-1)])\n        \n        # Final convolve layer\n        self.output = nn.Conv2d(up_channels[-1], image_channels, 1)\n\n    def forward(self, x, timestep):# U shape analogy, downsampling(encoder)-> get more abstract feature out of the image through smaller and deeper sample\n                                            # upsampling(decoder)  -> after undertanding the features, capturing/adding details back\n                                            # Skip Connections     -> refering high-res images while upsampling images to get OG texture and edges\n        # embed time\n        t_emb = self.time_mlp(timestep)\n        \n        # initial convolution\n        x = self.conv0(x)\n        \n        # downsampling path + save skip connections\n        residual_inputs = []  # saving high-res residual images for later(skip connections)\n        for down_block in self.downs:\n            x = down_block(x, t_emb)\n            residual_inputs.append(x)\n        \n        # upsampling path + use skip connections\n        for up_block in self.ups:\n            residual_x = residual_inputs.pop()   #get the corresponding skip connection from the downsampling path\n\n            # Concatenate the skip connection with the current feature map\n            x = torch.cat((x, residual_x), dim=1)           \n            x = up_block(x, t_emb)\n            \n        return self.output(x)\n\nmodel = UnetArch().to(DEVICE)\nprint(f\"Model architecture defined. Number of parameters: {sum(p.numel() for p in model.parameters()):,}\")\n\n\n\n","metadata":{},"outputs":[],"execution_count":null},{"id":"e1d47bad","cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom tqdm import tqdm\n\nLEARNING_RATE = 1e-4\nEPOCHS = 10 \n\noptimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)\ncriterion = nn.L1Loss() \n\nprint(\"Starting Training...\")\nmodel.train() \n\nfor epoch in range(EPOCHS):\n    pbar = tqdm(dataloader, desc=f\"Epoch {epoch+1}/{EPOCHS}\")\n    \n    for step, (images, _) in enumerate(pbar):\n        optimizer.zero_grad()\n\n        # sample a random timestep for each image in the batch.....We use a LongTensor because timesteps are integer indices.\n        t = torch.randint(0, timestamps, (batch_size,), device=DEVICE).long()\n        \n        # add noise to the images to get x_t and the added noise (our target)\n        x_t, noise = forward_diff_sample(images, t, DEVICE)\n        \n        # predict the noise using the U-Net......model is on the DEVICE, and x_t and t are already on the DEVICE\n        predicted_noise = model(x_t, t)\n        \n        # calculate the loss between the predicted noise and the actual noise\n        loss = criterion(noise, predicted_noise)\n        \n        # backpropagate and update weights\n        loss.backward()\n        optimizer.step()\n        \n        # update the progress bar description with the current loss\n        pbar.set_postfix(loss=f\"{loss.item():.4f}\")\n\nprint(\"Training finished.\")\n\nmodel_path = \"ddpm_cifar10_10_epochs.pth\"\ntorch.save(model.state_dict(), model_path)\nprint(f\"Model saved to {model_path}\")","metadata":{},"outputs":[],"execution_count":null},{"id":"1e4d55e4","cell_type":"code","source":"import torch\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nimport torchvision\n\n@torch.no_grad() # we are not training, so we don't need to calculate gradients\ndef sample_timestep(x, t, model):\n    betas_t = get_i_from_list(betas, t, x.shape)\n    sqrt_one_minus_alphas_cumprod_t = get_i_from_list(\n        sqrt_one_minus_alphas_cumprod, t, x.shape\n    )\n    sqrt_recip_alphas_t = get_i_from_list(sqrt_recip_alphas, t, x.shape)\n    \n    model_mean = sqrt_recip_alphas_t * (\n        x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t\n    )\n    posterior_variance_t = get_i_from_list(posterior_variance, t, x.shape)\n    \n    if t.all() == 0:\n        return model_mean\n    else:\n        noise = torch.randn_like(x)\n        return model_mean + torch.sqrt(posterior_variance_t) * noise\n\n\n@torch.no_grad()\ndef sample_plot_images(model, num_images=16):\n    print(\"Generating new images...\")\n    img = torch.randn((num_images, 3, img_size, img_size), device=DEVICE)\n    \n    for i in tqdm(reversed(range(0, timestamps)), desc='Sampling loop', total=timestamps):\n        t = torch.full((num_images,), i, device=DEVICE, dtype=torch.long)\n        img = sample_timestep(img, t, model)\n\n    show_img(img, \"Generated Images\")\n\n\nmodel = UnetArch().to(DEVICE)\n\nmodel_path = \"ddpm_cifar10_10_epochs.pth\"\nmodel.load_state_dict(torch.load(model_path, map_location=DEVICE))\n\nmodel.eval()\n\nsample_plot_images(model)","metadata":{},"outputs":[],"execution_count":null},{"id":"698d72a4","cell_type":"code","source":"","metadata":{},"outputs":[],"execution_count":null}]}