Skip to content

Commit d462457

Browse files
chaircbotchaircbot
committed
refactor: Add the sampling loop method to reduce the method code repetition rate.
1 parent 9580110 commit d462457

4 files changed

Lines changed: 134 additions & 124 deletions

File tree

‎iddm/model/samples/base.py‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,17 @@
77
"""
88
import math
99
from typing import Optional, List, Union, Tuple
10+
import logging
11+
import coloredlogs
1012

1113
import torch
1214
from torch import nn
1315

1416
from iddm.config.setting import DEFAULT_IMAGE_SIZE, IMAGE_CHANNEL
1517

18+
logger = logging.getLogger(__name__)
19+
coloredlogs.install(level="INFO")
20+
1621

1722
class BaseDiffusion:
1823
"""
@@ -176,15 +181,42 @@ def sample(
176181
cfg_scale: Optional[float] = None
177182
) -> torch.Tensor:
178183
"""
179-
Sample method, this method should be implemented in the subclass
184+
Sample method, it is inherited by subclasses
180185
:param model: Model
181186
:param x: Input image tensor, if provided, will be used as the starting point for sampling
182187
:param n: Number of sample images, x priority is greater than n
183188
:param labels: Labels
184189
:param cfg_scale: classifier-free guidance interpolation weight, users can better generate model effect.
185190
:return: Sample images
186191
"""
187-
pass
192+
x, n = self._get_input_image(n=n, x=x)
193+
logger.info(msg=f"{self.__class__.__name__} Sampling {n} images...")
194+
model.eval()
195+
with torch.no_grad():
196+
# Subclasses implement specific iterative logic
197+
x = self._sample_loop(model, x, n, labels, cfg_scale)
198+
x = self.post_process(x)
199+
model.train()
200+
return x
201+
202+
def _sample_loop(
203+
self,
204+
model: nn.Module,
205+
x: Optional[torch.Tensor] = None,
206+
n: int = 1,
207+
labels: Optional[torch.Tensor] = None,
208+
cfg_scale: Optional[float] = None
209+
) -> torch.Tensor:
210+
"""
211+
Sampling iteration logic, to be implemented in subclasses
212+
:param model: Model
213+
:param x: Input image tensor
214+
:param n: Number of sample images
215+
:param labels: Labels
216+
:param cfg_scale: Classifier-free guidance scale
217+
:return: None
218+
"""
219+
raise NotImplementedError("Subclasses need to implement sampling iteration logic")
188220

