Support multi-image edits for gpt-image-2

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

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,30 +264,46 @@ 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",
         count,
         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 []
 

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"},
     )
 

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.