AnyImage/src/snapform/server/model_manager.py

import gc
import time
from pathlib import Path
from types import SimpleNamespace

from . import model_catalog
from .model_download import download_model


class ModelSessionCache:
    """Cache ONNX sessions shared by Jobs in one Server process."""

    def __init__(self):
        self.ben2_model = None
        self.ben2_key = None
        self.da3_model = None
        self.da3_key = None
        self.moge_model = None
        self.moge_key = None

    def get_ben2(self, model_dir, device):
        requested_device = getattr(device, "type", device)
        key = (str(Path(model_dir).resolve()), requested_device, "onnx")
        if self.ben2_model is not None and self.ben2_key == key:
            return self.ben2_model, 0.0
        self.release_ben2()
        return self._load_ben2(model_dir, requested_device, key)

    def _load_ben2(self, model_dir, requested_device, key):
        from .models.onnx_ben2 import create_session

        started = time.perf_counter()
        model = create_session(model_dir, requested_device)
        self.ben2_model = model
        self.ben2_key = key
        elapsed_ms = (time.perf_counter() - started) * 1000.0
        return model, elapsed_ms

    def get_da3(self, model_dir, device):
        requested_device = getattr(device, "type", device)
        key = (str(Path(model_dir).resolve()), requested_device, "onnx")
        if self.moge_model is not None:
            self.release_moge()
        if self.da3_model is not None and self.da3_key == key:
            return self.da3_model
        self.release_da3()
        return self._load_da3(model_dir, requested_device, key)

    def _load_da3(self, model_dir, requested_device, key):
        from .models.onnx_da3 import create_session

        model = create_session(model_dir, requested_device)
        self.da3_model = model
        self.da3_key = key
        return model

    def get_moge2(self, model_dir, requested_device):
        from .models.onnx_moge2 import create_session

        key = (str(Path(model_dir).resolve()), requested_device, "onnx")
        if self.da3_model is not None:
            self.release_da3()
        if self.moge_model is not None and self.moge_key == key:
            return self.moge_model
        self.release_moge()
        model = create_session(model_dir, requested_device)
        self.moge_model = model
        self.moge_key = key
        return model

    def release_ben2(self):
        self.ben2_model = None
        self.ben2_key = None
        self._release_runtime_cache()

    def release_da3(self):
        self.da3_model = None
        self.da3_key = None
        self._release_runtime_cache()

    def release_moge(self):
        self.moge_model = None
        self.moge_key = None
        self._release_runtime_cache()

    def close(self):
        self.ben2_model = None
        self.da3_model = None
        self.moge_model = None
        self.ben2_key = None
        self.da3_key = None
        self.moge_key = None
        self._release_runtime_cache()

    def snapshot(self):
        return {
            "ben2": Path(self.ben2_key[0]).name if self.ben2_key else "",
            "da3": Path(self.da3_key[0]).name if self.da3_key else "",
            "moge2": Path(self.moge_key[0]).name if self.moge_key else "",
        }

    def _release_runtime_cache(self):
        gc.collect()


class ModelManager:
    """Own downloaded model files and loaded inference sessions."""

    def __init__(self, models_directory):
        self.models_directory = Path(models_directory)
        self.sessions = ModelSessionCache()

    def directory(self, model_key):
        model = model_catalog.get_downloadable_model(model_key)
        directory = self.models_directory / model.directory_name
        old_moge_directory = self.models_directory / "moge2_base_normal"
        if (
            model_key == "MOGE2_VITB_NORMAL"
            and not directory.is_dir()
            and old_moge_directory.is_dir()
        ):
            old_moge_directory.rename(directory)
        legacy_directory = self.models_directory / model_key.lower()
        if (
            not directory.is_dir()
            and legacy_directory.is_dir()
            and legacy_directory != directory
        ):
            legacy_directory.rename(directory)
        return directory

    def ready(self, model_key):
        model = model_catalog.get_downloadable_model(model_key)
        directory = self.directory(model_key)
        return all(any(directory.glob(pattern)) for pattern in model.ready_patterns)

    def get_ben2(self, model_dir, device):
        return self.sessions.get_ben2(model_dir, device)

    def get_da3(self, model_dir, device):
        return self.sessions.get_da3(model_dir, device)

    def get_moge2(self, model_dir, device):
        return self.sessions.get_moge2(model_dir, device)

    def snapshot(self):
        return {"loaded": self.sessions.snapshot()}

    def clear(self):
        self.sessions.close()

    def close(self):
        self.sessions.close()

    def download(self, context, model_key):
        return self._download(context, model_key, context.progress)

    def download_missing(self, context, model_keys):
        missing = [key for key in model_keys if not self.ready(key)]
        if not missing:
            context.progress(1.0, "Required models are ready")
            return {"downloaded": []}

        total = len(missing)
        for index, model_key in enumerate(missing):
            def report(value, message, *, offset=index):
                context.progress((offset + value) / total, message)

            self._download(context, model_key, report)
        return {"downloaded": missing}

    def _download(self, context, model_key, progress):
        model = model_catalog.get_downloadable_model(model_key)
        args = SimpleNamespace(
            repository=model.repository,
            model_dir=str(self.directory(model_key)),
            progress=progress,
            cancel_check=context.check_cancelled,
            label=model.label,
            allow_pattern=list(model.download_patterns),
            required_pattern=list(model.ready_patterns),
            r2_directory=model.r2_directory,
            r2_file=[
                f"{filename}|{size}|{checksum}"
                for filename, size, checksum in model.r2_files
            ],
        )
        context.check_cancelled()
        download_model(args)
        context.check_cancelled()