add forward diffusion process

Author: marjanfamili

exploratory/test_diffusion.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 1,
    "id": "7683c4c0",
    "metadata": {},
    "outputs": [],
@@ -27,9 +27,17 @@
     "schedule = LogLinearSchedule(sigma_min=1e-3, sigma_max=1e3)\n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "f08c9060",
+   "metadata": {},
+   "source": [
+    "# test model initiation "
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 2,
    "id": "8a6a6d9d",
    "metadata": {},
    "outputs": [],
@@ -47,9 +55,17 @@
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "f5a9df7e",
+   "metadata": {},
+   "source": [
+    "# test forward pass "
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 3,
    "id": "229ccb16",
    "metadata": {},
    "outputs": [
@@ -67,13 +83,35 @@
     "print(output.shape)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "7e7ae519",
+   "metadata": {},
+   "source": [
+    "# test sampler"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "id": "f5b4269c",
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([4, 3, 32, 32])\n"
+     ]
+    }
+   ],
+   "source": [
+    "x_0 = torch.randn(4, 3, 32, 32) # (B, C, H, W)\n",
+    "t = torch.rand(4)  # Random times\n",
+    "\n",
+    "x_t = processor.q_sample(x_0, t)\n",
+    "print(x_t.shape)"
+   ]
   },
   {
    "cell_type": "code",

src/auto_cast/processors/diffusion.py

@@ -56,6 +56,24 @@ def map(self, x: Tensor) -> Tensor:
     def forward(self, x: Tensor) -> Tensor:
         return self.map(x)
 
-
-
-
+    def q_sample(self, x_0: Tensor, t:Tensor) -> Tensor:
+        """Forward diffusion q(x_t | x_0).
+
+        Sample from q(x_t|x_0) = N(alpha_t * x_0, Sigma_t^2*I)
+        where alpha_t and sigma_t are obtained from the noise schedule.
+        
+        Args:
+        x_0: clean data (B, C, H, W)
+        t: time (B,)
+
+        Returns
+        -------
+        x_t: noised data at t (B, C, H, W)
+        """
+        alpha_t, sigma_t = self.schedule(t)
+        # Reshape (B,) to (B, 1, 1, 1) for broadcasting with (B, C, H, W)
+        alpha_t = alpha_t.view(-1, 1, 1, 1)
+        sigma_t = sigma_t.view(-1, 1, 1, 1)
+        noise = torch.randn_like(x_0)
+        x_t = alpha_t * x_0 + sigma_t * noise
+        return x_t