(self, cond, shape, verbose=True, callback=None, quantize_denoised=False,
img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0.,
score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None,
log_every_t=None)
| 991 | |
| 992 | @torch.no_grad() |
| 993 | def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False, |
| 994 | img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0., |
| 995 | score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None, |
| 996 | log_every_t=None): |
| 997 | if not log_every_t: |
| 998 | log_every_t = self.log_every_t |
| 999 | timesteps = self.num_timesteps |
| 1000 | if batch_size is not None: |
| 1001 | b = batch_size if batch_size is not None else shape[0] |
| 1002 | shape = [batch_size] + list(shape) |
| 1003 | else: |
| 1004 | b = batch_size = shape[0] |
| 1005 | if x_T is None: |
| 1006 | img = torch.randn(shape, device=self.device) |
| 1007 | else: |
| 1008 | img = x_T |
| 1009 | intermediates = [] |
| 1010 | if cond is not None: |
| 1011 | if isinstance(cond, dict): |
| 1012 | cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else |
| 1013 | list(map(lambda x: x[:batch_size], cond[key])) for key in cond} |
| 1014 | else: |
| 1015 | cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size] |
| 1016 | |
| 1017 | if start_T is not None: |
| 1018 | timesteps = min(timesteps, start_T) |
| 1019 | iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation', |
| 1020 | total=timesteps) if verbose else reversed( |
| 1021 | range(0, timesteps)) |
| 1022 | if type(temperature) == float: |
| 1023 | temperature = [temperature] * timesteps |
| 1024 | |
| 1025 | for i in iterator: |
| 1026 | ts = torch.full((b,), i, device=self.device, dtype=torch.long) |
| 1027 | if self.shorten_cond_schedule: |
| 1028 | assert self.model.conditioning_key != 'hybrid' |
| 1029 | tc = self.cond_ids[ts].to(cond.device) |
| 1030 | cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) |
| 1031 | |
| 1032 | img, x0_partial = self.p_sample(img, cond, ts, |
| 1033 | clip_denoised=self.clip_denoised, |
| 1034 | quantize_denoised=quantize_denoised, return_x0=True, |
| 1035 | temperature=temperature[i], noise_dropout=noise_dropout, |
| 1036 | score_corrector=score_corrector, corrector_kwargs=corrector_kwargs) |
| 1037 | if mask is not None: |
| 1038 | assert x0 is not None |
| 1039 | img_orig = self.q_sample(x0, ts) |
| 1040 | img = img_orig * mask + (1. - mask) * img |
| 1041 | |
| 1042 | if i % log_every_t == 0 or i == timesteps - 1: |
| 1043 | intermediates.append(x0_partial) |
| 1044 | if callback: callback(i) |
| 1045 | if img_callback: img_callback(img, i) |
| 1046 | return img, intermediates |
| 1047 | |
| 1048 | @torch.no_grad() |
| 1049 | def p_sample_loop(self, cond, shape, return_intermediates=False, |
no test coverage detected