| 508 | |
| 509 | @torch.no_grad() |
| 510 | def interpolate(self, x1, x2, t=None, lam=0.5): |
| 511 | b, *_, device = *x1.shape, x1.device |
| 512 | t = default(t, self.num_timesteps - 1) |
| 513 | |
| 514 | assert x1.shape == x2.shape |
| 515 | |
| 516 | t_batched = torch.stack([torch.tensor(t, device=device)] * b) |
| 517 | xt1, xt2 = map(lambda x: self.q_sample(x, t=t_batched), (x1, x2)) |
| 518 | |
| 519 | img = (1 - lam) * xt1 + lam * xt2 |
| 520 | for i in tqdm( |
| 521 | reversed(range(0, t)), desc="interpolation sample time step", total=t |
| 522 | ): |
| 523 | img = self.p_sample( |
| 524 | img, torch.full((b,), i, device=device, dtype=torch.long) |
| 525 | ) |
| 526 | |
| 527 | return img |
| 528 | |
| 529 | @staticmethod |
| 530 | def space_timesteps(num_timesteps, section_counts): |