@mcoda/mswarm 0.1.104 → 0.1.106

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/README.md CHANGED
@@ -8,6 +8,22 @@ Owner-run mswarm node for exposing local `mcoda` agents through the mswarm self-
8
8
  npm install -g @mcoda/mswarm
9
9
  ```
10
10
 
11
+ The package also ships hardened, standard-library Python bridges for
12
+ owner-local generative runtimes under its `scripts/` directory:
13
+
14
+ - `stable-diffusion-cpp-openai-server.py` adds effective, validated
15
+ `negative_prompt`, `seed`, and `steps` controls to the OpenAI image surface,
16
+ refuses redirects and proxy routing, and bounds accepted-client threads and
17
+ read time;
18
+ - `stable-audio-3-openai-server.py` wraps the approved Stable Audio 3 TensorRT
19
+ CLI profile;
20
+ - `wan22-openai-server.py` maps OpenAI-shaped video requests to native
21
+ stable-diffusion.cpp async jobs.
22
+
23
+ All three are loopback-only deployment components, not public listeners. See
24
+ `docs/ops/mswarm-self-hosted-generative-agents-runbook.md` in the mcoda
25
+ repository for pinned service, registration, readiness, and rollback guidance.
26
+
11
27
  ## Quick Setup
12
28
 
13
29
  ```sh
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@mcoda/mswarm",
3
- "version": "0.1.104",
3
+ "version": "0.1.106",
4
4
  "description": "Owner-run mswarm node for exposing local mcoda agents through the mswarm self-hosted gateway.",
5
5
  "type": "module",
6
6
  "main": "dist/server.js",
@@ -47,8 +47,8 @@
47
47
  },
48
48
  "dependencies": {
49
49
  "fastify": "^5.7.4",
50
- "@mcoda/db": "0.1.104",
51
- "@mcoda/shared": "0.1.104"
50
+ "@mcoda/shared": "0.1.106",
51
+ "@mcoda/db": "0.1.106"
52
52
  },
