ls_mlkit.diffuser.model_interface module

class ls_mlkit.diffuser.model_interface.Model4DiffuserInterface[source]

Bases: ABC

abstractmethod get_model_device() → device[source]

Get the device of the model

Returns:

the device of the model

Return type:

torch.device

abstractmethod prepare_batch_data_for_input(batch: dict[str, Any]) → dict[str, Any][source]

Prepare batch data for input

Parameters:

batch (dict[str, Any]) – the batch of data

Returns:

the prepared batch of data

Return type:

dict[str, Any]