otx.core.exporter.diffusion#

Exporter for diffusion models that uses native torch and OpenVINO conversion tools.

Classes

DiffusionOTXModelExporter(...[, mean, std, ...])

Exporter for diffusion models that uses native torch and OpenVINO conversion tools.
class otx.core.exporter.diffusion.DiffusionOTXModelExporter(task_level_export_parameters: TaskLevelExportParameters, input_size: tuple[int, ...], mean: tuple[float, float, float] = (0.0, 0.0, 0.0), std: tuple[float, float, float] = (1.0, 1.0, 1.0), resize_mode: Literal['crop', 'standard', 'fit_to_window', 'fit_to_window_letterbox'] = 'standard', pad_value: int = 0, swap_rgb: bool = False, via_onnx: bool = False, onnx_export_configuration: dict[str, Any] | None = None, output_names: list[str] | None = None, input_names: list[str] | None = None)[source]#

Bases: OTXNativeModelExporter

Exporter for diffusion models that uses native torch and OpenVINO conversion tools.

export(model: OTXModel, output_dir: Path, base_model_name: str = 'exported_model', export_format: OTXExportFormatType = OTXExportFormatType.OPENVINO, precision: OTXPrecisionType = OTXPrecisionType.FP32, to_exportable_code: bool = False) → Path[source]#

Exports input model to the specified deployable format, such as OpenVINO IR or ONNX.

Parameters:

  • model (OTXModel) – OTXModel to be exported

  • output_dir (Path) – path to the directory to store export artifacts

  • base_model_name (str, optional) – exported model name

  • format (OTXExportFormatType) – final format of the exported model

  • precision (OTXExportPrecisionType, optional) – precision of the exported model’s weights

  • to_exportable_code (bool, optional) – whether to generate exportable code. Currently not supported by Diffusion task.

Returns:

path to the exported model

Return type:

Path

to_onnx(model: OTXModel | Module, output_dir: Path, base_model_name: str = 'exported_model', precision: OTXPrecisionType = OTXPrecisionType.FP32, embed_metadata: bool = True, model_type: str = 'stable_diffusion') → Path[source]#

Export the given PyTorch model to ONNX format and save it to the specified output directory.

Parameters:

  • model (OTXModel) – OTXModel to be exported.

  • output_dir (Path) – The directory where the ONNX model will be saved.

  • base_model_name (str, optional) – The base name for the exported model. Defaults to “exported_model”.

  • precision (OTXPrecisionType, optional) – The precision type for the exported model.

  • OTXPrecisionType.FP32. (Defaults to) –

  • embed_metadata (bool, optional) – Whether to embed metadata in the ONNX model. Defaults to True.

Returns:

The path to the saved ONNX model.

Return type:

Path

to_openvino(model: OTXModel | Module, output_dir: Path, base_model_name: str = 'exported_model', precision: OTXPrecisionType = OTXPrecisionType.FP32) → Path[source]#

Export to OpenVINO Intermediate Representation format.

In this implementation the export is done only via standard OV/ONNX tools.