feat: add cache

src/infra/diffusers/_autoencoder_repository.py

@@ -18,6 +18,7 @@ from omegaconf import OmegaConf
 from safetensors import safe_open
 
 from ...domain.repositories import AutoencoderRepository
+from ._cache import is_vae_cached
 
 
 class DiffusersAutoencoderRepository(AutoencoderRepository):
@@ -30,6 +31,10 @@ class DiffusersAutoencoderRepository(AutoencoderRepository):
         checkpoint_name = os.path.basename(checkpoint_path)
         ckpt_cache_path = os.path.join(self.cache_dir, checkpoint_name)
 
+        # Check if already cached
+        if is_vae_cached(ckpt_cache_path):
+            return ckpt_cache_path
+
         self._vae_pt_to_vae_diffuser(checkpoint_path, ckpt_cache_path)
         return ckpt_cache_path
 

src/infra/diffusers/_cache.py

@@ -0,0 +1,35 @@
+import os
+from typing import Final
+
+# Centralized cache helpers for diffusers-related repositories
+
+_PIPELINE_INDEX_FILE: Final[str] = "model_index.json"
+_VAE_CONFIG_FILE: Final[str] = "config.json"
+_VAE_WEIGHT_CANDIDATES: Final[tuple[str, ...]] = (
+    "diffusion_pytorch_model.safetensors",
+    "diffusion_pytorch_model.bin",
+    "pytorch_model.bin",
+    "model.safetensors",
+)
+
+
+def is_pipeline_cached(cache_path: str) -> bool:
+    """Returns True if the Stable Diffusion pipeline has already been converted.
+
+    Criteria: directory exists and contains a `model_index.json`.
+    """
+    return os.path.isdir(cache_path) and os.path.isfile(
+        os.path.join(cache_path, _PIPELINE_INDEX_FILE)
+    )
+
+
+def is_vae_cached(cache_path: str) -> bool:
+    """Returns True if the VAE has already been converted.
+
+    Criteria: directory exists, contains `config.json`, and at least one known weight file.
+    """
+    if not (
+        os.path.isdir(cache_path) and os.path.isfile(os.path.join(cache_path, _VAE_CONFIG_FILE))
+    ):
+        return False
+    return any(os.path.isfile(os.path.join(cache_path, fname)) for fname in _VAE_WEIGHT_CANDIDATES)

src/infra/diffusers/_pipeline_repository.py

@@ -6,6 +6,7 @@ from comfy.model_management import get_torch_device  # pyright: ignore[reportMis
 from diffusers import StableDiffusionPipeline
 
 from ...domain.repositories import PipelineRepository
+from ._cache import is_pipeline_cached
 
 
 class DiffusersPipelineRepository(PipelineRepository):
@@ -17,6 +18,10 @@ class DiffusersPipelineRepository(PipelineRepository):
         checkpoint_name = os.path.basename(checkpoint_path)
         ckpt_cache_path = os.path.join(self.cache_dir, checkpoint_name)
 
+        # Check if already cached
+        if is_pipeline_cached(ckpt_cache_path):
+            return ckpt_cache_path
+
         StableDiffusionPipeline.from_single_file(
             pretrained_model_link_or_path=checkpoint_path,
             torch_dtype=dtype,