mirror of
https://github.com/Laurent2916/Diffusers.jl.git
synced 2024-09-19 11:05:28 +00:00
✨ add docs using Documenter.jl
This commit is contained in:
parent
3c9841ffd9
commit
3c455c8b2e
5
docs/Project.toml
Normal file
5
docs/Project.toml
Normal file
|
@ -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"
|
20
docs/make.jl
Normal file
20
docs/make.jl
Normal file
|
@ -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",
|
||||||
|
]
|
||||||
|
)
|
51
docs/src/index.md
Normal file
51
docs/src/index.md
Normal file
|
@ -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]
|
||||||
|
```
|
|
@ -1,4 +1,6 @@
|
||||||
import Diffusers
|
import Diffusers
|
||||||
|
import Diffusers.Schedulers
|
||||||
|
import Diffusers.Schedulers: DDPM
|
||||||
import Diffusers.BetaSchedules: cosine_beta_schedule, rescale_zero_terminal_snr
|
import Diffusers.BetaSchedules: cosine_beta_schedule, rescale_zero_terminal_snr
|
||||||
using Flux
|
using Flux
|
||||||
using Random
|
using Random
|
||||||
|
@ -41,7 +43,7 @@ scatter(data[1, :], data[2, :],
|
||||||
)
|
)
|
||||||
|
|
||||||
num_timesteps = 100
|
num_timesteps = 100
|
||||||
scheduler = Diffusers.DDPM(
|
scheduler = DDPM(
|
||||||
Vector{Float64},
|
Vector{Float64},
|
||||||
rescale_zero_terminal_snr(
|
rescale_zero_terminal_snr(
|
||||||
cosine_beta_schedule(num_timesteps, 0.999f0, 0.001f0)
|
cosine_beta_schedule(num_timesteps, 0.999f0, 0.001f0)
|
||||||
|
@ -51,7 +53,7 @@ scheduler = Diffusers.DDPM(
|
||||||
noise = randn(size(data))
|
noise = randn(size(data))
|
||||||
|
|
||||||
anim = @animate for i in cat(1:num_timesteps, repeat([num_timesteps], 50), dims=1)
|
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, :],
|
scatter(noise[1, :], noise[2, :],
|
||||||
alpha=0.1,
|
alpha=0.1,
|
||||||
aspectratio=:equal,
|
aspectratio=:equal,
|
||||||
|
@ -117,10 +119,10 @@ for epoch = 1:num_epochs
|
||||||
for data in dataloader
|
for data in dataloader
|
||||||
noise = randn(size(data))
|
noise = randn(size(data))
|
||||||
timesteps = rand(2:num_timesteps, size(data)[2]) # TODO: fix start at timestep=2, bruh
|
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
|
grads = Flux.gradient(model) do m
|
||||||
model_output = m(noisy_data, timesteps)
|
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)
|
loss(noise, noise_prediction)
|
||||||
end
|
end
|
||||||
Flux.update!(opt, params, grads)
|
Flux.update!(opt, params, grads)
|
||||||
|
@ -131,7 +133,7 @@ end
|
||||||
# sampling animation
|
# sampling animation
|
||||||
anim = @animate for timestep in num_timesteps:-1:2
|
anim = @animate for timestep in num_timesteps:-1:2
|
||||||
model_output = model(data, [timestep])
|
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, :],
|
p1 = scatter(sampled_data[1, :], sampled_data[2, :],
|
||||||
alpha=0.5,
|
alpha=0.5,
|
||||||
|
|
|
@ -1,3 +1,7 @@
|
||||||
|
"""
|
||||||
|
Abstract type for schedulers.
|
||||||
|
|
||||||
|
"""
|
||||||
abstract type Scheduler end
|
abstract type Scheduler end
|
||||||
|
|
||||||
"""
|
"""
|
||||||
|
@ -69,7 +73,7 @@ end
|
||||||
Add noise to clean data using the forward diffusion process.
|
Add noise to clean data using the forward diffusion process.
|
||||||
|
|
||||||
## Input
|
## Input
|
||||||
* `scheduler::Scheduler`: scheduler to use
|
* `scheduler::DDPM`: scheduler to use
|
||||||
* `x₀::AbstractArray`: clean data to add noise to
|
* `x₀::AbstractArray`: clean data to add noise to
|
||||||
* `ϵ::AbstractArray`: noise to add to clean data
|
* `ϵ::AbstractArray`: noise to add to clean data
|
||||||
* `t::AbstractArray`: timesteps used to weight the noise
|
* `t::AbstractArray`: timesteps used to weight the noise
|
||||||
|
|
Loading…
Reference in a new issue