Fixed issues #10, #9, #7, #4, and #3

README.md

@@ -67,7 +67,7 @@ Check out the [Getting Started Guide](https://cabralpinto.github.io/modular-diff
 
 ## Contributing
 
-We appreciate your support and welcome your contributions! Please fell free to submit pull requests if you found a bug or typo you want to fix. If you want to contribute a new prebuilt module or feature, please start by opening an issue and discussing it with us. If you don't know where to begin, take a look at the [open issues](https://github.com/cabralpinto/modular-diffusion/issues). Please read our [Contributing Guide](https://github.com/cabralpinto/modular-diffusion/blob/main/CONTRIBUTING.md) for more details.
+We appreciate your support and welcome your contributions! Please feel free to submit pull requests if you found a bug or typo you want to fix. If you want to contribute a new prebuilt module or feature, please start by opening an issue and discussing it with us. If you don't know where to begin, take a look at the [open issues](https://github.com/cabralpinto/modular-diffusion/issues). Please read our [Contributing Guide](https://github.com/cabralpinto/modular-diffusion/blob/main/CONTRIBUTING.md) for more details.
 
 ## License
 

diffusion/schedule.py

@@ -35,7 +35,6 @@ def compute(self) -> Tensor:
         delta = ((t / (self.steps + 2) + self.offset) / (1 + self.offset) * torch.pi /
                    2).cos()**self.exponent
         alpha = torch.clip(delta[1:] / delta[:-1], 1e-3, 1)
-        # TODO check if this is correct
         return alpha
 
 
@@ -45,6 +44,6 @@ class Sqrt(Schedule):
 
     def compute(self) -> Tensor:
         t = torch.arange(self.steps + 2)
-        delta = 1 - torch.sqrt(t / self.steps + self.offset)
+        delta = 1 - torch.sqrt(t / (self.steps + 1) + self.offset)
         alpha = torch.clip(delta[1:] / delta[:-1], 0, 0.999)
         return alpha

docs/src/components/Article.astro

@@ -1,5 +1,5 @@
 <article
-  class="grow min-w-0 max-w-[43rem] p-10 sm:p-16 lg:blur-0 transition-[filter]"
+  class="grow min-w-0 max-w-[43rem] p-10 sm:p-16 lg:blur-0 transition-[filter] bg-white dark:bg-gray-900 text-gray-900 dark:text-gray-100"
 >
   <slot />
 </article>
@@ -23,49 +23,46 @@
       }
     }
     a {
-      @apply text-green-600 hover:underline;
+      @apply text-green-600 dark:text-green-400 hover:underline dark:hover:text-green-300;
     }
     li {
       @apply mb-3 ml-6 -indent-3;
       &::before {
-        @apply content-["\2022"] text-gray-300 inline-block w-3;
+        @apply content-["\2022"] text-gray-300 dark:text-gray-500 inline-block w-3;
       }
     }
     blockquote {
-      @apply my-5 sm:my-7 border border-green-600 bg-gray-50 rounded-md sm:rounded-lg p-5;
+      @apply my-5 sm:my-7 border border-green-600 dark:border-green-400 bg-gray-50 dark:bg-gray-800 rounded-md sm:rounded-lg p-5;
       > p {
         @apply mb-1.5 sm:mb-2;
         &:not(:first-child) {
-          @apply text-gray-600 text-xs/[1.15rem] sm:text-sm/[1.36rem] mb-0;
+          @apply text-gray-600 dark:text-gray-400 text-xs/[1.15rem] sm:text-sm/[1.36rem] mb-0;
           > code {
             @apply text-[0.625rem]/[0.75rem] sm:text-xs px-1 py-0.5;
           }
         }
       }
     }
     code {
-      @apply text-xs sm:text-sm no-contextual;
+      @apply text-xs sm:text-sm no-contextual dark:text-gray-300;
       &:not(pre code) {
-        @apply px-1 sm:px-1 py-[0.0625rem] sm:py-0.5 bg-gray-200 rounded-md font-mono;
+        @apply px-1 sm:px-1 py-[0.0625rem] sm:py-0.5 bg-gray-200 dark:bg-gray-700 rounded-md font-mono;
       }
     }
     pre {
-      @apply my-5 sm:my-7 rounded-md sm:rounded-lg p-4 sm:p-5;
+      @apply my-5 sm:my-7 rounded-md sm:rounded-lg p-4 sm:p-5 bg-gray-200 dark:bg-gray-800 text-gray-900 dark:text-gray-100;
     }
     img {
       @apply mx-auto my-5 sm:my-7 rounded-md sm:rounded-lg;
     }
     hr {
-      @apply my-6 sm:my-8 border-gray-200;
+      @apply my-6 sm:my-8 border-gray-200 dark:border-gray-700;
     }
     strong {
       @apply font-medium;
     }
     em {
-      @apply text-gray-500;
+      @apply text-gray-500 dark:text-gray-400;
     }
