(self, cond, struct_cond, shape, return_intermediates=False,
x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False,
mask=None, x0=None, img_callback=None, start_T=None,
log_every_t=None, time_replace=None, adain_fea=None, interfea_path=None,
unconditional_conditioning=None,
unconditional_guidance_scale=None,
reference_sr=None, reference_lr=0.05, reference_step=1, reference_range=[100, 1000])
| 2854 | |
| 2855 | @torch.no_grad() |
| 2856 | def p_sample_loop(self, cond, struct_cond, shape, return_intermediates=False, |
| 2857 | x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False, |
| 2858 | mask=None, x0=None, img_callback=None, start_T=None, |
| 2859 | log_every_t=None, time_replace=None, adain_fea=None, interfea_path=None, |
| 2860 | unconditional_conditioning=None, |
| 2861 | unconditional_guidance_scale=None, |
| 2862 | reference_sr=None, reference_lr=0.05, reference_step=1, reference_range=[100, 1000]): |
| 2863 | |
| 2864 | if not log_every_t: |
| 2865 | log_every_t = self.log_every_t |
| 2866 | device = self.betas.device |
| 2867 | b = shape[0] |
| 2868 | if x_T is None: |
| 2869 | img = torch.randn(shape, device=device) |
| 2870 | else: |
| 2871 | img = x_T |
| 2872 | |
| 2873 | intermediates = [img] |
| 2874 | if timesteps is None: |
| 2875 | timesteps = self.num_timesteps |
| 2876 | |
| 2877 | iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed( |
| 2878 | range(0, timesteps)) |
| 2879 | |
| 2880 | if mask is not None: |
| 2881 | assert x0 is not None |
| 2882 | assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match |
| 2883 | |
| 2884 | batch_list = [] |
| 2885 | for i in iterator: |
| 2886 | if time_replace is None or time_replace == 1000: |
| 2887 | ts = torch.full((b,), i, device=device, dtype=torch.long) |
| 2888 | t_replace=None |
| 2889 | else: |
| 2890 | ts = torch.full((b,), i, device=device, dtype=torch.long) |
| 2891 | t_replace = repeat(torch.tensor([self.ori_timesteps[i]]), '1 -> b', b=img.size(0)) |
| 2892 | t_replace = t_replace.long().to(device) |
| 2893 | if self.shorten_cond_schedule: |
| 2894 | assert self.model.conditioning_key != 'hybrid' |
| 2895 | tc = self.cond_ids[ts].to(cond.device) |
| 2896 | cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) |
| 2897 | |
| 2898 | if t_replace is not None: |
| 2899 | if start_T is not None: |
| 2900 | if self.ori_timesteps[i] > start_T: |
| 2901 | continue |
| 2902 | struct_cond_input = self.structcond_stage_model(struct_cond, t_replace) |
| 2903 | else: |
| 2904 | if start_T is not None: |
| 2905 | if i > start_T: |
| 2906 | continue |
| 2907 | struct_cond_input = self.structcond_stage_model(struct_cond, ts) |
| 2908 | |
| 2909 | if interfea_path is not None: |
| 2910 | batch_list.append(struct_cond_input) |
| 2911 | |
| 2912 | img = self.p_sample(img, cond, struct_cond_input, ts, |
| 2913 | clip_denoised=self.clip_denoised, |
no test coverage detected