Skip to content

Variance Exploding

Variance Exploding diffusion model implementation.

This module implements the Variance Exploding diffusion model and its corresponding noise schedule, which is particularly effective for image generation tasks.

VarianceExploding

Bases: BaseDiffusion

Variance Exploding diffusion model implementation.

This model implements diffusion using a variance exploding process, where the noise increases exponentially with time.

Attributes:

Name Type Description
NEEDS_NOISE_SCHEDULE

Class constant indicating if a custom noise schedule is required.
Source code in image_gen\diffusion\ve.py 
 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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
class VarianceExploding(BaseDiffusion):
    """Variance Exploding diffusion model implementation.

    This model implements diffusion using a variance exploding process,
    where the noise increases exponentially with time.

    Attributes:
        NEEDS_NOISE_SCHEDULE: Class constant indicating if a custom noise
            schedule is required.
    """

    NEEDS_NOISE_SCHEDULE = False

    def __init__(self, *_, sigma: float = 25.0, **__):
        """Initialize the variance exploding diffusion model.

        Args:
            sigma: Base sigma value for variance control. Defaults to 25.0.
        """
        super().__init__(VarianceExplodingSchedule(sigma))

    def forward_sde(self, x: Tensor, t: Tensor, *_, **__) -> Tuple[
            Tensor, Tensor]:
        """Calculate drift and diffusion for the forward SDE.

        Args:
            x: Input tensor representing the current state.
            t: Time steps tensor.

        Returns:
            Tuple of (drift, diffusion) tensors.
        """
        drift = torch.zeros_like(x)
        diffusion = (self.schedule.sigma ** t).view(-1, 1, 1, 1)
        return drift, diffusion

    def forward_process(self, x0: Tensor, t: Tensor, *args: Any, **kwargs: Any) -> Tuple[
            Tensor, Tensor]:
        """Apply the forward diffusion process.

        Args:
            x0: Input tensor representing initial state.
            t: Time steps tensor.
            *args: Additional positional arguments.
            **kwargs: Additional keyword arguments.

        Returns:
            Tuple of (noisy_sample, noise) tensors.
        """
        sigma_t = self.schedule(t, *args, **kwargs)
        sigma = sigma_t.view(x0.shape[0], *([1] * (x0.dim() - 1)))
        noise = torch.randn_like(x0)
        return x0 + sigma * noise, noise

    def compute_loss(self, score: Tensor, noise: Tensor, t: Tensor,
                     *args: Any, **kwargs: Any) -> Tensor:
        """Compute loss between predicted score and actual noise.

        Args:
            score: Predicted score tensor.
            noise: Actual noise tensor.
            t: Time steps tensor.
            *args: Additional positional arguments.
            **kwargs: Additional keyword arguments.

        Returns:
            Loss tensor.
        """
        sigma_t = self.schedule(t, *args, **kwargs)
        sigma_t = sigma_t.view(score.shape[0], *([1] * (score.dim() - 1)))
        loss = (sigma_t * score + noise) ** 2
        return loss.sum(dim=tuple(range(1, loss.dim())))

    def config(self) -> dict:
        """Get configuration parameters for the diffusion model.

        Returns:
            Dictionary containing configuration parameters.
        """
        return self.schedule.config()

__init__(*_, sigma=25.0, **__)

Initialize the variance exploding diffusion model.

Parameters:

Name Type Description Default
sigma float

Base sigma value for variance control. Defaults to 25.0.

 25.0
Source code in image_gen\diffusion\ve.py 
83
84
85
86
87
88
89
def __init__(self, *_, sigma: float = 25.0, **__):
    """Initialize the variance exploding diffusion model.

    Args:
        sigma: Base sigma value for variance control. Defaults to 25.0.
    """
    super().__init__(VarianceExplodingSchedule(sigma))

compute_loss(score, noise, t, *args, **kwargs)

Compute loss between predicted score and actual noise.

Parameters:

Name Type Description Default
score Tensor

Predicted score tensor.

 required
noise Tensor

Actual noise tensor.

 required
t Tensor

Time steps tensor.

 required
*args Any

Additional positional arguments.

 ()
**kwargs Any

Additional keyword arguments.

 {}

Returns:

Type Description
Tensor

Loss tensor.
Source code in image_gen\diffusion\ve.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def compute_loss(self, score: Tensor, noise: Tensor, t: Tensor,
                 *args: Any, **kwargs: Any) -> Tensor:
    """Compute loss between predicted score and actual noise.

    Args:
        score: Predicted score tensor.
        noise: Actual noise tensor.
        t: Time steps tensor.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.

    Returns:
        Loss tensor.
    """
    sigma_t = self.schedule(t, *args, **kwargs)
    sigma_t = sigma_t.view(score.shape[0], *([1] * (score.dim() - 1)))
    loss = (sigma_t * score + noise) ** 2
    return loss.sum(dim=tuple(range(1, loss.dim())))

config()

Get configuration parameters for the diffusion model.

Returns:

Type Description
dict

Dictionary containing configuration parameters.
Source code in image_gen\diffusion\ve.py 
143
144
145
146
147
148
149
def config(self) -> dict:
    """Get configuration parameters for the diffusion model.

    Returns:
        Dictionary containing configuration parameters.
    """
    return self.schedule.config()

forward_process(x0, t, *args, **kwargs)

Apply the forward diffusion process.

Parameters:

Name Type Description Default
x0 Tensor

Input tensor representing initial state.

 required
t Tensor

Time steps tensor.

 required
*args Any

Additional positional arguments.

 ()
**kwargs Any

Additional keyword arguments.

 {}

Returns:

Type Description
Tuple[Tensor, Tensor]

