i2i_contract.py

"""CPU-only request and response contract for the image edit API.

This module deliberately performs technical validation only. It does not
classify, filter, or otherwise inspect image or prompt content for NSFW or
other semantic categories.
"""

from __future__ import annotations

import base64
import binascii
import json
import math
import warnings
from dataclasses import dataclass
from io import BytesIO
from typing import Callable, Collection, Sequence

from PIL import Image, ImageOps, UnidentifiedImageError


API_VERSION = "2026-08-12.3"
CONTENT_MODERATION_ENABLED = False
BASE_EDIT_MODE_ID = "__base__"

MIN_INPUT_IMAGES = 1
MAX_INPUT_IMAGES = 2
MAX_IMAGE_BYTES = 10 * 1024 * 1024
MAX_TOTAL_IMAGE_BYTES = MAX_IMAGE_BYTES * MAX_INPUT_IMAGES
# Bound the JSON string before parsing or allocating decoded byte buffers. The
# fixed allowance covers two data-URL headers, JSON quoting/separators, and a
# small amount of harmless whitespace while each encoded image remains subject
# to its own exact limit below.
MAX_REQUEST_JSON_CHARS = MAX_INPUT_IMAGES * (4 * math.ceil(MAX_IMAGE_BYTES / 3) + 64) + 256
MAX_IMAGE_PIXELS = 16_000_000
MAX_TOTAL_IMAGE_PIXELS = MAX_IMAGE_PIXELS * MAX_INPUT_IMAGES
MIN_IMAGE_SIDE = 64
MAX_OUTPUT_SIDE = 1024
MAX_PROMPT_CHARS = 1000
MAX_SEED = (2**31) - 1
MIN_STEPS = 1
MAX_STEPS = 50
MIN_GUIDANCE = 1.0
MAX_GUIDANCE = 10.0

ALLOWED_MIME_FORMATS = {
    "image/png": "PNG",
    "image/jpeg": "JPEG",
    "image/webp": "WEBP",
}

PUBLIC_ERROR_MESSAGES = {
    "I2I_BAD_REQUEST": "The request is invalid.",
    "I2I_UNSUPPORTED_TYPE": "Use a PNG, JPEG, or WebP image.",
    "I2I_FILE_TOO_LARGE": "The image exceeds the 10 MiB limit.",
    "I2I_IMAGE_DECODE_FAILED": "The image could not be decoded.",
    "I2I_IMAGE_DIMENSIONS_INVALID": "The image dimensions are not supported.",
    "I2I_PROMPT_REQUIRED": "Enter an edit prompt.",
    "I2I_PROMPT_TOO_LONG": "The edit prompt is too long.",
    "I2I_LORA_NOT_ALLOWED": "The requested edit mode is not available.",
    "I2I_PROVIDER_BUSY": "The image editor is busy. Try again later.",
    "I2I_INFERENCE_FAILED": "The image could not be generated.",
    "I2I_INVALID_OUTPUT": "The image editor returned an invalid result.",
}


class ContractError(ValueError):
    """A stable, low-cardinality API contract error."""

    def __init__(self, code: str):
        if code not in PUBLIC_ERROR_MESSAGES:
            raise ValueError(f"Unknown public error code: {code}")
        self.code = code
        super().__init__(code)


@dataclass(frozen=True)
class ValidatedInputImage:
    data: bytes
    mime: str
    width: int
    height: int


@dataclass(frozen=True)
class ValidatedEditRequest:
    images: tuple[ValidatedInputImage, ...]
    prompt: str
    edit_mode: str
    seed: int
    randomize_seed: bool
    guidance_scale: float
    steps: int


