diffengine.models.utils.timesteps

Module Contents

Classes

TimeSteps	Time Steps module.
LaterTimeSteps	Later biased Time Steps module.
EarlierTimeSteps	Earlier biased Time Steps module.
RangeTimeSteps	Range biased Time Steps module.
CubicSamplingTimeSteps	Cubic Sampling Time Steps module.
WuerstchenRandomTimeSteps	Wuerstchen Random Time Steps module.
DDIMTimeSteps	DDIM Time Steps module.

class diffengine.models.utils.timesteps.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.timesteps.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.timesteps.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.timesteps.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.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.timesteps.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

class diffengine.models.utils.timesteps.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) –