diff options
Diffstat (limited to 'src/mcp_sd15')
| -rw-r--r-- | src/mcp_sd15/pipeline_factory.py | 12 |
1 files changed, 11 insertions, 1 deletions
diff --git a/src/mcp_sd15/pipeline_factory.py b/src/mcp_sd15/pipeline_factory.py index cd92690..12d95ba 100644 --- a/src/mcp_sd15/pipeline_factory.py +++ b/src/mcp_sd15/pipeline_factory.py @@ -28,11 +28,13 @@ class PipelineFactory: def create_base(self) -> Any: torch, diffusers = self._runtime() + cache_dir = str(self.config.model.root / "huggingface") try: pipeline = diffusers.StableDiffusionPipeline.from_single_file( str(self.config.model.checkpoint_path()), torch_dtype=torch.float16, local_files_only=True, + cache_dir=cache_dir, ) vae_path = self.config.model.vae_path() if vae_path is not None: @@ -40,6 +42,11 @@ class PipelineFactory: str(vae_path), torch_dtype=torch.float16, local_files_only=True, + cache_dir=cache_dir, + ) + if getattr(pipeline, "vae", None) is None: + raise ConfigError( + "the checkpoint does not provide a VAE and no external VAE was loaded" ) adapter_names: list[str] = [] adapter_weights: list[float] = [] @@ -69,7 +76,10 @@ class PipelineFactory: torch, diffusers = self._runtime() try: return diffusers.ControlNetModel.from_single_file( - str(path), torch_dtype=torch.float16, local_files_only=True + str(path), + torch_dtype=torch.float16, + local_files_only=True, + cache_dir=str(self.config.model.root / "huggingface"), ).to("cuda") except Exception as exc: raise ConfigError(f"failed to load ControlNet model {path}: {exc}") from exc |