189221
def _get_input_image(
190222
self,

‎iddm/model/samples/ddim.py‎

Lines changed: 24 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ def _init_time_step(self):
7373
self.time_step = reversed(torch.cat((torch.tensor([0], dtype=torch.long), self.time_step)))
7474
self.time_step = list(zip(self.time_step[:-1], self.time_step[1:]))
7575

76-
def sample(
76+
def _sample_loop(
7777
self,
7878
model: nn.Module,
7979
x: Optional[torch.Tensor] = None,
@@ -82,7 +82,7 @@ def sample(
8282
cfg_scale: Optional[float] = None
8383
) -> torch.Tensor:
8484
"""
85-
DDIM sample method
85+
DDPM sample loop method
8686
:param model: Model
8787
:param x: Input image tensor, if provided, will be used as the starting point for sampling
8888
:param n: Number of sample images, x priority is greater than n
@@ -91,35 +91,28 @@ def sample(
9191
Avoiding the posterior collapse problem, Reference paper: 'Classifier-Free Diffusion Guidance'
9292
:return: Sample images
9393
"""
94-
# Input dim: [n, img_channel, img_size_h, img_size_w]
95-
x, n = self._get_input_image(n=n, x=x)
96-
logger.info(msg=f"DDIM Sampling {n} new images....")
97-
model.eval()
98-
with torch.no_grad():
99-
# The list of current time and previous time
100-
for i, p_i in tqdm(self.time_step):
101-
# Time step, creating a tensor of size n
102-
t = (torch.ones(n) * i).long().to(self.device)
103-
# Previous time step, creating a tensor of size n
104-
p_t = (torch.ones(n) * p_i).long().to(self.device)
105-
# Expand to a 4-dimensional tensor, and get the value according to the time step t
106-
alpha_t = self.alpha_hat[t][:, None, None, None]
107-
alpha_prev = self.alpha_hat[p_t][:, None, None, None]
108-
noise = torch.randn_like(x) if i > 1 else torch.zeros_like(x)
94+
# The list of current time and previous time
95+
for i, p_i in tqdm(self.time_step):
96+
# Time step, creating a tensor of size n
97+
t = (torch.ones(n) * i).long().to(self.device)
98+
# Previous time step, creating a tensor of size n
99+
p_t = (torch.ones(n) * p_i).long().to(self.device)
100+
# Expand to a 4-dimensional tensor, and get the value according to the time step t
101+
alpha_t = self.alpha_hat[t][:, None, None, None]
102+
alpha_prev = self.alpha_hat[p_t][:, None, None, None]
103+
noise = torch.randn_like(x) if i > 1 else torch.zeros_like(x)
109104

110-
# Predict noise
111-
predicted_noise = self._get_predicted_noise(model, x, t, labels, cfg_scale)
105+
# Predict noise
106+
predicted_noise = self._get_predicted_noise(model, x, t, labels, cfg_scale)
107+
108+
# Calculation formula
109+
# Division would cause the value to be too large or too small, and it needs to be constrained
110+
# https://github.com/ermongroup/ddim/blob/main/functions/denoising.py#L54C12-L54C54
111+
x0_t = torch.clamp((x - (predicted_noise * torch.sqrt((1 - alpha_t)))) / torch.sqrt(alpha_t), -1, 1)
112+
# Sigma
113+
c1 = self.eta * torch.sqrt((1 - alpha_t / alpha_prev) * (1 - alpha_prev) / (1 - alpha_t))
114+
c2 = torch.sqrt((1 - alpha_prev) - c1 ** 2)
115+
# Predicted x0 + direction pointing to xt + sigma * predicted noise
116+
x = torch.sqrt(alpha_prev) * x0_t + c2 * predicted_noise + c1 * noise
112117

113-
# Calculation formula
114-
# Division would cause the value to be too large or too small, and it needs to be constrained
115-
# https://github.com/ermongroup/ddim/blob/main/functions/denoising.py#L54C12-L54C54
116-
x0_t = torch.clamp((x - (predicted_noise * torch.sqrt((1 - alpha_t)))) / torch.sqrt(alpha_t), -1, 1)
117-
# Sigma
118-
c1 = self.eta * torch.sqrt((1 - alpha_t / alpha_prev) * (1 - alpha_prev) / (1 - alpha_t))
119-
c2 = torch.sqrt((1 - alpha_prev) - c1 ** 2)
120-
# Predicted x0 + direction pointing to xt + sigma * predicted noise
121-
x = torch.sqrt(alpha_prev) * x0_t + c2 * predicted_noise + c1 * noise
122-
# Post process
123-
x = self.post_process(x=x)
124-
model.train()
125118
return x

‎iddm/model/samples/ddpm.py‎

Lines changed: 25 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ def __init__(
5656
device=device, schedule_name=schedule_name, latent=latent, latent_channel=latent_channel,
5757
autoencoder=autoencoder)
5858

59-
def sample(
59+
def _sample_loop(
6060
self,
6161
model: nn.Module,
6262
x: Optional[torch.Tensor] = None,
@@ -65,7 +65,7 @@ def sample(
6565
cfg_scale: Optional[float] = None
6666
) -> torch.Tensor:
6767
"""
68-
DDPM sample method
68+
DDPM sample loop method
6969
:param model: Model
7070
:param x: Input image tensor, if provided, will be used as the starting point for sampling
7171
:param n: Number of sample images, x priority is greater than n
@@ -74,42 +74,32 @@ def sample(
7474
Avoiding the posterior collapse problem, Reference paper: 'Classifier-Free Diffusion Guidance'
7575
:return: Sample images
7676
"""
77-
# Input dim: [n, img_channel, img_size_h, img_size_w]
78-
x, n = self._get_input_image(n=n, x=x)
79-
logger.info(msg=f"DDPM Sampling {n} new images....")
80-
model.eval()
81-
with torch.no_grad():
82-
# 'reversed(range(1, self.noise_steps)' iterates over a sequence of integers in reverse
83-
for i in tqdm(reversed(range(1, self.noise_steps)), position=0, total=self.noise_steps - 1):
84-
# Time step, creating a tensor of size n
85-
t = (torch.ones(n) * i).long().to(self.device)
77+
# 'reversed(range(1, self.noise_steps)' iterates over a sequence of integers in reverse
78+
for i in tqdm(reversed(range(1, self.noise_steps)), position=0, total=self.noise_steps - 1):
79+
# Time step, creating a tensor of size n
80+
t = (torch.ones(n) * i).long().to(self.device)
8681

87-
# Whether the network has conditional input, such as multiple category input
88-
# Predict noise
89-
predicted_noise = self._get_predicted_noise(model, x, t, labels, cfg_scale)
82+
# Whether the network has conditional input, such as multiple category input
83+
# Predict noise
84+
predicted_noise = self._get_predicted_noise(model, x, t, labels, cfg_scale)
9085

91-
# Expand to a 4-dimensional tensor, and get the value according to the time step t
92-
alpha = self.alpha[t][:, None, None, None]
93-
alpha_hat = self.alpha_hat[t][:, None, None, None]
94-
beta = self.beta[t][:, None, None, None]
95-
# Only noise with a step size greater than 1 is required.
96-
# For details, refer to line 3 of Algorithm 2 on page 4 of the paper
97-
noise = torch.randn_like(x) if i > 1 else torch.zeros_like(x)
86+
# Expand to a 4-dimensional tensor, and get the value according to the time step t
87+
alpha = self.alpha[t][:, None, None, None]
88+
alpha_hat = self.alpha_hat[t][:, None, None, None]
89+
beta = self.beta[t][:, None, None, None]
90+
# Only noise with a step size greater than 1 is required.
91+
# For details, refer to line 3 of Algorithm 2 on page 4 of the paper
92+
noise = torch.randn_like(x) if i > 1 else torch.zeros_like(x)
9893

99-
# Fix latent diffusion explosion problem
100-
if self.latent:
101-
x = x.clamp(-1, 1)
102-
103-
# In each epoch, use x to calculate t - 1 of x
104-
# For details, refer to line 4 of Algorithm 2 on page 4 of the paper
105-
x = 1 / torch.sqrt(alpha) * (
106-
x - ((1 - alpha) / (torch.sqrt(1 - alpha_hat))) * predicted_noise) + torch.sqrt(
107-
beta) * noise
108-
# Post process, output of constraint x
94+
# Fix latent diffusion explosion problem
10995
if self.latent:
110-
x = self.post_process(x=x)
111-
else:
112-
x = self.post_process(x=x.clamp(-1, 1))
96+
x = x.clamp(-1, 1)
97+
98+
# In each epoch, use x to calculate t - 1 of x
99+
# For details, refer to line 4 of Algorithm 2 on page 4 of the paper
100+
x = 1 / torch.sqrt(alpha) * (
101+
x - ((1 - alpha) / (torch.sqrt(1 - alpha_hat))) * predicted_noise) + torch.sqrt(
102+
beta) * noise
113103

114-
model.train()
104+
x = x if self.latent else x.clamp(-1, 1)
115105
return x

‎iddm/model/samples/plms.py‎

Lines changed: 51 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ def __init__(
5757
img_size=img_size, device=device, schedule_name=schedule_name, latent=latent,
5858
latent_channel=latent_channel, autoencoder=autoencoder)
5959

60-
def sample(
60+
def _sample_loop(
6161
self,
6262
model: nn.Module,
6363
x: Optional[torch.Tensor] = None,
@@ -66,7 +66,7 @@ def sample(
6666
cfg_scale: Optional[float] = None
6767
) -> torch.Tensor:
6868
"""
69-
PLMS sample method
69+
PLMS sample loop method
7070
:param model: Model
7171
:param x: Input image tensor, if provided, will be used as the starting point for sampling
7272
:param n: Number of sample images, x priority is greater than n
@@ -75,62 +75,57 @@ def sample(
7575
Avoiding the posterior collapse problem, Reference paper: 'Classifier-Free Diffusion Guidance'
7676
:return: Sample images
7777
"""
78-
# Input dim: [n, img_channel, img_size_h, img_size_w]
79-
x, n = self._get_input_image(n=n, x=x)
80-
logger.info(msg=f"PLMS Sampling {n} new images....")
81-
model.eval()
82-
with torch.no_grad():
83-
# Old eps
84-
old_eps = []
85-
# The list of current time and previous time
86-
for i, p_i in tqdm(self.time_step):
87-
# Time step, creating a tensor of size n
88-
t = (torch.ones(n) * i).long().to(self.device)
89-
# Previous time step, creating a tensor of size n
90-
p_t = (torch.ones(n) * p_i).long().to(self.device)
91-
# Expand to a 4-dimensional tensor, and get the value according to the time step t
92-
alpha_t = self.alpha_hat[t][:, None, None, None]
93-
alpha_prev = self.alpha_hat[p_t][:, None, None, None]
94-
noise = torch.randn_like(x) if i > 1 else torch.zeros_like(x)
78+
# Old eps
79+
old_eps = []
80+
# The list of current time and previous time
81+
for i, p_i in tqdm(self.time_step):
82+
# Time step, creating a tensor of size n
83+
t = (torch.ones(n) * i).long().to(self.device)
84+
# Previous time step, creating a tensor of size n
85+
p_t = (torch.ones(n) * p_i).long().to(self.device)
86+
# Expand to a 4-dimensional tensor, and get the value according to the time step t
87+
alpha_t = self.alpha_hat[t][:, None, None, None]
88+
alpha_prev = self.alpha_hat[p_t][:, None, None, None]
89+
noise = torch.randn_like(x) if i > 1 else torch.zeros_like(x)
9590

96-
# Predict noise
97-
predicted_noise = self._get_predicted_noise(model, x, t, labels, cfg_scale)
91+
# Predict noise
92+
predicted_noise = self._get_predicted_noise(model, x, t, labels, cfg_scale)
9893

99-
# Calculation formula
100-
if len(old_eps) == 0:
101-
# Pseudo Improved Euler (2nd order)
102-
x0_t = torch.clamp((x - (predicted_noise * torch.sqrt((1 - alpha_t)))) / torch.sqrt(alpha_t), -1, 1)
103-
c1 = self.eta * torch.sqrt((1 - alpha_t / alpha_prev) * (1 - alpha_prev) / (1 - alpha_t))
104-
c2 = torch.sqrt((1 - alpha_prev) - c1 ** 2)
105-
p_x = torch.sqrt(alpha_prev) * x0_t + c2 * predicted_noise + c1 * noise
106-
if labels is None and cfg_scale is None:
107-
# Images and time steps input into the model
108-
predicted_noise_next = model(p_x, p_t)
109-
else:
110-
predicted_noise_next = model(p_x, p_t, labels)
111-
predicted_noise_prime = (predicted_noise + predicted_noise_next) / 2
112-
elif len(old_eps) == 1:
113-
# 2nd order Pseudo Linear Multistep (Adams-Bashforth)
114-
predicted_noise_prime = (3 * predicted_noise - old_eps[-1]) / 2
115-
elif len(old_eps) == 2:
116-
# 3rd order Pseudo Linear Multistep (Adams-Bashforth)
117-
predicted_noise_prime = (23 * predicted_noise - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
118-
elif len(old_eps) >= 3:
119-
# 4th order Pseudo Linear Multistep (Adams-Bashforth)
120-
predicted_noise_prime = (55 * predicted_noise - 59 * old_eps[-1] + 37 * old_eps[-2] -
121-
9 * old_eps[-3]) / 24
122-
123-
x0_t = torch.clamp((x - (predicted_noise_prime * torch.sqrt((1 - alpha_t)))) / torch.sqrt(alpha_t),
124-
-1, 1)
94+
# Calculation formula
95+
if len(old_eps) == 0:
96+
# Pseudo Improved Euler (2nd order)
97+
x0_t = torch.clamp((x - (predicted_noise * torch.sqrt((1 - alpha_t)))) / torch.sqrt(alpha_t), -1, 1)
12598
c1 = self.eta * torch.sqrt((1 - alpha_t / alpha_prev) * (1 - alpha_prev) / (1 - alpha_t))
12699
c2 = torch.sqrt((1 - alpha_prev) - c1 ** 2)
127-
x = torch.sqrt(alpha_prev) * x0_t + c2 * predicted_noise_prime + c1 * noise
128-
# Save old predicted_noise
129-
old_eps.append(predicted_noise)
130-
# Only the last 3 historical values are retained to save memory
131-
if len(old_eps) > 3:
132-
old_eps.pop(0)
133-
# Post process
134-
x = self.post_process(x=x)
135-
model.train()
100+
p_x = torch.sqrt(alpha_prev) * x0_t + c2 * predicted_noise + c1 * noise
101+
if labels is None and cfg_scale is None:
102+
# Images and time steps input into the model
103+
predicted_noise_next = model(p_x, p_t)
104+
else:
105+
predicted_noise_next = model(p_x, p_t, labels)
106+
predicted_noise_prime = (predicted_noise + predicted_noise_next) / 2
107+
elif len(old_eps) == 1:
108+
# 2nd order Pseudo Linear Multistep (Adams-Bashforth)
109+
predicted_noise_prime = (3 * predicted_noise - old_eps[-1]) / 2
110+
elif len(old_eps) == 2:
111+
# 3rd order Pseudo Linear Multistep (Adams-Bashforth)
112+
predicted_noise_prime = (23 * predicted_noise - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
113+
elif len(old_eps) >= 3:
114+
# 4th order Pseudo Linear Multistep (Adams-Bashforth)
115+
predicted_noise_prime = (55 * predicted_noise - 59 * old_eps[-1] + 37 * old_eps[-2] -
116+
9 * old_eps[-3]) / 24
117+
else:
118+
raise ValueError(f"Unexpected number of old_eps: {len(old_eps)}")
119+
120+
x0_t = torch.clamp((x - (predicted_noise_prime * torch.sqrt((1 - alpha_t)))) / torch.sqrt(alpha_t),
121+
-1, 1)
122+
c1 = self.eta * torch.sqrt((1 - alpha_t / alpha_prev) * (1 - alpha_prev) / (1 - alpha_t))
123+
c2 = torch.sqrt((1 - alpha_prev) - c1 ** 2)
124+
x = torch.sqrt(alpha_prev) * x0_t + c2 * predicted_noise_prime + c1 * noise
125+
# Save old predicted_noise
126+
old_eps.append(predicted_noise)
127+
# Only the last 3 historical values are retained to save memory
128+
if len(old_eps) > 3:
129+
old_eps.pop(0)
130+
136131
return x

0 commit comments

Comments
 (0)