"""Flask server construction for API profile mode."""
from __future__ import annotations
from dataclasses import dataclass
from types import SimpleNamespace
from flask import Flask, request
from diffusion_cli.api_profiles import API_PROFILES
from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.generation_service import (
MODEL_RESIDENCY_VALUES,
GenerationService,
)
from diffusion_cli.ui import registerUiRoutes
BINARY_JSON_FIELDS = ("init_images", "mask", "extra_images")
@dataclass(frozen=True)
class ServerConfig:
"""Validated settings for HTTP server mode."""
api_profile: str
host: str
port: int
model_residency: str
def validateServerConfig(
api_profile: str | None,
host: str,
port: int,
model_residency: str,
) -> ServerConfig:
"""Validate server mode command line settings."""
if api_profile is None:
raise DiffusionCliError("--api-profile is required")
if api_profile not in API_PROFILES:
raise DiffusionCliError(f"Unknown API profile: {api_profile}")
if not host:
raise DiffusionCliError("--host must not be empty")
if port < 1 or port > 65535:
raise DiffusionCliError(f"--port must be from 1 through 65535: {port}")
if model_residency not in MODEL_RESIDENCY_VALUES:
raise DiffusionCliError(
"--model-residency must be one of staged, cpu-cache"
)
return ServerConfig(
api_profile=api_profile,
host=host,
port=port,
model_residency=model_residency,
)
def _summarizeJsonBody(value):
if isinstance(value, dict):
result = {}
for key, item in value.items():
if key in BINARY_JSON_FIELDS and isinstance(item, str):
result[key] = f"<{len(item)} bytes>"
elif key in BINARY_JSON_FIELDS and isinstance(item, list):
result[key] = f"<{len(item)} items>"
else:
result[key] = _summarizeJsonBody(item)
return result
if isinstance(value, list):
return [_summarizeJsonBody(item) for item in value]
return value
def createApp(server_config: ServerConfig, generation_service) -> Flask:
"""Create a Flask app for the selected API profile."""
app = Flask(__name__)
context = SimpleNamespace(
server_config=server_config,
generation_service=generation_service,
)
@app.before_request
def logRequest() -> None:
body = None
if request.is_json:
body = _summarizeJsonBody(request.get_json(silent=True))
app.logger.debug(
"request method=%s path=%s query=%s json=%r",
request.method,
request.path,
request.query_string.decode("utf-8"),
body,
)
registerUiRoutes(app, context)
API_PROFILES[server_config.api_profile].register_routes(app, context)
return app
def serve(server_config: ServerConfig, user_config) -> None:
"""Run the configured Flask development server."""
generation_service = GenerationService(
user_config,
model_residency=server_config.model_residency,
)
app = createApp(server_config, generation_service)
app.run(host=server_config.host, port=server_config.port)