added noise scheduler and a notebook to make sure the input runs through the model

Author: marjanfamili

exploratory/test_diffusion.ipynb

@@ -0,0 +1,108 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "7683c4c0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from auto_cast.processors.diffusion import DiffusionProcessor, LogLinearSchedule\n",
+    "\n",
+    "class SimpleUNet(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.net = nn.Sequential(\n",
+    "            nn.Conv2d(3,64,3,padding=1),\n",
+    "            nn.ReLU(),\n",
+    "            nn.Conv2d(64,3,3,padding=1)\n",
+    "        )\n",
+    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
+    "        return self.net(x)\n",
+    "# Create components\n",
+    "denoiser_nn = SimpleUNet()\n",
+    "loss = nn.MSELoss()\n",
+    "schedule = LogLinearSchedule(sigma_min=1e-3, sigma_max=1e3)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "8a6a6d9d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from auto_cast.processors.diffusion import DiffusionProcessor, LogLinearSchedule\n",
+    "\n",
+    "# Create noise schedule\n",
+    "schedule = LogLinearSchedule(sigma_min=1e-3, sigma_max=1e3)\n",
+    "\n",
+    "# Create processor\n",
+    "processor = DiffusionProcessor(\n",
+    "    denoiser_nn=denoiser_nn,\n",
+    "    loss=loss,\n",
+    "    schedule=schedule\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "229ccb16",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([4, 3, 32, 32])\n"
+     ]
+    }
+   ],
+   "source": [
+    "x = torch.randn(4, 3, 32, 32)  # batch=4, channels=3, height=32, width=32\n",
+    "output = processor(x)\n",
+    "print(output.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f5b4269c",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "231244a0",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "auto-cast",
+   "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.12.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/auto_cast/processors/diffusion.py

@@ -1,16 +1,42 @@
+import math
+
 import torch
 import torch.nn as nn
-from auto_cast.types import Batch, EncodedBatch, RolloutOutput, Tensor
+
 from auto_cast.processors.base import Processor
+from auto_cast.types import Batch, EncodedBatch, RolloutOutput, Tensor
 
-from azula.noise import Schedule
 
+class NoiseSchedule(nn.Module):
+    """Noise Schedule Module."""
+
+    def forward(self, t: Tensor) -> tuple[Tensor, Tensor]:
+        """Get alpha and sigma for given time steps t."""
+        msg = "Subclasses should implement this method."
+        raise NotImplementedError(msg)
+    
+
+class LogLinearSchedule(NoiseSchedule):
+    """Log-Linear Noise Schedule.
+    
+    Implements a log-linear schedule for alpha and sigma.
+    """
+
+    def __init__(self, sigma_min: float = 0.002, sigma_max: float = 80.0):
+        super().__init__()
+        self.log_sigma_min = math.log(sigma_min)
+        self.log_sigma_max = math.log(sigma_max)
+
+    def forward(self, t: Tensor) -> tuple[Tensor, Tensor]:
+        alpha = torch.ones_like(t)
+        sigma = torch.exp(self.log_sigma_min * (1 - t) + self.log_sigma_max * t)
+        return alpha, sigma
 
 
 class DiffusionProcessor(Processor):
     """Diffusion Processor."""
 
-    def __init__(self, denoiser_nn, loss, schedule: Schedule, **kwargs):
+    def __init__(self, denoiser_nn, loss, schedule: NoiseSchedule, **kwargs):
         """Initialize the DiffusionProcessor.
 
         denoiser_nn: The neural network used for denoising.
@@ -26,3 +52,10 @@ def __init__(self, denoiser_nn, loss, schedule: Schedule, **kwargs):
     def map(self, x: Tensor) -> Tensor:
         """Map input window of states/times to output window using denoiser."""
         return self.denoiser_nn(x)
+    
+    def forward(self, x: Tensor) -> Tensor:
+        return self.map(x)
+
+
+
+