feat: completed adding DDIM, DDPM and Euler samplers + training

AshishKumar4 · Jul 8, 2024 · ea41d03 · ea41d03
1 parent 69490a1
commit ea41d03
Show file tree

Hide file tree

Showing 4 changed files with 15,773 additions and 402 deletions.
diff --git a/Diffusion flax linen.ipynb b/Diffusion flax linen.ipynb
diff --git a/README.md b/README.md
@@ -1,6 +1,6 @@
 # ![](images/logo.jpeg "FlaxDiff")
 
-## A Versatile and Easy-to-Understand Diffusion Library
+## A Versatile and simple Diffusion Library
 
 In recent years, diffusion and score-based multi-step models have revolutionized the generative AI domain. However, the latest research in this field has become highly math-intensive, making it challenging to understand how state-of-the-art diffusion models work and generate such impressive images. Replicating this research in code can be daunting.
 

diff --git a/example notebooks/simple diffusion flax.ipynb b/example notebooks/simple diffusion flax.ipynb
diff --git a/flaxdiff/models/simple_unet.py b/flaxdiff/models/simple_unet.py
@@ -69,14 +69,11 @@ def __call__(self, x):
 
 class TimeEmbedding(nn.Module):
     features:int
-    max_timesteps:int=10000
+    nax_positions:int=10000
 
     def setup(self):
-        # self.embeddings = nn.Embed(
-        #     num_embeddings=max_timesteps, features=out_features
-        # )
         half_dim = self.features // 2
-        emb = jnp.log(self.max_timesteps) / (half_dim - 1)
+        emb = jnp.log(self.nax_positions) / (half_dim - 1)
         emb = jnp.exp(-emb * jnp.arange(half_dim, dtype=jnp.float32))
         self.embeddings = emb