Class DiffusionFlow

diff2flow/diffusion.py:71–190

class DiffusionFlow(nn.Module):
    def __init__(
            self,
            net_cfg: dict,
            timesteps: int = 1000,
            beta_schedule: str = 'linear',
            loss_type: str = 'l2',
            parameterization: str = 'v',
            linear_start: float = 0.00085,
            linear_end: float = 0.0120,
            ddim_steps: int = 10,           # 10 timesteps default for FlowModel
    ):
        super().__init__()
        self.net = instantiate_from_config(net_cfg)

        self.diffusion_cfg = dict(
            timesteps=timesteps,
            beta_schedule=beta_schedule,
            zero_terminal_snr=False,
            loss_type=loss_type,
            parameterization=parameterization,
            linear_start=linear_start,
            linear_end=linear_end,
            cosine_s=8e-3,
            original_elbo_weight=0.,
            v_posterior=0.,
            l_simple_weight=1.0
        )
        self.diffusion = GaussianDiffusion(**self.diffusion_cfg)

        self.ddim_steps = ddim_steps
        self.ddim_sampler = DDIMSampler(self.diffusion)

    def forward(self, x: torch.Tensor, t: torch.Tensor, **kwargs):
        return self.net(x, t, **kwargs)

    def training_losses(self, x1: torch.Tensor, x0: torch.Tensor = None, **cond_kwargs):
        loss, _ = self.diffusion.training_losses(
            model=self.net,
            x_start=x1,
            model_kwargs=cond_kwargs,
            x_noise=x0,
        )
        return loss
    
    def generate(self, x: torch.Tensor, sample_kwargs=None, reverse=False, return_intermediates=False, **kwargs):
        """
        Args:
            x: source minibatch (bs, *dim)
            sample_kwargs: dict, additional sampling arguments for the solver
                progress: bool, whether to show a progress bar
                clip_denoised: bool, whether to clip the denoised images to [-1, 1]
                use_ddpm: bool, whether to use DDPM sampling instead of DDIM
                intermediate_key: str, key to use for intermediate outputs
                    (DDIM: 'x_inter', 'pred_x0' | DDPM: 'sample' or 'pred_xstart')
                intermediate_freq: int, frequency of intermediate outputs
                __ DDIM only __:
                    num_steps: int, number of DDIM steps to take