Method encode

diff2flow/ddim.py:224–281

(
        self,
        model,
        x0,
        t_enc=None,
        ddim_steps=100,
        use_original_steps=True,
        n_intermediates=0,
        model_kwargs=None,
        progress=True
    )

    
    @torch.no_grad()
    def encode(
        self,
        model,
        x0,
        t_enc=None,
        ddim_steps=100,
        use_original_steps=True,
        n_intermediates=0,
        model_kwargs=None,
        progress=True
    ):
        assert self.ddpm.parameterization == "eps", "Only works with eps parameterization"
        if not use_original_steps:
            warnings.warn('Using DDIM for encoding is not recommended, as it is not fully debugged.')
        
        bs, dev = x0.shape[0], x0.device
        model_kwargs = model_kwargs or {}
        return_intermediates = n_intermediates > 0

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

        # steps
        num_reference_steps = self.ddpm_num_timesteps if use_original_steps else self.ddim_timesteps.shape[0]
        if t_enc is None:
            t_enc = num_reference_steps
        assert t_enc <= num_reference_steps
        num_steps = t_enc

        if use_original_steps:
            alphas_next = self.ddpm.alphas_cumprod[:num_steps]
            alphas = self.ddpm.alphas_cumprod_prev[:num_steps]
        else:
            alphas_next = self.ddim_alphas[:num_steps]
            alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])

        x_next = x0
        intermediates = []
        inter_steps = []
        for i in tqdm(range(num_steps), desc='Encoding x0', disable=not progress):
            t = torch.full((bs,), i, device=dev, dtype=torch.long)
            
            noise_pred = model(x_next, t, **model_kwargs)

            xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next
            weighted_noise_pred = alphas_next[i].sqrt() * (
                    (1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt()) * noise_pred
            x_next = xt_weighted + weighted_noise_pred
            if return_intermediates and i % (num_steps // n_intermediates) == 0 and i < num_steps - 1:
                intermediates.append(x_next)
                inter_steps.append(i)
            elif return_intermediates and i >= num_steps - 2:
                intermediates.append(x_next)
                inter_steps.append(i)

        out = {'x_encoded': x_next, 'intermediate_steps': inter_steps}
        if return_intermediates:
            out.update({'intermediates': intermediates})
        return x_next, out