"""HTTP API profile registry and route implementations."""
from __future__ import annotations
from base64 import b64encode
from dataclasses import dataclass
import json
from typing import Callable
from flask import Response, jsonify, request
from diffusion_cli.config import (
ImageGenerationRequest,
normalizeOutputExtension,
validateDimensions,
validateOutputQuality,
)
from diffusion_cli.errors import DiffusionCliError
UNSUPPORTED_IMAGE_FIELDS = ("init_images", "mask", "extra_images")
@dataclass(frozen=True)
class ApiProfile:
"""HTTP API contract exposed by server mode."""
name: str
description: str
register_routes: Callable[[object, object], None]
def _optionalString(data: dict, key: str) -> str | None:
value = data.get(key)
if value is None:
return None
if not isinstance(value, str):
raise DiffusionCliError(f"{key} must be a string")
return value
def _optionalInt(data: dict, key: str) -> int | None:
value = data.get(key)
if value is None:
return None
if not isinstance(value, int) or isinstance(value, bool):
raise DiffusionCliError(f"{key} must be an integer")
return value
def _optionalFloat(data: dict, key: str) -> float | None:
value = data.get(key)
if value is None:
return None
if not isinstance(value, int | float) or isinstance(value, bool):
raise DiffusionCliError(f"{key} must be a number")
return float(value)
def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
for field in UNSUPPORTED_IMAGE_FIELDS:
if field in data:
raise DiffusionCliError(f"{field} is not supported")
prompt = data.get("prompt")
if not isinstance(prompt, str) or not prompt.strip():
raise DiffusionCliError("prompt required")
seed = _optionalInt(data, "seed")
if seed == -1:
seed = None
width = _optionalInt(data, "width")
height = _optionalInt(data, "height")
if width is not None and height is not None:
validateDimensions(width, height)
elif width is not None and (width <= 0 or width % 8 != 0):
raise DiffusionCliError(
f"Width must be a positive multiple of 8: got {width}"
)
elif height is not None and (height <= 0 or height % 8 != 0):
raise DiffusionCliError(
f"Height must be a positive multiple of 8: got {height}"
)
steps = _optionalInt(data, "steps")
batch_size = _optionalInt(data, "batch_size")
cfg = _optionalFloat(data, "cfg_scale")
output_extension = _optionalString(data, "output_extension")
output_quality = _optionalInt(data, "output_quality")
if output_extension is not None:
output_extension = normalizeOutputExtension(output_extension)
if output_quality is not None:
output_quality = validateOutputQuality(output_quality)
if steps is not None and steps < 1:
raise DiffusionCliError(f"Steps must be at least 1: got {steps}")
if batch_size is not None and batch_size < 1:
raise DiffusionCliError(
f"Batch size must be at least 1: got {batch_size}"
)
if cfg is not None and cfg < 0:
raise DiffusionCliError(f"CFG must be non-negative: got {cfg}")
return ImageGenerationRequest(
prompt=prompt,
negative_prompt=_optionalString(data, "negative_prompt"),
seed=seed,
width=width,
height=height,
batch_size=batch_size,
steps=steps,
cfg=cfg,
output_extension=output_extension,
output_quality=output_quality,
)
def _badRequest(message: str):
return jsonify({"error": "bad_request", "message": message}), 400
def _generationFailed(message: str):
return jsonify({"error": "generation_failed", "message": message}), 500
def registerSillyTavernSdcppRoutes(app, context) -> None:
"""Register SillyTavern stable-diffusion.cpp compatibility routes."""
@app.route("/health", methods=["GET"])
def health():
return jsonify({
"status": "ok",
"api_profile": context.server_config.api_profile,
})
@app.route("/v1/images/generations", methods=["OPTIONS"])
def imageGenerationsOptions():
response = Response(status=204)
response.headers["Allow"] = "OPTIONS, POST"
return response
@app.route("/v1/models", methods=["GET"])
def models():
return jsonify({
"data": [
{
"id": "z-image-local",
"object": "model",
"owned_by": "local",
}
]
})
@app.route("/sdapi/v1/txt2img", methods=["POST"])
def txt2img():
data = request.get_json(silent=True)
if not isinstance(data, dict) or not data:
return _badRequest("JSON object body required")
try:
generation_request = _txt2imgRequest(data)
except DiffusionCliError as exc:
return _badRequest(str(exc))
try:
images = context.generation_service.generateImages(
generation_request,
)
except DiffusionCliError as exc:
app.logger.exception("Generation failed")
return _generationFailed(str(exc))
except Exception:
app.logger.exception("Generation failed")
return _generationFailed("generation failed")
encoded_images = [
b64encode(image.data).decode("ascii") for image in images
]
info = {
"prompt": generation_request.prompt,
"negative_prompt": generation_request.negative_prompt,
"seed": images[0].seed if images else generation_request.seed,
"width": generation_request.width,
"height": generation_request.height,
"steps": generation_request.steps,
"cfg_scale": generation_request.cfg,
}
return jsonify({
"images": encoded_images,
"parameters": data,
"info": json.dumps(info, separators=(",", ":")),
})
SILLYTAVERN_SDCPP_PROFILE = ApiProfile(
name="sillytavern-sdcpp",
description="SillyTavern stable-diffusion.cpp adapter compatibility.",
register_routes=registerSillyTavernSdcppRoutes,
)
API_PROFILES = {
SILLYTAVERN_SDCPP_PROFILE.name: SILLYTAVERN_SDCPP_PROFILE,
}