53
53
  "scripts": {
54
54
  "build": "pnpm --filter @mcoda/shared run build && pnpm --filter @mcoda/codali run build && tsc -p tsconfig.json && node ../../scripts/copy-mswarm-codali-vendor.js",
@@ -0,0 +1,1157 @@
1
+ #!/usr/bin/env python3
2
+ """Loopback-only OpenAI image bridge for stable-diffusion.cpp."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import argparse
7
+ import base64
8
+ import binascii
9
+ import ipaddress
10
+ import json
11
+ import logging
12
+ import math
13
+ import re
14
+ import secrets
15
+ import signal
16
+ import socket
17
+ import struct
18
+ import threading
19
+ import time
20
+ import urllib.error
21
+ import urllib.parse
22
+ import urllib.request
23
+ import zlib
24
+ from dataclasses import dataclass, field
25
+ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
26
+ from socketserver import TCPServer
27
+ from typing import Any
28
+
29
+
30
+ DEFAULT_HOST = "127.0.0.1"
31
+ DEFAULT_PORT = 11449
32
+ DEFAULT_UPSTREAM_BASE_URL = "http://127.0.0.1:11445"
33
+ DEFAULT_MODEL_ID = "sd-cpp-local"
34
+ DEFAULT_CLIENT_READ_TIMEOUT_SECONDS = 10.0
35
+ DEFAULT_MAX_HANDLER_THREADS = 8
36
+ DEFAULT_MAX_REQUEST_BYTES = 131_072
37
+ OVERLOAD_HEADER_ALLOWANCE_BYTES = 65_536
38
+ OVERLOAD_DRAIN_TIMEOUT_SECONDS = 0.1
39
+ PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
40
+ ALLOWED_FIELDS = {
41
+ "model",
42
+ "prompt",
43
+ "negative_prompt",
44
+ "n",
45
+ "size",
46
+ "response_format",
47
+ "seed",
48
+ "steps",
49
+ }
50
+ RESERVED_CONTROL_TAG = re.compile(r"<\s*/?\s*sd_cpp_extra_args\b", re.IGNORECASE)
51
+ MODEL_ID_PATTERN = re.compile(r"[A-Za-z0-9][A-Za-z0-9._:/-]{0,127}")
52
+
53
+
54
+ class NoRedirectHandler(urllib.request.HTTPRedirectHandler):
55
+ def redirect_request(
56
+ self,
57
+ request: urllib.request.Request,
58
+ file_pointer: Any,
59
+ status: int,
60
+ message: str,
61
+ headers: Any,
62
+ new_url: str,
63
+ ) -> None:
64
+ del request, file_pointer, status, message, headers, new_url
65
+ return None
66
+
67
+
68
+ DIRECT_HTTP_OPENER = urllib.request.build_opener(
69
+ urllib.request.ProxyHandler({}),
70
+ NoRedirectHandler(),
71
+ )
72
+ OVERLOAD_BODY = json.dumps(
73
+ {
74
+ "error": {
75
+ "message": "the image bridge has too many active clients",
76
+ "type": "server_error",
77
+ "param": None,
78
+ "code": "server_busy",
79
+ }
80
+ },
81
+ separators=(",", ":"),
82
+ ).encode("utf-8")
83
+ OVERLOAD_RESPONSE = (
84
+ b"HTTP/1.1 503 Service Unavailable\r\n"
85
+ b"content-type: application/json\r\n"
86
+ + f"content-length: {len(OVERLOAD_BODY)}\r\n".encode("ascii")
87
+ + b"cache-control: no-store\r\n"
88
+ b"x-content-type-options: nosniff\r\n"
89
+ b"connection: close\r\n"
90
+ b"\r\n"
91
+ + OVERLOAD_BODY
92
+ )
93
+
94
+
95
+ class ApiError(Exception):
96
+ def __init__(
97
+ self,
98
+ status: int,
99
+ message: str,
100
+ *,
101
+ error_type: str = "invalid_request_error",
102
+ param: str | None = None,
103
+ code: str | None = None,
104
+ ) -> None:
105
+ super().__init__(message)
106
+ self.status = status
107
+ self.message = message
108
+ self.error_type = error_type
109
+ self.param = param
110
+ self.code = code
111
+
112
+
113
+ class UpstreamError(Exception):
114
+ pass
115
+
116
+
117
+ class UpstreamTimeout(UpstreamError):
118
+ pass
119
+
120
+
121
+ @dataclass
122
+ class ServerState:
123
+ upstream_base_url: str
124
+ model_id: str
125
+ upstream_model_id: str
126
+ max_request_bytes: int
127
+ max_output_bytes: int
128
+ max_prompt_chars: int
129
+ max_negative_prompt_chars: int
130
+ min_width: int
131
+ max_width: int
132
+ min_height: int
133
+ max_height: int
134
+ max_pixels: int
135
+ default_size: tuple[int, int]
136
+ default_steps: int
137
+ max_steps: int
138
+ request_timeout_seconds: float
139
+ inference_timeout_seconds: float
140
+ shutdown_requested: threading.Event = field(default_factory=threading.Event)
141
+ inference_slots: threading.BoundedSemaphore = field(
142
+ default_factory=lambda: threading.BoundedSemaphore(1)
143
+ )
144
+
145
+
146
+ class ImageHttpServer(ThreadingHTTPServer):
147
+ daemon_threads = True
148
+ request_queue_size = 8
149
+
150
+ def __init__(
151
+ self,
152
+ server_address: tuple[str, int],
153
+ request_handler_class: type[BaseHTTPRequestHandler],
154
+ *,
155
+ client_read_timeout_seconds: float = DEFAULT_CLIENT_READ_TIMEOUT_SECONDS,
156
+ max_handler_threads: int = DEFAULT_MAX_HANDLER_THREADS,
157
+ max_request_bytes: int = DEFAULT_MAX_REQUEST_BYTES,
158
+ bind_and_activate: bool = True,
159
+ ) -> None:
160
+ self.client_read_timeout_seconds = client_read_timeout_seconds
161
+ self.max_handler_threads = max_handler_threads
162
+ self.max_request_bytes = max_request_bytes
163
+ self._handler_slots = threading.BoundedSemaphore(max_handler_threads)
164
+ super().__init__(server_address, request_handler_class, bind_and_activate)
165
+
166
+ def server_bind(self) -> None:
167
+ # HTTPServer otherwise performs an unnecessary reverse-DNS lookup.
168
+ TCPServer.server_bind(self)
169
+ host, port = self.server_address[:2]
170
+ self.server_name = str(host)
171
+ self.server_port = int(port)
172
+
173
+ def get_request(self) -> tuple[socket.socket, Any]:
174
+ request, client_address = super().get_request()
175
+ request.settimeout(self.client_read_timeout_seconds)
176
+ return request, client_address
177
+
178
+ def process_request(self, request: socket.socket, client_address: Any) -> None:
179
+ if not self._handler_slots.acquire(blocking=False):
180
+ self._reject_overloaded_request(request)
181
+ return
182
+ try:
183
+ super().process_request(request, client_address)
184
+ except BaseException:
185
+ self._handler_slots.release()
186
+ raise
187
+
188
+ def _reject_overloaded_request(self, request: socket.socket) -> None:
189
+ drain_deadline = time.monotonic() + min(
190
+ self.client_read_timeout_seconds,
191
+ OVERLOAD_DRAIN_TIMEOUT_SECONDS,
192
+ )
193
+ drain_remaining = self.max_request_bytes + OVERLOAD_HEADER_ALLOWANCE_BYTES
194
+ try:
195
+ request.settimeout(max(0.001, drain_deadline - time.monotonic()))
196
+ request.sendall(OVERLOAD_RESPONSE)
197
+ try:
198
+ request.shutdown(socket.SHUT_WR)
199
+ except OSError:
200
+ return
201
+ while drain_remaining > 0:
202
+ timeout_remaining = drain_deadline - time.monotonic()
203
+ if timeout_remaining <= 0:
204
+ break
205
+ request.settimeout(timeout_remaining)
206
+ try:
207
+ chunk = request.recv(min(65_536, drain_remaining))
208
+ except OSError:
209
+ break
210
+ if not chunk:
211
+ break
212
+ drain_remaining -= len(chunk)
213
+ except OSError:
214
+ pass
215
+ finally:
216
+ self.close_request(request)
217
+
218
+ def process_request_thread(
219
+ self, request: socket.socket, client_address: Any
220
+ ) -> None:
221
+ try:
222
+ super().process_request_thread(request, client_address)
223
+ finally:
224
+ self._handler_slots.release()
225
+
226
+
227
+ def _is_int(value: Any) -> bool:
228
+ return isinstance(value, int) and not isinstance(value, bool)
229
+
230
+
231
+ def _contains_reserved_control_tag(value: str) -> bool:
232
+ return RESERVED_CONTROL_TAG.search(value) is not None
233
+
234
+
235
+ def _validate_text(
236
+ value: Any,
237
+ *,
238
+ field_name: str,
239
+ max_chars: int,
240
+ allow_empty: bool,
241
+ ) -> str:
242
+ if not isinstance(value, str) or (not allow_empty and not value.strip()):
243
+ qualifier = "a string" if allow_empty else "a non-empty string"
244
+ raise ApiError(
245
+ 400,
246
+ f"{field_name} must be {qualifier}",
247
+ param=field_name,
248
+ code=f"invalid_{field_name}",
249
+ )
250
+ if len(value) > max_chars:
251
+ raise ApiError(
252
+ 400,
253
+ f"{field_name} exceeds the configured limit",
254
+ param=field_name,
255
+ code=f"{field_name}_limit_exceeded",
256
+ )
257
+ if "\x00" in value:
258
+ raise ApiError(
259
+ 400,
260
+ f"{field_name} must not contain a null character",
261
+ param=field_name,
262
+ code=f"invalid_{field_name}",
263
+ )
264
+ try:
265
+ value.encode("utf-8", errors="strict")
266
+ except UnicodeEncodeError as error:
267
+ raise ApiError(
268
+ 400,
269
+ f"{field_name} must be valid UTF-8 text",
270
+ param=field_name,
271
+ code=f"invalid_{field_name}",
272
+ ) from error
273
+ if _contains_reserved_control_tag(value):
274
+ raise ApiError(
275
+ 400,
276
+ f"{field_name} contains the reserved stable-diffusion.cpp control tag",
277
+ param=field_name,
278
+ code="reserved_control_tag",
279
+ )
280
+ return value
281
+
282
+
283
+ def _parse_size_text(value: Any, *, field_name: str = "size") -> tuple[int, int]:
284
+ if not isinstance(value, str):
285
+ raise ApiError(
286
+ 400,
287
+ "size must use WIDTHxHEIGHT",
288
+ param=field_name,
289
+ code="invalid_size",
290
+ )
291
+ match = re.fullmatch(r"([1-9][0-9]*)x([1-9][0-9]*)", value)
292
+ if match is None:
293
+ raise ApiError(
294
+ 400,
295
+ "size must use WIDTHxHEIGHT",
296
+ param=field_name,
297
+ code="invalid_size",
298
+ )
299
+ return int(match.group(1)), int(match.group(2))
300
+
301
+
302
+ def _parse_generation_request(payload: Any, state: ServerState) -> dict[str, Any]:
303
+ if not isinstance(payload, dict):
304
+ raise ApiError(400, "request body must be a JSON object", code="invalid_json")
305
+ unknown = sorted(set(payload) - ALLOWED_FIELDS)
306
+ if unknown:
307
+ raise ApiError(
308
+ 400,
309
+ f"unsupported request parameter: {unknown[0]}",
310
+ param=unknown[0],
311
+ code="unsupported_parameter",
312
+ )
313
+ if payload.get("model") != state.model_id:
314
+ raise ApiError(
315
+ 400,
316
+ f"model must be {state.model_id}",
317
+ param="model",
318
+ code="model_not_found",
319
+ )
320
+
321
+ prompt = _validate_text(
322
+ payload.get("prompt"),
323
+ field_name="prompt",
324
+ max_chars=state.max_prompt_chars,
325
+ allow_empty=False,
326
+ )
327
+ negative_prompt = _validate_text(
328
+ payload.get("negative_prompt", ""),
329
+ field_name="negative_prompt",
330
+ max_chars=state.max_negative_prompt_chars,
331
+ allow_empty=True,
332
+ )
333
+
334
+ n = payload.get("n", 1)
335
+ if not _is_int(n) or n != 1:
336
+ raise ApiError(400, "n must be 1", param="n", code="invalid_n")
337
+
338
+ default_width, default_height = state.default_size
339
+ width, height = _parse_size_text(
340
+ payload.get("size", f"{default_width}x{default_height}")
341
+ )
342
+ if width < state.min_width or height < state.min_height:
343
+ raise ApiError(
344
+ 400,
345
+ "size is below the configured minimum",
346
+ param="size",
347
+ code="unsupported_size",
348
+ )
349
+ if (
350
+ width > state.max_width
351
+ or height > state.max_height
352
+ or width * height > state.max_pixels
353
+ ):
354
+ raise ApiError(
355
+ 400,
356
+ "size exceeds the configured limit",
357
+ param="size",
358
+ code="unsupported_size",
359
+ )
360
+
361
+ response_format = payload.get("response_format", "b64_json")
362
+ if response_format != "b64_json":
363
+ raise ApiError(
364
+ 400,
365
+ "response_format must be b64_json",
366
+ param="response_format",
367
+ code="unsupported_response_format",
368
+ )
369
+
370
+ seed = payload.get("seed")
371
+ if seed is None or seed == -1:
372
+ seed = secrets.randbelow(2**31)
373
+ if not _is_int(seed) or not 0 <= seed <= 2**31 - 1:
374
+ raise ApiError(
375
+ 400,
376
+ "seed must be -1 or an integer between 0 and 2147483647",
377
+ param="seed",
378
+ code="invalid_seed",
379
+ )
380
+
381
+ steps = payload.get("steps", state.default_steps)
382
+ if not _is_int(steps) or not 1 <= steps <= state.max_steps:
383
+ raise ApiError(
384
+ 400,
385
+ f"steps must be an integer between 1 and {state.max_steps}",
386
+ param="steps",
387
+ code="invalid_steps",
388
+ )
389
+
390
+ return {
391
+ "prompt": prompt,
392
+ "negative_prompt": negative_prompt,
393
+ "n": 1,
394
+ "width": width,
395
+ "height": height,
396
+ "response_format": "b64_json",
397
+ "seed": seed,
398
+ "steps": steps,
399
+ }
400
+
401
+
402
+ def _controlled_prompt(request: dict[str, Any]) -> str:
403
+ controlled_args = {
404
+ "negative_prompt": request["negative_prompt"],
405
+ "seed": request["seed"],
406
+ "sample_params": {"sample_steps": request["steps"]},
407
+ }
408
+ serialized = json.dumps(
409
+ controlled_args,
410
+ ensure_ascii=True,
411
+ separators=(",", ":"),
412
+ )
413
+ return (
414
+ f"{request['prompt']}\n"
415
+ f"<sd_cpp_extra_args>{serialized}</sd_cpp_extra_args>"
416
+ )
417
+
418
+
419
+ def _is_timeout_error(error: BaseException) -> bool:
420
+ if isinstance(error, (TimeoutError, socket.timeout)):
421
+ return True
422
+ if isinstance(error, urllib.error.URLError):
423
+ return isinstance(error.reason, (TimeoutError, socket.timeout))
424
+ return False
425
+
426
+
427
+ def _bounded_json_request(
428
+ url: str,
429
+ *,
430
+ method: str,
431
+ timeout: float | None,
432
+ max_bytes: int,
433
+ payload: dict[str, Any] | None = None,
434
+ ) -> tuple[int, dict[str, Any]]:
435
+ body = None
436
+ headers = {"accept": "application/json"}
437
+ if payload is not None:
438
+ body = json.dumps(payload, separators=(",", ":")).encode("utf-8")
439
+ headers["content-type"] = "application/json"
440
+ request = urllib.request.Request(url, data=body, headers=headers, method=method)
441
+ content_type = ""
442
+ try:
443
+ with DIRECT_HTTP_OPENER.open(request, timeout=timeout) as response:
444
+ raw = response.read(max_bytes + 1)
445
+ status = response.status
446
+ content_type = response.headers.get_content_type()
447
+ if response.geturl() != url:
448
+ raise UpstreamError("stable-diffusion.cpp changed the upstream URL")
449
+ except urllib.error.HTTPError as error:
450
+ if 300 <= error.code < 400:
451
+ error.close()
452
+ raise UpstreamError("stable-diffusion.cpp redirects are forbidden") from error
453
+ raw = error.read(max_bytes + 1)
454
+ status = error.code
455
+ content_type = error.headers.get_content_type()
456
+ except (OSError, urllib.error.URLError) as error:
457
+ if _is_timeout_error(error):
458
+ raise UpstreamTimeout("stable-diffusion.cpp request timed out") from error
459
+ raise UpstreamError("stable-diffusion.cpp is unreachable") from error
460
+ if len(raw) > max_bytes:
461
+ raise UpstreamError("stable-diffusion.cpp response exceeds the configured limit")
462
+ if content_type != "application/json":
463
+ raise UpstreamError("stable-diffusion.cpp returned the wrong content type")
464
+ try:
465
+ parsed = json.loads(raw)
466
+ except (UnicodeDecodeError, json.JSONDecodeError) as error:
467
+ raise UpstreamError("stable-diffusion.cpp returned invalid JSON") from error
468
+ if not isinstance(parsed, dict):
469
+ raise UpstreamError("stable-diffusion.cpp returned an invalid object")
470
+ return status, parsed
471
+
472
+
473
+ def _check_upstream_model(state: ServerState) -> None:
474
+ status, payload = _bounded_json_request(
475
+ f"{state.upstream_base_url}/v1/models",
476
+ method="GET",
477
+ timeout=state.request_timeout_seconds,
478
+ max_bytes=1_048_576,
479
+ )
480
+ data = payload.get("data")
481
+ if (
482
+ status != 200
483
+ or ("object" in payload and payload.get("object") != "list")
484
+ or not isinstance(data, list)
485
+ or len(data) != 1
486
+ or not isinstance(data[0], dict)
487
+ or data[0].get("id") != state.upstream_model_id
488
+ or data[0].get("object") != "model"
489
+ ):
490
+ raise UpstreamError("stable-diffusion.cpp model catalog does not match")
491
+
492
+
493
+ def _native_request(request: dict[str, Any]) -> dict[str, Any]:
494
+ return {
495
+ "prompt": _controlled_prompt(request),
496
+ "n": 1,
497
+ "size": f"{request['width']}x{request['height']}",
498
+ "output_format": "png",
499
+ }
500
+
501
+
502
+ def _generate_image(state: ServerState, request: dict[str, Any]) -> dict[str, Any]:
503
+ max_json_bytes = math.ceil(state.max_output_bytes / 3) * 4 + 131_072
504
+ status, payload = _bounded_json_request(
505
+ f"{state.upstream_base_url}/v1/images/generations",
506
+ method="POST",
507
+ # The handler enforces the client-facing generation deadline. This
508
+ # transport remains open so the sole inference slot is held until the
509
+ # native server has actually completed and replied.
510
+ timeout=None,
511
+ max_bytes=max_json_bytes,
512
+ payload=_native_request(request),
513
+ )
514
+ if status == 429:
515
+ raise ApiError(
516
+ 429,
517
+ "the local image generator queue is full",
518
+ error_type="rate_limit_error",
519
+ code="upstream_queue_saturated",
520
+ )
521
+ if status != 200:
522
+ raise UpstreamError(f"stable-diffusion.cpp rejected image generation ({status})")
523
+ return payload
524
+
525
+
526
+ def _validate_png(
527
+ decoded: bytes,
528
+ *,
529
+ expected_width: int,
530
+ expected_height: int,
531
+ ) -> None:
532
+ if not decoded.startswith(PNG_SIGNATURE):
533
+ raise UpstreamError("stable-diffusion.cpp returned an invalid PNG signature")
534
+ offset = len(PNG_SIGNATURE)
535
+ chunk_count = 0
536
+ seen_ihdr = False
537
+ seen_idat = False
538
+ seen_iend = False
539
+ seen_plte = False
540
+ idat_ended = False
541
+ idat_parts: list[bytes] = []
542
+ bit_depth = 0
543
+ color_type = -1
544
+ interlace_method = -1
545
+ while offset < len(decoded):
546
+ if offset + 12 > len(decoded):
547
+ raise UpstreamError("stable-diffusion.cpp returned a truncated PNG")
548
+ chunk_count += 1
549
+ if chunk_count > 4_096:
550
+ raise UpstreamError("stable-diffusion.cpp returned too many PNG chunks")
551
+ chunk_length = struct.unpack(">I", decoded[offset : offset + 4])[0]
552
+ chunk_type = decoded[offset + 4 : offset + 8]
553
+ if chunk_length > 0x7FFFFFFF or any(
554
+ not (65 <= byte <= 90 or 97 <= byte <= 122) for byte in chunk_type
555
+ ):
556
+ raise UpstreamError("stable-diffusion.cpp returned an invalid PNG chunk")
557
+ chunk_end = offset + 12 + chunk_length
558
+ if chunk_end > len(decoded):
559
+ raise UpstreamError("stable-diffusion.cpp returned a truncated PNG")
560
+ chunk_data = decoded[offset + 8 : offset + 8 + chunk_length]
561
+ expected_crc = struct.unpack(">I", decoded[chunk_end - 4 : chunk_end])[0]
562
+ actual_crc = zlib.crc32(chunk_data, zlib.crc32(chunk_type)) & 0xFFFFFFFF
563
+ if actual_crc != expected_crc:
564
+ raise UpstreamError("stable-diffusion.cpp returned an invalid PNG checksum")
565
+ if not seen_ihdr:
566
+ if chunk_type != b"IHDR" or chunk_length != 13:
567
+ raise UpstreamError("stable-diffusion.cpp returned an invalid PNG header")
568
+ width, height = struct.unpack(">II", chunk_data[:8])
569
+ if width != expected_width or height != expected_height:
570
+ raise UpstreamError("stable-diffusion.cpp returned unexpected image dimensions")
571
+ if width <= 0 or height <= 0:
572
+ raise UpstreamError("stable-diffusion.cpp returned invalid image dimensions")
573
+ bit_depth = chunk_data[8]
574
+ color_type = chunk_data[9]
575
+ compression_method = chunk_data[10]
576
+ filter_method = chunk_data[11]
577
+ interlace_method = chunk_data[12]
578
+ valid_depths = {
579
+ 0: {1, 2, 4, 8, 16},
580
+ 2: {8, 16},
581
+ 3: {1, 2, 4, 8},
582
+ 4: {8, 16},
583
+ 6: {8, 16},
584
+ }
585
+ if color_type not in valid_depths or bit_depth not in valid_depths[color_type]:
586
+ raise UpstreamError("stable-diffusion.cpp returned an invalid PNG color mode")
587
+ if compression_method != 0 or filter_method != 0:
588
+ raise UpstreamError("stable-diffusion.cpp returned unsupported PNG methods")
589
+ if interlace_method not in (0, 1):
590
+ raise UpstreamError("stable-diffusion.cpp returned invalid PNG interlacing")
591
+ seen_ihdr = True
592
+ elif chunk_type == b"IHDR":
593
+ raise UpstreamError("stable-diffusion.cpp returned duplicate PNG headers")
594
+ elif not seen_ihdr:
595
+ raise UpstreamError("stable-diffusion.cpp returned an invalid PNG header")
596
+ if chunk_type == b"PLTE":
597
+ if seen_idat or seen_plte or chunk_length < 3 or chunk_length > 768 or chunk_length % 3:
598
+ raise UpstreamError("stable-diffusion.cpp returned an invalid PNG palette")
599
+ if color_type in (0, 4):
600
+ raise UpstreamError("stable-diffusion.cpp returned a forbidden PNG palette")
601
+ if color_type == 3 and chunk_length // 3 > 2**bit_depth:
602
+ raise UpstreamError("stable-diffusion.cpp returned an oversized PNG palette")
603
+ seen_plte = True
604
+ if chunk_type == b"IDAT":
605
+ if idat_ended or (color_type == 3 and not seen_plte):
606
+ raise UpstreamError("stable-diffusion.cpp returned invalid PNG image data")
607
+ seen_idat = True
608
+ idat_parts.append(chunk_data)
609
+ elif seen_idat:
610
+ idat_ended = True
611
+ if chunk_type == b"IEND":
612
+ if chunk_length != 0 or chunk_end != len(decoded):
613
+ raise UpstreamError("stable-diffusion.cpp returned an invalid PNG terminator")
614
+ seen_iend = True
615
+ elif chunk_type not in (b"IHDR", b"PLTE", b"IDAT") and chunk_type[0] & 0x20 == 0:
616
+ raise UpstreamError("stable-diffusion.cpp returned an unknown critical PNG chunk")
617
+ offset = chunk_end
618
+ if not seen_ihdr or not seen_idat or not seen_iend:
619
+ raise UpstreamError("stable-diffusion.cpp returned an incomplete PNG")
620
+
621
+ channels = {0: 1, 2: 3, 3: 1, 4: 2, 6: 4}[color_type]
622
+ bits_per_pixel = channels * bit_depth
623
+ passes = [(0, 0, 1, 1)]
624
+ if interlace_method == 1:
625
+ passes = [
626
+ (0, 0, 8, 8),
627
+ (4, 0, 8, 8),
628
+ (0, 4, 4, 8),
629
+ (2, 0, 4, 4),
630
+ (0, 2, 2, 4),
631
+ (1, 0, 2, 2),
632
+ (0, 1, 1, 2),
633
+ ]
634
+ scanline_layout: list[tuple[int, int]] = []
635
+ expected_scanline_bytes = 0
636
+ for start_x, start_y, step_x, step_y in passes:
637
+ pass_width = max(0, (expected_width - start_x + step_x - 1) // step_x)
638
+ pass_height = max(0, (expected_height - start_y + step_y - 1) // step_y)
639
+ if pass_width == 0 or pass_height == 0:
640
+ continue
641
+ row_bytes = (pass_width * bits_per_pixel + 7) // 8
642
+ scanline_layout.append((pass_height, row_bytes))
643
+ expected_scanline_bytes += pass_height * (row_bytes + 1)
644
+
645
+ decompressor = zlib.decompressobj()
646
+ try:
647
+ scanlines = decompressor.decompress(
648
+ b"".join(idat_parts), expected_scanline_bytes + 1
649
+ )
650
+ except zlib.error as error:
651
+ raise UpstreamError("stable-diffusion.cpp returned invalid PNG compression") from error
652
+ if (
653
+ len(scanlines) != expected_scanline_bytes
654
+ or not decompressor.eof
655
+ or decompressor.unconsumed_tail
656
+ or decompressor.unused_data
657
+ ):
658
+ raise UpstreamError("stable-diffusion.cpp returned invalid PNG scanlines")
659
+ scanline_offset = 0
660
+ for pass_height, row_bytes in scanline_layout:
661
+ for _ in range(pass_height):
662
+ if scanlines[scanline_offset] > 4:
663
+ raise UpstreamError("stable-diffusion.cpp returned an invalid PNG filter")
664
+ scanline_offset += row_bytes + 1
665
+
666
+
667
+ def _validate_generation_result(
668
+ payload: dict[str, Any],
669
+ request: dict[str, Any],
670
+ state: ServerState,
671
+ ) -> dict[str, Any]:
672
+ created = payload.get("created")
673
+ data = payload.get("data")
674
+ if not _is_int(created) or created < 0:
675
+ raise UpstreamError("stable-diffusion.cpp returned an invalid creation time")
676
+ if payload.get("output_format") != "png":
677
+ raise UpstreamError("stable-diffusion.cpp returned the wrong image format")
678
+ if (
679
+ not isinstance(data, list)
680
+ or len(data) != 1
681
+ or not isinstance(data[0], dict)
682
+ ):
683
+ raise UpstreamError("stable-diffusion.cpp returned an invalid image list")
684
+ encoded = data[0].get("b64_json")
685
+ if not isinstance(encoded, str) or not encoded:
686
+ raise UpstreamError("stable-diffusion.cpp returned no encoded image")
687
+ max_encoded_bytes = math.ceil(state.max_output_bytes / 3) * 4 + 4
688
+ if len(encoded) > max_encoded_bytes:
689
+ raise ApiError(
690
+ 413,
691
+ "generated image exceeds the configured output limit",
692
+ error_type="server_error",
693
+ code="output_too_large",
694
+ )
695
+ try:
696
+ decoded = base64.b64decode(encoded, validate=True)
697
+ except (ValueError, binascii.Error) as error:
698
+ raise UpstreamError("stable-diffusion.cpp returned invalid image base64") from error
699
+ if not decoded or len(decoded) > state.max_output_bytes:
700
+ raise ApiError(
701
+ 413,
702
+ "generated image exceeds the configured output limit",
703
+ error_type="server_error",
704
+ code="output_too_large",
705
+ )
706
+ _validate_png(
707
+ decoded,
708
+ expected_width=request["width"],
709
+ expected_height=request["height"],
710
+ )
711
+ return {
712
+ "created": created,
713
+ "data": [{"b64_json": base64.b64encode(decoded).decode("ascii")}],
714
+ }
715
+
716
+
717
+ def _handler_for(state: ServerState) -> type[BaseHTTPRequestHandler]:
718
+ class Handler(BaseHTTPRequestHandler):
719
+ server_version = "stable-diffusion-cpp-openai/1"
720
+
721
+ def log_message(self, format_string: str, *args: Any) -> None:
722
+ del format_string, args
723
+
724
+ def _send_json(self, status: int, payload: dict[str, Any]) -> None:
725
+ encoded = json.dumps(payload, separators=(",", ":")).encode("utf-8")
726
+ try:
727
+ self.send_response(status)
728
+ self.send_header("content-type", "application/json")
729
+ self.send_header("content-length", str(len(encoded)))
730
+ self.send_header("cache-control", "no-store")
731
+ self.send_header("x-content-type-options", "nosniff")
732
+ self.end_headers()
733
+ self.wfile.write(encoded)
734
+ except OSError:
735
+ self.close_connection = True
736
+
737
+ def _send_error(self, error: ApiError) -> None:
738
+ self._send_json(
739
+ error.status,
740
+ {
741
+ "error": {
742
+ "message": error.message,
743
+ "type": error.error_type,
744
+ "param": error.param,
745
+ "code": error.code,
746
+ }
747
+ },
748
+ )
749
+
750
+ def _method_not_allowed(self) -> None:
751
+ self._send_error(
752
+ ApiError(
753
+ 405,
754
+ "method not allowed",
755
+ error_type="invalid_request_error",
756
+ code="method_not_allowed",
757
+ )
758
+ )
759
+
760
+ do_DELETE = _method_not_allowed
761
+ do_OPTIONS = _method_not_allowed
762
+ do_PATCH = _method_not_allowed
763
+ do_PUT = _method_not_allowed
764
+
765
+ def do_GET(self) -> None:
766
+ if self.path not in ("/healthz", "/v1/models"):
767
+ self._send_error(
768
+ ApiError(404, "not found", error_type="not_found_error")
769
+ )
770
+ return
771
+ try:
772
+ _check_upstream_model(state)
773
+ except UpstreamError:
774
+ if self.path == "/healthz":
775
+ self._send_json(
776
+ 503,
777
+ {"status": "unhealthy", "model": state.model_id},
778
+ )
779
+ else:
780
+ self._send_error(
781
+ ApiError(
782
+ 503,
783
+ "image model is unavailable",
784
+ error_type="server_error",
785
+ code="upstream_unavailable",
786
+ )
787
+ )
788
+ return
789
+ if self.path == "/healthz":
790
+ self._send_json(
791
+ 200,
792
+ {
793
+ "status": "ok",
794
+ "model": state.model_id,
795
+ "upstream_model": state.upstream_model_id,
796
+ "upstream": "stable-diffusion.cpp",
797
+ },
798
+ )
799
+ return
800
+ self._send_json(
801
+ 200,
802
+ {
803
+ "object": "list",
804
+ "data": [
805
+ {
806
+ "id": state.model_id,
807
+ "object": "model",
808
+ "owned_by": "local",
809
+ }
810
+ ],
811
+ },
812
+ )
813
+
814
+ def do_POST(self) -> None:
815
+ if self.path != "/v1/images/generations":
816
+ self._send_error(
817
+ ApiError(404, "not found", error_type="not_found_error")
818
+ )
819
+ return
820
+ try:
821
+ if self.headers.get("transfer-encoding"):
822
+ raise ApiError(
823
+ 400,
824
+ "chunked request bodies are not supported",
825
+ code="unsupported_transfer_encoding",
826
+ )
827
+ content_type = self.headers.get("content-type", "")
828
+ if content_type.split(";", 1)[0].strip().lower() != "application/json":
829
+ raise ApiError(
830
+ 400,
831
+ "content-type must be application/json",
832
+ code="invalid_content_type",
833
+ )
834
+ length_header = self.headers.get("content-length")
835
+ if length_header is None:
836
+ raise ApiError(411, "content-length is required", code="length_required")
837
+ try:
838
+ content_length = int(length_header)
839
+ except ValueError as error:
840
+ raise ApiError(
841
+ 400,
842
+ "content-length is invalid",
843
+ code="invalid_content_length",
844
+ ) from error
845
+ if content_length < 1:
846
+ raise ApiError(400, "request body is empty", code="invalid_json")
847
+ if content_length > state.max_request_bytes:
848
+ raise ApiError(
849
+ 413,
850
+ "request body exceeds the configured limit",
851
+ code="request_too_large",
852
+ )
853
+ try:
854
+ raw = self.rfile.read(content_length)
855
+ except (TimeoutError, socket.timeout) as error:
856
+ raise ApiError(
857
+ 408,
858
+ "request body read timed out",
859
+ code="request_timeout",
860
+ ) from error
861
+ if len(raw) != content_length:
862
+ raise ApiError(400, "request body is incomplete", code="invalid_json")
863
+ try:
864
+ payload = json.loads(raw)
865
+ except (UnicodeDecodeError, json.JSONDecodeError) as error:
866
+ raise ApiError(
867
+ 400,
868
+ "request body is invalid JSON",
869
+ code="invalid_json",
870
+ ) from error
871
+ parsed = _parse_generation_request(payload, state)
872
+ if state.shutdown_requested.is_set():
873
+ raise ApiError(
874
+ 503,
875
+ "image generator is shutting down",
876
+ error_type="server_error",
877
+ code="service_shutting_down",
878
+ )
879
+ except ApiError as error:
880
+ self._send_error(error)
881
+ return
882
+
883
+ if not state.inference_slots.acquire(blocking=False):
884
+ self._send_error(
885
+ ApiError(
886
+ 429,
887
+ "the local image generator is busy",
888
+ error_type="rate_limit_error",
889
+ code="local_concurrency_saturated",
890
+ )
891
+ )
892
+ return
893
+ started_at = time.monotonic()
894
+ completed = threading.Event()
895
+ outcome: dict[str, Any] = {}
896
+
897
+ def run_generation() -> None:
898
+ try:
899
+ _check_upstream_model(state)
900
+ upstream_result = _generate_image(state, parsed)
901
+ outcome["response"] = _validate_generation_result(
902
+ upstream_result, parsed, state
903
+ )
904
+ except BaseException as error:
905
+ outcome["error"] = error
906
+ finally:
907
+ # The slot belongs to the native request, not the client
908
+ # handler. In particular, a client-facing timeout must not
909
+ # admit overlapping GPU work while sd-server is still busy.
910
+ state.inference_slots.release()
911
+ completed.set()
912
+
913
+ threading.Thread(
914
+ target=run_generation,
915
+ name="sd-cpp-image-generation",
916
+ daemon=True,
917
+ ).start()
918
+ if not completed.wait(state.inference_timeout_seconds):
919
+ logging.error("generation exceeded client deadline; native request still active")
920
+ self._send_error(
921
+ ApiError(
922
+ 504,
923
+ "image generation timed out",
924
+ error_type="server_error",
925
+ code="generation_timeout",
926
+ )
927
+ )
928
+ return
929
+
930
+ error = outcome.get("error")
931
+ if error is None:
932
+ self._send_json(200, outcome["response"])
933
+ logging.info(
934
+ "generation completed width=%d height=%d steps=%d elapsed_ms=%d",
935
+ parsed["width"],
936
+ parsed["height"],
937
+ parsed["steps"],
938
+ int((time.monotonic() - started_at) * 1000),
939
+ )
940
+ elif isinstance(error, ApiError):
941
+ self._send_error(error)
942
+ elif isinstance(error, UpstreamTimeout):
943
+ logging.error("generation failed reason=upstream timeout")
944
+ self._send_error(
945
+ ApiError(
946
+ 504,
947
+ "image generation timed out",
948
+ error_type="server_error",
949
+ code="generation_timeout",
950
+ )
951
+ )
952
+ elif isinstance(error, UpstreamError):
953
+ logging.error("generation failed reason=%s", str(error))
954
+ self._send_error(
955
+ ApiError(
956
+ 502,
957
+ "image generation failed",
958
+ error_type="server_error",
959
+ code="upstream_generation_failed",
960
+ )
961
+ )
962
+ else:
963
+ logging.exception("generation failed with an unexpected worker error", exc_info=error)
964
+ self._send_error(
965
+ ApiError(
966
+ 500,
967
+ "image generation failed",
968
+ error_type="server_error",
969
+ code="internal_error",
970
+ )
971
+ )
972
+
973
+ return Handler
974
+
975
+
976
+ def _arg_size(value: str) -> tuple[int, int]:
977
+ try:
978
+ return _parse_size_text(value)
979
+ except ApiError as error:
980
+ raise argparse.ArgumentTypeError("size must use WIDTHxHEIGHT") from error
981
+
982
+
983
+ def _parse_args() -> argparse.Namespace:
984
+ parser = argparse.ArgumentParser(description=__doc__)
985
+ parser.add_argument("--host", default=DEFAULT_HOST)
986
+ parser.add_argument("--port", type=int, default=DEFAULT_PORT)
987
+ parser.add_argument("--upstream-base-url", default=DEFAULT_UPSTREAM_BASE_URL)
988
+ parser.add_argument("--model-id", default=DEFAULT_MODEL_ID)
989
+ parser.add_argument("--upstream-model-id", default=DEFAULT_MODEL_ID)
990
+ parser.add_argument(
991
+ "--max-request-bytes", type=int, default=DEFAULT_MAX_REQUEST_BYTES
992
+ )
993
+ parser.add_argument("--max-output-bytes", type=int, default=67_108_864)
994
+ parser.add_argument("--max-prompt-chars", type=int, default=8_192)
995
+ parser.add_argument("--max-negative-prompt-chars", type=int, default=8_192)
996
+ parser.add_argument("--min-width", type=int, default=64)
997
+ parser.add_argument("--max-width", type=int, default=2_048)
998
+ parser.add_argument("--min-height", type=int, default=64)
999
+ parser.add_argument("--max-height", type=int, default=2_048)
1000
+ parser.add_argument("--max-pixels", type=int, default=4_194_304)
1001
+ parser.add_argument("--default-size", type=_arg_size, default=(768, 768))
1002
+ parser.add_argument("--default-steps", type=int, default=20)
1003
+ parser.add_argument("--max-steps", type=int, default=100)
1004
+ parser.add_argument(
1005
+ "--client-read-timeout-seconds",
1006
+ type=float,
1007
+ default=DEFAULT_CLIENT_READ_TIMEOUT_SECONDS,
1008
+ )
1009
+ parser.add_argument(
1010
+ "--max-handler-threads",
1011
+ type=int,
1012
+ default=DEFAULT_MAX_HANDLER_THREADS,
1013
+ )
1014
+ parser.add_argument("--request-timeout-seconds", type=float, default=5.0)
1015
+ parser.add_argument("--inference-timeout-seconds", type=float, default=900.0)
1016
+ return parser.parse_args()
1017
+
1018
+
1019
+ def _validate_loopback_origin(value: str) -> str:
1020
+ parsed = urllib.parse.urlsplit(value)
1021
+ try:
1022
+ address = ipaddress.ip_address(parsed.hostname or "")
1023
+ port = parsed.port
1024
+ except ValueError as error:
1025
+ raise SystemExit("--upstream-base-url must be a loopback HTTP origin") from error
1026
+ if (
1027
+ parsed.scheme != "http"
1028
+ or not address.is_loopback
1029
+ or port is None
1030
+ or parsed.username
1031
+ or parsed.password
1032
+ or parsed.query
1033
+ or parsed.fragment
1034
+ or parsed.path not in ("", "/")
1035
+ ):
1036
+ raise SystemExit("--upstream-base-url must be a loopback HTTP origin")
1037
+ return value.rstrip("/")
1038
+
1039
+
1040
+ def main() -> int:
1041
+ args = _parse_args()
1042
+ try:
1043
+ address = ipaddress.ip_address(args.host)
1044
+ except ValueError as error:
1045
+ raise SystemExit("--host must be a loopback IP address") from error
1046
+ if not address.is_loopback:
1047
+ raise SystemExit("--host must be a loopback IP address")
1048
+ if not 1 <= args.port <= 65_535 and args.port != 0:
1049
+ raise SystemExit("--port must be 0 or between 1 and 65535")
1050
+ upstream_base_url = _validate_loopback_origin(args.upstream_base_url)
1051
+ if MODEL_ID_PATTERN.fullmatch(args.model_id) is None:
1052
+ raise SystemExit("--model-id is invalid")
1053
+ if MODEL_ID_PATTERN.fullmatch(args.upstream_model_id) is None:
1054
+ raise SystemExit("--upstream-model-id is invalid")
1055
+
1056
+ integer_limits = (
1057
+ args.max_request_bytes,
1058
+ args.max_output_bytes,
1059
+ args.max_prompt_chars,
1060
+ args.max_negative_prompt_chars,
1061
+ args.min_width,
1062
+ args.max_width,
1063
+ args.min_height,
1064
+ args.max_height,
1065
+ args.max_pixels,
1066
+ args.default_steps,
1067
+ args.max_steps,
1068
+ args.max_handler_threads,
1069
+ )
1070
+ if any(value <= 0 for value in integer_limits):
1071
+ raise SystemExit("request and generation limits must be positive")
1072
+ if args.min_width > args.max_width or args.min_height > args.max_height:
1073
+ raise SystemExit("minimum image dimensions must not exceed maximums")
1074
+ default_width, default_height = args.default_size
1075
+ if (
1076
+ not args.min_width <= default_width <= args.max_width
1077
+ or not args.min_height <= default_height <= args.max_height
1078
+ or default_width * default_height > args.max_pixels
1079
+ ):
1080
+ raise SystemExit("default size is outside the configured image limits")
1081
+ if not 1 <= args.default_steps <= args.max_steps:
1082
+ raise SystemExit("default steps are outside the configured range")
1083
+ for value, name in (
1084
+ (args.client_read_timeout_seconds, "client read timeout"),
1085
+ (args.request_timeout_seconds, "request timeout"),
1086
+ (args.inference_timeout_seconds, "inference timeout"),
1087
+ ):
1088
+ if not math.isfinite(value) or value <= 0:
1089
+ raise SystemExit(f"{name} must be positive")
1090
+
1091
+ state = ServerState(
1092
+ upstream_base_url=upstream_base_url,
1093
+ model_id=args.model_id,
1094
+ upstream_model_id=args.upstream_model_id,
1095
+ max_request_bytes=args.max_request_bytes,
1096
+ max_output_bytes=args.max_output_bytes,
1097
+ max_prompt_chars=args.max_prompt_chars,
1098
+ max_negative_prompt_chars=args.max_negative_prompt_chars,
1099
+ min_width=args.min_width,
1100
+ max_width=args.max_width,
1101
+ min_height=args.min_height,
1102
+ max_height=args.max_height,
1103
+ max_pixels=args.max_pixels,
1104
+ default_size=args.default_size,
1105
+ default_steps=args.default_steps,
1106
+ max_steps=args.max_steps,
1107
+ request_timeout_seconds=args.request_timeout_seconds,
1108
+ inference_timeout_seconds=args.inference_timeout_seconds,
1109
+ )
1110
+ logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
1111
+ logging.info("checking stable-diffusion.cpp model catalog")
1112
+ try:
1113
+ _check_upstream_model(state)
1114
+ except UpstreamError as error:
1115
+ raise SystemExit("stable-diffusion.cpp model preflight failed") from error
1116
+ server = ImageHttpServer(
1117
+ (args.host, args.port),
1118
+ _handler_for(state),
1119
+ client_read_timeout_seconds=args.client_read_timeout_seconds,
1120
+ max_handler_threads=args.max_handler_threads,
1121
+ max_request_bytes=args.max_request_bytes,
1122
+ )
1123
+
1124
+ def request_shutdown(signum: int, frame: Any) -> None:
1125
+ del signum, frame
1126
+ state.shutdown_requested.set()
1127
+ threading.Thread(target=server.shutdown, daemon=True).start()
1128
+
1129
+ signal.signal(signal.SIGTERM, request_shutdown)
1130
+ signal.signal(signal.SIGINT, request_shutdown)
1131
+ host, port = server.server_address[:2]
1132
+ print(
1133
+ json.dumps(
1134
+ {
1135
+ "event": "listening",
1136
+ "host": host,
1137
+ "port": port,
1138
+ "model": state.model_id,
1139
+ "upstream_model": state.upstream_model_id,
1140
+ "upstream": state.upstream_base_url,
1141
+ "client_read_timeout_seconds": args.client_read_timeout_seconds,
1142
+ "max_handler_threads": args.max_handler_threads,
1143
+ "inference_slots": 1,
1144
+ },
1145
+ separators=(",", ":"),
1146
+ ),
1147
+ flush=True,
1148
+ )
1149
+ try:
1150
+ server.serve_forever(poll_interval=0.5)
1151
+ finally:
1152
+ server.server_close()
1153
+ return 0
1154
+
1155
+
1156
+ if __name__ == "__main__":
1157
+ raise SystemExit(main())