Class ContinuousTimeGaussianDiffusion

models/diffusion/continuous_time.py:63–449

Continuous-time Gaussian diffusion https://arxiv.org/pdf/2107.00630.pdf

class ContinuousTimeGaussianDiffusion(base.GaussianDiffusion):
    """
    Continuous-time Gaussian diffusion
    https://arxiv.org/pdf/2107.00630.pdf
    """

    def __init__(
        self,
        denoiser: nn.Module,
        criterion: Literal["l2", "l1", "huber"] | nn.Module = "l2",
        objective: Literal["eps", "v", "x0"] = "eps",
        prediction: Literal["eps", "v", "x0"] = "eps",
        beta_schedule: Literal[
            "linear", "cosine", "cosine_shifted", "cosine_interpolated"
        ] = "cosine",
        min_snr_loss_weight: bool = True,
        min_snr_gamma: float = 5.0,
        sampling_resolution: tuple[int, int] | None = None,
        clip_sample: bool = True,
        clip_sample_range: float = 1,
        image_d: float = None,
        noise_d_low: float = None,
        noise_d_high: float = None,
        num_training_steps: int = 1000,
        classifier_free_scale: float = 7.5,
        sampling: str= "ddpm",
        use_seg: bool = False,
        use_guidence_net: bool = False,
        guidence_step_interval: int = -1,
        guidence_weights: tuple = None,
        use_control_net = False,
        upsampling = False,
        up_rate=2,
        base_up_rate=0.25,
        downsampling =False,
        down_rate=0.5,
        base_down_rate=1.0,
        interpolate_rate=0.5,
        inference=False
    ):
        super().__init__(
            denoiser=denoiser,
            sampling=sampling,
            criterion=criterion,
            num_training_steps=num_training_steps,
            objective=objective,
            prediction=prediction,
            beta_schedule=beta_schedule,
            min_snr_loss_weight=min_snr_loss_weight,
            min_snr_gamma=min_snr_gamma,
            sampling_resolution=sampling_resolution,
            clip_sample=clip_sample,
            clip_sample_range=clip_sample_range,
            classifier_free_scale=classifier_free_scale,
            use_seg=use_seg,
            use_guidence_net=use_guidence_net,
            guidence_step_interval=guidence_step_interval,
            guidence_weights=guidence_weights,