144 lines
5.0 KiB
Python
144 lines
5.0 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Expose global quant cache status and safe inactive-artifact clearing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from collections.abc import Callable, Coroutine
|
|
from typing import Any, Protocol, cast
|
|
|
|
from aiohttp import web
|
|
|
|
from ..services.quant_cache_service import (
|
|
QuantCacheEvictionResult,
|
|
QuantCacheService,
|
|
QuantCacheStatus,
|
|
)
|
|
from ..services.quantized_model_boundaries import QuantCacheLimitProvider
|
|
from .quant_cache_settings import SettingsQuantCacheLimitProvider
|
|
|
|
QUANT_CACHE_ROUTE = "/simple-syrup/quant-cache"
|
|
Handler = Callable[[Any], Coroutine[Any, Any, web.Response]]
|
|
_REGISTERED_PROMPT_SERVERS: set[int] = set()
|
|
|
|
|
|
class QuantCacheRoutesProtocol(Protocol):
|
|
"""Describe the route decorators used by quant cache endpoints."""
|
|
|
|
def get(self, path: str) -> Callable[[Handler], Handler]:
|
|
"""Return a GET route decorator."""
|
|
|
|
def delete(self, path: str) -> Callable[[Handler], Handler]:
|
|
"""Return a DELETE route decorator."""
|
|
|
|
def post(self, path: str) -> Callable[[Handler], Handler]:
|
|
"""Return a POST route decorator."""
|
|
|
|
|
|
class QuantCachePromptServerProtocol(Protocol):
|
|
"""Describe the PromptServer state required by cache routes."""
|
|
|
|
routes: QuantCacheRoutesProtocol
|
|
|
|
|
|
class QuantCacheServiceBoundary(Protocol):
|
|
"""Describe global cache operations exposed through HTTP."""
|
|
|
|
def status(self) -> QuantCacheStatus:
|
|
"""Return current cache state."""
|
|
|
|
def clear_inactive(self) -> QuantCacheEvictionResult:
|
|
"""Remove every inactive managed artifact."""
|
|
|
|
def enforce_limit(self, limit_bytes: int) -> QuantCacheEvictionResult:
|
|
"""Apply the current global LRU budget."""
|
|
|
|
|
|
class QuantCacheHandlers:
|
|
"""Serve global quant cache state and explicit clear requests."""
|
|
|
|
def __init__(
|
|
self,
|
|
cache_service: QuantCacheServiceBoundary,
|
|
limit_provider: QuantCacheLimitProvider,
|
|
) -> None:
|
|
"""Create handlers with explicit authoritative collaborators."""
|
|
|
|
self._cache_service = cache_service
|
|
self._limit_provider = limit_provider
|
|
|
|
async def get_status(self, _request: Any) -> web.Response:
|
|
"""Return current global cache usage and configured limit."""
|
|
|
|
status = self._cache_service.status()
|
|
return web.json_response(status.to_payload(self._limit_provider.limit_bytes()))
|
|
|
|
async def clear_inactive(self, _request: Any) -> web.Response:
|
|
"""Clear inactive artifacts and return updated global cache status."""
|
|
|
|
eviction = self._cache_service.clear_inactive()
|
|
status = self._cache_service.status()
|
|
payload = status.to_payload(self._limit_provider.limit_bytes())
|
|
payload.update(
|
|
{
|
|
"removed_artifacts": eviction.removed_artifacts,
|
|
"removed_bytes": eviction.removed_bytes,
|
|
}
|
|
)
|
|
return web.json_response(payload)
|
|
|
|
async def enforce_limit(self, _request: Any) -> web.Response:
|
|
"""Apply the persisted limit and return updated global cache status."""
|
|
|
|
limit_bytes = self._limit_provider.limit_bytes()
|
|
eviction = self._cache_service.enforce_limit(limit_bytes)
|
|
status = self._cache_service.status()
|
|
payload = status.to_payload(limit_bytes)
|
|
payload.update(
|
|
{
|
|
"removed_artifacts": eviction.removed_artifacts,
|
|
"removed_bytes": eviction.removed_bytes,
|
|
}
|
|
)
|
|
return web.json_response(payload)
|
|
|
|
|
|
def register_quant_cache_routes(
|
|
cache_service: QuantCacheServiceBoundary | None = None,
|
|
limit_provider: QuantCacheLimitProvider | None = None,
|
|
prompt_server: QuantCachePromptServerProtocol | None = None,
|
|
) -> bool:
|
|
"""Register global cache routes with ComfyUI when PromptServer is available."""
|
|
|
|
server_instance = prompt_server or _prompt_server_instance()
|
|
if server_instance is None:
|
|
return False
|
|
server_key = id(server_instance)
|
|
if prompt_server is None and server_key in _REGISTERED_PROMPT_SERVERS:
|
|
return True
|
|
handlers = QuantCacheHandlers(
|
|
cache_service or QuantCacheService(),
|
|
limit_provider or SettingsQuantCacheLimitProvider(),
|
|
)
|
|
server_instance.routes.get(QUANT_CACHE_ROUTE)(handlers.get_status)
|
|
server_instance.routes.post(QUANT_CACHE_ROUTE)(handlers.enforce_limit)
|
|
server_instance.routes.delete(QUANT_CACHE_ROUTE)(handlers.clear_inactive)
|
|
if prompt_server is None:
|
|
_REGISTERED_PROMPT_SERVERS.add(server_key)
|
|
return True
|
|
|
|
|
|
def _prompt_server_instance() -> QuantCachePromptServerProtocol | None:
|
|
"""Return ComfyUI's PromptServer instance when available."""
|
|
|
|
try:
|
|
server_module = sys.modules["server"]
|
|
prompt_server = server_module.PromptServer
|
|
instance = prompt_server.instance
|
|
except (KeyError, AttributeError):
|
|
return None
|
|
return cast(QuantCachePromptServerProtocol, instance)
|