Shaman
63 lines
1.7 KiB
Python
63 lines
1.7 KiB
Python
"""Pydantic models for the OpenAI images contract."""
|
|
import base64
|
|
import binascii
|
|
|
|
from pydantic import BaseModel, field_validator
|
|
|
|
DEFAULT_SIZE = "1024x1024"
|
|
|
|
|
|
class ImageGenerationRequest(BaseModel):
|
|
prompt: str
|
|
model: str | None = None
|
|
n: int | None = 1
|
|
size: str = DEFAULT_SIZE
|
|
seed: int | None = None
|
|
response_format: str | None = None # accepted; only b64_json is produced
|
|
|
|
@field_validator("size")
|
|
@classmethod
|
|
def _size_shape(cls, v: str) -> str:
|
|
w_str, sep, h_str = v.lower().partition("x")
|
|
if not sep or not w_str.isdigit() or not h_str.isdigit() or int(w_str) <= 0 or int(h_str) <= 0:
|
|
raise ValueError(f"invalid size '{v}', expected format 'WIDTHxHEIGHT'")
|
|
return v
|
|
|
|
|
|
class ImageEditJsonRequest(BaseModel):
|
|
image: str # base64-encoded PNG/JPEG (data-URL prefix allowed)
|
|
prompt: str
|
|
model: str | None = None
|
|
size: str | None = None # accepted but ignored: edits size comes from the input image
|
|
seed: int | None = None
|
|
|
|
@field_validator("image")
|
|
@classmethod
|
|
def _decode_b64(cls, v: str) -> str:
|
|
raw = v.split(",", 1)[1] if v.startswith("data:") else v
|
|
try:
|
|
base64.b64decode(raw, validate=True)
|
|
except (binascii.Error, ValueError) as e:
|
|
raise ValueError("image must be valid base64") from e
|
|
return raw
|
|
|
|
|
|
class ImageItem(BaseModel):
|
|
b64_json: str
|
|
|
|
|
|
class ImagesResponse(BaseModel):
|
|
created: int
|
|
data: list[ImageItem]
|
|
|
|
|
|
class ModelCard(BaseModel):
|
|
id: str
|
|
object: str = "model"
|
|
owned_by: str = "comfy-adapter"
|
|
|
|
|
|
class ModelsResponse(BaseModel):
|
|
object: str = "list"
|
|
data: list[ModelCard]
|