1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
|
"""
Diffusion process utilities.
Contains forward and reverse diffusion functions.
"""
import torch
from .config import DEVICE, ALPHA, ALPHA_BAR, BETA, DIFFU_STEPS
def q_xt_xt_1(xt_1: torch.Tensor, t: int | torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Forward diffusion step: q(x_t | x_{t-1}).
Adds noise to image at step t-1 to get image at step t.
Args:
xt_1: Image at timestep t-1 [B, C, H, W]
t: Timestep (int or tensor)
Returns:
xt: Noisy image at timestep t
eps: The noise that was added
"""
if isinstance(t, int):
t_ind = torch.tensor(t, dtype=torch.long, device=DEVICE)
else:
t_ind = t.to(dtype=torch.long, device=DEVICE)
alpha = ALPHA[t_ind]
mean = torch.sqrt(alpha) * xt_1
std = torch.sqrt(1 - alpha)
eps = torch.randn(xt_1.shape, device=DEVICE)
xt = mean + std * eps
return xt, eps
def q_xt_x0(x0: torch.Tensor, t: int | torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Forward diffusion: q(x_t | x_0).
Directly compute noisy image at any timestep t from clean image x0.
Args:
x0: Clean image [B, C, H, W] or [C, H, W]
t: Timestep tensor [B, 1] or [B]
Returns:
xt: Noisy image at timestep t
eps: The noise that was added
"""
if isinstance(t, int):
t_ind = torch.tensor(t, dtype=torch.long, device=DEVICE)
else:
t_ind = t.to(dtype=torch.long, device=DEVICE)
# Handle both batched and single images
reshaped = len(x0.shape) == 3
if reshaped:
c, w, h = x0.shape
b = 1
x0 = x0.view(b, c, w, h)
else:
b, c, w, h = x0.shape
# Reshape t for broadcasting
t_ind = t_ind.view(b, 1, 1, 1)
t_ind = t_ind.expand(b, c, w, h)
alpha_bar = ALPHA_BAR[t_ind]
mean = torch.sqrt(alpha_bar) * x0
std = torch.sqrt(1 - alpha_bar)
eps = torch.randn(x0.shape, device=DEVICE)
xt = mean + std * eps
if reshaped:
xt = xt.view(c, w, h)
return xt, eps
def p_xt_1_xt(model: torch.nn.Module, xt: torch.Tensor, t: torch.Tensor,
vec: torch.Tensor) -> torch.Tensor:
"""
Reverse diffusion step: p(x_{t-1} | x_t).
Denoise image at step t to get image at step t-1.
Model predicts the noise.
Args:
model: UNet model that predicts noise
xt: Noisy image at timestep t [B, C, H, W]
t: Timestep tensor [B, 1]
vec: Label one-hot vector [B, NB_LABEL]
Returns:
xt_1: Denoised image at timestep t-1
"""
t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t
alpha_bar_t = ALPHA_BAR[t_ind].view(-1, 1, 1, 1)
alpha_bar_t_1 = ALPHA_BAR[t_ind - 1].view(-1, 1, 1, 1)
alpha_t = ALPHA[t_ind].view(-1, 1, 1, 1)
beta_t = BETA[t_ind].view(-1, 1, 1, 1)
beta_tilde = (1 - alpha_bar_t_1) / (1 - alpha_bar_t) * beta_t
# Model predicts the noise
epsilon_theta = model(xt, t, vec)
sigma_theta = torch.sqrt(beta_tilde)
mu_theta = (xt - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * epsilon_theta) / torch.sqrt(alpha_t)
# Don't add noise at t=1
mask_t0 = (t > 1).to(dtype=torch.float32).view(-1, 1, 1, 1)
noise = torch.randn(xt.shape, device=DEVICE) * mask_t0
return mu_theta + sigma_theta * noise
def p_xt_1_xt_x0_pred(model: torch.nn.Module, xt: torch.Tensor, t: torch.Tensor,
vec: torch.Tensor) -> torch.Tensor:
"""
Reverse diffusion step where model predicts x0 directly.
Args:
model: UNet model that predicts clean image x0
xt: Noisy image at timestep t [B, C, H, W]
t: Timestep tensor [B, 1]
vec: Label one-hot vector [B, NB_LABEL]
Returns:
xt_1: Denoised image at timestep t-1
"""
x0_pred = model(xt, t, vec)
xt_1, _ = q_xt_x0(x0_pred, t - 1)
return xt_1
def forward_diffusion(x0: torch.Tensor) -> list[torch.Tensor]:
"""
Run full forward diffusion process.
Args:
x0: Clean image [C, H, W] or [B, C, H, W]
Returns:
List of images at each timestep [x0, x1, ..., xT]
"""
x = x0.clone()
xs = [x]
for t in range(1, DIFFU_STEPS + 1):
x, _ = q_xt_xt_1(x, t)
xs.append(x)
return xs
|