-    // .text {
-    //     @apply font-sans text-sm leading-6 sm:text-base sm:leading-7;
-    // }
   }
 </style>

docs/src/components/Header.astro

@@ -1,48 +1,92 @@
 <header
-  class="fixed w-full bg-white flex items-center border-b h-12 sm:h-14 px-7 sm:px-10 z-50 text-sm sm:text-base overflow-x-auto"
+  class="fixed w-full bg-white dark:bg-gray-900 flex items-center border-b border-gray-200 dark:border-gray-700 h-12 sm:h-14 px-7 sm:px-10 z-50 text-sm sm:text-base overflow-x-auto"
 >
-  <div class="sm:hidden bg-gradient-to-r from-white to-transparent h-11 w-7 fixed left-0"></div>
+  <div class="sm:hidden bg-gradient-to-r from-white dark:from-gray-900 to-transparent h-11 w-7 fixed left-0"></div>
   <img
     src="/modular-diffusion/images/menu.svg"
-    class="toggle lg:hidden h-7 hover:scale-y-110 transition-[transform] -ml-1.5 mr-4 sm:mr-6 cursor-pointer"
+    class="toggle lg:hidden h-7 hover:scale-y-110 -ml-1.5 mr-4 sm:mr-6 cursor-pointer dark:filter dark:invert"
     alt="Menu"
   />
   <a class="flex items-center" href="/modular-diffusion/">
     <img
       class="mr-3 rounded-md"
       src="/modular-diffusion/images/logo.png"
       alt="Logo"
+      class="dark:filter dark:invert"
     />
-    <p class="whitespace-nowrap">Modular Diffusion</p>
+    <p class="whitespace-nowrap text-gray-900 dark:text-gray-100">Modular Diffusion</p>
   </a>
   <div class="grow px-8"></div>
   <a
-    class="hover:underline"
-    href="https://github.com/cabralpinto/modular-diffusion">GitHub</a
+    class="hover:underline text-gray-900 dark:text-gray-100 dark:hover:text-gray-300"
+    href="https://github.com/cabralpinto/modular-diffusion"
   >
+    GitHub
+  </a>
   <a
-    class="ml-4 sm:ml-6 hover:underline"
-    href="https://discord.com/invite/4uBXpwg6">Discord</a
+    class="ml-4 sm:ml-6 hover:underline text-gray-900 dark:text-gray-100 dark:hover:text-gray-300"
+    href="https://discord.com/invite/4uBXpwg6"
   >
+    Discord
+  </a>
   <a
-    class="ml-4 sm:ml-6 hover:underline"
-    href="https://pypi.org/project/modular-diffusion/">PyPI</a
+    class="ml-4 sm:ml-6 hover:underline text-gray-900 dark:text-gray-100 dark:hover:text-gray-300"
+    href="https://pypi.org/project/modular-diffusion"
   >
