From b4ed6e3c99fbaef0d3d4295ead8745582bb5ee17 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Laure=CE=B7t?= <laurent@fainsin.bzh>
Date: Tue, 25 Jul 2023 20:08:35 +0200
Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=9A=20move=20DDPM's=20step=20from=20Sc?=
 =?UTF-8?q?hedulers.jl=20to=20DDPM.jl?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/DDPM.jl       | 49 ++++++++++++++++++++++++++++++++++++++++++++---
 src/Schedulers.jl | 48 ----------------------------------------------
 2 files changed, 46 insertions(+), 51 deletions(-)

diff --git a/src/DDPM.jl b/src/DDPM.jl
index ebb3fd3..322af2f 100644
--- a/src/DDPM.jl
+++ b/src/DDPM.jl
@@ -19,7 +19,6 @@ struct DDPM{V<:AbstractVector} <: Scheduler
   sqrt_one_minus_α_cumprods::V
 end
 
-
 function DDPM(V::DataType, βs::AbstractVector)
   αs = 1 .- βs
   α_cumprods = cumprod(αs)
@@ -39,6 +38,50 @@ function DDPM(V::DataType, βs::AbstractVector)
   )
 end
 
-function DDPM(V::DataType, beta_scheduler)
-  DDPM(V, beta_scheduler)
+function step(
+  scheduler::DDPM,
+  sample::AbstractArray,
+  model_output::AbstractArray,
+  timesteps::AbstractArray,
+)
+  """
+  Remove noise from model output using the backward diffusion process.
+
+  Args:
+    scheduler (`Scheduler`): scheduler object.
+    sample (`AbstractArray`): sample to remove noise from, i.e. model_input.
+    model_output (`AbstractArray`): predicted noise from the model.
+    timesteps (`AbstractArray`): timesteps to remove noise from.
+
+  Returns:
+    `AbstractArray`: denoised model output at the given timestep.
+  """
+  # 1. compute alphas, betas
+  α_cumprod_t = scheduler.α_cumprods[timesteps]
+  α_cumprod_t_prev = scheduler.α_cumprods[timesteps.-1]
+  β_cumprod_t = 1 .- α_cumprod_t
+  β_cumprod_t_prev = 1 .- α_cumprod_t_prev
+  current_α_t = α_cumprod_t ./ α_cumprod_t_prev
+  current_β_t = 1 .- current_α_t
+
+  # 2. compute predicted original sample from predicted noise also called
+  # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
+  # epsilon prediction type
+  # print shapes of thingies
+  x_0_pred = (sample - sqrt.(β_cumprod_t)' .* model_output) ./ sqrt.(α_cumprod_t)'
+
+  # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
+  # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
+  pred_original_sample_coeff = (sqrt.(α_cumprod_t_prev) .* current_β_t) ./ β_cumprod_t
+  current_sample_coeff = sqrt.(current_α_t) .* β_cumprod_t_prev ./ β_cumprod_t
+
+  # 5. Compute predicted previous sample µ_t
+  # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
+  pred_prev_sample = pred_original_sample_coeff' .* x_0_pred + current_sample_coeff' .* sample
+
+  # 6. Add noise
+  variance = sqrt.(scheduler.βs[timesteps])' .* randn(size(model_output))
+  pred_prev_sample = pred_prev_sample + variance
+
+  return pred_prev_sample, x_0_pred
 end
diff --git a/src/Schedulers.jl b/src/Schedulers.jl
index d3ecd99..1ba7869 100644
--- a/src/Schedulers.jl
+++ b/src/Schedulers.jl
@@ -25,51 +25,3 @@ function add_noise(
 
   return noisy_data
 end
-
-function step(
-  scheduler::Scheduler,
-  sample::AbstractArray,
-  model_output::AbstractArray,
-  timesteps::AbstractArray,
-)
-  """
-  Remove noise from model output using the backward diffusion process.
-
-  Args:
-    scheduler (`Scheduler`): scheduler object.
-    sample (`AbstractArray`): sample to remove noise from, i.e. model_input.
-    model_output (`AbstractArray`): predicted noise from the model.
-    timesteps (`AbstractArray`): timesteps to remove noise from.
-
-  Returns:
-    `AbstractArray`: denoised model output at the given timestep.
-  """
-  # 1. compute alphas, betas
-  α_cumprod_t = scheduler.α_cumprods[timesteps]
-  α_cumprod_t_prev = scheduler.α_cumprods[timesteps.-1]
-  β_cumprod_t = 1 .- α_cumprod_t
-  β_cumprod_t_prev = 1 .- α_cumprod_t_prev
-  current_α_t = α_cumprod_t ./ α_cumprod_t_prev
-  current_β_t = 1 .- current_α_t
-
-  # 2. compute predicted original sample from predicted noise also called
-  # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
-  # epsilon prediction type
-  # print shapes of thingies
-  x_0_pred = (sample - sqrt.(β_cumprod_t)' .* model_output) ./ sqrt.(α_cumprod_t)'
-
-  # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
-  # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
-  pred_original_sample_coeff = (sqrt.(α_cumprod_t_prev) .* current_β_t) ./ β_cumprod_t
-  current_sample_coeff = sqrt.(current_α_t) .* β_cumprod_t_prev ./ β_cumprod_t
-
-  # 5. Compute predicted previous sample µ_t
-  # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
-  pred_prev_sample = pred_original_sample_coeff' .* x_0_pred + current_sample_coeff' .* sample
-
-  # 6. Add noise
-  variance = sqrt.(scheduler.βs[timesteps])' .* randn(size(model_output))
-  pred_prev_sample = pred_prev_sample + variance
-
-  return pred_prev_sample
-end