Skip to content

Base Sampler

Base sampler class for diffusion models.

This module provides a base abstract class for all samplers used in diffusion models. It defines the common interface that all samplers should implement.

BaseSampler

Bases: ABC

Abstract base class for all diffusion model samplers.

All samplers inherit from this class and must implement the call method which performs the actual sampling process.

Attributes:

Name Type Description
diffusion

The diffusion model to sample from.
verbose

Whether to print progress information during sampling.
Source code in image_gen\samplers\base.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
class BaseSampler(ABC):
    """Abstract base class for all diffusion model samplers.

    All samplers inherit from this class and must implement the call method
    which performs the actual sampling process.

    Attributes:
        diffusion: The diffusion model to sample from.
        verbose: Whether to print progress information during sampling.
    """

    def __init__(self, diffusion: BaseDiffusion, *_, verbose: bool = True, **__):
        """Initialize the sampler.

        Args:
            diffusion: The diffusion model to sample from.
            verbose: Whether to print progress information during sampling.
                Defaults to True.
        """
        self.diffusion = diffusion
        self.verbose = verbose

    @abstractmethod
    def __call__(
            self,
            x_T: Tensor,
            score_model: Callable,
            *args: Any,
            n_steps: int = 500,
            seed: Optional[int] = None,
            callback: Optional[Callable[[Tensor, int], None]] = None,
            callback_frequency: int = 50,
            guidance: Optional[Callable[[
                Tensor, Tensor, Tensor], Tensor]] = None,
            **kwargs: Any
    ) -> Tensor:
        """Perform the sampling process.

        Args:
            x_T: The initial noise tensor to start sampling from.
            score_model: The score model function that predicts the score.
            *args: Additional positional arguments.
            n_steps: Number of sampling steps. Defaults to 500.
            seed: Random seed for reproducibility. Defaults to None.
            callback: Optional function called during sampling to monitor progress.
                It takes the current sample and step number as inputs.
                Defaults to None.
            callback_frequency: How often to call the callback function.
                Defaults to 50.
            guidance: Optional guidance function for conditional sampling.
                Defaults to None.
            **kwargs: Additional keyword arguments.

        Returns:
            A tuple containing the final sample and the sequence of all samples.
        """
        ...

    def config(self) -> dict:
        """Return the configuration of the sampler.

        Returns:
            A dictionary with the sampler's configuration parameters.
        """
        return {}

    def __str__(self) -> str:
        """Return a string representation of the sampler.

        Returns:
            A string with the sampler's class name and its configuration.
        """
        config = self.config()
        params = ", ".join(f"{k}: {v}" for k, v in config.items())
        return f"{self._class_name}({params})"

    @property
    def _class_name(self) -> str:
        """Get the class name of the sampler.

        This property is automatically overridden in custom classes made by users.

        Returns:
            The name of the sampler class.
        """
        return self.__class__.__name__

__call__(x_T, score_model, *args, n_steps=500, seed=None, callback=None, callback_frequency=50, guidance=None, **kwargs) abstractmethod

Perform the sampling process.

Parameters:

Name Type Description Default
x_T Tensor

The initial noise tensor to start sampling from.

 required
score_model Callable

The score model function that predicts the score.

 required
*args Any

Additional positional arguments.

 ()
n_steps int

Number of sampling steps. Defaults to 500.

 500
seed Optional[int]

Random seed for reproducibility. Defaults to None.

 None
callback Optional[Callable[[Tensor, int], None]]

Optional function called during sampling to monitor progress. It takes the current sample and step number as inputs. Defaults to None.

 None
callback_frequency int

How often to call the callback function. Defaults to 50.

 50
guidance Optional[Callable[[Tensor, Tensor, Tensor], Tensor]]

Optional guidance function for conditional sampling. Defaults to None.

 None
**kwargs Any

Additional keyword arguments.

 {}

Returns:

Type Description
Tensor

A tuple containing the final sample and the sequence of all samples.
Source code in image_gen\samplers\base.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
@abstractmethod
def __call__(
        self,
        x_T: Tensor,
        score_model: Callable,
        *args: Any,
        n_steps: int = 500,
        seed: Optional[int] = None,
        callback: Optional[Callable[[Tensor, int], None]] = None,
        callback_frequency: int = 50,
        guidance: Optional[Callable[[
            Tensor, Tensor, Tensor], Tensor]] = None,
        **kwargs: Any
) -> Tensor:
    """Perform the sampling process.

    Args:
        x_T: The initial noise tensor to start sampling from.
        score_model: The score model function that predicts the score.
        *args: Additional positional arguments.
        n_steps: Number of sampling steps. Defaults to 500.
        seed: Random seed for reproducibility. Defaults to None.
        callback: Optional function called during sampling to monitor progress.
            It takes the current sample and step number as inputs.
            Defaults to None.
        callback_frequency: How often to call the callback function.
            Defaults to 50.
        guidance: Optional guidance function for conditional sampling.
            Defaults to None.
        **kwargs: Additional keyword arguments.

    Returns:
        A tuple containing the final sample and the sequence of all samples.
    """
    ...

__init__(diffusion, *_, verbose=True, **__)

Initialize the sampler.

Parameters:

Name Type Description Default
diffusion BaseDiffusion

The diffusion model to sample from.

 required
verbose bool

Whether to print progress information during sampling. Defaults to True.

 True
Source code in image_gen\samplers\base.py 
26
27
28
29
30
31
32
33
34
35
def __init__(self, diffusion: BaseDiffusion, *_, verbose: bool = True, **__):
    """Initialize the sampler.

    Args:
        diffusion: The diffusion model to sample from.
        verbose: Whether to print progress information during sampling.
            Defaults to True.
    """
    self.diffusion = diffusion
    self.verbose = verbose

__str__()

Return a string representation of the sampler.

Returns:

Type Description
str

A string with the sampler's class name and its configuration.
Source code in image_gen\samplers\base.py 
81
82
83
84
85
86
87
88
89
def __str__(self) -> str:
    """Return a string representation of the sampler.

    Returns:
        A string with the sampler's class name and its configuration.
    """
    config = self.config()
    params = ", ".join(f"{k}: {v}" for k, v in config.items())
    return f"{self._class_name}({params})"

config()

Return the configuration of the sampler.

Returns:

Type Description
dict

A dictionary with the sampler's configuration parameters.
Source code in image_gen\samplers\base.py 
73
74
75
76
77
78
79
def config(self) -> dict:
    """Return the configuration of the sampler.

    Returns:
        A dictionary with the sampler's configuration parameters.
    """
    return {}