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
|
"""Serialized generation service and persistence policy."""
from __future__ import annotations
import asyncio
import logging
import secrets
import time
import uuid
from datetime import UTC, datetime
from typing import Protocol
from PIL import Image
from .config import AppConfig
from .errors import GenerationError, InvalidRequestError, SD15Error
from .images import decode_image, encode_png
from .schema import (
GenerationResult,
GuidedImageRequest,
InpaintImageRequest,
TextToImageRequest,
)
logger = logging.getLogger(__name__)
class GenerationBackend(Protocol):
"""Synchronous GPU boundary; the service runs it in a worker thread."""
def text_to_image(self, request: TextToImageRequest, seed: int) -> Image.Image: ...
def guided_image(
self, request: GuidedImageRequest, image: Image.Image, seed: int
) -> Image.Image: ...
def inpaint_image(
self,
request: InpaintImageRequest,
image: Image.Image,
mask: Image.Image,
seed: int,
) -> Image.Image: ...
class ImageGenerationService:
"""Validate/decode requests and serialize access to the single GPU."""
def __init__(self, config: AppConfig, backend: GenerationBackend) -> None:
self.config = config
self.backend = backend
self._generation_lock = asyncio.Lock()
async def generate_image(self, request: TextToImageRequest) -> GenerationResult:
return await self._run("text_to_image", request)
async def generate_guided_image(self, request: GuidedImageRequest) -> GenerationResult:
control = self.config.controlnet.for_mode(request.control_mode)
if not control.enabled:
raise InvalidRequestError(f"ControlNet mode is disabled: {request.control_mode}")
image = decode_image(request.input_image)
return await self._run("guided_image", request, image)
async def inpaint_image(self, request: InpaintImageRequest) -> GenerationResult:
if not self.config.controlnet.inpaint.enabled:
raise InvalidRequestError("ControlNet mode is disabled: inpaint")
image = decode_image(request.input_image)
mask = decode_image(request.mask_image, mask=True)
return await self._run("inpaint_image", request, image, mask)
async def _run(
self,
method: str,
request: TextToImageRequest | GuidedImageRequest | InpaintImageRequest,
*images: Image.Image,
) -> GenerationResult:
seed = secrets.randbelow(2**63)
started = time.monotonic()
async with self._generation_lock:
logger.info("generation started mode=%s size=%s seed=%d", method, request.size, seed)
try:
function = getattr(self.backend, method)
output = await asyncio.to_thread(function, request, *images, seed)
if not isinstance(output, Image.Image):
raise GenerationError("generation backend returned a non-image result")
except SD15Error:
raise
except Exception as exc:
logger.exception("generation failed mode=%s", method)
if "out of memory" in str(exc).lower():
self._recover_cuda()
raise GenerationError("CUDA out of memory") from exc
raise GenerationError(f"generation failed: {exc}") from exc
metadata: dict[str, str | int | float] = {
"seed": seed,
"size": str(request.size),
"mode": method,
"scheduler": self.config.generation.scheduler,
"timestamp": datetime.now(UTC).isoformat(),
"elapsed_seconds": round(time.monotonic() - started, 3),
}
png_data = encode_png(output, metadata)
if self.config.output.enabled:
self._save(png_data)
logger.info("generation completed mode=%s seed=%d", method, seed)
return GenerationResult(png_data, metadata)
def _save(self, png_data: bytes) -> None:
directory = self.config.output.directory
try:
directory.mkdir(parents=True, exist_ok=True)
destination = directory / f"{uuid.uuid4()}.png"
destination.write_bytes(png_data)
except OSError as exc:
raise GenerationError(f"could not save generated image: {exc}") from exc
@staticmethod
def _recover_cuda() -> None:
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except ImportError:
pass
|
