| 882 | |
| 883 | @torch.no_grad() |
| 884 | def p_sample_loop(self, cond, shape, return_intermediates=False, x_T=None, verbose=True, callback=None, \ |
| 885 | timesteps=None, mask=None, x0=None, img_callback=None, start_T=None, log_every_t=None, **kwargs): |
| 886 | |
| 887 | if not log_every_t: |
| 888 | log_every_t = self.log_every_t |
| 889 | device = self.betas.device |
| 890 | b = shape[0] |
| 891 | # sample an initial noise |
| 892 | if x_T is None: |
| 893 | img = torch.randn(shape, device=device) |
| 894 | else: |
| 895 | img = x_T |
| 896 | |
| 897 | intermediates = [img] |
| 898 | if timesteps is None: |
| 899 | timesteps = self.num_timesteps |
| 900 | if start_T is not None: |
| 901 | timesteps = min(timesteps, start_T) |
| 902 | |
| 903 | iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed(range(0, timesteps)) |
| 904 | |
| 905 | if mask is not None: |
| 906 | assert x0 is not None |
| 907 | assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match |
| 908 | |
| 909 | for i in iterator: |
| 910 | ts = torch.full((b,), i, device=device, dtype=torch.long) |
| 911 | if self.shorten_cond_schedule: |
| 912 | assert self.model.conditioning_key != 'hybrid' |
| 913 | tc = self.cond_ids[ts].to(cond.device) |
| 914 | cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) |
| 915 | |
| 916 | img = self.p_sample(img, cond, ts, clip_denoised=self.clip_denoised, **kwargs) |
| 917 | if mask is not None: |
| 918 | img_orig = self.q_sample(x0, ts) |
| 919 | img = img_orig * mask + (1. - mask) * img |
| 920 | |
| 921 | if i % log_every_t == 0 or i == timesteps - 1: |
| 922 | intermediates.append(img) |
| 923 | if callback: callback(i) |
| 924 | if img_callback: img_callback(img, i) |
| 925 | |
| 926 | if return_intermediates: |
| 927 | return img, intermediates |
| 928 | return img |
| 929 | |
| 930 | @torch.no_grad() |
| 931 | def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None, \ |