ls_mlkit.diffuser.time_scheduler module

Time Scheduler for Diffusion

class ls_mlkit.diffuser.time_scheduler.TimeScheduler(continuous_time_start: float = 0.0, continuous_time_end: float = 1.0, num_train_timesteps: int = 1000, num_inference_steps: int = None)[source]

Bases: ABC

Time Scheduler for Diffusion

continuous_time_to_discrete_time(continuous_time: Tensor) → Tensor[source]

Convert a continuous time to a discrete time.

Parameters:

continuous_time (Tensor) – the continuous time

Returns:

the discrete time

Return type:

Tensor

discrete_time_to_continuous_time(discrete_time: Tensor) → Tensor[source]

Convert a discrete time to a continuous time.

Parameters:

discrete_time (Tensor) – the discrete time

Returns:

the continuous time

Return type:

Tensor

get_continuous_timesteps_schedule() → Tensor[source]

Get the continuous one-dimensional timesteps for sampling/inference.

Returns:

the continuous one-dimensional timesteps for sampling/inference

Return type:

Tensor

get_discrete_timesteps_schedule() → Tensor[source]

Get the discrete one-dimensional timesteps for sampling/inference.

Returns:

the discrete one-dimensional timesteps for sampling/inference

Return type:

Tensor

initialize_timesteps_schedule()[source]

Initialize timesteps schedule for DDPM sampling. For DDPM, we should use ALL timesteps in reverse order for best quality. For DDIM-style sampling, use linear spacing.

sample_a_continuous_time_step_uniformly(macro_shape: Tuple[int, ...], same_timesteps_for_all_samples: bool = False) → Tensor[source]

Sample a continuous timestep in the range of [continuous_time_start, continuous_time_end] uniformly.

Parameters:

  • macro_shape (Tuple[int, ...]) – the macro shape of the samples

  • same_timesteps_for_all_samples (bool, optional) – whether to use the same timestep for all samples. Defaults to False.

Returns:

the continuous timestep

Return type:

Tensor

sample_a_discrete_time_step_uniformly(macro_shape: Tuple[int, ...], same_timesteps_for_all_samples: bool = False) → Tensor[source]

Sample a discrete timestep in the range of [0, num_train_timesteps - 1] uniformly.

Parameters:

  • macro_shape (Tuple[int, ...]) – the macro shape of the samples

  • same_timesteps_for_all_samples (bool, optional) – whether to use the same timestep for all samples. Defaults to False.

Returns:

the discrete timestep

Return type:

Tensor

set_continuous_timesteps_schedule(continuous_timesteps: list[float])[source]

Set the continuous one-dimensional timesteps for sampling/inference.

Parameters:

continuous_timesteps (list[float]) – the continuous one-dimensional timesteps for sampling/inference

set_discrete_timesteps_schedule(discrete_timesteps: list[int])[source]

Set the discrete one-dimensional timesteps for sampling/inference.

Parameters:

discrete_timesteps (list[int]) – the discrete one-dimensional timesteps for sampling/inference