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

cleanup again

Browse Source
This commit is contained in:
Phil Wang
2022-06-28 13:42:41 -07:00
parent c3d1607019
commit be2bd8d320
1 changed files with 3 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
denoising_diffusion_pytorch/elucidated_diffusion.py
+3 -10
View File
@@ -117,26 +117,19 @@ class ElucidatedDiffusion(nn.Module):
# sampling
@torch.no_grad()
def sample_one_timestep(self, x, time, time_next):
batch, *_, device = *x.shape, x.device
return x
def sample(self, batch_size = 16):
shape = (batch_size, self.channels, self.image_size, self.image_size)
@torch.no_grad()
def sample_all_timesteps(self, shape):
images = torch.randn(shape, device = self.device)
steps = torch.linspace(1., 0., 100 + 1, device = self.device)
for i in tqdm(range(100), desc = 'sampling loop time step', total = 100):
times = steps[i]
times_next = steps[i + 1]
images = self.sample_one_timestep(images, times, times_next)
images = images
return unnormalize_to_zero_to_one(images)
@torch.no_grad()
def sample(self, batch_size = 16):
return self.sample_all_timesteps((batch_size, self.channels, self.image_size, self.image_size))
# training
def forward(self, images):