class SingleResidentAdapterManager:
    """Maintain a fail-closed, at-most-one-resident LoRA state machine.

    The caller must serialize ``activate_mode`` with the same lock used for
    inference. Marking a load dirty before mutation means a partial adapter
    load is never treated as a clean base-model state.
    """

    def __init__(
        self,
        *,
        adapter_modes: Collection[str],
        load_adapter: Callable[[str], None],
        activate_adapter: Callable[[str], None],
        unload_adapters: Callable[[], None],
    ):
        self._adapter_modes = frozenset(adapter_modes)
        self._load_adapter = load_adapter
        self._activate_adapter = activate_adapter
        self._unload_adapters = unload_adapters
        self._resident_adapter_mode: str | None = None
        self._dirty = False

    @property
    def resident_adapter_mode(self) -> str | None:
        return self._resident_adapter_mode

    @property
    def dirty(self) -> bool:
        return self._dirty

    def _drop_resident_adapter(self) -> None:
        if self._resident_adapter_mode is None and not self._dirty:
            return
        try:
            self._unload_adapters()
        except Exception:
            self._resident_adapter_mode = None
            self._dirty = True
            raise
        self._resident_adapter_mode = None
        self._dirty = False

    def activate_mode(self, edit_mode: str) -> None:
        if edit_mode != BASE_EDIT_MODE_ID and edit_mode not in self._adapter_modes:
            raise ContractError("I2I_LORA_NOT_ALLOWED")
        if edit_mode == self._resident_adapter_mode and not self._dirty:
            return

        self._drop_resident_adapter()
        if edit_mode == BASE_EDIT_MODE_ID:
            return

        # Loading can mutate the pipeline before raising, so assume the state
        # is dirty until load + activation both complete.
        self._dirty = True
        try:
            self._load_adapter(edit_mode)
            self._activate_adapter(edit_mode)
        except Exception:
            try:
                self._unload_adapters()
            except Exception:
                self._dirty = True
            else:
                self._dirty = False
            self._resident_adapter_mode = None
            raise

        self._resident_adapter_mode = edit_mode
        self._dirty = False


def public_error_text(code: str) -> str:
    """Return a stable client-safe code and message."""

    return f"{code}: {PUBLIC_ERROR_MESSAGES[code]}"


def _detected_mime(data: bytes) -> str | None:
    if data.startswith(b"\x89PNG\r\n\x1a\n"):
        return "image/png"
    if data.startswith(b"\xff\xd8\xff"):
        return "image/jpeg"
    if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return "image/webp"
    return None


def _validate_dimensions(width: int, height: int, *, output: bool = False) -> None:
    if width <= 0 or height <= 0:
        raise ContractError("I2I_IMAGE_DIMENSIONS_INVALID")
    if width * height > MAX_IMAGE_PIXELS:
        raise ContractError("I2I_IMAGE_DIMENSIONS_INVALID")
    if output:
        if max(width, height) > MAX_OUTPUT_SIDE:
            raise ContractError("I2I_INVALID_OUTPUT")
        return
    if min(width, height) < MIN_IMAGE_SIDE:
        raise ContractError("I2I_IMAGE_DIMENSIONS_INVALID")
    # The model rounds output dimensions to a multiple of eight. Reject an
    # aspect ratio that would collapse the shorter side to zero.
    long_side = max(width, height)
    short_side = min(width, height)
    if int(MAX_OUTPUT_SIDE * short_side / long_side) < 8:
        raise ContractError("I2I_IMAGE_DIMENSIONS_INVALID")


def _decode_image_data_url(value: str) -> ValidatedInputImage:
    if not isinstance(value, str):
        raise ContractError("I2I_BAD_REQUEST")

    header, separator, encoded = value.partition(",")
    expected_headers = {f"data:{mime};base64": mime for mime in ALLOWED_MIME_FORMATS}
    mime = expected_headers.get(header)
    if not separator or mime is None:
        raise ContractError("I2I_UNSUPPORTED_TYPE")
    if not encoded:
        raise ContractError("I2I_IMAGE_DECODE_FAILED")

    # Reject oversized payloads before allocating the decoded byte buffer.
    max_encoded_chars = 4 * math.ceil(MAX_IMAGE_BYTES / 3)
    if len(encoded) > max_encoded_chars:
        raise ContractError("I2I_FILE_TOO_LARGE")
    try:
        data = base64.b64decode(encoded, validate=True)
    except (binascii.Error, ValueError):
        raise ContractError("I2I_IMAGE_DECODE_FAILED") from None
    if len(data) > MAX_IMAGE_BYTES:
        raise ContractError("I2I_FILE_TOO_LARGE")
    if _detected_mime(data) != mime:
        raise ContractError("I2I_UNSUPPORTED_TYPE")
    if mime == "image/png" and len(data) >= 24 and data[12:16] == b"IHDR":
        _validate_dimensions(
            int.from_bytes(data[16:20], "big"),
            int.from_bytes(data[20:24], "big"),
        )

    try:
        with warnings.catch_warnings():
            warnings.simplefilter("error", Image.DecompressionBombWarning)
            with Image.open(BytesIO(data)) as probe:
                if getattr(probe, "is_animated", False) and getattr(probe, "n_frames", 1) != 1:
                    raise ContractError("I2I_IMAGE_DECODE_FAILED")
                if probe.format != ALLOWED_MIME_FORMATS[mime]:
                    raise ContractError("I2I_UNSUPPORTED_TYPE")
                _validate_dimensions(*probe.size)
                probe.verify()

            with Image.open(BytesIO(data)) as decoded:
                decoded.load()
                image = ImageOps.exif_transpose(decoded).convert("RGB")
                image.load()
    except ContractError:
        raise
    except (Image.DecompressionBombError, Image.DecompressionBombWarning):
        raise ContractError("I2I_IMAGE_DIMENSIONS_INVALID") from None
    except (UnidentifiedImageError, OSError, SyntaxError, ValueError):
        raise ContractError("I2I_IMAGE_DECODE_FAILED") from None

    _validate_dimensions(*image.size)
    return ValidatedInputImage(
        data=data,
        mime=mime,
        width=image.width,
        height=image.height,
    )


