Method p_sample_loop

diffusion/diffusion_ddpm_pan.py:445–507

(self, x_in, continous=False, get_interm_fm=False)

    @torch.no_grad()
    def p_sample_loop(self, x_in, continous=False, get_interm_fm=False):
        device = self.betas.device
        clip_noise = True if exists(self.clamp_range) else False
        sample_inter = 1 | (self.num_timesteps // 10)

        # self-conditioning
        x_start = None

        if not self.conditional:
            shape = x_in
            b = shape[0]
            img = torch.randn(shape, device=device)
            ret_img = img
            for i in tqdm(
                reversed(range(0, self.num_timesteps)),
                desc="ddpm sampling loop time step",
                total=self.num_timesteps,
            ):
                self_cond = x_start if self.self_condition else None
                img = self.p_sample(
                    img,
                    torch.full((b,), i, device=device, dtype=torch.long),
                    self_cond=self_cond,
                    clip_denoised=clip_noise,
                    get_interm_fm=get_interm_fm,
                )

                if i % sample_inter == 0 and continous:
                    ret_img = torch.cat([ret_img, img], dim=0)

                x_start = img
            return img
        else:
            x = x_in  # sr
            shape = x_in.shape[-2:]
            if isinstance(x, list):
                b = x[0].shape[0]
            elif isinstance(x, Tensor):
                b = x.shape[0]
            img = torch.randn((b, self.channels, *shape), device=device)
            ret_img = img
            for i in tqdm(
                reversed(range(0, self.num_timesteps)),
                desc="sampling loop time step",
                total=self.num_timesteps,
            ):
                self_cond = x_start if self.self_condition else None
                img = self.p_sample(
                    img,
                    torch.full((b,), i, device=device, dtype=torch.long),
                    condition_x=x,
                    self_cond=self_cond,
                    clip_denoised=clip_noise,
                    get_interm_fm=get_interm_fm,
                )
                if i % sample_inter == 0 and continous:
                    ret_img = torch.cat([ret_img, img], dim=0)
                x_start = img