diff --git a/src/infra/diffusers/_autoencoder_repository.py b/src/infra/diffusers/_autoencoder_repository.py index c85b0af..7403291 100644 --- a/src/infra/diffusers/_autoencoder_repository.py +++ b/src/infra/diffusers/_autoencoder_repository.py @@ -5,8 +5,7 @@ from comfy.model_management import get_torch_device # pyright: ignore[reportMis from diffusers import AutoencoderKL from ...domain.repositories import AutoencoderRepository -from ...utils import get_cache_path -from ._cache import is_vae_cached +from ._cache import get_cache_path, is_vae_cached class DiffusersAutoencoderRepository(AutoencoderRepository): diff --git a/src/infra/diffusers/_cache.py b/src/infra/diffusers/_cache.py index 204ca66..ed566c0 100644 --- a/src/infra/diffusers/_cache.py +++ b/src/infra/diffusers/_cache.py @@ -1,6 +1,9 @@ import os +from pathlib import Path from typing import Final +from ...config import get_cache_dir, get_project_root + # Centralized cache helpers for diffusers-related repositories _PIPELINE_INDEX_FILE: Final[str] = "model_index.json" @@ -13,6 +16,26 @@ _VAE_WEIGHT_CANDIDATES: Final[tuple[str, ...]] = ( ) +def setup_cache_path() -> Path: + env_cache_str = get_cache_dir() + env_cache_path = Path(env_cache_str) + if env_cache_path.is_absolute(): + cache_dir = env_cache_path + else: + project_root = get_project_root() + cache_dir = (project_root / env_cache_path).resolve() + if not cache_dir.exists(): + cache_dir.mkdir(parents=True, exist_ok=True) + return cache_dir + + +def get_cache_path(filename: str | None = None) -> Path: + cache_dir = setup_cache_path() + if filename is None: + return cache_dir + return cache_dir / filename + + def is_pipeline_cached(cache_path: str) -> bool: """Returns True if the Stable Diffusion pipeline has already been converted. diff --git a/src/infra/diffusers/_pipeline_repository.py b/src/infra/diffusers/_pipeline_repository.py index 6485e8c..d2828f2 100644 --- a/src/infra/diffusers/_pipeline_repository.py +++ b/src/infra/diffusers/_pipeline_repository.py @@ -5,8 +5,7 @@ from comfy.model_management import get_torch_device # pyright: ignore[reportMis from diffusers import StableDiffusionPipeline from ...domain.repositories import PipelineRepository -from ...utils import get_cache_path -from ._cache import is_pipeline_cached +from ._cache import get_cache_path, is_pipeline_cached class DiffusersPipelineRepository(PipelineRepository): diff --git a/src/infra/diffusers/_scheduler_repository.py b/src/infra/diffusers/_scheduler_repository.py index b82e530..c99ad65 100644 --- a/src/infra/diffusers/_scheduler_repository.py +++ b/src/infra/diffusers/_scheduler_repository.py @@ -3,7 +3,7 @@ from diffusers.schedulers.scheduling_utils import SchedulerMixin from ...domain.model import Scheduler from ...domain.repositories import SchedulerRepository -from ...utils import get_cache_path +from ._cache import get_cache_path class DiffusersSchedulerRepository(SchedulerRepository): diff --git a/src/utils/__init__.py b/src/utils/__init__.py deleted file mode 100644 index 8e67c85..0000000 --- a/src/utils/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from ._cache import ensure_cache_dir, get_cache_path - -__all__ = ["ensure_cache_dir", "get_cache_path"] diff --git a/src/utils/_cache.py b/src/utils/_cache.py deleted file mode 100644 index 08cda45..0000000 --- a/src/utils/_cache.py +++ /dev/null @@ -1,23 +0,0 @@ -from pathlib import Path - -from ..config import get_cache_dir, get_project_root - - -def ensure_cache_dir() -> Path: - env_cache_str = get_cache_dir() - env_cache_path = Path(env_cache_str) - if env_cache_path.is_absolute(): - cache_dir = env_cache_path - else: - project_root = get_project_root() - cache_dir = (project_root / env_cache_path).resolve() - if not cache_dir.exists(): - cache_dir.mkdir(parents=True, exist_ok=True) - return cache_dir - - -def get_cache_path(filename: str | None = None) -> Path: - cache_dir = ensure_cache_dir() - if filename is None: - return cache_dir - return cache_dir / filename