141 lines
3.2 KiB
Python
141 lines
3.2 KiB
Python
"""
|
|
Lilith Model Loader - Unified ML model loading and caching.
|
|
|
|
Provides a consistent interface for loading models across different frameworks:
|
|
- HuggingFace Transformers (text, classification, QA, etc.)
|
|
- Diffusers (Stable Diffusion, SDXL, Flux)
|
|
- GGUF (llama-cpp-python for quantized models)
|
|
|
|
Features:
|
|
- Remote model fetching via rsync/scp
|
|
- Local caching with size verification
|
|
- GPU/CPU device management
|
|
- Manifest-based model registry
|
|
- Async and sync loading APIs
|
|
|
|
Example:
|
|
>>> from lilith_model_loader import get_loader, ensure_model
|
|
|
|
>>> # Quick path resolution
|
|
>>> path = ensure_model("ministral-3b-instruct")
|
|
|
|
>>> # Load with framework-specific loader
|
|
>>> loader = get_loader("hf")
|
|
>>> model = await loader.load("Marqo/nsfw-image-detection-384", task="image-classification")
|
|
|
|
>>> # Load SDXL
|
|
>>> loader = get_loader("diffusers")
|
|
>>> pipeline = await loader.load("stabilityai/stable-diffusion-xl-base-1.0")
|
|
|
|
>>> # Load GGUF model
|
|
>>> loader = get_loader("gguf")
|
|
>>> llm = await loader.load("ministral-3b-instruct", n_gpu_layers=-1)
|
|
"""
|
|
|
|
__version__ = "1.1.0"
|
|
|
|
# Base classes and types
|
|
from .base import (
|
|
BaseModelLoader,
|
|
ModelInfo,
|
|
ModelLoadError,
|
|
ModelNotFoundError,
|
|
DeviceNotAvailableError,
|
|
)
|
|
|
|
from .types import (
|
|
ModelEntry,
|
|
LoadResult,
|
|
TransferProgress,
|
|
CachedModel,
|
|
CacheStats,
|
|
)
|
|
|
|
# Device management
|
|
from .device import (
|
|
DeviceManager,
|
|
DeviceInfo,
|
|
get_best_device,
|
|
get_device_count,
|
|
allocate_device,
|
|
clear_gpu_cache,
|
|
)
|
|
|
|
# Registry
|
|
from .registry import (
|
|
register_loader,
|
|
get_loader,
|
|
get_loader_class,
|
|
list_loaders,
|
|
list_aliases,
|
|
is_loader_registered,
|
|
)
|
|
|
|
# Original loader (path resolution and caching)
|
|
from .loader import (
|
|
ModelLoader as PathModelLoader,
|
|
ensure_model,
|
|
ensure_model_sync,
|
|
get_loader as get_path_loader,
|
|
)
|
|
|
|
# Backwards compatibility
|
|
ModelLoader = PathModelLoader
|
|
|
|
# Framework-specific loaders (imported lazily to avoid dependency issues)
|
|
# These will auto-register with the registry when imported
|
|
try:
|
|
from .hf_loader import HFModelLoader
|
|
except ImportError:
|
|
HFModelLoader = None # type: ignore
|
|
|
|
try:
|
|
from .diffusers_loader import DiffusersLoader
|
|
except ImportError:
|
|
DiffusersLoader = None # type: ignore
|
|
|
|
try:
|
|
from .gguf_loader import GGUFModelLoader
|
|
except ImportError:
|
|
GGUFModelLoader = None # type: ignore
|
|
|
|
__all__ = [
|
|
# Version
|
|
"__version__",
|
|
# Base classes
|
|
"BaseModelLoader",
|
|
"ModelInfo",
|
|
"ModelLoadError",
|
|
"ModelNotFoundError",
|
|
"DeviceNotAvailableError",
|
|
# Types
|
|
"ModelEntry",
|
|
"LoadResult",
|
|
"TransferProgress",
|
|
"CachedModel",
|
|
"CacheStats",
|
|
# Device management
|
|
"DeviceManager",
|
|
"DeviceInfo",
|
|
"get_best_device",
|
|
"get_device_count",
|
|
"allocate_device",
|
|
"clear_gpu_cache",
|
|
# Registry
|
|
"register_loader",
|
|
"get_loader",
|
|
"get_loader_class",
|
|
"list_loaders",
|
|
"list_aliases",
|
|
"is_loader_registered",
|
|
# Path resolution (original API)
|
|
"ModelLoader",
|
|
"PathModelLoader",
|
|
"ensure_model",
|
|
"ensure_model_sync",
|
|
"get_path_loader",
|
|
# Framework loaders
|
|
"HFModelLoader",
|
|
"DiffusersLoader",
|
|
"GGUFModelLoader",
|
|
]
|