feat: migrate DDIM to Reactant

Author: avik-pal

Diff for: examples/DDIM/Project.toml

@@ -1,10 +1,10 @@
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
-CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Comonicon = "863f3e99-da2a-4334-8734-de3dacbe5542"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DataAugmentation = "88a5189c-e7ff-4f85-ac6b-e6158070f02e"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 ImageCore = "a09fc81d-aa75-5fe9-8630-4744c3626534"
 ImageIO = "82e4d734-157c-48bb-816b-45c225c6df19"
@@ -14,8 +14,9 @@ LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 ParameterSchedulers = "d7d3b36b-41b8-4d0d-a2bf-768c6151755e"
-ProgressBars = "49802e3a-d2f1-5c88-81d8-b72133a6f568"
+ProgressTables = "e0b4b9f6-8cc7-451e-9c86-94c5316e9f73"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -24,7 +25,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ArgCheck = "2.3.0"
-CairoMakie = "0.12"
 Comonicon = "1"
 ConcreteStructs = "0.2.3"
 DataAugmentation = "0.3"
@@ -38,7 +38,6 @@ LuxCUDA = "0.3"
 MLUtils = "0.4"
 Optimisers = "0.4.1"
 ParameterSchedulers = "0.4.1"
-ProgressBars = "1"
 Random = "1.10"
 Setfield = "1"
 StableRNGs = "1.0.2"

Diff for: examples/DDIM/main.jl

@@ -6,9 +6,9 @@
 
 # ## Package Imports
 
-using ArgCheck, CairoMakie, ConcreteStructs, Comonicon, DataAugmentation, DataDeps, FileIO,
-      ImageCore, JLD2, Lux, LuxCUDA, MLUtils, Optimisers, ParameterSchedulers, ProgressBars,
-      Random, Setfield, StableRNGs, Statistics, Zygote
+using ArgCheck, ConcreteStructs, Comonicon, DataAugmentation, DataDeps, Enzyme, FileIO,
+      ImageCore, JLD2, Lux, LuxCUDA, MLUtils, Optimisers, ParameterSchedulers,
+      ProgressTables, Random, Reactant, Setfield, StableRNGs, Statistics, Zygote
 using TensorBoardLogger: TBLogger, log_value, log_images
 
 CUDA.allowscalar(false)
@@ -19,16 +19,15 @@ CUDA.allowscalar(false)
 # [the Keras example](https://keras.io/examples/generative/ddim/). Embed noise variances to
 # embedding.
 
-function sinusoidal_embedding(x::AbstractArray{T, 4}, min_freq::T, max_freq::T,
-        embedding_dims::Int) where {T <: AbstractFloat}
-    size(x)[1:3] != (1, 1, 1) &&
+function sinusoidal_embedding(
+        x::AbstractArray{T, 4}, min_freq, max_freq, embedding_dims::Int) where {T}
+    if size(x)[1:3] != (1, 1, 1)
         throw(DimensionMismatch("Input shape must be (1, 1, 1, batch)"))
+    end
 
-    lower, upper = log(min_freq), log(max_freq)
+    lower, upper = T(log(min_freq)), T(log(max_freq))
     n = embedding_dims ÷ 2
-    d = (upper - lower) / (n - 1)
-    freqs = reshape(exp.(lower:d:upper) |> get_device(x), 1, 1, n, 1)
-    x_ = 2 .* x .* freqs
+    x_ = 2 .* x .* exp.(reshape(range(lower, upper; length=n), 1, 1, n, 1))
     return cat(sinpi.(x_), cospi.(x_); dims=Val(3))
 end
 
@@ -97,7 +96,7 @@ function unet_model(image_size::Tuple{Int, Int}; channels=[32, 64, 96, 128],
     return @compact(;
         upsample, conv_in, conv_out, down_blocks, residual_blocks, up_blocks,
         min_freq, max_freq, embedding_dims,
-        num_blocks=(length(channels) - 1)) do x::Tuple{AbstractArray{<:Real, 4}, AbstractArray{<:Real, 4}}
+        num_blocks=(length(channels) - 1)) do x::Tuple{<:AbstractArray, <:AbstractArray}
     #! format: on
         noisy_images, noise_variances = x
 
@@ -125,8 +124,7 @@ function ddim(rng::AbstractRNG, args...; min_signal_rate=0.02f0,
     unet = unet_model(args...; kwargs...)
     bn = BatchNorm(3; affine=false, track_stats=true)
 
-    return @compact(; unet, bn, rng, min_signal_rate,
-        max_signal_rate, dispatch=:DDIM) do x::AbstractArray{<:Real, 4}
+    return @compact(; unet, bn, rng, min_signal_rate, max_signal_rate, dispatch=:DDIM) do x
         images = bn(x)
         rng = Lux.replicate(rng)
 
@@ -144,10 +142,10 @@ function ddim(rng::AbstractRNG, args...; min_signal_rate=0.02f0,
     end
 end
 
-function diffusion_schedules(diffusion_times::AbstractArray{T, 4}, min_signal_rate::T,
-        max_signal_rate::T) where {T <: Real}
-    start_angle = acos(max_signal_rate)
-    end_angle = acos(min_signal_rate)
+function diffusion_schedules(
+        diffusion_times::AbstractArray{T, 4}, min_signal_rate, max_signal_rate) where {T}
+    start_angle = T(acos(max_signal_rate))
+    end_angle = T(acos(min_signal_rate))
 
     diffusion_angles = @. start_angle + (end_angle - start_angle) * diffusion_times
 
@@ -157,8 +155,14 @@ function diffusion_schedules(diffusion_times::AbstractArray{T, 4}, min_signal_ra
     return noise_rates, signal_rates
 end
 
-function denoise(unet, noisy_images::AbstractArray{T, 4}, noise_rates::AbstractArray{T, 4},
-        signal_rates::AbstractArray{T, 4}) where {T <: Real}
+function denoise(
+        unet, noisy_images::AbstractArray{T1, 4}, noise_rates::AbstractArray{T2, 4},
+        signal_rates::AbstractArray{T3, 4}) where {T1, T2, T3}
+    T = promote_type(T1, T2, T3)
+    noisy_images = T.(noisy_images)
+    noise_rates = T.(noise_rates)
+    signal_rates = T.(signal_rates)
+
     pred_noises = unet((noisy_images, noise_rates .^ 2))
     pred_images = @. (noisy_images - pred_noises * noise_rates) / signal_rates
     return pred_noises, pred_images