mirror of
https://github.com/Laurent2916/Diffusers.jl.git
synced 2024-09-18 18:45:28 +00:00
🧪 (Schedulers) add get_velocity tests
This commit is contained in:
parent
b90b50de27
commit
184b9665ec
|
@ -1,8 +1,9 @@
|
||||||
import Diffusers: reverse, forward, DDPM, cosine_beta_schedule
|
import Diffusers: reverse, forward, get_velocity, DDPM, cosine_beta_schedule, VELOCITY
|
||||||
import Statistics: mean, std
|
import Statistics: mean, std
|
||||||
using Test
|
using Test
|
||||||
|
|
||||||
@testset "Schedulers tests" begin
|
@testset "Schedulers tests" begin
|
||||||
|
|
||||||
@testset "check `reverse` correctness" begin
|
@testset "check `reverse` correctness" begin
|
||||||
T = 10
|
T = 10
|
||||||
batch_size = 8
|
batch_size = 8
|
||||||
|
@ -15,17 +16,38 @@ using Test
|
||||||
x₀ = ones(Float32, size, size, batch_size)
|
x₀ = ones(Float32, size, size, batch_size)
|
||||||
ϵ = randn(Float32, size, size, batch_size)
|
ϵ = randn(Float32, size, size, batch_size)
|
||||||
|
|
||||||
|
@testset "PredictionType == EPSILON" begin
|
||||||
for t in 1:T
|
for t in 1:T
|
||||||
t = ones(UInt32, batch_size) .* t
|
t = ones(UInt32, batch_size) .* t
|
||||||
# corrupt x₀ with noise
|
# get xₜ from forward diffusion process
|
||||||
xₜ = forward(ddpm, x₀, ϵ, t)
|
xₜ = forward(ddpm, x₀, ϵ, t)
|
||||||
# suppose a model predicted ϵ perfectly
|
# suppose a model predicted ϵ perfectly
|
||||||
_, x̂₀ = reverse(ddpm, xₜ, ϵ, t)
|
ϵᵧ = ϵ
|
||||||
|
# use reverse diffusion process to retreive x̂₀
|
||||||
|
_, x̂₀ = reverse(ddpm, xₜ, ϵᵧ, t)
|
||||||
# test that we recover x₀
|
# test that we recover x₀
|
||||||
@test x̂₀ ≈ x₀
|
@test x̂₀ ≈ x₀
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
|
||||||
|
@testset "PredictionType == VELOCITY" begin
|
||||||
|
for t in 1:T
|
||||||
|
t = ones(UInt32, batch_size) .* t
|
||||||
|
# get xₜ from forward diffusion process
|
||||||
|
xₜ = forward(ddpm, x₀, ϵ, t)
|
||||||
|
# compute vₜ to train model
|
||||||
|
vₜ = get_velocity(ddpm, x₀, ϵ, t)
|
||||||
|
# suppose a model predicted vₜ perfectly
|
||||||
|
vᵧ = vₜ
|
||||||
|
# use reverse diffusion process to retreive x̂₀
|
||||||
|
_, x̂₀ = Diffusers.reverse(ddpm, xₜ, vₜ, t, VELOCITY)
|
||||||
|
# test that we recover x₀
|
||||||
|
@test x̂₀ ≈ x₀
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
end
|
||||||
|
|
||||||
@testset "check `forward` terminal SNR" begin
|
@testset "check `forward` terminal SNR" begin
|
||||||
T = 10
|
T = 10
|
||||||
batch_size = 1
|
batch_size = 1
|
||||||
|
@ -45,4 +67,5 @@ using Test
|
||||||
@test std(xₜ) ≈ 1.0 atol = 1.0f-3
|
@test std(xₜ) ≈ 1.0 atol = 1.0f-3
|
||||||
@test mean(xₜ) ≈ 0.0 atol = 1.0f-2
|
@test mean(xₜ) ≈ 0.0 atol = 1.0f-2
|
||||||
end
|
end
|
||||||
|
|
||||||
end
|
end
|
||||||
|
|
Loading…
Reference in a new issue