def _decode_images(images_b64_json: str) -> tuple[ValidatedInputImage, ...]:
    if not isinstance(images_b64_json, str):
        raise ContractError("I2I_BAD_REQUEST")
    if len(images_b64_json) > MAX_REQUEST_JSON_CHARS:
        raise ContractError("I2I_FILE_TOO_LARGE")
    try:
        values = json.loads(images_b64_json)
    except (json.JSONDecodeError, TypeError):
        raise ContractError("I2I_BAD_REQUEST") from None
    if not isinstance(values, list) or not MIN_INPUT_IMAGES <= len(values) <= MAX_INPUT_IMAGES:
        raise ContractError("I2I_BAD_REQUEST")

    decoded_images: list[ValidatedInputImage] = []
    total_bytes = 0
    total_pixels = 0
    for value in values:
        image = _decode_image_data_url(value)
        total_bytes += len(image.data)
        if total_bytes > MAX_TOTAL_IMAGE_BYTES:
            raise ContractError("I2I_FILE_TOO_LARGE")
        total_pixels += image.width * image.height
        if total_pixels > MAX_TOTAL_IMAGE_PIXELS:
            raise ContractError("I2I_IMAGE_DIMENSIONS_INVALID")
        decoded_images.append(image)
    return tuple(decoded_images)


def parse_example_index(value: object, total_examples: int) -> int | None:
    """Return a bounded integer example index, rejecting coercive values."""

    if isinstance(total_examples, bool) or not isinstance(total_examples, int) or total_examples < 0:
        raise ValueError("total_examples must be a non-negative integer")
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return None
    if isinstance(value, float) and (not math.isfinite(value) or not value.is_integer()):
        return None
    index = int(value)
    return index if 0 <= index < total_examples else None


def validate_edit_request(
    *,
    images_b64_json: str,
    prompt: str,
    lora_adapter: str,
    seed: int,
    randomize_seed: bool,
    guidance_scale: float,
    steps: int,
    allowed_edit_modes: Collection[str],
) -> ValidatedEditRequest:
    """Validate the seven-field public API contract without semantic review."""

    images = _decode_images(images_b64_json)

    if not isinstance(prompt, str) or not prompt.strip():
        raise ContractError("I2I_PROMPT_REQUIRED")
    normalized_prompt = prompt.strip()
    if len(normalized_prompt) > MAX_PROMPT_CHARS:
        raise ContractError("I2I_PROMPT_TOO_LONG")

    if not isinstance(lora_adapter, str) or lora_adapter not in allowed_edit_modes:
        raise ContractError("I2I_LORA_NOT_ALLOWED")
    if isinstance(seed, bool) or not isinstance(seed, int) or not 0 <= seed <= MAX_SEED:
        raise ContractError("I2I_BAD_REQUEST")
    if not isinstance(randomize_seed, bool):
        raise ContractError("I2I_BAD_REQUEST")
    if (
        isinstance(guidance_scale, bool)
        or not isinstance(guidance_scale, (int, float))
        or not math.isfinite(float(guidance_scale))
        or not MIN_GUIDANCE <= float(guidance_scale) <= MAX_GUIDANCE
    ):
        raise ContractError("I2I_BAD_REQUEST")
    if isinstance(steps, bool) or not isinstance(steps, int) or not MIN_STEPS <= steps <= MAX_STEPS:
        raise ContractError("I2I_BAD_REQUEST")

    return ValidatedEditRequest(
        images=images,
        prompt=normalized_prompt,
        edit_mode=lora_adapter,
        seed=seed,
        randomize_seed=randomize_seed,
        guidance_scale=float(guidance_scale),
        steps=steps,
    )


