Method decode

diff2flow/ddim.py:284–324

(
        self,
        model,
        x_latent,
        t_start=None,
        ddim_steps=100,
        use_original_steps=True,
        model_kwargs=None,
        progress=True,
        clip_denoised=False
    )

    @torch.no_grad()
    def decode(
        self,
        model,
        x_latent,
        t_start=None,
        ddim_steps=100,
        use_original_steps=True,
        model_kwargs=None,
        progress=True,
        clip_denoised=False
    ):
        bs, dev = x_latent.shape[0], x_latent.device
        model_kwargs = model_kwargs or {}

        self.make_schedule(ddim_num_steps=ddim_steps, device=dev, ddim_eta=0.0, verbose=False)

        timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
        if t_start is not None:
            timesteps = timesteps[:t_start]

        time_range = np.flip(timesteps)
        total_steps = timesteps.shape[0]

        iterator = tqdm(time_range, desc='Decoding image', total=total_steps, disable=not progress)
        
        x_dec = x_latent
        for i, step in enumerate(iterator):
            index = total_steps - i - 1
            ts = torch.full((bs,), step, device=dev, dtype=torch.long)

            x_dec, _ = self.p_sample_ddim(
                model=model,
                x=x_dec,
                t=ts,
                index=index,
                model_kwargs=model_kwargs,
                use_original_steps=use_original_steps,
                clip_denoised=clip_denoised
            )

        return x_dec