diff --git a/src/infra/diffusers/_autoencoder_repository.py b/src/infra/diffusers/_autoencoder_repository.py index e6958dc..2337a84 100644 --- a/src/infra/diffusers/_autoencoder_repository.py +++ b/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 diff --git a/src/infra/diffusers/_cache.py b/src/infra/diffusers/_cache.py new file mode 100644 index 0000000..204ca66 --- /dev/null +++ b/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) diff --git a/src/infra/diffusers/_pipeline_repository.py b/src/infra/diffusers/_pipeline_repository.py index a024f75..8322176 100644 --- a/src/infra/diffusers/_pipeline_repository.py +++ b/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,