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
|
from __future__ import annotations
import asyncio
import io
import threading
import time
from dataclasses import replace
import pytest
from PIL import Image
from mcp_sd15.config import ControlModelConfig, ControlNetConfig
from mcp_sd15.errors import GenerationError, InvalidRequestError
from mcp_sd15.schema import GuidedImageRequest, InpaintImageRequest, TextToImageRequest
from mcp_sd15.service import ImageGenerationService
class FakeBackend:
def __init__(self, delay: float = 0) -> None:
self.calls: list[str] = []
self.delay = delay
self.active = 0
self.max_active = 0
self.guard = threading.Lock()
def _result(self, name: str) -> Image.Image:
with self.guard:
self.active += 1
self.max_active = max(self.max_active, self.active)
time.sleep(self.delay)
with self.guard:
self.active -= 1
self.calls.append(name)
return Image.new("RGB", (4, 4), "red")
def text_to_image(self, request: object, seed: int) -> Image.Image:
return self._result("text")
def guided_image(self, request: object, image: Image.Image, seed: int) -> Image.Image:
assert image.mode == "RGB"
return self._result("guided")
def inpaint_image(
self, request: object, image: Image.Image, mask: Image.Image, seed: int
) -> Image.Image:
assert mask.mode == "L"
return self._result("inpaint")
async def test_routes_all_generation_types(app_config, image_base64: str) -> None:
backend = FakeBackend()
service = ImageGenerationService(app_config, backend)
text = await service.generate_image(TextToImageRequest.create("cat"))
guided = await service.generate_guided_image(
GuidedImageRequest.create("cat", "depth", image_base64)
)
inpaint = await service.inpaint_image(
InpaintImageRequest.create("cat", image_base64, image_base64)
)
assert backend.calls == ["text", "guided", "inpaint"]
for result in (text, guided, inpaint):
assert result.png_data.startswith(b"\x89PNG")
assert isinstance(result.metadata["seed"], int)
async def test_generation_is_strictly_serialized(app_config) -> None:
backend = FakeBackend(delay=0.03)
service = ImageGenerationService(app_config, backend)
await asyncio.gather(
*(service.generate_image(TextToImageRequest.create(f"image {i}")) for i in range(4))
)
assert backend.max_active == 1
async def test_disabled_controlnet_fails_before_backend(app_config, image_base64: str) -> None:
disabled = replace(
app_config,
controlnet=ControlNetConfig(
depth=ControlModelConfig(False, None, "midas"),
openpose=app_config.controlnet.openpose,
inpaint=app_config.controlnet.inpaint,
),
)
service = ImageGenerationService(disabled, FakeBackend())
with pytest.raises(InvalidRequestError, match="disabled"):
await service.generate_guided_image(GuidedImageRequest.create("cat", "depth", image_base64))
async def test_backend_error_is_sanitized(app_config) -> None:
class BrokenBackend(FakeBackend):
def text_to_image(self, request: object, seed: int) -> Image.Image:
raise RuntimeError("boom")
service = ImageGenerationService(app_config, BrokenBackend())
with pytest.raises(GenerationError, match="boom"):
await service.generate_image(TextToImageRequest.create("cat"))
async def test_cuda_oom_is_classified(app_config) -> None:
class OomBackend(FakeBackend):
def text_to_image(self, request: object, seed: int) -> Image.Image:
raise RuntimeError("CUDA out of memory while allocating")
service = ImageGenerationService(app_config, OomBackend())
with pytest.raises(GenerationError, match="CUDA out of memory"):
await service.generate_image(TextToImageRequest.create("cat"))
async def test_optional_output_is_saved(app_config) -> None:
enabled = replace(app_config, output=replace(app_config.output, enabled=True))
service = ImageGenerationService(enabled, FakeBackend())
await service.generate_image(TextToImageRequest.create("cat"))
files = list(enabled.output.directory.glob("*.png"))
assert len(files) == 1
with Image.open(io.BytesIO(files[0].read_bytes())) as image:
assert image.format == "PNG"
|
