diff --git a/.gitignore b/.gitignore index 3dc7007..ce15f48 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ __pycache__/ *.jpeg *.webp .DS_Store +.pytest_cache/ diff --git a/README.md b/README.md index d938e33..f5e35e9 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,23 @@ Generate images using Google's Gemini / Imagen models. Defaults to **Nano Banana | `thinking_level` | `minimal` · `high` | none | NB2 only. Improves composition at the cost of latency. | | `output_path` | string | auto-generated | Full path e.g. `/tmp/my-image.png` | +### `edit_image_gemini` +Image-**to**-image editing with Gemini (Nano Banana 2). The image-to-image counterpart to `generate_image_gemini`: it takes one or more **reference images** as input so you can edit a single image, compose/fuse several, or keep a character consistent across renders. Unlike the generation tools, it **returns the edited image back to the caller as a viewable image** (not just a file path), closing the generate → see → edit loop. + +| Parameter | Options | Default | Notes | +|---|---|---|---| +| `prompt` | string | required | The edit/composition instruction | +| `input_images` | list of paths | required | 1–14 reference images (PNG/JPEG/WebP/GIF). Use up to 5 of the same subject for character consistency | +| `model` | `gemini-3.1-flash-image-preview` · `gemini-3-pro-image-preview` | `gemini-3.1-flash-image-preview` | | +| `aspect_ratio` | same set as `generate_image_gemini` | `1:1` | | +| `image_size` | `512px` · `1K` · `2K` · `4K` | model default | | +| `thinking_level` | `minimal` · `high` | none | NB2 only | +| `output_path` | string | auto-generated | The edited image is also saved here | + +``` +Edit /tmp/room.png: replace the sofa with a green velvet one, keep everything else. +``` + ### `generate_image_openai` Generate images using OpenAI's GPT Image models. diff --git a/demo_edit.py b/demo_edit.py new file mode 100644 index 0000000..2d49c23 --- /dev/null +++ b/demo_edit.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +""" +Offline proof-of-concept for neoimage's image-to-image editing + closed visual loop. + +Runs with NO network and NO API key: a stubbed Gemini endpoint (httpx.MockTransport) +stands in for the real API. It proves three things end-to-end: + + 1. The edit request actually carries the input image(s) back to the model + (inlineData parts) -- i.e. this is genuine image-to-image, not text-to-image. + 2. The MCP tool returns the edited image as a VIEWABLE image content block + (mime image/png), so the calling model can see what it made and iterate. + 3. The edited bytes are also persisted to disk, and structured metadata + (reference_count, model, file_path, text_response) comes back alongside. + +Exit code 0 == all assertions passed. + + ./venv/bin/python demo_edit.py +""" + +import asyncio +import base64 +import json +import os +import sys +import tempfile +from unittest.mock import patch + +import httpx + +import image_edit + +# A 1x1 PNG used both as the "input" image and as the stubbed model "output". +PNG_1X1 = bytes.fromhex( + "89504e470d0a1a0a0000000d4948445200000001000000010802000000907753" + "de0000000c4944415408d763f8cfc0f01f0005010102a0bb3d630000000049454e44ae426082" +) + +# Distinct bytes for the "edited" result so we can prove we got the model's +# output back, not just echoed our input. +EDITED_PNG = bytes.fromhex( + "89504e470d0a1a0a0000000d4948445200000001000000010802000000907753" + "de0000000c4944415408d76368606000000005000157a4d2b30000000049454e44ae426082" +) + +captured_request: dict = {} + + +def gemini_mock_handler(request: httpx.Request) -> httpx.Response: + """Stand-in for generativelanguage.googleapis.com generateContent.""" + body = json.loads(request.content) + captured_request["body"] = body + parts = body["contents"][0]["parts"] + ref_count = sum(1 for p in parts if "inlineData" in p) + # Refuse to "edit" if no image was sent -- this stub only honors image-to-image. + if ref_count == 0: + return httpx.Response(400, json={"error": "no input image"}) + response = { + "candidates": [ + { + "content": { + "parts": [ + {"text": f"Applied edit using {ref_count} reference image(s)."}, + { + "inlineData": { + "mimeType": "image/png", + "data": base64.b64encode(EDITED_PNG).decode("ascii"), + } + }, + ] + } + } + ] + } + return httpx.Response(200, json=response) + + +def check(label: str, condition: bool) -> None: + status = "ok " if condition else "FAIL" + print(f" [{status}] {label}") + if not condition: + raise SystemExit(f"assertion failed: {label}") + + +async def main() -> None: + # --- Part A: pure-logic guards (no network at all) ------------------------ + print("Part A -- request construction & validation (pure):") + ref = image_edit.ReferenceImage(mime_type="image/png", data=PNG_1X1) + body = image_edit.build_gemini_edit_body("make it pop", [ref]) + parts = body["contents"][0]["parts"] + check("prompt is the first part", parts[0] == {"text": "make it pop"}) + check("input image attached as inlineData", "inlineData" in parts[1]) + check("mime detected from bytes, not extension", image_edit.detect_mime(PNG_1X1) == "image/png") + + raised = False + try: + image_edit.build_gemini_edit_body("x", []) + except ValueError: + raised = True + check("editing with zero references is rejected", raised) + + raised = False + try: + image_edit.build_gemini_edit_body("x", [ref] * (image_edit.MAX_REFERENCE_IMAGES + 1)) + except ValueError: + raised = True + check(f"more than {image_edit.MAX_REFERENCE_IMAGES} references is rejected", raised) + + # --- Part B: full edit round-trip via mocked transport -------------------- + print("Part B -- image-to-image round-trip (mocked Gemini, no network):") + with tempfile.TemporaryDirectory() as tmp: + in_path = os.path.join(tmp, "input.png") + with open(in_path, "wb") as f: + f.write(PNG_1X1) + + async with httpx.AsyncClient( + transport=httpx.MockTransport(gemini_mock_handler) + ) as mock_client: + result = await image_edit.gemini_edit( + prompt="swap the background to a beach", + reference_paths=[in_path], + api_key="offline-test-key", + http_client=mock_client, + ) + check("model received the input image", captured_request["body"] + ["contents"][0]["parts"][1].get("inlineData") is not None) + check("edited bytes are the model output (not the input)", result.image_bytes == EDITED_PNG) + check("reference_count reported", result.reference_count == 1) + check("model text response captured", "reference image" in (result.text_response or "")) + + # --- Part C: closed visual loop through the real MCP tool ----------------- + print("Part C -- MCP tool returns a VIEWABLE image (closed loop):") + # Inject the mock transport into the tool's HTTP factory so the real + # edit_image_gemini tool runs fully offline. Scoped via context managers so + # neither the factory patch nor the fake API key leaks past this block. + def _mock_factory(timeout: float) -> httpx.AsyncClient: + return httpx.AsyncClient( + transport=httpx.MockTransport(gemini_mock_handler), timeout=timeout + ) + + from fastmcp import Client + import server + + with tempfile.TemporaryDirectory() as tmp, \ + patch.object(image_edit, "_make_client", _mock_factory), \ + patch.dict(os.environ, {"GOOGLE_API_KEY": "offline-test-key"}): + in_path = os.path.join(tmp, "portrait.png") + out_path = os.path.join(tmp, "edited.png") + with open(in_path, "wb") as f: + f.write(PNG_1X1) + + async with Client(server.mcp) as client: + tools = {t.name for t in await client.list_tools()} + check("edit_image_gemini is registered as an MCP tool", "edit_image_gemini" in tools) + + res = await client.call_tool( + "edit_image_gemini", + { + "prompt": "give the character a red jacket, keep the face identical", + "input_images": [in_path], + "output_path": out_path, + }, + ) + + image_blocks = [b for b in res.content if getattr(b, "type", None) == "image"] + check("tool returned a viewable image content block", len(image_blocks) == 1) + check("returned image is a PNG", image_blocks[0].mimeType == "image/png") + decoded = base64.b64decode(image_blocks[0].data) + check("viewable image is the edited output", decoded == EDITED_PNG) + check("edited image also written to disk", os.path.exists(out_path)) + with open(out_path, "rb") as f: + check("file on disk matches edited bytes", f.read() == EDITED_PNG) + check("structured metadata reports reference_count", + res.structured_content.get("reference_count") == 1) + + print("\nAll checks passed -- image-to-image editing + closed visual loop proven offline.") + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except SystemExit as exc: + print(exc, file=sys.stderr) + sys.exit(1) + sys.exit(0) diff --git a/image_edit.py b/image_edit.py new file mode 100644 index 0000000..0acc0ba --- /dev/null +++ b/image_edit.py @@ -0,0 +1,235 @@ +""" +Image-to-image editing for Google Gemini (Nano Banana 2). + +This is the capability the generation-only `server.py` never built: feeding one or +more *reference images* back into the model so it can edit, compose, or hold a +character consistent across renders -- the headline feature of the Gemini image +models that the README advertises ("character consistency for up to 5 characters") +but the existing tools cannot actually exercise, because they only ever send text. + +The functions here are deliberately split into: + * pure, network-free helpers (mime detection, request-body construction, + response parsing) that can be unit-tested offline with no API key, and + * a single async `gemini_edit()` with an injectable HTTP client so the whole + request/response round-trip is exercisable against `httpx.MockTransport`. + +Nothing here imports FastMCP -- the MCP tool wrapper lives in `server.py`. +""" + +from __future__ import annotations + +import base64 +from dataclasses import dataclass + +import httpx + +# Gemini Nano Banana 2 accepts up to 14 reference images in a single request +# (and can hold up to 5 of them consistent as "characters"). +MAX_REFERENCE_IMAGES = 14 + +DEFAULT_MODEL = "gemini-3.1-flash-image-preview" +GEMINI_ENDPOINT = ( + "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent" +) + +# Magic-byte signatures -> MIME type. Order matters only in that each check is +# unambiguous; we never guess. +_MAGIC = ( + (b"\x89PNG\r\n\x1a\n", "image/png"), + (b"\xff\xd8\xff", "image/jpeg"), + (b"GIF87a", "image/gif"), + (b"GIF89a", "image/gif"), +) + + +@dataclass(frozen=True) +class ReferenceImage: + """A decoded input image ready to be attached to a request.""" + + mime_type: str + data: bytes + + @property + def as_inline_part(self) -> dict: + return { + "inlineData": { + "mimeType": self.mime_type, + "data": base64.b64encode(self.data).decode("ascii"), + } + } + + +@dataclass(frozen=True) +class EditResult: + """The product of an edit: raw image bytes plus provenance metadata.""" + + image_bytes: bytes + mime_type: str + text_response: str | None + model: str + reference_count: int + + +def detect_mime(data: bytes) -> str: + """Detect an image MIME type from its magic bytes. Trusts no file extension. + + Raises ValueError for unrecognized/empty input rather than guessing -- the + request must declare a real mimeType or the model will reject it. + """ + if not data: + raise ValueError("empty image data") + for signature, mime in _MAGIC: + if data.startswith(signature): + return mime + # WebP is RIFFWEBP, so the marker sits at byte 8, not byte 0. + if data[:4] == b"RIFF" and data[8:12] == b"WEBP": + return "image/webp" + raise ValueError("unrecognized image format (not PNG/JPEG/WebP/GIF)") + + +def load_reference_image(path: str) -> ReferenceImage: + """Read an image file from disk and sniff its MIME type from content.""" + with open(path, "rb") as f: + data = f.read() + return ReferenceImage(mime_type=detect_mime(data), data=data) + + +def build_gemini_edit_body( + prompt: str, + references: list[ReferenceImage], + *, + model: str = DEFAULT_MODEL, + aspect_ratio: str = "1:1", + image_size: str | None = None, + thinking_level: str | None = None, +) -> dict: + """Construct the `generateContent` request body for an image edit. + + The shape mirrors the text-only body in `server.py` but interleaves the + reference images as `inlineData` parts alongside the prompt text -- which is + exactly what turns "generate" into "edit/compose/keep-consistent". + """ + if not prompt or not prompt.strip(): + raise ValueError("prompt must be a non-empty string") + if not references: + raise ValueError( + "edit requires at least one reference image; " + "use generate_image_gemini for text-only generation" + ) + if len(references) > MAX_REFERENCE_IMAGES: + raise ValueError( + f"too many reference images: {len(references)} " + f"(Gemini accepts at most {MAX_REFERENCE_IMAGES})" + ) + + parts: list[dict] = [{"text": prompt}] + parts.extend(ref.as_inline_part for ref in references) + + body: dict = { + "contents": [{"parts": parts}], + "generationConfig": {"responseModalities": ["IMAGE", "TEXT"]}, + } + + image_config: dict = {} + if aspect_ratio != "1:1": + image_config["aspectRatio"] = aspect_ratio + if image_size: + image_config["imageSize"] = image_size + if image_config: + body["generationConfig"]["imageConfig"] = image_config + + # thinking mode is Nano Banana 2 only; silently irrelevant elsewhere. + if thinking_level and model == "gemini-3.1-flash-image-preview": + body["generationConfig"]["thinkingConfig"] = { + "thinkingLevel": thinking_level, + "includeThoughts": False, + } + + return body + + +def parse_gemini_image_response(result: dict) -> tuple[bytes, str | None]: + """Pull the image bytes (and any text) out of a generateContent response. + + Same envelope the existing generator parses, factored out so the edit path + and a future refactor of the generate path can share one parser. + """ + candidates = result.get("candidates", []) + if not candidates: + raise ValueError("no candidates returned from Gemini API") + + parts = candidates[0].get("content", {}).get("parts", []) + # Take the FIRST image/text part, not the last. Gemini can return multiple + # inlineData parts (multi-candidate / multi-turn); a last-wins loop would + # silently discard the earlier image. + image_data: str | None = next( + (p["inlineData"]["data"] for p in parts if "inlineData" in p), None + ) + text_response: str | None = next( + (p["text"] for p in parts if "text" in p), None + ) + + if not image_data: + raise ValueError( + f"no image data in response. Text response: {text_response}" + ) + return base64.b64decode(image_data), text_response + + +def _make_client(timeout: float) -> httpx.AsyncClient: + """HTTP client factory. Overridable in tests/demos to inject a MockTransport + so the full edit round-trip runs with zero network and no API key.""" + return httpx.AsyncClient(timeout=timeout) + + +async def gemini_edit( + *, + prompt: str, + reference_paths: list[str], + api_key: str, + model: str = DEFAULT_MODEL, + aspect_ratio: str = "1:1", + image_size: str | None = None, + thinking_level: str | None = None, + http_client: httpx.AsyncClient | None = None, +) -> EditResult: + """Edit/compose an image from a prompt plus one or more reference images. + + Pass `http_client` to inject a transport (offline testing); otherwise a real + client is created and closed for you. + """ + references = [load_reference_image(p) for p in reference_paths] + body = build_gemini_edit_body( + prompt, + references, + model=model, + aspect_ratio=aspect_ratio, + image_size=image_size, + thinking_level=thinking_level, + ) + + url = GEMINI_ENDPOINT.format(model=model) + headers = {"Content-Type": "application/json", "x-goog-api-key": api_key} + timeout = 300.0 if image_size == "4K" else 120.0 + + client = http_client or _make_client(timeout) + owns_client = http_client is None + try: + response = await client.post(url, headers=headers, json=body) + if response.status_code != 200: + raise ValueError( + f"Gemini API error ({response.status_code}): {response.text}" + ) + result = response.json() + finally: + if owns_client: + await client.aclose() + + image_bytes, text_response = parse_gemini_image_response(result) + return EditResult( + image_bytes=image_bytes, + mime_type=detect_mime(image_bytes), + text_response=text_response, + model=model, + reference_count=len(references), + ) diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..2f4c80e --- /dev/null +++ b/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +asyncio_mode = auto diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..bf51dd4 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,3 @@ +-r requirements.txt +pytest +pytest-asyncio diff --git a/server.py b/server.py index 7a36aa3..75e8a6f 100644 --- a/server.py +++ b/server.py @@ -12,6 +12,10 @@ import httpx from fastmcp import FastMCP +from fastmcp.tools.tool import ToolResult +from fastmcp.utilities.types import Image + +import image_edit mcp = FastMCP("neoimage") @@ -387,5 +391,98 @@ async def generate_image_grok( } +@mcp.tool() +async def edit_image_gemini( + prompt: str, + input_images: list[str], + model: Literal[ + "gemini-3.1-flash-image-preview", + "gemini-3-pro-image-preview", + ] = "gemini-3.1-flash-image-preview", + aspect_ratio: Literal[ + "1:1", + "2:3", "3:2", + "3:4", "4:3", + "4:5", "5:4", + "9:16", "16:9", + "21:9", + "4:1", "1:4", + "8:1", "1:8", + ] = "1:1", + image_size: Literal["512px", "1K", "2K", "4K"] | None = None, + thinking_level: Literal["minimal", "high"] | None = None, + output_path: str | None = None, +) -> ToolResult: + """ + Edit, compose, or character-consistently re-render images with Google Gemini + (Nano Banana 2) using one or more REFERENCE IMAGES as input. + + This is the image-to-image counterpart to generate_image_gemini. Where the + generator only takes text, this tool feeds existing images back into the model + so you can: + - EDIT a single image ("make the sky stormy", "remove the background") + - COMPOSE/FUSE several images ("put the product from image 1 onto the desk in image 2") + - keep a CHARACTER CONSISTENT across renders (pass up to 5 portraits of the same person) + + Unlike the generation tools, this one RETURNS THE EDITED IMAGE to you directly + as a viewable image (not just a file path), so you can see the result and decide + whether to iterate -- closing the generate -> see -> edit loop. + + Args: + prompt: Instruction describing the edit/composition to perform. + input_images: One or more paths to reference images (PNG/JPEG/WebP/GIF). + At least 1, at most 14. Use up to 5 images of the same subject for + character consistency. + model: Gemini image model. Defaults to Nano Banana 2. + aspect_ratio: Output dimensions ratio. + image_size: Output resolution (512px/1K/2K/4K). None uses the model default. + thinking_level: Enable thinking mode (Nano Banana 2 only). + output_path: Where to save the edited image. Auto-generated if not provided. + + Returns: + A ToolResult whose content includes the edited image (viewable inline) and + whose structured metadata carries file_path, model, reference_count, the + prompt, and any text_response from the model. + """ + api_key = os.environ.get("GOOGLE_API_KEY") + if not api_key: + raise ValueError("GOOGLE_API_KEY environment variable is not set") + + result = await image_edit.gemini_edit( + prompt=prompt, + reference_paths=input_images, + api_key=api_key, + model=model, + aspect_ratio=aspect_ratio, + image_size=image_size, + thinking_level=thinking_level, + ) + + if output_path is None: + output_path = generate_output_path("gemini_edit", "png") + + try: + with open(output_path, "wb") as f: + f.write(result.image_bytes) + except OSError as e: + # Surface disk errors as ValueError, consistent with how every other + # failure in this server is reported to the MCP caller. + raise ValueError(f"could not write edited image to {output_path}: {e}") from e + + metadata = { + "file_path": os.path.abspath(output_path), + "model": result.model, + "aspect_ratio": aspect_ratio, + "image_size": image_size, + "thinking_level": thinking_level, + "prompt": prompt, + "reference_count": result.reference_count, + "text_response": result.text_response, + } + + image = Image(data=result.image_bytes, format="png") + return ToolResult(content=[image.to_image_content()], structured_content=metadata) + + if __name__ == "__main__": mcp.run() diff --git a/tests/test_image_edit.py b/tests/test_image_edit.py new file mode 100644 index 0000000..0336b05 --- /dev/null +++ b/tests/test_image_edit.py @@ -0,0 +1,184 @@ +"""Offline unit/integration tests for the Gemini image-edit path. + +No network and no API key required -- the HTTP layer is stubbed with +httpx.MockTransport. Run with: ./venv/bin/python -m pytest tests/ -q +""" + +import base64 +import json +import os +import sys + +import httpx +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import image_edit # noqa: E402 + +PNG_1X1 = bytes.fromhex( + "89504e470d0a1a0a0000000d4948445200000001000000010802000000907753" + "de0000000c4944415408d763f8cfc0f01f0005010102a0bb3d630000000049454e44ae426082" +) +JPEG_HEADER = b"\xff\xd8\xff\xe0\x00\x10JFIF" +WEBP_HEADER = b"RIFF\x00\x00\x00\x00WEBPVP8 " +EDITED_PNG = bytes.fromhex( + "89504e470d0a1a0a0000000d4948445200000001000000010802000000907753" + "de0000000c4944415408d76368606000000005000157a4d2b30000000049454e44ae426082" +) + + +# --- pure helpers ---------------------------------------------------------- + +@pytest.mark.parametrize( + "data,expected", + [ + (PNG_1X1, "image/png"), + (JPEG_HEADER, "image/jpeg"), + (WEBP_HEADER, "image/webp"), + (b"GIF89a....", "image/gif"), + ], +) +def test_detect_mime_from_magic_bytes(data, expected): + assert image_edit.detect_mime(data) == expected + + +def test_detect_mime_rejects_unknown(): + with pytest.raises(ValueError): + image_edit.detect_mime(b"not an image") + with pytest.raises(ValueError): + image_edit.detect_mime(b"") + + +def test_build_body_interleaves_prompt_and_images(): + ref = image_edit.ReferenceImage(mime_type="image/png", data=PNG_1X1) + body = image_edit.build_gemini_edit_body("edit me", [ref, ref]) + parts = body["contents"][0]["parts"] + assert parts[0] == {"text": "edit me"} + assert sum(1 for p in parts if "inlineData" in p) == 2 + # image bytes are base64-encoded into the request + assert parts[1]["inlineData"]["data"] == base64.b64encode(PNG_1X1).decode() + assert parts[1]["inlineData"]["mimeType"] == "image/png" + assert body["generationConfig"]["responseModalities"] == ["IMAGE", "TEXT"] + + +def test_build_body_requires_at_least_one_reference(): + with pytest.raises(ValueError, match="at least one reference"): + image_edit.build_gemini_edit_body("x", []) + + +def test_build_body_caps_reference_count(): + ref = image_edit.ReferenceImage(mime_type="image/png", data=PNG_1X1) + too_many = [ref] * (image_edit.MAX_REFERENCE_IMAGES + 1) + with pytest.raises(ValueError, match="too many reference images"): + image_edit.build_gemini_edit_body("x", too_many) + + +def test_build_body_rejects_empty_prompt(): + ref = image_edit.ReferenceImage(mime_type="image/png", data=PNG_1X1) + with pytest.raises(ValueError, match="non-empty"): + image_edit.build_gemini_edit_body(" ", [ref]) + + +def test_build_body_thinking_only_for_nb2(): + ref = image_edit.ReferenceImage(mime_type="image/png", data=PNG_1X1) + nb2 = image_edit.build_gemini_edit_body( + "x", [ref], model="gemini-3.1-flash-image-preview", thinking_level="high" + ) + assert nb2["generationConfig"]["thinkingConfig"]["thinkingLevel"] == "high" + pro = image_edit.build_gemini_edit_body( + "x", [ref], model="gemini-3-pro-image-preview", thinking_level="high" + ) + assert "thinkingConfig" not in pro["generationConfig"] + + +def test_parse_response_extracts_image_and_text(): + payload = { + "candidates": [ + { + "content": { + "parts": [ + {"text": "done"}, + {"inlineData": {"data": base64.b64encode(EDITED_PNG).decode()}}, + ] + } + } + ] + } + img, text = image_edit.parse_gemini_image_response(payload) + assert img == EDITED_PNG + assert text == "done" + + +def test_parse_response_raises_without_image(): + with pytest.raises(ValueError): + image_edit.parse_gemini_image_response({"candidates": [{"content": {"parts": []}}]}) + + +# --- full round-trip via mocked transport ---------------------------------- + +def _mock_client(handler): + return httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + +@pytest.mark.asyncio +async def test_gemini_edit_round_trip(tmp_path): + seen = {} + + def handler(request): + seen["body"] = json.loads(request.content) + seen["api_key"] = request.headers.get("x-goog-api-key") + return httpx.Response( + 200, + json={ + "candidates": [ + { + "content": { + "parts": [ + {"text": "edited"}, + {"inlineData": {"data": base64.b64encode(EDITED_PNG).decode()}}, + ] + } + } + ] + }, + ) + + in_path = tmp_path / "in.png" + in_path.write_bytes(PNG_1X1) + + result = await image_edit.gemini_edit( + prompt="recolor", + reference_paths=[str(in_path)], + api_key="k-123", + http_client=_mock_client(handler), + ) + + # the input image really went to the model + assert "inlineData" in seen["body"]["contents"][0]["parts"][1] + assert seen["api_key"] == "k-123" + assert result.image_bytes == EDITED_PNG + assert result.reference_count == 1 + assert result.text_response == "edited" + + +@pytest.mark.asyncio +async def test_gemini_edit_surfaces_api_error(): + def handler(request): + return httpx.Response(403, text="permission denied") + + import tempfile + + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f: + f.write(PNG_1X1) + path = f.name + try: + with pytest.raises(ValueError, match="403"): + await image_edit.gemini_edit( + prompt="x", + reference_paths=[path], + api_key="k", + http_client=_mock_client(handler), + ) + finally: + os.unlink(path)