5a2de30066
Use the module singleton pattern to share global state. This can also be a modified version of the Global Object Pattern. The main reason this pattern is used is for ease of use when handling global state rather than adding extra dependencies for a DI parameter. Signed-off-by: kingbri <bdashore3@proton.me>
76 lines
2.1 KiB
Python
76 lines
2.1 KiB
Python
"""
|
|
Manages the storage and utility of model containers.
|
|
|
|
Containers exist as a common interface for backends.
|
|
"""
|
|
|
|
import pathlib
|
|
from loguru import logger
|
|
from typing import Optional
|
|
|
|
from backends.exllamav2.model import ExllamaV2Container
|
|
from common.logger import get_loading_progress_bar
|
|
from common.utils import load_progress
|
|
|
|
|
|
# Global model container
|
|
container: Optional[ExllamaV2Container] = None
|
|
|
|
|
|
async def unload_model():
|
|
"""Unloads a model"""
|
|
global container
|
|
|
|
container.unload()
|
|
container = None
|
|
|
|
|
|
async def load_model_gen(model_path: pathlib.Path, **kwargs):
|
|
"""Generator to load a model"""
|
|
global container
|
|
|
|
# Check if the model is already loaded
|
|
if container and container.model:
|
|
loaded_model_name = container.get_model_path().name
|
|
|
|
if loaded_model_name == model_path.name:
|
|
raise ValueError(
|
|
f'Model "{loaded_model_name}" is already loaded! Aborting.'
|
|
)
|
|
|
|
# Unload the existing model
|
|
if container and container.model:
|
|
logger.info("Unloading existing model.")
|
|
await unload_model()
|
|
|
|
container = ExllamaV2Container(model_path.resolve(), False, **kwargs)
|
|
|
|
model_type = "draft" if container.draft_config else "model"
|
|
load_status = container.load_gen(load_progress)
|
|
|
|
progress = get_loading_progress_bar()
|
|
progress.start()
|
|
|
|
try:
|
|
for module, modules in load_status:
|
|
if module == 0:
|
|
loading_task = progress.add_task(
|
|
f"[cyan]Loading {model_type} modules", total=modules
|
|
)
|
|
else:
|
|
progress.advance(loading_task)
|
|
if module == modules:
|
|
# Switch to model progress if the draft model is loaded
|
|
if model_type == "draft":
|
|
model_type = "model"
|
|
else:
|
|
progress.stop()
|
|
|
|
yield module, modules, model_type
|
|
finally:
|
|
progress.stop()
|
|
|
|
|
|
async def load_model(model_path: pathlib.Path, **kwargs):
|
|
async for _, _, _ in load_model_gen(model_path, **kwargs):
|
|
pass
|