| 1047 | |
| 1048 | @torch.no_grad() |
| 1049 | def p_sample_loop(self, cond, shape, return_intermediates=False, |
| 1050 | x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False, |
| 1051 | mask=None, x0=None, img_callback=None, start_T=None, |
| 1052 | log_every_t=None): |
| 1053 | |
| 1054 | if not log_every_t: |
| 1055 | log_every_t = self.log_every_t |
| 1056 | device = self.betas.device |
| 1057 | b = shape[0] |
| 1058 | if x_T is None: |
| 1059 | img = torch.randn(shape, device=device) |
| 1060 | else: |
| 1061 | img = x_T |
| 1062 | |
| 1063 | intermediates = [img] |
| 1064 | if timesteps is None: |
| 1065 | timesteps = self.num_timesteps |
| 1066 | |
| 1067 | if start_T is not None: |
| 1068 | timesteps = min(timesteps, start_T) |
| 1069 | iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed( |
| 1070 | range(0, timesteps)) |
| 1071 | |
| 1072 | if mask is not None: |
| 1073 | assert x0 is not None |
| 1074 | assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match |
| 1075 | |
| 1076 | for i in iterator: |
| 1077 | ts = torch.full((b,), i, device=device, dtype=torch.long) |
| 1078 | if self.shorten_cond_schedule: |
| 1079 | assert self.model.conditioning_key != 'hybrid' |
| 1080 | tc = self.cond_ids[ts].to(cond.device) |
| 1081 | cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) |
| 1082 | |
| 1083 | img = self.p_sample(img, cond, ts, |
| 1084 | clip_denoised=self.clip_denoised, |
| 1085 | quantize_denoised=quantize_denoised) |
| 1086 | if mask is not None: |
| 1087 | img_orig = self.q_sample(x0, ts) |
| 1088 | img = img_orig * mask + (1. - mask) * img |
| 1089 | |
| 1090 | if i % log_every_t == 0 or i == timesteps - 1: |
| 1091 | intermediates.append(img) |
| 1092 | if callback: callback(i) |
| 1093 | if img_callback: img_callback(img, i) |
| 1094 | |
| 1095 | if return_intermediates: |
| 1096 | return img, intermediates |
| 1097 | return img |
| 1098 | |
| 1099 | @torch.no_grad() |
| 1100 | def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None, |