-  <div class="sm:hidden bg-gradient-to-l from-white to-transparent h-11 w-7 fixed right-0"></div>
+    PyPI
+  </a>
+  <button id="dark-mode-toggle" class="hidden sm:flex ml-4 p-2 rounded bg-gray-200 dark:bg-gray-700 hover:scale-110 transition-transform duration-300 relative items-center">
+    <i id="dark-mode-icon" class="fas fa-sun text-yellow-500"></i>
+  </button>
+  <div class="sm:hidden bg-gradient-to-l from-white dark:from-gray-900 to-transparent h-11 w-7 fixed right-0"></div>
 </header>
 
 <script>
   const toggle = document.querySelector(".toggle");
+  const darkModeToggle = document.getElementById('dark-mode-toggle');
+  const darkModeIcon = document.getElementById('dark-mode-icon');
   const nav = document.querySelector("nav");
   const article = document.querySelector("article");
+
   toggle?.addEventListener("click", () => {
     nav?.classList.toggle("-ml-72");
     nav?.classList.toggle("ml-0");
     article?.classList.toggle("blur-sm");
   });
+
   article?.addEventListener("click", () => {
     nav?.classList.add("-ml-72");
     nav?.classList.remove("ml-0");
     article?.classList.remove("blur-sm");
   });
+
+  darkModeToggle.addEventListener('click', () => {
+    document.documentElement.classList.toggle('dark');
+    const isDark = document.documentElement.classList.contains('dark');
+    const theme = isDark ? 'dark' : 'light';
+    localStorage.setItem('theme', theme);
+    updateDarkModeIcon(isDark);
+  });
+
+  // Function to update icon based on the mode
+  function updateDarkModeIcon(isDark) {
+    if (isDark) {
+      darkModeIcon.classList.remove('fa-sun');
+      darkModeIcon.classList.add('fa-moon');
+      darkModeIcon.classList.remove('text-yellow-500');
+      darkModeIcon.classList.add('text-gray-200');
+    } else {
+      darkModeIcon.classList.remove('fa-moon');
+      darkModeIcon.classList.add('fa-sun');
+      darkModeIcon.classList.remove('text-gray-200');
+      darkModeIcon.classList.add('text-yellow-500');
+    }
+  }
+
+  // Apply the saved theme on load
+  const savedTheme = localStorage.getItem('theme');
+  if (savedTheme) {
+    document.documentElement.classList.add(savedTheme);
+    updateDarkModeIcon(savedTheme === 'dark');
+  }
 </script>

docs/src/components/Index.astro

@@ -10,13 +10,13 @@ headings[0].text = "Overview";
 ---
 
 <div
-  class="hidden xl:block sticky top-8 w-64 pr-16 pt-16 h-full text-sm min-w-fit"
+  class="hidden xl:block sticky top-8 w-64 pr-16 pt-16 h-full text-sm min-w-fit bg-white dark:bg-gray-900 text-gray-900 dark:text-gray-100"
 >
