wassname/denoising-diffusion-pytorch
1
0
Fork 0
mirror of https://github.com/wassname/denoising-diffusion-pytorch.git synced 2026-09-12 12:22:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

release elucidating diffusion

Browse Source
This commit is contained in:
Phil Wang
2022-06-28 17:39:48 -07:00
parent c535d31fc5
commit 86fcb6785b
2 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
denoising_diffusion_pytorch/elucidated_diffusion.py
+5 -3
View File
@@ -140,16 +140,18 @@ class ElucidatedDiffusion(nn.Module):
# sampling
@torch.no_grad()
def sample(self, batch_size = 16):
def sample(self, batch_size = 16, num_sample_steps = None):
num_sample_steps = default(num_sample_steps, self.num_sample_steps)
shape = (batch_size, self.channels, self.image_size, self.image_size)
# get the schedule, which is returned as (sigma, gamma) tuple, and pair up with the next sigma and gamma
sigmas = self.sample_schedule()
sigmas = self.sample_schedule(num_sample_steps)
gammas = torch.where(
(sigmas >= self.S_tmin) & (sigmas <= self.S_tmax),
min(self.S_churn / self.num_sample_steps, sqrt(2) - 1),
min(self.S_churn / num_sample_steps, sqrt(2) - 1),
0.
)