| 37 | return torch.randint(low=1, high=self.noise_steps, size=(n,)) |
| 38 | |
| 39 | def sample(self, model, n): |
| 40 | logging.info(f"Sampling {n} new images....") |
| 41 | model.eval() |
| 42 | with torch.no_grad(): |
| 43 | x = torch.randn((n, 3, self.img_size, self.img_size)).to(self.device) |
| 44 | for i in tqdm(reversed(range(1, self.noise_steps)), position=0): |
| 45 | t = (torch.ones(n) * i).long().to(self.device) |
| 46 | predicted_noise = model(x, t) |
| 47 | alpha = self.alpha[t][:, None, None, None] |
| 48 | alpha_hat = self.alpha_hat[t][:, None, None, None] |
| 49 | beta = self.beta[t][:, None, None, None] |
| 50 | if i > 1: |
| 51 | noise = torch.randn_like(x) |
| 52 | else: |
| 53 | noise = torch.zeros_like(x) |
| 54 | x = 1 / torch.sqrt(alpha) * (x - ((1 - alpha) / (torch.sqrt(1 - alpha_hat))) * predicted_noise) + torch.sqrt(beta) * noise |
| 55 | model.train() |
| 56 | x = (x.clamp(-1, 1) + 1) / 2 |
| 57 | x = (x * 255).type(torch.uint8) |
| 58 | return x |
| 59 | |
| 60 | |
| 61 | def train(args): |