diff options
Diffstat (limited to 'tests/test_pipeline_factory.py')
| -rw-r--r-- | tests/test_pipeline_factory.py | 99 |
1 files changed, 99 insertions, 0 deletions
diff --git a/tests/test_pipeline_factory.py b/tests/test_pipeline_factory.py new file mode 100644 index 0000000..454a5b1 --- /dev/null +++ b/tests/test_pipeline_factory.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from dataclasses import replace +from types import SimpleNamespace + +import pytest + +from mcp_sd15.config import LoRAConfig +from mcp_sd15.errors import ConfigError +from mcp_sd15.pipeline_factory import PipelineFactory + + +class FakePipeline: + def __init__(self, *, with_vae: bool = True) -> None: + self.vae = object() if with_vae else None + self.scheduler = SimpleNamespace(config={}) + self.lora_loads: list[dict[str, object]] = [] + self.adapters: tuple[list[str], list[float]] | None = None + self.device: str | None = None + + def load_lora_weights(self, directory: str, **kwargs: object) -> None: + self.lora_loads.append({"directory": directory, **kwargs}) + + def set_adapters(self, names: list[str], *, adapter_weights: list[float]) -> None: + self.adapters = (names, adapter_weights) + + def set_progress_bar_config(self, *, disable: bool) -> None: + assert disable is True + + def to(self, device: str) -> FakePipeline: + self.device = device + return self + + +class FakeScheduler: + @classmethod + def from_config(cls, config: dict[str, object]) -> object: + return {"scheduler": cls.__name__, **config} + + +def fake_diffusers(pipeline: FakePipeline, external_vae: object) -> SimpleNamespace: + class StableDiffusionPipeline: + @classmethod + def from_single_file(cls, path: str, **kwargs: object) -> FakePipeline: + pipeline.base_load = {"path": path, **kwargs} # type: ignore[attr-defined] + return pipeline + + class AutoencoderKL: + @classmethod + def from_single_file(cls, path: str, **kwargs: object) -> object: + pipeline.vae_load = {"path": path, **kwargs} # type: ignore[attr-defined] + return external_vae + + return SimpleNamespace( + StableDiffusionPipeline=StableDiffusionPipeline, + AutoencoderKL=AutoencoderKL, + DPMSolverMultistepScheduler=FakeScheduler, + DDIMScheduler=FakeScheduler, + EulerDiscreteScheduler=FakeScheduler, + EulerAncestralDiscreteScheduler=FakeScheduler, + ) + + +def test_base_pipeline_requires_vae_and_applies_all_loras(app_config, monkeypatch) -> None: + model = replace( + app_config.model, + vae="external.safetensors", + loras=( + LoRAConfig("detail.safetensors", 0.8), + LoRAConfig("style.safetensors", 0.5), + ), + ) + factory = PipelineFactory(replace(app_config, model=model)) + pipeline = FakePipeline() + external_vae = object() + runtime = (SimpleNamespace(float16="float16"), fake_diffusers(pipeline, external_vae)) + monkeypatch.setattr(factory, "_runtime", lambda: runtime) + + result = factory.create_base() + + assert result is pipeline + assert pipeline.vae is external_vae + assert pipeline.adapters == (["configured_lora_0", "configured_lora_1"], [0.8, 0.5]) + assert [item["weight_name"] for item in pipeline.lora_loads] == [ + "detail.safetensors", + "style.safetensors", + ] + assert all(item["local_files_only"] is True for item in pipeline.lora_loads) + assert pipeline.device == "cuda" + + +def test_checkpoint_without_vae_fails_pipeline_initialization(app_config, monkeypatch) -> None: + factory = PipelineFactory(app_config) + pipeline = FakePipeline(with_vae=False) + runtime = (SimpleNamespace(float16="float16"), fake_diffusers(pipeline, object())) + monkeypatch.setattr(factory, "_runtime", lambda: runtime) + + with pytest.raises(ConfigError, match="does not provide a VAE"): + factory.create_base() |
