diffengine.models.utils

Submodules

• diffengine.models.utils.noise
• diffengine.models.utils.timesteps

Package Contents

Classes

OffsetNoise	Offset noise module.
PyramidNoise	Pyramid noise module.
WhiteNoise	White noise module.
CubicSamplingTimeSteps	Cubic Sampling Time Steps module.
DDIMTimeSteps	DDIM Time Steps module.
EarlierTimeSteps	Earlier biased Time Steps module.
LaterTimeSteps	Later biased Time Steps module.
RangeTimeSteps	Range biased Time Steps module.
TimeSteps	Time Steps module.
WuerstchenRandomTimeSteps	Wuerstchen Random Time Steps module.

class diffengine.models.utils.OffsetNoise(offset_weight=0.05)

Bases: torch.nn.Module

Offset noise module.

https://www.crosslabs.org/blog/diffusion-with-offset-noise

Args:

offset_weight (float): Noise offset weight. Defaults to 0.05.

forward(latents)

Forward pass.

Generates noise for the given latents.

Args:

latents (torch.Tensor): Latent vectors.

Parameters:

latents (torch.Tensor) –

Return type:

torch.Tensor

Parameters:

offset_weight (float) –

class diffengine.models.utils.PyramidNoise(discount=0.9, *, random_multiplier=True)

Bases: torch.nn.Module

Pyramid noise module.

https://wandb.ai/johnowhitaker/multires_noise/reports/ Multi-Resolution-Noise-for-Diffusion-Model-Training–VmlldzozNjYyOTU2

Args:

discount (float): Noise offset weight. Defaults to 0.9. random_multiplier (bool): Whether to use random multiplier.

Defaults to True.

forward(latents)

Forward pass.

Generates noise for the given latents.

Args:

latents (torch.Tensor): Latent vectors.

Parameters:

latents (torch.Tensor) –

Return type:

torch.Tensor

Parameters:

• discount (float) –

• random_multiplier (bool) –

class diffengine.models.utils.WhiteNoise(*args, **kwargs)

Bases: torch.nn.Module

White noise module.

forward(latents)

Forward pass.

Generates noise for the given latents.

Args:

latents (torch.Tensor): Latent vectors.

Parameters:

latents (torch.Tensor) –

Return type:

torch.Tensor

class diffengine.models.utils.CubicSamplingTimeSteps(*args, **kwargs)

Bases: torch.nn.Module

Cubic Sampling Time Steps module.

For more details about why cubic sampling is used, refer to section 3.4 of https://arxiv.org/abs/2302.08453

forward(scheduler, num_batches, device)

Forward pass.

Generates time steps for the given batches.

Args:

scheduler (DDPMScheduler): Scheduler for training diffusion model. num_batches (int): Batch size. device (str): Device.

Parameters:

• scheduler (diffusers.DDPMScheduler) –

• num_batches (int) –

• device (str) –

Return type:

torch.Tensor

class diffengine.models.utils.DDIMTimeSteps(num_ddim_timesteps=50)

Bases: torch.nn.Module

DDIM Time Steps module.

Args:

num_ddim_timesteps (int): Number of DDIM timesteps. Defaults to 50.

forward(scheduler, num_batches, device)

Forward pass.

Generates time steps for the given batches.

Args:

scheduler (DDPMScheduler): Scheduler for training diffusion model. num_batches (int): Batch size. device (str): Device.

Parameters:

• scheduler (diffusers.DDPMScheduler) –

• num_batches (int) –

• device (str) –

Return type:

torch.Tensor

Parameters:

num_ddim_timesteps (int) –

class diffengine.models.utils.EarlierTimeSteps(bias_multiplier=5.0, bias_portion=0.25)

Bases: torch.nn.Module

Earlier biased Time Steps module.

Args:

bias_multiplier (float): Bias multiplier. Defaults to 10. bias_portion (float): Portion of earlier time steps to bias.

Defaults to 0.25.

forward(scheduler, num_batches, device)

Forward pass.

Generates time steps for the given batches.

Args:

scheduler (DDPMScheduler): Scheduler for training diffusion model. num_batches (int): Batch size. device (str): Device.

Parameters:

• scheduler (diffusers.DDPMScheduler) –

• num_batches (int) –

• device (str) –

Return type:

torch.Tensor

Parameters:

• bias_multiplier (float) –

• bias_portion (float) –

class diffengine.models.utils.LaterTimeSteps(bias_multiplier=5.0, bias_portion=0.25)

Bases: torch.nn.Module

Later biased Time Steps module.

Args:

bias_multiplier (float): Bias multiplier. Defaults to 10. bias_portion (float): Portion of later time steps to bias.

Defaults to 0.25.

forward(scheduler, num_batches, device)

Forward pass.

Generates time steps for the given batches.

Args:

scheduler (DDPMScheduler): Scheduler for training diffusion model. num_batches (int): Batch size. device (str): Device.

Parameters:

• scheduler (diffusers.DDPMScheduler) –

• num_batches (int) –

• device (str) –

Return type:

torch.Tensor

Parameters:

• bias_multiplier (float) –

• bias_portion (float) –

class diffengine.models.utils.RangeTimeSteps(bias_multiplier=5.0, bias_begin=0.25, bias_end=0.75)

Bases: torch.nn.Module

Range biased Time Steps module.

Args:

bias_multiplier (float): Bias multiplier. Defaults to 10. bias_begin (float): Portion of begin time steps to bias.

Defaults to 0.25.

bias_end (float): Portion of end time steps to bias.

Defaults to 0.75.

forward(scheduler, num_batches, device)

Forward pass.

Generates time steps for the given batches.

Args:

scheduler (DDPMScheduler): Scheduler for training diffusion model. num_batches (int): Batch size. device (str): Device.

Parameters:

• scheduler (diffusers.DDPMScheduler) –

• num_batches (int) –

• device (str) –

Return type:

torch.Tensor

Parameters:

• bias_multiplier (float) –

• bias_begin (float) –

• bias_end (float) –

class diffengine.models.utils.TimeSteps(*args, **kwargs)

Bases: torch.nn.Module

Time Steps module.

forward(scheduler, num_batches, device)

Forward pass.

Generates time steps for the given batches.

Args:

scheduler (DDPMScheduler): Scheduler for training diffusion model. num_batches (int): Batch size. device (str): Device.

Parameters:

• scheduler (diffusers.DDPMScheduler) –

• num_batches (int) –

• device (str) –

Return type:

torch.Tensor

class diffengine.models.utils.WuerstchenRandomTimeSteps(*args, **kwargs)

Bases: torch.nn.Module

Wuerstchen Random Time Steps module.

forward(num_batches, device)

Forward pass.

Generates time steps for the given batches.

Args:

scheduler (DDPMScheduler): Scheduler for training diffusion model. num_batches (int): Batch size. device (str): Device.

Parameters:

• num_batches (int) –

• device (str) –

Return type:

torch.Tensor