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