initial comm

.gitignore

@@ -0,0 +1,126 @@
+noise-game/*
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyderworkspace
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "files.exclude": {
+        "logseq/*": true
+    }
+}

README.md

@@ -1,3 +1,98 @@
----
-license: mit
----
+# DDPM Project
+
+This repository contains the implementation of Denoising Diffusion Probabilistic Models (DDPM).
+
+## Table of Contents
+- [Introduction](#introduction)
+- [Installation](#installation)
+- [Usage](#usage)
+- [Contributing](#contributing)
+
+## Introduction
+Denoising Diffusion Probabilistic Models (DDPM) are a class of generative models that learn to generate data by reversing a diffusion process. This repository provides a comprehensive implementation of DDPM.
+
+## Installation
+To install the necessary dependencies, run:
+```bash
+pip install -r requirements.txt
+```
+
+## Usage
+To train the model, use the following command:
+```bash
+python train.py
+```
+To generate samples, use:
+```bash
+python generate.py
+```
+
+## Game
+To understand the model and it's workings, we're working on a cool cute little game where the user is the UNET reverser/diffusion model and is tasked to denoise the images with noise made of grids of lines.
+
+Use [learndiffusion.vercel.app](learndiffusion.vercel.app) to access the primitive version of the game. You can also contribute to the game by checking out at the diffusion_game branch. A new model showcase will also be added such that the model's weights are loaded from the internet, model's files are installed and loaded into a gradio interface for direct use/inference on the vercel. Feel free to make changes for the same, issue is opened.
+
+## Explanations and Mathematics
+- slides from presentation : 
+- notes/explanations : [HERE](slides\notes)
+- a cute lab talk ppt: 
+- plato's allegory : \<link to REPUBLIC>
+
+## Resources
+- Original Paper : https://arxiv.org/pdf/2006.11239
+- Improvement Paper : https://arxiv.org/abs/2102.09672
+- Improvement by OpenAI : https://arxiv.org/pdf/2105.05233
+- Stable Diffusion Paper : https://arxiv.org/abs/2112.10752
+- 
+
+### Papers for background
+- UNET Paper for Biomedical Segmentation
+- Autoencooder
+- Variational Autoencoder
+- Markov Hierarchical VAE
+- Introductory Lectures on Diffusion Process
+
+### Youtube videos and courses
+#### Mathematics
+- Outliers
+- Omar Jahil
+
+#### Pytorch Implementation
+- [Deep Findr](https://www.youtube.com/watch?v=a4Yfz2FxXiY)
+- [Notebook from Deep Findr](https://colab.research.google.com/drive/1sjy9odlSSy0RBVgMTgP7s99NXsqglsUL?usp=sharing)
+
+## Pretrained Weights
+weights from the model can be found in [pretrained_weights](https://drive.google.com/drive/folders/1NiQDI3e67I9FITVnrzNPP2Az0LABRpic?usp=sharing)
+
+For loading the pretrained weights:
+```
+model2 = SimpleUnet()
+model2.load_state_dict(torch.load("/content/drive/MyDrive/Research Work/mlsa/DDPM/model_weights.pth"))
+model2.eval()
+```
+
+For making inferences
+TODO: Errors in the sampling function, boolean errors and etc. Will open issues for solving by others as exercise if needed.
+```
+num_samples = 8  # Number of images to generate
+image_size = (3, 32, 32)  # Example for CIFAR10
+noise = torch.randn(num_samples, *image_size).to("cuda")
+
+model2.to("cuda")
+# Generate images by denoising
+with torch.no_grad():
+    generated_images = model2.sample(noise)
+
+# Save the generated images
+save_image(generated_images, "generated_images.png", nrow=4, normalize=True)
+```
+
+
+## Contributing
+Contributions are welcome! Please open an issue or submit a pull request.
+
+
+## Future Ideas
+- Make the model onnx compatible for training and inferencing on Intel GPUs
+- Build a Stable Diffusion model Text2Img using CLIP implementationnnnn !!!
+- Train the current model for a much larger dataset with more generalizations and nuances

dataloader/dataloader_cifar10.py

@@ -0,0 +1,39 @@
+from torchvision import transforms 
+from torch.utils.data import Subset, DataLoader
+import numpy as np
+import torch
+import torchvision
+import matplotlib.pyplot as plt
+from sklearn.model_selection import train_test_split
+
+def load_transformed_dataset(IMG_SIZE=64):
+    # Define the transformation pipeline
+    data_transforms = [
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),  # Scales data into [0,1]
+        transforms.Lambda(lambda t: (t * 2) - 1)  # Scale between [-1, 1]
+    ]
+    data_transform = transforms.Compose(data_transforms)
+
+    # Load CIFAR10 dataset without splitting
+    cifar10_dataset = torchvision.datasets.CIFAR10(root=".", download=True, transform=data_transform)
+
+    # Split indices into train and test using sklearn's train_test_split
+    dataset_size = len(cifar10_dataset)
+    indices = list(range(dataset_size))
+    train_indices, test_indices = train_test_split(indices, test_size=0.2, random_state=42)
+
+    # Create train and test subsets
+    train_subset = Subset(cifar10_dataset, train_indices)
+    test_subset = Subset(cifar10_dataset, test_indices)
+
+    # Combine train and test subsets into a single ConcatDataset
+    combined_dataset = torch.utils.data.ConcatDataset([train_subset, test_subset])
+
+    return combined_dataset
+
+def load_dataloader(combined_dataset, batch_size=64):
+    # Create dataloaders
+    dataloader = DataLoader(combined_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
+    return dataloader

journals/2025_01_29.md

@@ -0,0 +1,63 @@
+<!-- i am gonna make a cute lil game for the first years to understand this ddpm paper i think
+for 1st task, i am thinking of making a game where user gets an image and behind the scene there are 15-20 "noise" patterns made of lines (diagonals and straight) that get transposed over the image at our ease procedurally.
+
+
+now at any given level, any of the N noise patterns can be on the image and the noisy image as well as the original image will be provided to the user.
+the user's task therefore is to choose which out of the mcq noise options are there in the image.
+
+In task 2, the user will be shown a simulation. infront of the user, starting from 1 to N, each individual noise pattern will be transposed on the image infront of the user. the user needs to memorize the pattern and after the simulation, choose the patterns first forwards and then backwards.
+then in the final boss level, i'll have actual gaussian normal distributions as the noise 😂 and the user will decipher that -->
+
+
+### **Game Structure**
+1. **Task 1: Noise Pattern Identification**
+   - User sees a noisy image and the original image side by side.
+   - User selects which noise patterns (from MCQ options) are applied to the noisy image.
+   - Feedback is provided on correctness.
+
+2. **Task 2: Noise Pattern Memorization**
+   - User watches a simulation where noise patterns are sequentially applied to an image.
+   - After the simulation, the user must recall and select the patterns in the correct order (forwards and backwards).
+   - Feedback is provided on accuracy.
+
+3. **Boss Level: Gaussian Noise Deciphering**
+   - User is shown an image with Gaussian noise applied.
+   - User must identify the noise characteristics (e.g., intensity, distribution) or match it to a reference.
+
+---
+
+#### 1. **Image and Noise Generator**
+   - Create a `SquareImage` component that displays the original and noisy images.
+   - Create a `NoisePattern` component that generates the noise patterns (grids with lines).
+   - Use a `NoiseGenerator` utility to procedurally apply noise patterns to the image.
+
+#### 2. **Task 1: Noise Identification**
+   - Create a `Task1` component that:
+     - Displays the original and noisy images.
+     - Provides MCQ options for noise patterns.
+     - Handles user input and provides feedback.
+
+#### 3. **Task 2: Noise Memorization**
+   - Create a `Task2` component that:
+     - Plays a simulation of noise patterns being applied sequentially.
+     - Provides an interface for the user to select patterns in order.
+     - Validates the user’s input and provides feedback.
+
+#### 4. **Boss Level: Gaussian Noise**
+   - Create a `BossLevel` component that:
+     - Applies Gaussian noise to an image.
+     - Provides tools for the user to analyze and decipher the noise.
+
+#### 5. **Game Manager**
+   - Create a `GameManager` component that:
+     - Tracks the user’s progress through tasks.
+     - Handles transitions between tasks.
+     - Displays scores and feedback.
+
+
+### **Next Steps**
+1. Implement the `NoiseGenerator` utility to create and apply noise patterns.
+2. Add feedback mechanisms for user selections.
+3. Build the simulation for Task 2.
+4. Integrate Gaussian noise for the Boss Level.
+5. Style the app to make it visually appealing.

model/attn_utils.py

@@ -0,0 +1,97 @@
+import torch
+from torch import nn
+import math
+import torch.nn.functional as F
+
+
+class SelfAttention(nn.Module):
+  def __init__(self, chs, num_heads=1, ffn_expansion=4, dropout=0.1):
+    super().__init__()
+    self.norm = nn.LayerNorm(chs)
+    self.attn = nn.MultiheadAttention(embed_dim=chs, num_heads=num_heads, batch_first=True)
+
+    self.ffn = nn.Sequential(
+        nn.Linear(chs, chs*ffn_expansion),
+        nn.GELU(),
+        nn.Linear(chs*ffn_expansion, chs)
+    )
+
+    self.norm2 = nn.LayerNorm(chs)
+    self.dropout = nn.Dropout(dropout)
+  
+  def forward(self, x):
+    b,c,h,w = x.shape
+    x_reshaped = x.view(b,c,h*w).transpose(1,2)
+    
+    attn_out, _ = self.attn(self.norm(x_reshaped), self.norm(x_reshaped), self.norm(x_reshaped))
+    x_attn = x_reshaped + self.dropout(attn_out)
+
+    ffn_out = self.ffn(self.norm2(x_attn))
+    x_out = x_attn + self.dropout(ffn_out)
+
+
+    x_out = x_out.transpose(1,2).view(b,c,h,w)
+    return x_out
+  
+
+
+
+class CBAM(nn.Module):
+  def __init__(self,chs, reduction=16):
+    super().__init__()
+
+    self.channel_attn = nn.Sequential(
+        nn.AdaptiveAvgPool2d(1),
+        nn.Conv2d(chs, chs//reduction, 1),
+        nn.ReLU(),
+        nn.Conv2d(chs//reduction, chs, 1),
+        nn.Sigmoid()
+    )
+
+    self.spatial_attn = nn.Sequential(
+        nn.Conv2d(2,1,kernel_size=7,padding=3),
+        nn.Sigmoid()
+    )
+
+  def forward(self,x):
+    ch_wt = self.channel_attn(x)
+    x = x*ch_wt
+
+    avg_pool = torch.mean(x, dim=1, keepdim=True)
+    max_pool, _ = torch.max(x, dim=1, keepdim=True)
+    sp_wt = self.spatial_attn(torch.cat([avg_pool, max_pool], dim=1))
+    x = x* sp_wt
+
+    return x
+
+
+
+
+class Block_CBAM(nn.Module):
+  def __init__(self, in_ch, out_ch, time_emb_dim, up=False):
+    super().__init__()
+    self.time_mlp = nn.Linear(time_emb_dim, out_ch)
+    if up:
+      ## up channel - go big big big bigg from smol smol smol with 3x3 kernel
+      self.conv1 = nn.Conv2d(2*in_ch, out_ch, 3, padding=1)
+      self.transform = nn.ConvTranspose2d(out_ch, out_ch, 4, 2, 1)
+    else:
+      self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
+      self.transform = nn.Conv2d(out_ch, out_ch, 4,2,1)
+
+    self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
+    self.relu = nn.ReLU()
+    self.batch_norm1 = nn.BatchNorm2d(out_ch)
+    self.batch_norm2 = nn.BatchNorm2d(out_ch)
+
+    self.cbam = CBAM(out_ch)
+
+  def forward(self, x, t, ):
+    h = self.batch_norm1(self.relu(self.conv1(x)))
+    time_emb = self.relu(self.time_mlp(t))
+    time_emb = time_emb[(..., ) + (None, ) * 2]
+    h = h + time_emb
+    h = self.batch_norm2(self.relu(self.conv2(h)))
+    
+    h = self.cbam(h)
+    return self.transform(h)

model/model.py

@@ -0,0 +1,320 @@
+'''
+THis gile is to contain the DDPM implementation modularized for loading, prediciton and training.
+'''
+
+from torch import nn
+import math
+import torch
+from utils import forward_diffusion_sample, sample_timestep, sample_plot_image
+import torch.nn.functional as F
+from attn_utils import SelfAttention, CBAM, Block_CBAM
+
+class Block(nn.Module):
+  def __init__(self, in_ch, out_ch, time_emb_dim, up=False):
+    super().__init__()
+    self.time_mlp = nn.Linear(time_emb_dim, out_ch)
+    if up:
+      ## up channel - go big big big bigg from smol smol smol with 3x3 kernel
+      self.conv1 = nn.Conv2d(2*in_ch, out_ch, 3, padding=1)
+      self.transform = nn.ConvTranspose2d(out_ch, out_ch, 4, 2, 1)
+    else:
+      self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
+      self.transform = nn.Conv2d(out_ch, out_ch, 4,2,1)
+    self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
+    self.relu = nn.ReLU()
+    self.batch_norm1 = nn.BatchNorm2d(out_ch)
+    self.batch_norm2 = nn.BatchNorm2d(out_ch)
+
+  def forward(self, x, t, ):
+    h = self.batch_norm1(self.relu(self.conv1(x)))
+    time_emb = self.relu(self.time_mlp(t))
+    time_emb = time_emb[(..., ) + (None, ) * 2]
+    h = h + time_emb
+    h = self.batch_norm2(self.relu(self.conv2(h)))
+    return self.transform(h)
+  
+class PositionEmbeddings(nn.Module):
+  def __init__(self,dim):
+    super().__init__()
+    self.dim = dim
+
+  def forward(self, time):
+    device = time.device
+    half_dim = self.dim // 2
+    embeddings = math.log(10000) / (half_dim - 1)
+    embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
+    embeddings = time[:, None] * embeddings[None, :]
+    embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
+    return embeddings
+  
+
+
+class SimpleUnet(nn.Module):
+  def __init__(self):
+    super().__init__()
+    image_channels = 3
+    down_channels = (64, 128, 256, 512, 1024)
+    up_channels = (1024, 512, 256, 128, 64)
+    self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    out_dim = 3
+    time_emb_dim = 32
+
+    ## timestep stored as positional encoding in terms of sine
+    self.time_mlp = nn.Sequential(
+        PositionEmbeddings(time_emb_dim),
+        nn.Linear(time_emb_dim, time_emb_dim),
+        nn.ReLU()
+    )
+
+    
+    self.conv0 = nn.Conv2d(image_channels, down_channels[0], 3, padding=1)
+    self.down_blocks = nn.ModuleList([
+        Block(down_channels[i], down_channels[i+1], time_emb_dim)
+        for i in range(len(down_channels)-1)
+    ])
+    self.up_blocks = nn.ModuleList([
+        Block(up_channels[i], up_channels[i+1], time_emb_dim, up=True)
+        for i in range(len(up_channels)-1)
+    ])
+
+    ## readout layer
+    self.output = nn.Conv2d(up_channels[-1], out_dim, 1)
+
+  def forward(self, x, timestep):
+    t = self.time_mlp(timestep)
+    x = self.conv0(x)
+    residual_inputs = []
+    for down in self.down_blocks:
+      x = down(x, t)
+      residual_inputs.append(x)
+    for up in self.up_blocks:
+      residual_x = residual_inputs.pop()
+      x = torch.cat((x, residual_x), dim=1)
+      x = up(x, t)
+    return self.output(x)
+  
+  @torch.no_grad()
+  def sample(self, noise):
+      """
+      Generate an image by denoising a given noise tensor using the reverse diffusion process.
+
+      Args:
+          noise (torch.Tensor): Initial noise tensor (e.g., sampled from a Gaussian distribution).
+      
+      Returns:
+          torch.Tensor: Denoised image.
+      """
+      img = noise  # Start with the provided noise tensor
+      T = self.num_timesteps  # Total timesteps for diffusion
+      stepsize = 1  # You can adjust if needed
+
+      # Iterate through the timesteps in reverse order
+      for i in range(0, T)[::-1]:
+          t = torch.full((noise.size(0),), i, device=noise.device, dtype=torch.long)  # Current timestep
+          img = sample_timestep(self, img, t)  # Perform one reverse diffusion step
+          img = torch.clamp(img, -1.0, 1.0)  # Clamp the image to ensure values stay in [-1, 1]
+
+      return img
+
+  def get_loss(self, x_0, t):
+      x_noisy, noise = forward_diffusion_sample(x_0, t, self.device)
+      noise_pred = self(x_noisy, t)
+      return F.l1_loss(noise, noise_pred)
+
+  def train(self, dataloader, BATCH_SIZE=64,T=300, EPOCHS=50, verbose=True):
+      from torch.optim import Adam
+
+      device = "cuda" if torch.cuda.is_available() else "cpu"
+      self.to(device)
+      optimizer = Adam(self.parameters(), lr=0.001)
+      epochs = EPOCHS
+
+      for epoch in range(epochs):
+          for step, batch in enumerate(dataloader):
+              optimizer.zero_grad()
+
+              t = torch.randint(0, T, (BATCH_SIZE,), device=device).long()
+              loss = self.get_loss(self, batch[0], t)
+              loss.backward()
+              optimizer.step()
+
+              if verbose:
+                if epoch % 5 == 0 and step % 150 == 0:
+                    print(f"Epoch {epoch} | step {step:03d} Loss: {loss.item()} ")
+                    sample_plot_image(self)
+
+  def test():
+    ## TODO: add the testing loop here
+    pass
+
+
+
+
+################################################################################################
+####################### ATTENTION LAYERS ADDEDD TO THE MODEL ###################################
+################################################################################################
+
+class SimpleUnetWSelfAttn(nn.Module):
+  def __init__(self):
+    super().__init__()
+    image_channels = 3
+    down_channels = (64, 128, 256, 512, 1024)
+    up_channels = (1024, 512, 256, 128, 64)
+
+    out_dim = 3
+    time_emb_dim = 32
+
+    ## timestep stored as positional encoding in terms of sine
+    self.time_mlp = nn.Sequential(
+        PositionEmbeddings(time_emb_dim),
+        nn.Linear(time_emb_dim, time_emb_dim),
+        nn.ReLU()
+    )
+    self.num_timesteps = 300
+
+
+    self.conv0 = nn.Conv2d(image_channels, down_channels[0], 3, padding=1)
+    self.down_blocks = nn.ModuleList([
+        Block(down_channels[i], down_channels[i+1], time_emb_dim)
+        for i in range(len(down_channels)-1)
+    ])
+    self.up_blocks = nn.ModuleList([
+        Block(up_channels[i], up_channels[i+1], time_emb_dim, up=True)
+        for i in range(len(up_channels)-1)
+    ])
+
+    self.self_attention = SelfAttention(down_channels[-1])
+
+
+    ## readout layer
+    self.output = nn.Conv2d(up_channels[-1], out_dim, 1)
+
+
+  # def settimestep()
+
+  def forward(self, x, timestep):
+    self.num_timesteps = timestep
+    t = self.time_mlp(timestep)
+    x = self.conv0(x)
+    residual_inputs = []
+    for down in self.down_blocks:
+      x = down(x, t)
+      residual_inputs.append(x)
+
+    x = self.self_attention(x)
+    
+    for up in self.up_blocks:
+      residual_x = residual_inputs.pop()
+      x = torch.cat((x, residual_x), dim=1)
+      x = up(x, t)
+    return self.output(x)
+
+  @torch.no_grad()
+  def sample(self, noise):
+      """
+      Generate an image by denoising a given noise tensor using the reverse diffusion process.
+
+      Args:
+          noise (torch.Tensor): Initial noise tensor (e.g., sampled from a Gaussian distribution).
+      
+      Returns:
+          torch.Tensor: Denoised image.
+      """
+      img = noise  # Start with the provided noise tensor
+      T = self.num_timesteps  # Total timesteps for diffusion
+      stepsize = 1  # You can adjust if needed
+      print(noise.device)
+
+      # Iterate through the timesteps in reverse order
+      for i in range(T - 1, -1, -1):
+          t = torch.full((noise.size(0),), i, device=noise.device, dtype=torch.long)  # Current timestep
+          img = sample_timestep(self, img, t)  # Perform one reverse diffusion step
+          img = torch.clamp(img, -1.0, 1.0)  # Clamp the image to ensure values stay in [-1, 1]
+
+      return img
+  
+
+
+################################################################################################
+#################### Convolutional Block Attention Module ADDED TO THE MODEL ###################
+################################################################################################
+
+class SimpleUnetWCBAM(nn.Module):
+  def __init__(self):
+    super().__init__()
+    image_channels = 3
+    down_channels = (64, 128, 256, 512, 1024)
+    up_channels = (1024, 512, 256, 128, 64)
+
+    out_dim = 3
+    time_emb_dim = 32
+
+    ## timestep stored as positional encoding in terms of sine
+    self.time_mlp = nn.Sequential(
+        PositionEmbeddings(time_emb_dim),
+        nn.Linear(time_emb_dim, time_emb_dim),
+        nn.ReLU()
+    )
+    self.num_timesteps = 300
+
+
+    self.conv0 = nn.Conv2d(image_channels, down_channels[0], 3, padding=1)
+    self.down_blocks = nn.ModuleList([
+        Block_CBAM(down_channels[i], down_channels[i+1], time_emb_dim)
+        for i in range(len(down_channels)-1)
+    ])
+    self.up_blocks = nn.ModuleList([
+        Block_CBAM(up_channels[i], up_channels[i+1], time_emb_dim, up=True)
+        for i in range(len(up_channels)-1)
+    ])
+
+    self.self_attention = SelfAttention(down_channels[-1])
+
+
+    ## readout layer
+    self.output = nn.Conv2d(up_channels[-1], out_dim, 1)
+
+
+  # def settimestep()
+
+  def forward(self, x, timestep):
+    self.num_timesteps = timestep
+    t = self.time_mlp(timestep)
+    x = self.conv0(x)
+    residual_inputs = []
+    for down in self.down_blocks:
+      x = down(x, t)
+      residual_inputs.append(x)
+
+    x = self.self_attention(x)
+    
+    for up in self.up_blocks:
+      residual_x = residual_inputs.pop()
+      x = torch.cat((x, residual_x), dim=1)
+      x = up(x, t)
+    return self.output(x)
+
+  @torch.no_grad()
+  def sample(self, noise):
+      """
+      Generate an image by denoising a given noise tensor using the reverse diffusion process.
+
+      Args:
+          noise (torch.Tensor): Initial noise tensor (e.g., sampled from a Gaussian distribution).
+      
+      Returns:
+          torch.Tensor: Denoised image.
+      """
+      img = noise  # Start with the provided noise tensor
+      T = self.num_timesteps  # Total timesteps for diffusion
+      stepsize = 1  # You can adjust if needed
+      print(noise.device)
+
+      # Iterate through the timesteps in reverse order
+      for i in range(T - 1, -1, -1):
+          t = torch.full((noise.size(0),), i, device=noise.device, dtype=torch.long)  # Current timestep
+          img = sample_timestep(self, img, t)  # Perform one reverse diffusion step
+          img = torch.clamp(img, -1.0, 1.0)  # Clamp the image to ensure values stay in [-1, 1]
+
+      return img

model/precomputes.py

@@ -0,0 +1,14 @@
+import torch
+from torch.nn import functional as F
+
+T = 300 ## according to the paper
+
+### SOO MMANNYY PRECOMPUTEDD VALUESS TO TRACKKKK
+betas = torch.linspace(1e-4, 0.02, T)
+alphas = 1. - betas
+alphas_cumulative_products = torch.cumprod(alphas, axis=0)
+alphas_cumulative_products_prev = F.pad(alphas_cumulative_products[:-1], (1, 0), value=1.0)
+sqrt_recip_alphas = torch.sqrt(1.0 / alphas)
+sqrt_alphas_cumulative_products = torch.sqrt(alphas_cumulative_products)
+sqrt_one_minus_alphas_cumulative_products = torch.sqrt(1. - alphas_cumulative_products)
+posterior_variance = betas * (1. - alphas_cumulative_products_prev) / (1. - alphas_cumulative_products)

model/utils.py

@@ -0,0 +1,90 @@
+## Scheduler
+'''
+sequentially add noise to images
+precomputed values used
+'''
+
+import torch.nn.functional as F
+import torch
+from precomputes import betas, sqrt_recip_alphas, sqrt_alphas_cumulative_products, sqrt_one_minus_alphas_cumulative_products, posterior_variance
+# from model import model
+from torchvision import transforms
+from matplotlib import pyplot as plt
+import numpy as np
+
+def get_index_from_list(vals, t, x_shape):
+  batch_size = t.shape[0]
+  out = vals.gather(-1, t.cpu())
+  return out.reshape(batch_size, *((1,)* (len(x_shape) - 1))).to(t.device)
+
+def forward_diffusion_sample(x_0, t, device="cpu"):
+  noise = torch.randn_like(x_0)
+  sqrt_alphas_cumulative_products_t = get_index_from_list(sqrt_alphas_cumulative_products, t, x_0.shape)
+  sqrt_one_minus_alphas_cumulative_products_t = get_index_from_list(
+      sqrt_one_minus_alphas_cumulative_products, t, x_0.shape
+  )
+  ## formulae for image augged looks like sqrt(pi(alpha_t)) * x_t-1 * sqrt(pi(1-alpha_t)) * noise~N(0,1)
+  return sqrt_alphas_cumulative_products_t.to(device) * x_0.to(device) \
+      + sqrt_one_minus_alphas_cumulative_products_t.to(device) * noise.to(device), noise.to(device)
+
+
+@torch.no_grad()
+def sample_timestep(model, x, t):
+    """
+    Calls the model to predict the noise in the image and returns 
+    the denoised image. 
+    Applies noise to this image, if we are not in the last step yet.
+    """
+    betas_t = get_index_from_list(betas, t, x.shape)
+    sqrt_one_minus_alphas_cumulative_products_t = get_index_from_list(
+        sqrt_one_minus_alphas_cumulative_products, t, x.shape
+    )
+    sqrt_recip_alphas_t = get_index_from_list(sqrt_recip_alphas, t, x.shape)
+    
+    # Call model (current image - noise prediction)
+    model_mean = sqrt_recip_alphas_t * (
+        x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumulative_products_t
+    )
+    posterior_variance_t = get_index_from_list(posterior_variance, t, x.shape)
+    
+    if t == 0:
+        return model_mean
+    else:
+        noise = torch.randn_like(x)
+        return model_mean + torch.sqrt(posterior_variance_t) * noise 
+
+@torch.no_grad()
+def sample_plot_image(model, IMG_SIZE=64, device="cpu", T=300):
+    
+    # Sample noise
+    img_size = IMG_SIZE
+    img = torch.randn((1, 3, img_size, img_size), device=device)
+    plt.figure(figsize=(15,15))
+    plt.axis('off')
+    num_images = 10
+    stepsize = int(T/num_images)
+
+    for i in range(0,T)[::-1]:
+        t = torch.full((1,), i, device=device, dtype=torch.long)
+        img = sample_timestep(img, t)
+        # Edit: This is to maintain the natural range of the distribution
+        img = torch.clamp(img, -1.0, 1.0)
+        if i % stepsize == 0:
+            plt.subplot(1, num_images, int(i/stepsize)+1)
+            show_tensor_image(model, img.detach().cpu())
+    # plt.show()
+    return img
+
+def show_tensor_image(image):
+    reverse_transforms = transforms.Compose([
+        transforms.Lambda(lambda t: (t + 1) / 2),
+        transforms.Lambda(lambda t: t.permute(1, 2, 0)), # CHW to HWC
+        transforms.Lambda(lambda t: t * 255.),
+        transforms.Lambda(lambda t: t.numpy().astype(np.uint8)),
+        transforms.ToPILImage(),
+    ])
+
+    # Take first image of batch
+    if len(image.shape) == 4:
+        image = image[0, :, :, :] 
+    plt.imshow(reverse_transforms(image))

requirements.txt

@@ -0,0 +1,6 @@
+torch
+math
+torchvision
+matplotlib
+numpy
+scikit-learn

simple_game/ddpm.py

@@ -0,0 +1,73 @@
+# import pygame
+# import random
+
+# # Initialize Pygame
+# pygame.init()
+
+# # Screen dimensions
+# WIDTH, HEIGHT = 800, 600
+# screen = pygame.display.set_mode((WIDTH, HEIGHT))
+# pygame.display.set_caption("DDPM Noise Game")
+
+# # Colors
+# WHITE = (255, 255, 255)
+# BLACK = (0, 0, 0)
+
+# # Load image
+# image = pygame.image.load("estelle-peplum-top-tops-509.webp")
+# image = pygame.transform.scale(image, (300, 300))
+
+# # Generate noise patterns
+# def generate_noise_pattern():
+#     pattern = pygame.Surface((300, 300), pygame.SRCALPHA)
+#     for _ in range(50):  # Draw random lines
+#         start = (random.randint(0, 300), random.randint(0, 300))
+#         end = (random.randint(0, 300), random.randint(0, 300))
+#         pygame.draw.line(pattern, BLACK, start, end, 2)
+#     return pattern
+
+# noise_patterns = [generate_noise_pattern() for _ in range(20)]
+
+# # Task 1: Identify the Noise
+# def task1():
+#     noisy_image = image.copy()
+#     current_noise = random.choice(noise_patterns)
+#     noisy_image.blit(current_noise, (0, 0))
+
+#     screen.fill(WHITE)
+#     screen.blit(image, (50, 50))
+#     screen.blit(noisy_image, (400, 50))
+#     pygame.display.flip()
+
+#     # Wait for user input
+#     running = True
+#     while running:
+#         for event in pygame.event.get():
+#             if event.type == pygame.QUIT:
+#                 pygame.quit()
+#                 return
+#             if event.type == pygame.KEYDOWN:
+#                 if event.key == pygame.K_1:  # Example: User selects pattern 1
+#                     print("You selected Pattern 1")
+#                     running = False
+
+# # Task 2: Memorize the Sequence
+# def task2():
+#     sequence = random.sample(noise_patterns, 5)  # Random sequence of 5 patterns
+#     for pattern in sequence:
+#         screen.fill(WHITE)
+#         screen.blit(image, (50, 50))
+#         # noisy_image = image.copy()
+#         # screen.blit(pattern, (0, 0))
+#         screen.set_alpha(0)
+#         screen.blit(pattern, (400, 50))
+#         pygame.display.flip()
+#         pygame.time.wait(1000)  # Show each pattern for 1 second
+
+#     # Ask user to recall the sequence
+#     print("Recall the sequence of patterns!")
+
+# # Main loop
+# task1()
+# task2()
+# pygame.quit()

usage/generate.py

@@ -0,0 +1,3 @@
+'''
+TODO: Add code here for generating example results from the model
+'''

usage/train.py

@@ -0,0 +1,3 @@
+'''
+TODO: Add the training script here
+'''