ls_mlkit.diffusion.conditioner.conditioner module

class ls_mlkit.diffusion.conditioner.conditioner.Conditioner(guidance_scale: float = 1.0)

Bases: ABC

disable()

enable()

abstractmethod get_conditional_score(x_t: Tensor, t: Tensor, padding_mask: Tensor, *args: Any, **kwargs: Any) → Tensor

Get conditional score

Parameters:

• x_t (Tensor) – the input tensor

• t (Tensor) – the time tensor

• padding_mask (Tensor) – the padding mask

Returns:

the conditional score

Return type:

Tensor

get_guidance_scale()

property guidance_scale

is_enabled() → bool

abstractmethod prepare_condition_dict(train: bool = True, *args: list[Any], **kwargs: dict[Any, Any]) → dict[str, Any]

Prepare the condition dictionary

Parameters:

train (bool, optional) – whether the conditioner is used in training. Defaults to True.

Returns:

the condition dictionary

Return type:

dict[str, Any]

abstractmethod set_condition(*args: list[Any], **kwargs: dict[Any, Any]) → None

Set the condition

Parameters:

• *args – additional arguments

• **kwargs – additional keyword arguments

set_guidance_scale(guidance_scale: float)

class ls_mlkit.diffusion.conditioner.conditioner.LGDConditioner(guidance_scale: float = 1.0)

Bases: Conditioner

Loss Guidance Diffusion Conditioner

abstractmethod compute_conditional_loss(p_gt_data: Tensor, padding_mask: Tensor) → Tensor

Compute the conditional loss

Parameters:

• p_gt_data (Tensor) – predicted clean data.

• padding_mask (Tensor) – the padding mask

Returns:

the conditional loss

Return type:

Tensor

get_conditional_score(x_t: Tensor, t: Tensor, padding_mask: Tensor, *args: Any, **kwargs: Any) → Tensor

Get conditional score

Parameters:

• x_t (Tensor) – the input tensor

• t (Tensor) – the time tensor

• padding_mask (Tensor) – the padding mask

Returns:

the conditional score

Return type:

Tensor