1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
|
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], **kwargs: object) -> object:
return SimpleNamespace(implementation=cls.__name__, config=config, options=kwargs)
class FakeDPMSolverMultistepScheduler(FakeScheduler):
pass
class FakeDDIMScheduler(FakeScheduler):
pass
class FakeEulerAncestralDiscreteScheduler(FakeScheduler):
pass
class FakeLCMScheduler(FakeScheduler):
pass
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=FakeDPMSolverMultistepScheduler,
DDIMScheduler=FakeDDIMScheduler,
EulerAncestralDiscreteScheduler=FakeEulerAncestralDiscreteScheduler,
LCMScheduler=FakeLCMScheduler,
)
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)
monkeypatch.setattr(
PipelineFactory, "_vae_config_dir", staticmethod(lambda cache_dir: "fake_vae_config")
)
result = factory.create_base()
assert result is pipeline
assert pipeline.vae is external_vae
assert pipeline.vae_load["config"] == "fake_vae_config" # type: ignore[attr-defined]
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()
@pytest.mark.parametrize(
("name", "implementation", "options"),
[
(
"dpmpp_2m_karras",
"FakeDPMSolverMultistepScheduler",
{
"algorithm_type": "dpmsolver++",
"solver_order": 2,
"use_karras_sigmas": True,
},
),
(
"dpmpp_2m_sde_karras",
"FakeDPMSolverMultistepScheduler",
{
"algorithm_type": "sde-dpmsolver++",
"solver_order": 2,
"use_karras_sigmas": True,
},
),
("euler_a", "FakeEulerAncestralDiscreteScheduler", {}),
("ddim", "FakeDDIMScheduler", {}),
("lcm", "FakeLCMScheduler", {}),
],
)
def test_scheduler_mapping(app_config, name, implementation, options) -> None:
generation = replace(app_config.generation, scheduler=name)
factory = PipelineFactory(replace(app_config, generation=generation))
current_config = {"beta_start": 0.00085, "beta_end": 0.012}
current = SimpleNamespace(config=current_config)
diffusers = fake_diffusers(FakePipeline(), object())
result = factory._create_scheduler(current, diffusers)
assert result.implementation == implementation
assert result.config is current_config
assert result.options == options
|
