import asyncio import base64 import io import threading import time import pytest from PIL import Image from modelship.infer.infer_config import RawRequestProxy, StableDiffusionCppConfig from modelship.infer.stable_diffusion_cpp.openai.serving_image import ( _DEFAULT_EDIT_STRENGTH, _DEFAULT_VARIATION_STRENGTH, OpenAIServingImage, ) from modelship.openai.protocol import ( ErrorResponse, ImageEditRequest, ImageGenerationRequest, ImageGenerationResponse, ImageVariationRequest, ) class _StubSD: """Stand-in for a stable-diffusion.cpp handle. Records the kwargs each `generate_image` call received and returns `batch_count` solid PIL images.""" def __init__(self): self.calls: list[dict] = [] def generate_image(self, **kwargs): self.calls.append(kwargs) n = kwargs.get("batch_count", 1) return [Image.new("RGB", (8, 8), (10, 20, 30)) for _ in range(n)] class _ConcurrencyProbeSD: """Records the peak number of overlapping generate_image calls.""" def __init__(self): self.active = 0 self.max_active = 0 self._lock = threading.Lock() def generate_image(self, **kwargs): with self._lock: self.active += 1 self.max_active = max(self.max_active, self.active) time.sleep(0.05) # widen the window so unserialized calls would overlap with self._lock: self.active -= 1 n = kwargs.get("batch_count", 1) return [Image.new("RGB", (8, 8)) for _ in range(n)] def _png_bytes(w: int = 16, h: int = 16) -> bytes: buf = io.BytesIO() Image.new("RGB", (w, h), (200, 100, 50)).save(buf, format="PNG") return buf.getvalue() def _rgba_png(w: int = 16, h: int = 16, *, transparent: bool) -> bytes: img = Image.new("RGBA", (w, h), (200, 100, 50, 255)) if transparent: for x in range(w // 4, w * 3 // 4): for y in range(h // 4, h * 3 // 4): img.putpixel((x, y), (0, 0, 0, 0)) buf = io.BytesIO() img.save(buf, format="PNG") return buf.getvalue() def _serving(sd=None) -> OpenAIServingImage: return OpenAIServingImage( sd=sd if sd is not None else _StubSD(), # type: ignore[arg-type] config=StableDiffusionCppConfig( sample_steps=5, cfg_scale=3.0, sample_method="euler", scheduler="karras", vae_tiling=True ), ) def _proxy() -> RawRequestProxy: return RawRequestProxy(None, {}) def _assert_valid_b64_image(obj) -> None: raw = base64.b64decode(obj.b64_json) Image.open(io.BytesIO(raw)).verify() @pytest.mark.asyncio class TestImageGeneration: async def test_generation_returns_n_images_and_forwards_config(self): sd = _StubSD() serving = _serving(sd) request = ImageGenerationRequest(model="m", prompt="a red bicycle", n=2, size="64x32") result = await serving.create_image_generation(request, _proxy()) assert isinstance(result, ImageGenerationResponse) assert len(result.data) == 2 for obj in result.data: assert obj.revised_prompt == "a red bicycle" _assert_valid_b64_image(obj) call = sd.calls[0] assert call["prompt"] == "a red bicycle" assert call["width"] == 64 and call["height"] == 32 assert call["batch_count"] == 2 # Per-model config defaults are applied on every call. assert call["cfg_scale"] == 3.0 assert call["sample_steps"] == 5 assert call["sample_method"] == "euler" assert call["scheduler"] == "karras" assert call["vae_tiling"] is True assert "init_image" not in call # txt2img async def test_generation_invalid_size_errors(self): serving = _serving() request = ImageGenerationRequest(model="m", prompt="x", n=1, size="not-a-size") result = await serving.create_image_generation(request, _proxy()) assert isinstance(result, ErrorResponse) async def test_calls_are_serialized(self): probe = _ConcurrencyProbeSD() serving = _serving(probe) request = ImageGenerationRequest(model="m", prompt="x", n=1, size="16x16") await asyncio.gather(*(serving.create_image_generation(request, _proxy()) for _ in range(4))) assert probe.max_active == 1 @pytest.mark.asyncio class TestImageEdit: async def test_edit_without_mask_passes_init_image_no_mask(self): sd = _StubSD() serving = _serving(sd) request = ImageEditRequest.model_construct(model="m", prompt="add a hat", n=2, size="32x32", strength=None) result = await serving.create_image_edit(_png_bytes(), None, request, _proxy()) assert isinstance(result, ImageGenerationResponse) assert len(result.data) == 2 call = sd.calls[0] assert call["prompt"] == "add a hat" assert call["init_image"] is not None assert call["strength"] == _DEFAULT_EDIT_STRENGTH assert "mask_image" not in call async def test_edit_with_mask_passes_mask_image(self): sd = _StubSD() serving = _serving(sd) request = ImageEditRequest.model_construct(model="m", prompt="erase", n=1, size="16x16", strength=0.5) result = await serving.create_image_edit(_png_bytes(), _png_bytes(), request, _proxy()) assert isinstance(result, ImageGenerationResponse) call = sd.calls[0] assert call["strength"] == 0.5 assert "mask_image" in call async def test_edit_transparent_image_derives_mask_from_alpha(self): sd = _StubSD() serving = _serving(sd) request = ImageEditRequest.model_construct(model="m", prompt="fill", n=1, size="16x16", strength=None) result = await serving.create_image_edit(_rgba_png(transparent=True), None, request, _proxy()) assert isinstance(result, ImageGenerationResponse) call = sd.calls[0] assert "mask_image" in call # Transparent center -> repaint (white); opaque corner -> keep (black). mask = call["mask_image"] assert mask.getpixel((8, 8)) == 255 assert mask.getpixel((0, 0)) == 0 async def test_edit_opaque_rgba_no_mask_is_plain_img2img(self): sd = _StubSD() serving = _serving(sd) request = ImageEditRequest.model_construct(model="m", prompt="x", n=1, size="16x16", strength=None) result = await serving.create_image_edit(_rgba_png(transparent=False), None, request, _proxy()) assert isinstance(result, ImageGenerationResponse) assert "mask_image" not in sd.calls[0] async def test_edit_invalid_size_errors(self): serving = _serving() request = ImageEditRequest.model_construct(model="m", prompt="x", n=1, size="bad", strength=None) result = await serving.create_image_edit(_png_bytes(), None, request, _proxy()) assert isinstance(result, ErrorResponse) @pytest.mark.asyncio class TestImageVariation: async def test_variation_uses_empty_prompt_and_default_strength(self): sd = _StubSD() serving = _serving(sd) request = ImageVariationRequest.model_construct(model="m", n=3, size="16x16", strength=None) result = await serving.create_image_variation(_png_bytes(), request, _proxy()) assert isinstance(result, ImageGenerationResponse) assert len(result.data) == 3 for obj in result.data: assert obj.revised_prompt is None _assert_valid_b64_image(obj) call = sd.calls[0] assert call["prompt"] == "" assert call["init_image"] is not None assert call["strength"] == _DEFAULT_VARIATION_STRENGTH async def test_variation_invalid_size_errors(self): serving = _serving() request = ImageVariationRequest.model_construct(model="m", n=1, size="bad", strength=None) result = await serving.create_image_variation(_png_bytes(), request, _proxy()) assert isinstance(result, ErrorResponse)