-  <h2 class="font-medium mb-1">On this page</h2>
+  <h2 class="font-medium mb-1 text-gray-900 dark:text-gray-100">On this page</h2>
   <ul class="index">
     {
       headings.map((heading) => (
-        <li class="text-gray-600 hover:text-gray-800 hover:font-medium transition-colors w-full">
+        <li class="text-gray-600 dark:text-gray-400 hover:text-gray-800 dark:hover:text-gray-300 hover:font-medium transition-colors w-full">
           <a class="block py-1" href={`#${heading.slug}`}>
             {heading.text}
           </a>

docs/src/components/Sidebar.astro

@@ -16,15 +16,15 @@ const articles = (await Astro.glob("../pages/[!_]*/*.mdx"))
 ---
 
 <nav
-  class="-ml-72 lg:ml-0 transition-[margin-left] lg:transition-none flex fixed mt-12 sm:mt-14 flex-col bg-white border-r h-full w-60 sm:w-72 px-7 sm:px-10 py-5 sm:py-8 space-y-4 sm:space-y-6 z-50"
+  class="-ml-72 lg:ml-0 transition-[margin-left] lg:transition-none flex fixed mt-12 sm:mt-14 flex-col bg-white dark:bg-gray-900 border-r dark:border-gray-700 h-full w-60 sm:w-72 px-7 sm:px-10 py-5 sm:py-8 space-y-4 sm:space-y-6 z-50 "
 >
   {
     Object.entries(articles).map(([group, titles]) => (
-      <group class="text-sm">
-        <h2 class="font-medium mb-1">{group}</h2>
+      <group class="text-sm dark:text-gray-100">
+        <h2 class="font-medium mb-1 text-gray-900 dark:text-gray-100">{group}</h2>
         <ul>
           {titles.map(({ title, href }) => (
-            <li class="text-gray-600 hover:text-gray-800 hover:font-medium transition-colors w-full">
+            <li class="text-gray-600 dark:text-gray-400 hover:text-gray-800 dark:hover:text-gray-300 hover:font-medium w-full">
               <a class="block py-1" href={href}>
                 {title}
               </a>

docs/src/layouts/Layout.astro

@@ -5,6 +5,7 @@ import Header from "../components/Header.astro";
 import Index from "../components/Index.astro";
 import Sidebar from "../components/Sidebar.astro";
 import "../styles/fonts.css";
+import "../styles/global.css";
 
 type Props = MarkdownLayoutProps<{
   title: string;
@@ -19,27 +20,37 @@ const {
 
 <!DOCTYPE html>
 <html class="scroll-smooth scroll-pt-[5.5rem]" style="zoom: 105%" lang="en">
-  <head>
-    <meta charset="utf-8" />
-    <meta name="viewport" content="width=device-width, initial-scale=1" />
-    <link
-      rel="icon"
-      type="image/x-icon"
-      href="/modular-diffusion/images/logo.ico"
-    />
-    <title>{title} | Modular Diffusion</title>
-    <meta
-      name="description"
-      content="Learn about Modular Diffusion, an easy-to-use Python library for designing and training your own custom Generative Diffusion Models."
-    />
-    <link rel="sitemap" href="/sitemap-index.xml" />
-    <link
-      rel="stylesheet"
-      href="https://cdn.jsdelivr.net/npm/[email protected]/dist/katex.min.css"
-      integrity="sha384-MlJdn/WNKDGXveldHDdyRP1R4CTHr3FeuDNfhsLPYrq2t0UBkUdK2jyTnXPEK1NQ"
-      crossorigin="anonymous"
-    />
-  </head>
+<head>
+  <meta charset="utf-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1" />
+  <link
+    rel="icon"
+    type="image/x-icon"
+    href="/modular-diffusion/images/logo.ico"
+  />
+  <title>{title} | Modular Diffusion</title>
+  <meta
+    name="description"
+    content="Learn about Modular Diffusion, an easy-to-use Python library for designing and training your own custom Generative Diffusion Models."
+  />
+  <link rel="sitemap" href="/sitemap-index.xml" />
+  <link
+    rel="stylesheet"
+    href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0-beta3/css/all.min.css"
+    integrity="sha384-5e2ESR8Ycmos6g3gAKr1Jvwye8sW4U1u/cAKulfVJnkakCcMqhOudbtPnvJ+nbv7"
+    crossorigin="anonymous"
+  />
+  <script>
+    const theme = `(function() {
+      const savedTheme = localStorage.getItem('theme');
+      if (savedTheme) {
+        document.documentElement.classList.add(savedTheme);
+      }
+    })();`
+  </script>
+</head>
+
+
   <body>
     <div class="flex flex-col text-gray-800">
       <Header />

docs/src/pages/guides/custom-modules.mdx

@@ -16,7 +16,7 @@ When tinkering with Diffusion Models, the time will come when you need to ventur
 
 ## Data transform
 
-In many Diffusion Model applications, the diffusion process takes place in the dataset space. If this is your case, the prebuilt `Identity` data transform module will serve your purposes, leaving your data untouched before applying noise during training. However, a growing number of algorithms, like [Stable Diffusion](https://arxiv.org/abs/2112.10752) and [Diffusion-LM](https://arxiv.org/abs/2205.14217), project data onto a latent space before applying diffusion. 
+In many Diffusion Model applications, the diffusion process takes place in the dataset space. If this is your case, the prebuilt `Identity` data transform module will serve your purposes, leaving your data untouched before applying noise during training. However, a growing number of algorithms, like [Stable Diffusion](https://arxiv.org/abs/2112.10752) and [Diffusion-LM](https://arxiv.org/abs/2205.14217), project data onto a latent space before applying diffusion.
 
 In the case of Diffusion-LM, the dataset consists of sequences of word IDs, but the diffusion process happens in the word embedding space. This means you need a way of converting sequences of word IDs into sequences of embeddings, and train the embeddings along with the Diffusion Model. In Modular Diffusion, this can be achieved by extending the `Data` base class and implement its `encode` and `decode` methods. The former projects the data into the latent space and the latter retrieves it to the dataset space. Let's take a look at how you could implement the aforementioned transform:
 
@@ -60,6 +60,7 @@ class Linear(Schedule):
 ```
 
 Given that `steps` is already a parameter in the base class, all we need to do is define `start` and `end` parameters, and use them to compute the $a_t$ values. Then, we can initialize the schedule with the syntax `Linear(steps, start, end)`.
+
 ## Probability distribution
 
 In the diffusion process, the chosen probability distribution plays a crucial role in modeling the noise that guides the transition between different states. The library comes prepackaged with a growing set of commonly used distributions, such as the `Normal` distribution, but different applications or experimental setups might require you to implement your own.
@@ -96,7 +97,7 @@ In most Diffusion Model applications, the standard choice of noise is Gaussian,
 
 - `schedule(self, alpha: Tensor) -> None`: This method is intended for precomputing resources based on the noise schedule $\alpha_t$ for $t \in {0,\dots,T}$. This can be beneficial for performance reasons when some calculations can be done ahead of time. A common use is calculating $\bar{\alpha}_{t}=\prod_{t=1}^{T}\alpha_{t}$.
 - `stationary(self, shape: tuple[int, ...]) -> Distribution`: This method is tasked with computing the stationary distribution $q(x_T)$, i.e., the noise distribution at the final time step, given a target shape.
-- `prior(self, x: Tensor, t: Tensor) -> Distribution`: This method computes the prior distribution $q(x_t | x_0)$, i.e., the distribution of the noisy images $x_t$ or `z`  given the initial image $x_0$ or `x`.
+- `prior(self, x: Tensor, t: Tensor) -> Distribution`: This method computes the prior distribution $q(x_t | x_0)$, i.e., the distribution of the noisy images $x_t$ or `z` given the initial image $x_0$ or `x`.
 - `posterior(self, x: Tensor, z: Tensor, t: Tensor) -> Distribution`: This method computes the posterior distribution $q(x_{t-1} | x_t, x_0)$, i.e., the distribution of the less noisy images $x_{t-1}$ given the current noisy image $x_t$ or `z` and the initial image $x_0$ or `x`.
 - `approximate(self, z: Tensor, t: Tensor, hat: Tensor) -> Distribution`: This method computes the approximate posterior distribution $p_\theta(x_{t-1} | x_t)$, i.e., the distribution of the less noisy images $x_{t-1}$ given the current noisy image $x_t$ or `z`. This is an approximation to the true posterior distribution that is easier to sample from or compute. The tensor `hat` is the output of the denoiser network containing the predicted parameters -- named this way because predicted values often are denoted with a hat, e.g., $\hat{\epsilon}$.
 
@@ -206,4 +207,4 @@ loss = Simple(parameter="epsilon") + 0.001 * VLB()
 
 ## Guidance
 
-As of right now, `ClassifierFree` guidance is hardcoded into the diffusion process, and there is no way of extending the base `Guidance` class, unless you create your own custom `Model` class. You can expect this behavior to change in an upcoming release. Please refer to our official [Issue Tracker](https://github.com/cabralpinto/modular-diffusion/issues) for updates. 
+As of right now, `ClassifierFree` guidance is hardcoded into the diffusion process, and there is no way of extending the base `Guidance` class, unless you create your own custom `Model` class. You can expect this behavior to change in an upcoming release. Please refer to our official [Issue Tracker](https://github.com/cabralpinto/modular-diffusion/issues) for updates.