Tuple of (noisy_sample, noise) tensors.
Source code in image_gen\diffusion\ve.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def forward_process(self, x0: Tensor, t: Tensor, *args: Any, **kwargs: Any) -> Tuple[
        Tensor, Tensor]:
    """Apply the forward diffusion process.

    Args:
        x0: Input tensor representing initial state.
        t: Time steps tensor.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.

    Returns:
        Tuple of (noisy_sample, noise) tensors.
    """
    sigma_t = self.schedule(t, *args, **kwargs)
    sigma = sigma_t.view(x0.shape[0], *([1] * (x0.dim() - 1)))
    noise = torch.randn_like(x0)
    return x0 + sigma * noise, noise

forward_sde(x, t, *_, **__)

Calculate drift and diffusion for the forward SDE.

Parameters:

Name Type Description Default
x Tensor

Input tensor representing the current state.

 required
t Tensor

Time steps tensor.

 required

Returns:

Type Description
Tuple[Tensor, Tensor]

Tuple of (drift, diffusion) tensors.
Source code in image_gen\diffusion\ve.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def forward_sde(self, x: Tensor, t: Tensor, *_, **__) -> Tuple[
        Tensor, Tensor]:
    """Calculate drift and diffusion for the forward SDE.

    Args:
        x: Input tensor representing the current state.
        t: Time steps tensor.

    Returns:
        Tuple of (drift, diffusion) tensors.
    """
    drift = torch.zeros_like(x)
    diffusion = (self.schedule.sigma ** t).view(-1, 1, 1, 1)
    return drift, diffusion

VarianceExplodingSchedule

Bases: BaseNoiseSchedule

Variance Exploding noise schedule.

This schedule models noise that increases exponentially over time, creating a "variance exploding" effect.

Attributes:

Name Type Description
sigma

Base sigma value that controls the rate of variance explosion.
Source code in image_gen\diffusion\ve.py 
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
class VarianceExplodingSchedule(BaseNoiseSchedule):
    """Variance Exploding noise schedule.

    This schedule models noise that increases exponentially over time,
    creating a "variance exploding" effect.

    Attributes:
        sigma: Base sigma value that controls the rate of variance explosion.
    """

    def __init__(self, sigma: float, *_, **__):
        """Initialize the variance exploding noise schedule.

        Args:
            sigma: Base sigma value for the schedule.
        """
        self.sigma = sigma

    def __call__(self, t: Tensor, *_, **__) -> Tensor:
        """Calculate the noise magnitude at time t.

        Args:
            t: Time step tensor.

        Returns:
            Tensor containing noise magnitudes at time t.
        """
        log_sigma = torch.log(torch.tensor(
            self.sigma, dtype=torch.float32, device=t.device))
        return torch.sqrt(0.5 * (self.sigma ** (2 * t) - 1.0) / log_sigma)

    def integral_beta(self, t: Tensor, *_, **__) -> Tensor:
        """Calculate the integrated noise intensity up to time t.

        Args:
            t: Time step tensor.

        Returns:
            Tensor containing integrated noise values.
        """
        return 0.5 * (self.sigma ** (2 * t) - 1) / np.log(self.sigma)

    def config(self) -> dict:
        """Get configuration parameters for the schedule.

        Returns:
            Dictionary containing configuration parameters.
        """
        return {
            "sigma": self.sigma
        }

__call__(t, *_, **__)

Calculate the noise magnitude at time t.

Parameters:

Name Type Description Default
t Tensor

Time step tensor.

 required

Returns:

Type Description
Tensor

Tensor containing noise magnitudes at time t.
Source code in image_gen\diffusion\ve.py 
35
36
37
38
39
40
41
42
43
44
45
46
def __call__(self, t: Tensor, *_, **__) -> Tensor:
    """Calculate the noise magnitude at time t.

    Args:
        t: Time step tensor.

    Returns:
        Tensor containing noise magnitudes at time t.
    """
    log_sigma = torch.log(torch.tensor(
        self.sigma, dtype=torch.float32, device=t.device))
    return torch.sqrt(0.5 * (self.sigma ** (2 * t) - 1.0) / log_sigma)

__init__(sigma, *_, **__)

Initialize the variance exploding noise schedule.

Parameters:

Name Type Description Default
sigma float

Base sigma value for the schedule.

 required
Source code in image_gen\diffusion\ve.py 
27
28
29
30
31
32
33
def __init__(self, sigma: float, *_, **__):
    """Initialize the variance exploding noise schedule.

    Args:
        sigma: Base sigma value for the schedule.
    """
    self.sigma = sigma

config()

Get configuration parameters for the schedule.

Returns:

Type Description
dict

Dictionary containing configuration parameters.
Source code in image_gen\diffusion\ve.py 
59
60
61
62
63
64
65
66
67
def config(self) -> dict:
    """Get configuration parameters for the schedule.

    Returns:
        Dictionary containing configuration parameters.
    """
    return {
        "sigma": self.sigma
    }

integral_beta(t, *_, **__)

Calculate the integrated noise intensity up to time t.

Parameters:

Name Type Description Default
t Tensor

Time step tensor.

 required

Returns:

Type Description
Tensor

Tensor containing integrated noise values.
Source code in image_gen\diffusion\ve.py 
48
49
50
51
52
53
54
55
56
57
def integral_beta(self, t: Tensor, *_, **__) -> Tensor:
    """Calculate the integrated noise intensity up to time t.

    Args:
        t: Time step tensor.

    Returns:
        Tensor containing integrated noise values.
    """
    return 0.5 * (self.sigma ** (2 * t) - 1) / np.log(self.sigma)