Flow vs Diffusion

1
# x0: [B, C, H, W]  real image in [-1,1] or [0,1]
2
t = sample_uniform_t(batch_size)
3
alpha_t = get_alpha(t)                      # precomputed schedule
4
eps = torch.randn_like(x0)                  # Gaussian noise
5

6
# forward noising
7
xt = (alpha_t.sqrt() * x0 +
8
      (1 - alpha_t).sqrt() * eps)
9

10
pred_eps = model(xt, t)                     # predict noise
11

12
loss = ((pred_eps - eps)**2).mean()

1
# x0: [B, C, H, W]
2
t = sample_uniform_t(batch_size)
3
x1 = torch.randn_like(x0)                   # Gaussian prior
4

5
# forward interpolation
6
xt = (1 - t) * x0 + t * x1
7

8
# true velocity
9
v_true = x1 - x0
10

11
pred_v = model(xt, t)                       # predict velocity
12

13
loss = ((pred_v - v_true)**2).mean()