diff --git a/CLAUDE.md b/CLAUDE.md index cc6d02a..6bbc093 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -131,11 +131,17 @@ Gemini's inline-image models (images returned out-of-band under the message a hosted URL (Z.ai, Seedance, Seedream 5.0 Lite) are fetched inside the enclave and inlined as `data:` URIs (the fetch is guarded: http(s) only, non-public IP hosts rejected, redirects + size capped, and only ever called on provider- -response URLs, never client input). Image-to-image editing sends the prior image -back inline (a `data:` URI / `image_url` content part on the latest user turn), -forwarded to providers that support it via the endpoint's `image` field. -Per-provider request quirks (response format, `n`, size/watermark, reference -support) live in `model_registry.py`. These models are billed a flat per-image +response URLs, never client input). Image-to-image editing and multi-image +compositing ("add this logo to this photo") send the input images inline +(`data:` URIs / `image_url` content parts on the latest user turn, up to 10), +forwarded to providers that support it. Delivery is one of two per-model paths: +ByteDance carries the references inline in the JSON `image` field of +`/images/generations`; OpenAI gpt-image is routed to its separate +`/images/edits` endpoint, where the references ride as multipart `image[]` file +uploads (only inline `data:` references are uploaded — a plain-URL reference is +skipped rather than dereferenced in the enclave). Per-provider request quirks +(response format, `n`, size/watermark, reference support, edit endpoint) live in +`model_registry.py`. These models are billed a flat per-image price (see `per_image_price_usd`), not per token. ## Verification Examples diff --git a/tee_gateway/image_generation.py b/tee_gateway/image_generation.py index f4d5355..88f3ee9 100644 --- a/tee_gateway/image_generation.py +++ b/tee_gateway/image_generation.py @@ -7,7 +7,10 @@ * ``generate_images`` — shape and send the provider request, always returning inline ``data:`` URIs (a provider-hosted URL is fetched into the enclave so - the client never sees a raw URL). + the client never sees a raw URL). Reference images for image-to-image edits + are delivered one of two ways depending on the provider: inline in the JSON + ``image`` field (ByteDance), or as multipart ``image[]`` file uploads to a + dedicated edits endpoint (OpenAI gpt-image; see ``image_edit_endpoint``). * ``create_image_generation_response`` / ``create_image_generation_streaming_response`` — surface the result on ``/v1/chat/completions`` exactly like Gemini's inline-image models (images @@ -19,6 +22,7 @@ """ import base64 +import binascii import ipaddress import json import logging @@ -137,6 +141,87 @@ def _fetch_url_as_data_uri(url: str) -> str: return f"data:{mime or 'image/jpeg'};base64,{b64}" +# MIME -> filename extension for multipart reference uploads. The edits endpoint +# keys off the upload's filename/content-type; anything unmapped falls back to +# png (gpt-image accepts png/jpeg/webp). +_IMAGE_MIME_EXT = { + "image/png": "png", + "image/jpeg": "jpg", + "image/jpg": "jpg", + "image/webp": "webp", + "image/gif": "gif", +} + + +def _decode_data_uri(uri: str) -> Optional[tuple[bytes, str]]: + """Decode a base64 ``data:`` URI into ``(raw_bytes, mime)``. + + Returns ``None`` for anything that isn't a base64 ``data:`` URI (e.g. a plain + http(s) URL) or that fails to decode. Callers skip those rather than fetch + client-supplied URLs to upload — dereferencing client input inside the + enclave would be an SSRF vector (see ``_validate_fetch_url``). + """ + if not uri.startswith("data:"): + return None + header, sep, b64 = uri.partition(",") + if not sep or not b64: + return None + meta = header[len("data:") :] + params = meta.split(";") + if "base64" not in params[1:]: + return None # only base64 payloads are supported + mime = params[0].strip() or "image/png" + try: + raw = base64.b64decode(b64, validate=True) + except (ValueError, binascii.Error): + return None + if not raw: + return None + return raw, mime + + +def _build_reference_uploads( + reference_images: List[str], +) -> list[tuple[str, tuple[str, bytes, str]]]: + """Decode base64 ``data:`` references into httpx multipart ``files`` tuples. + + Each entry is ``("image[]", (filename, bytes, mime))`` — the repeated + ``image[]`` field OpenAI's edits endpoint expects for multi-image edits. + Non-``data:`` references (plain URLs) are skipped rather than dereferenced. + """ + uploads: list[tuple[str, tuple[str, bytes, str]]] = [] + for i, ref in enumerate(reference_images): + decoded = _decode_data_uri(ref) + if decoded is None: + logger.warning("Skipping non-inline reference image for edit upload") + continue + raw, mime = decoded + ext = _IMAGE_MIME_EXT.get(mime, "png") + uploads.append(("image[]", (f"image_{i}.{ext}", raw, mime))) + return uploads + + +def _build_generations_payload( + cfg: Any, prompt: str, count: int, refs: Optional[List[str]] +) -> dict[str, Any]: + """Build the JSON body for a ``/images/generations`` request. + + ``refs`` (only set for providers that carry references inline) go in the + ``image`` field — a bare string for a single reference, an array for several + (Seedream/Seedance accept either). + """ + payload: dict[str, Any] = {"model": cfg.api_name, "prompt": prompt} + if cfg.image_response_format is not None: + payload["response_format"] = cfg.image_response_format + if cfg.image_send_n: + payload["n"] = count + if cfg.image_extra_params: + payload.update(cfg.image_extra_params) + if refs: + payload["image"] = refs[0] if len(refs) == 1 else refs + return payload + + def generate_images( model: str, prompt: str, @@ -145,10 +230,21 @@ def generate_images( ) -> tuple[list[str], int]: """Generate images via a provider's OpenAI-compatible images endpoint. - ``reference_images`` carries input images for image-to-image editing (e.g. a - follow-up "add a hat" that builds on the previously generated image), sent on - the same endpoint via the ``image`` field (a URL or base64 ``data:`` URI, or - an array of up to 10). Models whose endpoint doesn't support it ignore them. + ``reference_images`` carries input images for image-to-image editing and + multi-image compositing (e.g. "add this logo to this photo", or a follow-up + "add a hat" that builds on the previous result). Up to 10 are forwarded to + providers that support them (``image_supports_reference``); others ignore + them. Delivery differs by provider and is a single choice per model: + + * ``image_edit_endpoint`` unset — inline in the JSON ``image`` field of + ``/images/generations`` (a string for one, an array for several). + Used by ByteDance Seedream/Seedance. + * ``image_edit_endpoint`` set — multipart ``image[]`` file uploads to that + endpoint (OpenAI gpt-image's ``/images/edits``). Only inline ``data:`` + references can be uploaded; a reference that is a plain URL is skipped + (we don't dereference client-supplied URLs inside the enclave), and if + that leaves nothing to upload we fall back to a plain text-to-image + generation rather than failing. Returns ``(data_uris, image_count)``. Every entry is a ``data:`` URI — when a provider returns a hosted URL instead of inline bytes, the gateway fetches it @@ -168,22 +264,11 @@ def generate_images( # n is clamped to the OpenAI-compatible providers' documented 1..10 range. count = max(1, min(int(n), 10)) - payload: dict[str, Any] = {"model": cfg.api_name, "prompt": prompt} - if cfg.image_response_format is not None: - payload["response_format"] = cfg.image_response_format - if cfg.image_send_n: - payload["n"] = count - if cfg.image_extra_params: - payload.update(cfg.image_extra_params) - # Image-to-image editing: forward reference images via the ``image`` field (a - # single string, or an array for multi-reference edits, up to 10). Without - # this a follow-up edit like "add a hat" would ignore the prior image and - # generate a fresh one from the prompt text alone. Filter to non-empty strings - # so a malformed entry can't break JSON serialization of the request payload. + # Filter references to non-empty strings and cap at the providers' 10-image + # limit; only kept for models that accept reference images. + refs: Optional[list[str]] = None if reference_images and cfg.image_supports_reference: - refs = [r for r in reference_images if isinstance(r, str) and r][:10] - if refs: - payload["image"] = refs[0] if len(refs) == 1 else refs + refs = [r for r in reference_images if isinstance(r, str) and r][:10] or None logger.info( "Generating %d image(s) - Provider: %s, Model: %s", @@ -191,7 +276,34 @@ def generate_images( provider, cfg.api_name, ) - resp = client.post(_IMAGE_GENERATION_PATH, json=payload) + + # Two delivery paths, picked by the model's config: multipart file uploads to + # a dedicated edits endpoint, or the JSON generations endpoint (with inline + # references in the ``image`` field when the provider carries them there). + uploads = ( + _build_reference_uploads(refs) if (refs and cfg.image_edit_endpoint) else [] + ) + if uploads: + # ``uploads`` is only non-empty when image_edit_endpoint is set. + edit_endpoint = cfg.image_edit_endpoint + assert edit_endpoint is not None + form: dict[str, str] = {"model": cfg.api_name, "prompt": prompt} + if cfg.image_send_n: + form["n"] = str(count) + if cfg.image_response_format is not None: + form["response_format"] = cfg.image_response_format + if cfg.image_extra_params: + form.update({k: str(v) for k, v in cfg.image_extra_params.items()}) + resp = client.post(edit_endpoint, data=form, files=uploads) + else: + # No edit endpoint (or nothing uploadable): JSON generations. Inline + # references only ride along for providers that carry them there and + # aren't routed to an edits endpoint. + json_refs = refs if not cfg.image_edit_endpoint else None + resp = client.post( + _IMAGE_GENERATION_PATH, + json=_build_generations_payload(cfg, prompt, count, json_refs), + ) resp.raise_for_status() data = resp.json().get("data", []) or [] diff --git a/tee_gateway/model_registry.py b/tee_gateway/model_registry.py index 2f3ea75..43cd176 100644 --- a/tee_gateway/model_registry.py +++ b/tee_gateway/model_registry.py @@ -44,9 +44,19 @@ class ModelConfig: # Whether to send the OpenAI-style ``n`` count. Some endpoints (Z.ai # GLM-Image, ByteDance Seedance) don't document it and reject/ignore it. image_send_n: bool = True - # Whether the endpoint accepts reference images for image-to-image editing, - # sent via the ``image`` field. Text-to-image-only endpoints reject it. + # Whether the endpoint accepts reference images for image-to-image editing. + # Text-to-image-only endpoints reject it. HOW the references are delivered + # depends on ``image_edit_endpoint`` below. image_supports_reference: bool = False + # Where/how reference images are sent (image_supports_reference models only): + # * ``None`` — inline them in the ``image`` field of the JSON + # ``/images/generations`` request (a string, or an array for multi-ref + # edits). Used by ByteDance Seedream/Seedance. + # * a path (e.g. ``"/images/edits"``) — POST multipart/form-data to that + # endpoint with each reference uploaded as a repeated ``image[]`` file. + # Used by OpenAI gpt-image, whose editing/compositing lives on a separate + # endpoint from text-to-image generation. + image_edit_endpoint: Optional[str] = None # Static extra params merged verbatim into the request payload (e.g. size, # watermark). Keyed by field name; values must be JSON-serializable. image_extra_params: Optional[Mapping[str, Any]] = None @@ -171,10 +181,13 @@ class SupportedModel(Enum): # Image generation via OpenAI's /images/generations endpoint (gpt-image). # Unlike DALL·E, gpt-image models always return base64 (``b64_json``) and # reject the ``response_format`` field, so it's omitted. Image-to-image - # editing is a separate ``/images/edits`` endpoint OpenAI-side, so reference - # images aren't forwarded here (text-to-image only). Size/quality are pinned - # so the flat per-image price stays predictable. Billed at a flat $0.05 per - # generated image; token prices unused. + # editing and multi-image compositing (e.g. "add this logo to this photo") + # live on OpenAI's separate ``/images/edits`` endpoint, which takes the + # reference images as multipart file uploads rather than a JSON ``image`` + # field — so reference turns are routed there via ``image_edit_endpoint`` + # (up to 10 references per request). Size/quality are pinned so the flat + # per-image price stays predictable. Billed at a flat $0.05 per generated + # image; token prices unused. GPT_IMAGE_2 = ModelConfig( provider="openai", api_name="gpt-image-2", @@ -183,6 +196,8 @@ class SupportedModel(Enum): image_generation=True, per_image_price_usd=Decimal("0.05"), image_response_format=None, + image_supports_reference=True, + image_edit_endpoint="/images/edits", image_extra_params={"size": "1024x1024", "quality": "medium"}, ) diff --git a/tee_gateway/test/test_image_generation.py b/tee_gateway/test/test_image_generation.py index 0aee7dc..c6c9c6b 100644 --- a/tee_gateway/test/test_image_generation.py +++ b/tee_gateway/test/test_image_generation.py @@ -227,6 +227,72 @@ def test_reference_images_clamped_to_ten(self): payload = client.post.call_args.kwargs["json"] self.assertEqual(len(payload["image"]), 10) + def test_gpt_image_edits_uploads_references_as_multipart(self): + # gpt-image reference edits go to /images/edits as multipart image[] + # file uploads (not the JSON generations path), with n/size/quality as + # form fields and no response_format (gpt-image rejects it). + client = MagicMock() + client.post.return_value = _mock_response([{"b64_json": "aGVsbG8="}]) + refs = [ + "data:image/png;base64,QUJD", # "ABC" + "data:image/jpeg;base64,REVG", # "DEF" + ] + with patch.object(llm_backend, "openai_http_client", client): + images, count = generate_images( + GPT_IMAGE, "add the logo to the photo", n=1, reference_images=refs + ) + + self.assertEqual(count, 1) + self.assertEqual(images, ["data:image/jpeg;base64,aGVsbG8="]) + + args, kwargs = client.post.call_args + self.assertEqual(args[0], "/images/edits") + # No JSON body on the multipart path. + self.assertNotIn("json", kwargs) + form = kwargs["data"] + self.assertEqual(form["model"], get_model_config(GPT_IMAGE).api_name) + self.assertEqual(form["prompt"], "add the logo to the photo") + self.assertEqual(form["n"], "1") + self.assertEqual(form["size"], "1024x1024") + self.assertEqual(form["quality"], "medium") + self.assertNotIn("response_format", form) + # Both references are uploaded under the repeated image[] field, decoded + # back to their raw bytes with a mime-appropriate filename. + uploads = kwargs["files"] + self.assertEqual([field for field, _ in uploads], ["image[]", "image[]"]) + self.assertEqual(uploads[0][1][0], "image_0.png") + self.assertEqual(uploads[0][1][1], b"ABC") + self.assertEqual(uploads[0][1][2], "image/png") + self.assertEqual(uploads[1][1][0], "image_1.jpg") + self.assertEqual(uploads[1][1][1], b"DEF") + + def test_gpt_image_without_references_uses_generations(self): + # No references -> plain text-to-image on the JSON generations endpoint. + client = MagicMock() + client.post.return_value = _mock_response([{"b64_json": "aGVsbG8="}]) + with patch.object(llm_backend, "openai_http_client", client): + generate_images(GPT_IMAGE, "a red cube", n=1) + + args, kwargs = client.post.call_args + self.assertEqual(args[0], "/images/generations") + self.assertIn("json", kwargs) + self.assertNotIn("files", kwargs) + self.assertNotIn("image", kwargs["json"]) + + def test_gpt_image_non_inline_references_fall_back_to_generation(self): + # A plain URL reference can't be uploaded (we won't dereference client + # URLs in the enclave); with nothing uploadable, fall back to a plain + # generation rather than sending an empty edit request. + client = MagicMock() + client.post.return_value = _mock_response([{"b64_json": "aGVsbG8="}]) + with patch.object(llm_backend, "openai_http_client", client): + generate_images(GPT_IMAGE, "p", n=1, reference_images=["https://cdn/x.jpg"]) + + args, kwargs = client.post.call_args + self.assertEqual(args[0], "/images/generations") + self.assertNotIn("files", kwargs) + self.assertNotIn("image", kwargs["json"]) + def test_reference_images_ignored_for_non_bytedance(self): # xAI/Z.ai text-to-image endpoints don't support image edit; the `image` # field must not leak into their payloads.