ls_mlkit.diffuser.euclidean_diffuser module

class ls_mlkit.diffuser.euclidean_diffuser.EuclideanDiffuser(config: EuclideanDiffuserConfig, time_scheduler: TimeScheduler, masker: MaskerInterface, conditioner_list: list[Conditioner] = [])[source]

Bases: BaseDiffuser

forward_process_n_step(x: Tensor, t: Tensor, next_t: Tensor, padding_mask: Tensor, *args: Any, **kwargs: Any) → Tensor[source]

Forward process n step, from t to next_t

Parameters:

  • x (Tensor) – the sample

  • t (Tensor) – the timestep

  • next_t (Tensor) – the next timestep

  • padding_mask (Tensor) – the padding mask

Returns:

the sample at the next timestep

Return type:

Tensor

forward_process_one_step(x: Tensor, t: Tensor, padding_mask: Tensor, *args: Any, **kwargs: Any) → Tensor[source]

Forward process one step

Parameters:

  • x (Tensor) – the sample

  • t (Tensor) – the timestep

  • padding_mask (Tensor) – the padding mask

Returns:

the sample at the next timestep

Return type:

Tensor

get_accumulated_conditional_score(x_t: Tensor, t: Tensor, padding_mask: Tensor, *args: Any, **kwargs: Any) → Tensor[source]

Get the accumulated conditional score

Parameters:

  • x_t (Tensor) – \(x_t\)

  • t (Tensor) – \(t\)

  • padding_mask (Tensor) – the padding mask

Returns:

the accumulated conditional score

Return type:

Tensor

inpainting_x0_unconditionally(x_0: Tensor, padding_mask: Tensor, inpainting_mask: Tensor, device, recovery_mode: Literal['x_0', 'x_t'] = 'x_t', n_repaint_steps: int = 1, x_init_posterior: Tensor = None, inpainting_mask_key='inpainting_mask', *args: Any, **kwargs: Any) → Tensor[source]

Inpaint \(x_0\) unconditionally

Parameters:

  • x_0 (Tensor) – \(x_0\)

  • padding_mask (Tensor) – the padding mask

  • inpainting_mask (Tensor) – the inpainting mask

  • device (device) – the device to use for sampling

  • recovery_mode (Literal["x_0", "x_t"]) – the recovery mode

  • n_repaint_steps (int) – the number of repaint steps, n_repaint_steps > 1 means use RePaint.

  • x_init_posterior (Tensor) – Use x_init_posterior as the initial posterior if provided, otherwise sample from prior.

  • *args – additional arguments

  • **kwargs – additional keyword arguments

Returns:

\(x_0\) inpainted

Return type:

Tensor

sample_x0_unconditionally(shape: Tuple[int, ...], device, x_init_posterior: Tensor = None, *args: Any, **kwargs: Any) → Tensor[source]

Sample \(x_0\) unconditionally

Parameters:

  • shape (Tuple[int, ...]) – the shape of the sample

  • device (device) – the device to use for sampling

  • x_init_posterior (Tensor) – Use x_init_posterior as the initial posterior if provided, otherwise sample from prior.

  • *args – additional arguments

  • **kwargs – additional keyword arguments

Returns:

\(x_0\)

Return type:

Tensor

class ls_mlkit.diffuser.euclidean_diffuser.EuclideanDiffuserConfig(n_discretization_steps: int = 1000, ndim_micro_shape: int = 2, *args, **kwargs)[source]

Bases: BaseDiffuserConfig