ls_mlkit.diffusion.euclidean_diffuser module

class ls_mlkit.diffusion.euclidean_diffuser.EuclideanDiffuser(config: EuclideanDiffuserConfig, time_scheduler: DiffusionTimeScheduler, masker: MaskerInterface)[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

inpainting(x: Tensor, padding_mask: Tensor, inpainting_mask: Tensor, device, x_init_posterior: Tensor = None, inpainting_mask_key='inpainting_mask', n_repaint_steps: int = 1, return_all=False, *args: Any, **kwargs: Any) → dict[source]

_summary_

Parameters:

  • x (_type_) – _description_

  • padding_mask (_type_) – _description_

  • inpainting_mask (_type_) – _description_

  • device (_type_) – _description_

  • x_init_posterior (_type_, optional) – _description_. Defaults to None.

  • inpainting_mask_key (str, optional) – _description_. Defaults to “inpainting_mask”.

  • sapmling_condition_key (str, optional) – _description_. Defaults to “sapmling_condition”.

  • return_all (bool, optional) – _description_. Defaults to False.

  • sampling_condition (_type_, optional) – _description_. Defaults to None.

Returns:

_description_

Return type:

dict

recover_bright_region(x_known, x_t, t, padding_mask, inpainting_mask, x_prior) → Tensor[source]
sampling(shape: Tuple[int, ...], device, x_init_posterior: Tensor = None, return_all=False, *args: Any, **kwargs: Any) → dict[source]

_summary_

Parameters:

  • shape (_type_) – _description_

  • device (_type_) – _description_

  • x_init_posterior (_type_, optional) – _description_. Defaults to None.

  • return_all (bool, optional) – _description_. Defaults to False.

  • sampling_condition (_type_, optional) – _description_. Defaults to None.

  • sapmling_condition_key (str, optional) – _description_. Defaults to “sapmling_condition”.

Returns:

_description_

Return type:

dict

class ls_mlkit.diffusion.euclidean_diffuser.EuclideanDiffuserConfig(n_discretization_steps: int = 1000, ndim_micro_shape: int = 2, n_inference_steps: int = None, use_batch_flattening: bool = False, *args, **kwargs)[source]

Bases: BaseDiffuserConfig