From 3c455c8b2e13a3dbff6ea4b3ebb7ddbee0712520 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Laure=CE=B7t?= <laurent@fainsin.bzh>
Date: Sat, 29 Jul 2023 17:52:36 +0200
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20add=20docs=20using=20Documenter.jl?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 docs/Project.toml      |  5 +++++
 docs/make.jl           | 20 +++++++++++++++++
 docs/src/index.md      | 51 ++++++++++++++++++++++++++++++++++++++++++
 examples/swissroll.jl  | 12 +++++-----
 src/Schedulers/DDPM.jl |  6 ++++-
 5 files changed, 88 insertions(+), 6 deletions(-)
 create mode 100644 docs/Project.toml
 create mode 100644 docs/make.jl
 create mode 100644 docs/src/index.md

diff --git a/docs/Project.toml b/docs/Project.toml
new file mode 100644
index 0000000..e8f458e
--- /dev/null
+++ b/docs/Project.toml
@@ -0,0 +1,5 @@
+[deps]
+Diffusers = "90edb7a8-79d7-49b2-b6b1-9322c3fdead8"
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+PlotlyJS = "f0f68f2c-4968-5e81-91da-67840de0976a"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
diff --git a/docs/make.jl b/docs/make.jl
new file mode 100644
index 0000000..4b41d9b
--- /dev/null
+++ b/docs/make.jl
@@ -0,0 +1,20 @@
+using Diffusers
+using Documenter
+using Plots
+
+DocMeta.setdocmeta!(Diffusers, :DocTestSetup, :(using Diffusers); recursive=true)
+
+makedocs(;
+  modules=[Diffusers],
+  authors="Laurent Fainsin <laurent@fainsin.bzh>",
+  repo="https://github.com/Laurent2916/Diffusers.jl/blob/{commit}{path}#{line}",
+  sitename="Diffusers.jl",
+  format=Documenter.HTML(;
+    prettyurls=get(ENV, "CI", "false") == "true",
+    edit_link="main",
+    assets=String[]
+  ),
+  pages=[
+    "Home" => "index.md",
+  ]
+)
diff --git a/docs/src/index.md b/docs/src/index.md
new file mode 100644
index 0000000..55ff327
--- /dev/null
+++ b/docs/src/index.md
@@ -0,0 +1,51 @@
+```@meta
+CurrentModule = Diffusers
+```
+
+# Diffusers
+
+Documentation for [Diffusers.jl](https://github.com/Laurent2916/Diffusers.jl).
+
+```@eval
+using Diffusers.BetaSchedules
+using Plots
+plotlyjs()
+
+T = 1000
+linear = linear_beta_schedule(T)
+scaled_linear = scaled_linear_beta_schedule(T)
+cosine = cosine_beta_schedule(T)
+sigmoid = sigmoid_beta_schedule(T)
+
+plot(
+  [linear, scaled_linear, cosine, sigmoid],
+  label=["linear" "scaled_linear" "cosine" "sigmoid"],
+  xlabel="t",
+  ylabel="β",
+  title="Beta schedules",
+  legend=:topleft,
+  yscale=:log10,
+)
+
+savefig("beta_schedules.html")
+nothing
+```
+
+```@raw html
+<object type="text/html" data="beta_schedules.html" style="width:100%;height:420px;"></object>
+```
+
+```@index
+```
+
+```@autodocs
+Modules = [Diffusers]
+```
+
+```@autodocs
+Modules = [Diffusers.BetaSchedules]
+```
+
+```@autodocs
+Modules = [Diffusers.Schedulers]
+```
diff --git a/examples/swissroll.jl b/examples/swissroll.jl
index 046c8a7..6de614f 100644
--- a/examples/swissroll.jl
+++ b/examples/swissroll.jl
@@ -1,4 +1,6 @@
 import Diffusers
+import Diffusers.Schedulers
+import Diffusers.Schedulers: DDPM
 import Diffusers.BetaSchedules: cosine_beta_schedule, rescale_zero_terminal_snr
 using Flux
 using Random
@@ -41,7 +43,7 @@ scatter(data[1, :], data[2, :],
 )
 
 num_timesteps = 100
-scheduler = Diffusers.DDPM(
+scheduler = DDPM(
   Vector{Float64},
   rescale_zero_terminal_snr(
     cosine_beta_schedule(num_timesteps, 0.999f0, 0.001f0)
@@ -51,7 +53,7 @@ scheduler = Diffusers.DDPM(
 noise = randn(size(data))
 
 anim = @animate for i in cat(1:num_timesteps, repeat([num_timesteps], 50), dims=1)
-  noisy_data = Diffusers.add_noise(scheduler, data, noise, [i])
+  noisy_data = Diffusers.Schedulers.add_noise(scheduler, data, noise, [i])
   scatter(noise[1, :], noise[2, :],
     alpha=0.1,
     aspectratio=:equal,
@@ -117,10 +119,10 @@ for epoch = 1:num_epochs
   for data in dataloader
     noise = randn(size(data))
     timesteps = rand(2:num_timesteps, size(data)[2]) # TODO: fix start at timestep=2, bruh
-    noisy_data = Diffusers.add_noise(scheduler, data, noise, timesteps)
+    noisy_data = Diffusers.Schedulers.add_noise(scheduler, data, noise, timesteps)
     grads = Flux.gradient(model) do m
       model_output = m(noisy_data, timesteps)
-      noise_prediction, _ = Diffusers.step(scheduler, noisy_data, model_output, timesteps)
+      noise_prediction, _ = Diffusers.Schedulers.step(scheduler, noisy_data, model_output, timesteps)
       loss(noise, noise_prediction)
     end
     Flux.update!(opt, params, grads)
@@ -131,7 +133,7 @@ end
 # sampling animation
 anim = @animate for timestep in num_timesteps:-1:2
   model_output = model(data, [timestep])
-  sampled_data, x0_pred = Diffusers.step(scheduler, data, model_output, [timestep])
+  sampled_data, x0_pred = Diffusers.Schedulers.step(scheduler, data, model_output, [timestep])
 
   p1 = scatter(sampled_data[1, :], sampled_data[2, :],
     alpha=0.5,
diff --git a/src/Schedulers/DDPM.jl b/src/Schedulers/DDPM.jl
index 1f655e7..bbc4ac6 100644
--- a/src/Schedulers/DDPM.jl
+++ b/src/Schedulers/DDPM.jl
@@ -1,3 +1,7 @@
+"""
+Abstract type for schedulers.
+
+"""
 abstract type Scheduler end
 
 """
@@ -69,7 +73,7 @@ end
 Add noise to clean data using the forward diffusion process.
 
 ## Input
-  * `scheduler::Scheduler`: scheduler to use
+  * `scheduler::DDPM`: scheduler to use
   * `x₀::AbstractArray`: clean data to add noise to
   * `ϵ::AbstractArray`: noise to add to clean data
   * `t::AbstractArray`: timesteps used to weight the noise