def decode_validated_images(images: Sequence[ValidatedInputImage]) -> list[Image.Image]:
    """Decode previously validated bytes inside the ZeroGPU worker."""

    decoded_images: list[Image.Image] = []
    try:
        for item in images:
            with Image.open(BytesIO(item.data)) as decoded:
                decoded.load()
                image = ImageOps.exif_transpose(decoded).convert("RGB")
                image.load()
                if image.size != (item.width, item.height):
                    raise ContractError("I2I_IMAGE_DECODE_FAILED")
                decoded_images.append(image)
    except ContractError:
        raise
    except (UnidentifiedImageError, OSError, SyntaxError, ValueError):
        raise ContractError("I2I_IMAGE_DECODE_FAILED") from None
    return decoded_images


def serialize_png_candidate(image: Image.Image) -> bytes:
    """Serialize one model image for the CPU response validator."""

    if not isinstance(image, Image.Image):
        raise ContractError("I2I_INVALID_OUTPUT")
    try:
        buffer = BytesIO()
        image.save(buffer, format="PNG")
        return buffer.getvalue()
    except (OSError, SyntaxError, ValueError):
        raise ContractError("I2I_INVALID_OUTPUT") from None


def encode_png_result(png: bytes, seed: int) -> dict[str, str | int]:
    """Validate and encode one PNG byte payload as a public data URL."""

    if not isinstance(png, bytes):
        raise ContractError("I2I_INVALID_OUTPUT")
    if isinstance(seed, bool) or not isinstance(seed, int) or not 0 <= seed <= MAX_SEED:
        raise ContractError("I2I_INVALID_OUTPUT")

    try:
        if not png.startswith(b"\x89PNG\r\n\x1a\n"):
            raise ContractError("I2I_INVALID_OUTPUT")
        with Image.open(BytesIO(png)) as decoded:
            if decoded.format != "PNG":
                raise ContractError("I2I_INVALID_OUTPUT")
            _validate_dimensions(*decoded.size, output=True)
            decoded.verify()
    except ContractError:
        raise
    except (OSError, SyntaxError, ValueError):
        raise ContractError("I2I_INVALID_OUTPUT") from None

    encoded = base64.b64encode(png).decode("ascii")
    return {"image": f"data:image/png;base64,{encoded}", "seed": seed}


def execute_validated_edit(
    *,
    images_b64_json: str,
    prompt: str,
    lora_adapter: str,
    seed: int,
    randomize_seed: bool,
    guidance_scale: float,
    steps: int,
    allowed_edit_modes: Collection[str],
    runner: Callable[[ValidatedEditRequest], tuple[bytes, int]],
) -> dict[str, str | int]:
    """Validate before invoking the supplied GPU runner, then validate output."""

    request = validate_edit_request(
        images_b64_json=images_b64_json,
        prompt=prompt,
        lora_adapter=lora_adapter,
        seed=seed,
        randomize_seed=randomize_seed,
        guidance_scale=guidance_scale,
        steps=steps,
        allowed_edit_modes=allowed_edit_modes,
    )
    png, used_seed = runner(request)
    return encode_png_result(png, used_seed)


def output_dimensions(input_image: ValidatedInputImage | Image.Image) -> tuple[int, int]:
    """Scale to a 1024px longest side and round both dimensions to 8px."""

    if isinstance(input_image, ValidatedInputImage):
        width, height = input_image.width, input_image.height
    else:
        width, height = input_image.size
    _validate_dimensions(width, height)
    if width >= height:
        scaled_width = MAX_OUTPUT_SIDE
        scaled_height = int(MAX_OUTPUT_SIDE * height / width)
    else:
        scaled_height = MAX_OUTPUT_SIDE
        scaled_width = int(MAX_OUTPUT_SIDE * width / height)
    return max(8, (scaled_width // 8) * 8), max(8, (scaled_height // 8) * 8)