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
|
"""Validated request and response types independent of MCP and Diffusers."""
from __future__ import annotations
from dataclasses import dataclass
from enum import StrEnum
from typing import Any
from .errors import InvalidRequestError
class ImageSize(StrEnum):
SQUARE_512 = "square_512"
LANDSCAPE_512 = "landscape_512"
PORTRAIT_512 = "portrait_512"
@property
def dimensions(self) -> tuple[int, int]:
return {
ImageSize.SQUARE_512: (512, 512),
ImageSize.LANDSCAPE_512: (512, 384),
ImageSize.PORTRAIT_512: (384, 512),
}[self]
class ControlMode(StrEnum):
DEPTH = "depth"
OPENPOSE = "openpose"
OPENPOSE_FULL = "openpose_full"
def _prompt(value: Any, label: str, *, required: bool) -> str:
if value is None and not required:
return ""
if not isinstance(value, str):
raise InvalidRequestError(f"{label} must be a string")
value = value.strip()
if required and not value:
raise InvalidRequestError(f"{label} must not be empty")
if len(value) > 10_000:
raise InvalidRequestError(f"{label} is too long")
return value
def _size(value: str | ImageSize | None, default: str) -> ImageSize:
try:
return ImageSize(value or default)
except ValueError as exc:
raise InvalidRequestError(f"unsupported size: {value}") from exc
@dataclass(frozen=True)
class TextToImageRequest:
positive_prompt: str
negative_prompt: str
size: ImageSize
@classmethod
def create(
cls,
positive_prompt: str,
negative_prompt: str | None = None,
size: str | ImageSize | None = None,
*,
default_size: str = ImageSize.SQUARE_512,
) -> TextToImageRequest:
return cls(
_prompt(positive_prompt, "positive_prompt", required=True),
_prompt(negative_prompt, "negative_prompt", required=False),
_size(size, default_size),
)
@dataclass(frozen=True)
class GuidedImageRequest:
positive_prompt: str
negative_prompt: str
size: ImageSize
control_mode: ControlMode
input_image: str
@classmethod
def create(
cls,
positive_prompt: str,
control_mode: str | ControlMode,
input_image: str,
negative_prompt: str | None = None,
size: str | ImageSize | None = None,
*,
default_size: str = ImageSize.SQUARE_512,
) -> GuidedImageRequest:
base = TextToImageRequest.create(
positive_prompt, negative_prompt, size, default_size=default_size
)
try:
mode = ControlMode(control_mode)
except ValueError as exc:
raise InvalidRequestError(f"unsupported control_mode: {control_mode}") from exc
if not isinstance(input_image, str) or not input_image.strip():
raise InvalidRequestError("input_image must be a non-empty Base64 string")
return cls(base.positive_prompt, base.negative_prompt, base.size, mode, input_image)
@dataclass(frozen=True)
class InpaintImageRequest:
positive_prompt: str
negative_prompt: str
size: ImageSize
input_image: str
mask_image: str
@classmethod
def create(
cls,
positive_prompt: str,
input_image: str,
mask_image: str,
negative_prompt: str | None = None,
size: str | ImageSize | None = None,
*,
default_size: str = ImageSize.SQUARE_512,
) -> InpaintImageRequest:
base = TextToImageRequest.create(
positive_prompt, negative_prompt, size, default_size=default_size
)
for label, value in (("input_image", input_image), ("mask_image", mask_image)):
if not isinstance(value, str) or not value.strip():
raise InvalidRequestError(f"{label} must be a non-empty Base64 string")
return cls(
base.positive_prompt,
base.negative_prompt,
base.size,
input_image,
mask_image,
)
@dataclass(frozen=True)
class GenerationResult:
png_data: bytes
metadata: dict[str, str | int | float]
|
