♻️ (Schedulers) rename methods names

add_noise -> forward
step -> reverse

src/Schedulers/Abstract.jl

@@ -16,7 +16,7 @@ Add noise to clean data using the forward diffusion process.
 ## Output
   * `xₜ::AbstractArray`: noisy data at the given timesteps
 """
-function add_noise(
+function forward(
   scheduler::Scheduler,
   x₀::AbstractArray,
   ϵ::AbstractArray,
@@ -36,7 +36,7 @@ Remove noise from model output using the backward diffusion process.
   * `xₜ₋₁::AbstractArray`: denoised sample at t=t-1
   * `x̂₀::AbstractArray`: denoised sample at t=0
 """
-function step(
+function reverse(
   scheduler::Scheduler,
   xₜ::AbstractArray,
   ϵᵧ::AbstractArray,

src/Schedulers/DDPM.jl

@@ -65,7 +65,7 @@ function DDPM(β::AbstractVector)
   )
 end
 
-function add_noise(
+function forward(
   scheduler::DDPM,
   x₀::AbstractArray,
   ϵ::AbstractArray,
@@ -86,7 +86,7 @@ function add_noise(
   return xₜ
 end
 
-function step(
+function reverse(
   scheduler::DDPM,
   xₜ::AbstractArray,
   ϵᵧ::AbstractArray,

src/Schedulers/Schedulers.jl

@@ -7,7 +7,7 @@ export
   DDPM,
 
   # Scheduler methods
-  add_noise,
+  forward,
-  step
+  reverse
 
 end # module Schedulers

test/BetaSchedules.jl

@@ -15,7 +15,7 @@ using Test
     α̅_cosine = cumprod(1 .- β_cosine)
     α̅_sigmoid = cumprod(1 .- β_sigmoid)
 
-    # arxiv:2208.11970 (eq. 109)
+    # arxiv:2208.11970 Eq. 109
     SNR_linear = α̅_linear ./ (1 .- α̅_linear)
     SNR_scaled_linear = α̅_scaled_linear ./ (1 .- α̅_scaled_linear)
     SNR_cosine = α̅_cosine ./ (1 .- α̅_cosine)

test/Schedulers.jl

@@ -1,9 +1,9 @@
-import Diffusers: step, add_noise, DDPM, cosine_beta_schedule, rescale_zero_terminal_snr
+import Diffusers: reverse, forward, DDPM, cosine_beta_schedule
 using Statistics
 using Test
 
 @testset "Schedulers tests" begin
-  @testset "check `step` correctness" begin
+  @testset "check `reverse` correctness" begin
     T = 10
     batch_size = 8
     size = 128
@@ -20,15 +20,15 @@ using Test
     for t in 1:T
       t = ones(UInt32, batch_size) .* t
       # corrupt x₀ with noise
-      xₜ = Diffusers.add_noise(ddpm, x₀, ϵ, t)
+      xₜ = forward(ddpm, x₀, ϵ, t)
       # suppose a model predicted ϵ perfectly
-      _, x̂₀ = Diffusers.step(ddpm, xₜ, ϵ, t)
+      _, x̂₀ = reverse(ddpm, xₜ, ϵ, t)
       # test that we recover x₀
       @test x̂₀ ≈ x₀
     end
   end
 
-  @testset "check `add_noise` terminal SNR" begin
+  @testset "check `forward` terminal SNR" begin
     T = 10
     batch_size = 1
     size = 2500
@@ -46,7 +46,7 @@ using Test
 
     t = ones(UInt32, batch_size) .* T
     # corrupt x₀ with noise
-    xₜ = Diffusers.add_noise(ddpm, x₀, ϵ, t)
+    xₜ = forward(ddpm, x₀, ϵ, t)
 
     @test std(xₜ) ≈ 1.0 atol = 1.0f-3
     @test mean(xₜ) ≈ 0.0 atol = 1.0f-2