From a8604a31720690ecbb341a7fa5eab34b6f0e695d Mon Sep 17 00:00:00 2001
From: Xubin Ren <52506698+Re-bin@users.noreply.github.com>
Date: Sun, 26 Jul 2026 17:39:36 +0800
Subject: [PATCH] fix(image): scope Gemini image sizes by model

---
 nanobot/providers/image_generation.py    | 30 +++++++++++++++---------
 tests/providers/test_image_generation.py | 28 ++++++++++++++++++++++
 2 files changed, 47 insertions(+), 11 deletions(-)

diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py
index f9d76e28..9e9332fd 100644
--- a/nanobot/providers/image_generation.py
+++ b/nanobot/providers/image_generation.py
@@ -40,9 +40,11 @@ _GEMINI_IMAGEN_ASPECT_RATIOS = {"1:1", "9:16", "16:9", "3:4", "4:3"}
 _GEMINI_FLASH_ASPECT_RATIOS = {
     "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9",
 }
-# Image-size tokens accepted by Gemini 3+ image models (earlier Flash image
-# models expose only a single fixed resolution).
-_GEMINI_FLASH_IMAGE_SIZES = {"512", "1K", "2K", "4K"}
+# Gemini 3 Pro image models accept these sizes. Gemini 3.1 Flash adds 512,
+# while Gemini 3.1 Flash Lite supports only 1K.
+_GEMINI_3_IMAGE_SIZES = {"1K", "2K", "4K"}
+_GEMINI_31_FLASH_IMAGE_SIZES = {"512", *_GEMINI_3_IMAGE_SIZES}
+_GEMINI_31_FLASH_LITE_IMAGE_SIZES = {"1K"}
 _OLLAMA_DEFAULT_SIDE = 1024
 _OLLAMA_SIZE_PRESETS = {
     "1K": 1024,
@@ -782,21 +784,27 @@ def _gemini_flash_image_config(
     config: dict[str, str] = {}
     if aspect_ratio and aspect_ratio in _GEMINI_FLASH_ASPECT_RATIOS:
         config["aspectRatio"] = aspect_ratio
-    if image_size and _gemini_flash_supports_image_size(model):
+    if image_size:
         normalized = image_size.strip().upper()
-        if normalized in _GEMINI_FLASH_IMAGE_SIZES:
+        if normalized in _gemini_flash_supported_image_sizes(model):
             config["imageSize"] = normalized
     return config
 
 
-def _gemini_flash_supports_image_size(model: str) -> bool:
-    """Return whether the model honors ``imageSize``.
+def _gemini_flash_supported_image_sizes(model: str) -> set[str]:
+    """Return the ``imageSize`` values documented for a Flash-path model.
 
-    Only Gemini 3+ image models expose a configurable image size; earlier Flash
-    image models (2.0, 2.5) generate at a single fixed resolution, so ``imageSize``
-    is identified positively rather than by excluding a single version.
+    Earlier Flash image models (2.0, 2.5) expose no configurable size. Gemini
+    3.1 Flash Lite is intentionally checked before the broader Flash match.
     """
-    return "gemini-3" in model.lower()
+    normalized = model.lower()
+    if "gemini-3.1-flash-lite-image" in normalized:
+        return _GEMINI_31_FLASH_LITE_IMAGE_SIZES
+    if "gemini-3.1-flash-image" in normalized:
+        return _GEMINI_31_FLASH_IMAGE_SIZES
+    if "gemini-3" in normalized:
+        return _GEMINI_3_IMAGE_SIZES
+    return set()
 
 
 async def _aihubmix_images_from_payload(
diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py
index 19f5eb4d..7fe63e43 100644
--- a/tests/providers/test_image_generation.py
+++ b/tests/providers/test_image_generation.py
@@ -480,6 +480,34 @@ async def test_gemini_flash_2_0_drops_image_size() -> None:
     assert image_config == {"aspectRatio": "16:9"}
 
 
+@pytest.mark.parametrize(
+    ("model", "image_size", "expected"),
+    [
+        ("gemini-3-pro-image", "512", None),
+        ("gemini-3.1-flash-lite-image", "2K", None),
+        ("gemini-3.1-flash-lite-image", "1K", {"imageSize": "1K"}),
+        ("gemini-3.1-flash-image", "512", {"imageSize": "512"}),
+    ],
+)
+@pytest.mark.asyncio
+async def test_gemini_flash_scopes_image_size_by_model(
+    model: str,
+    image_size: str,
+    expected: dict[str, str] | None,
+) -> None:
+    fake = FakeClient(_gemini_flash_image_response())
+    client = GeminiImageGenerationClient(api_key="AIza-test", client=fake)  # type: ignore[arg-type]
+
+    await client.generate(
+        prompt="draw a cat",
+        model=model,
+        image_size=image_size,
+    )
+
+    response_format = fake.calls[0]["json"]["generationConfig"].get("responseFormat")
+    assert response_format == ({"image": expected} if expected else None)
+
+
 @pytest.mark.asyncio
 async def test_gemini_flash_ignores_unsupported_hints() -> None:
     fake = FakeClient(_gemini_flash_image_response())