Auto3D.model_factory.ModelFactory

class Auto3D.model_factory.ModelFactory

Bases: object

Factory for creating and managing NNP model adapters.

Centralizes model creation logic to eliminate code duplication across batchopt.py, SPE.py, and thermo.py. Returns adapter instances that provide a consistent interface for all model types.

Includes model caching to avoid reloading models from disk repeatedly. Use clear_cache() to free memory when models are no longer needed.

Example

>>> factory = ModelFactory()
>>> model = factory.create("AIMNET", device=torch.device("cuda:0"))
>>> # Or use the convenience function
>>> model = create_model("ANI2x", device=torch.device("cpu"))
>>> # Clear cache when done
>>> ModelFactory.clear_cache()

__init__()

Methods

__init__()
available_models()	Return list of registered model names.
clear_cache()	Clear the model cache to free memory.
create(name[, device, compile_model, ...])	Create a model adapter by name.
get_cache_info()	Return cache statistics.

classmethod clear_cache() → None

Clear the model cache to free memory.

Call this at the end of a workflow to release GPU memory held by cached models.

classmethod get_cache_info() → dict[str, int]

Return cache statistics.

Returns:

Dictionary with cache size information.

classmethod create(name: str, device: device | None = None, compile_model: bool | None = None, use_ensemble: bool | None = None, use_cache: bool = True, **kwargs: Any) → BaseModelAdapter

Create a model adapter by name.

Parameters:

• name – Model name (‘AIMNET’, ‘ANI2x’, ‘ANI2xt’) or path to custom model.

• device – Target device for the model.

• compile_model – Whether to use torch.compile() for optimization. If None, checks AUTO3D_COMPILE_MODEL environment variable.

• use_ensemble – For AIMNET only: whether to use ensemble model (default False). Single model is ~35x faster and accurate enough for geometry optimization. Set True for highest accuracy. If None, checks AUTO3D_USE_ENSEMBLE env var.

• use_cache – Whether to cache and reuse model instances. Default True. Set False to force creating a new model instance.

• **kwargs – Additional arguments passed to the adapter constructor.

Returns:

Initialized model adapter on the specified device.

Raises:

ValueError – If model name is not recognized and not a valid path.

classmethod available_models() → list[str]

Return list of registered model names.