Expand migrated image API: model selection, edits and 2K export

README.md

@@ -244,6 +244,14 @@ You can also pass `flow_project_id` per project on `POST /api/projects`.
 | `FLOW_ALLOW_DEGRADED` | `0` | `1` lets scene chaining and r2v fall back to plain i2v instead of failing. |
 | `DEFAULT_PAYGATE_TIER` | `PAYGATE_TIER_TWO` | Carried for the DB and dashboard; no longer selects a model. |
 
+### Image API
+
+The migrated image path supports Nano Banana Pro, Nano Banana 2 and Nano Banana
+2 Lite, all five current aspect ratios, 1-4 outputs, true base-image editing and
+native 2K image export. Exact future Flow image model wire ids pass through
+without being silently replaced by the default model. See
+[`docs/IMAGE_API.md`](docs/IMAGE_API.md).
+
 ### What does not work on the new API yet
 
 Three capabilities have no captured payload, so they fail with

agent/api/flow.py

@@ -1,6 +1,6 @@
 """Direct Flow API endpoints — for manual operations outside the queue."""
-from fastapi import APIRouter, HTTPException
-from pydantic import BaseModel
+from fastapi import APIRouter, HTTPException, Response
+from pydantic import BaseModel, Field
 from typing import Literal, Optional
 
 from agent.config import USE_BATCH_RPC, FLOW_PROJECT_ID, FLOW_ALLOW_DEGRADED
@@ -20,6 +20,10 @@ class GenerateImageRequest(BaseModel):
     project_id: str
     aspect_ratio: str = "IMAGE_ASPECT_RATIO_PORTRAIT"
     user_paygate_tier: str = "PAYGATE_TIER_ONE"
+    image_model: Optional[str] = None
+    count: int = Field(default=1, ge=1, le=4)
+    seed: Optional[int] = Field(default=None, ge=1, le=1_000_000_000)
+    reference_media_ids: Optional[list[str]] = None
     character_media_ids: Optional[list[str]] = None
 
 
@@ -94,6 +98,16 @@ class EditImageRequest(BaseModel):
     project_id: str
     aspect_ratio: str = "IMAGE_ASPECT_RATIO_PORTRAIT"
     user_paygate_tier: str = "PAYGATE_TIER_ONE"
+    image_model: Optional[str] = None
+    count: int = Field(default=1, ge=1, le=4)
+    seed: Optional[int] = Field(default=None, ge=1, le=1_000_000_000)
+    reference_media_ids: Optional[list[str]] = None
+
+
+class UpscaleImageRequest(BaseModel):
+    media_id: str
+    project_id: str
+    quality: Literal["2k", "4k"] = "2k"
 
 
 @router.get("/status")
@@ -127,11 +141,14 @@ async def get_credits():
 
 @router.post("/generate-image")
 async def generate_image(body: GenerateImageRequest):
-    """Generate image directly (bypasses queue)."""
+    """Generate 1-4 images with an explicit Flow image model."""
     client = get_flow_client()
     if not client.connected:
         raise HTTPException(503, "Extension not connected")
-    result = await client.generate_images(**body.model_dump())
+    data = body.model_dump(exclude={"reference_media_ids"})
+    refs = list(dict.fromkeys((body.reference_media_ids or []) + (body.character_media_ids or [])))
+    data["character_media_ids"] = refs or None
+    result = await client.generate_images(**data)
     if result.get("error") or (isinstance(result.get("status"), int) and result["status"] >= 400):
         raise HTTPException(result.get("status", 502), result.get("error", result.get("data")))
     return result.get("data", result)
@@ -335,20 +352,59 @@ async def get_media(media_id: str):
 
 @router.post("/edit-image")
 async def edit_image(body: EditImageRequest):
-    """Edit an existing image using IMAGE_INPUT_TYPE_BASE_IMAGE (bypasses queue)."""
+    """Edit an existing image using the current Flow BASE_IMAGE wire input."""
     client = get_flow_client()
     if not client.connected:
         raise HTTPException(503, "Extension not connected")
     result = await client.edit_image(
-        body.prompt, body.source_media_id, body.project_id,
+        body.prompt,
+        body.source_media_id,
+        body.project_id,
         aspect_ratio=body.aspect_ratio,
         user_paygate_tier=body.user_paygate_tier,
+        character_media_ids=body.reference_media_ids,
+        image_model=body.image_model,
+        count=body.count,
+        seed=body.seed,
     )
     if result.get("error") or (isinstance(result.get("status"), int) and result["status"] >= 400):
         raise HTTPException(result.get("status", 502), result.get("error", result.get("data")))
     return result.get("data", result)
 
 
+@router.post("/export-image")
+@router.post("/upscale-image", include_in_schema=False)
+async def export_image(body: UpscaleImageRequest):
+    """Download a generated Flow image at 2K (or plan-gated 4K)."""
+    import base64
+    import binascii
+
+    client = get_flow_client()
+    if not client.connected:
+        raise HTTPException(503, "Extension not connected")
+    result = await client.upscale_image(
+        body.media_id,
+        body.project_id,
+        resolution=body.quality.upper(),
+    )
+    if result.get("error") or (isinstance(result.get("status"), int) and result["status"] >= 400):
+        raise HTTPException(result.get("status", 502), result.get("error", result.get("data")))
+    data = result.get("data", result)
+    try:
+        content = base64.b64decode(data["encodedImage"], validate=True)
+    except (KeyError, TypeError, binascii.Error) as exc:
+        raise HTTPException(502, "Flow image upscale returned invalid image data") from exc
+    quality = body.quality.lower()
+    return Response(
+        content=content,
+        media_type=data.get("contentType", "image/jpeg"),
+        headers={
+            "Content-Disposition": f'attachment; filename="flow-{body.media_id}-{quality}.jpg"',
+            "X-Flow-Image-Quality": quality,
+        },
+    )
+
+
 @router.post("/upload-image")
 async def upload_image(body: UploadImageRequest):
     """Upload a local image file to Google Flow and get a media_id."""

agent/config.py

@@ -61,8 +61,9 @@
 VIDEO_MODELS = _MODELS["video_models"]
 UPSCALE_MODELS = _MODELS["upscale_models"]
 IMAGE_MODELS = _MODELS["image_models"]
-# Nickname from image_models. The batch path accepts GEM_PIX_2 (Nano Banana Pro)
-# and NARWHAL (Banana 2) and rejects everything else.
+# Nickname from image_models. Known aliases live in models.json, while the
+# batch path also accepts syntactically valid Flow wire model ids directly so
+# newly introduced image models do not require a Flow Kit release.
 DEFAULT_IMAGE_MODEL = _MODELS.get("default_image_model", "NANO_BANANA_PRO")
 
 # ─── API Endpoints ───────────────────────────────────────────

agent/models.json

@@ -55,7 +55,8 @@
   },
   "image_models": {
     "NANO_BANANA_PRO": "GEM_PIX_2",
-    "NANO_BANANA_2": "NARWHAL"
+    "NANO_BANANA_2": "NARWHAL",
+    "NANO_BANANA_2_LITE": "HARBOR_SEAL"
   },
   "default_image_model": "NANO_BANANA_PRO",
   "batch_video_models": {

agent/services/flow_batch.py

@@ -40,6 +40,7 @@
 RPC_PROJECT_MEDIA = "Zzl0ze"
 RPC_MEDIA = "as29s"
 RPC_UPLOAD_IMAGE = "maseQ"
+RPC_UPSCALE_IMAGE = "SPrCad"
 
 CAPTCHA_IMAGE = "IMAGE_GENERATION"
 CAPTCHA_VIDEO = "VIDEO_GENERATION"
@@ -49,15 +50,21 @@
 #: happen in the page, moments before the request leaves.
 CAPTCHA_SLOT = "__CAPTCHA__"
 
-#: Wire names this path accepts. Everything else is rejected outright by Flow.
-#: ``GEM_PIX_2`` is Nano Banana Pro, ``NARWHAL`` is Banana 2. Flow Kit uses Pro
-#: by default (see agent/models.json), which is also what the new path defaults
-#: to; a caller that wants Banana 2 has to name it.
-IMAGE_MODELS = {"GEM_PIX_2", "NARWHAL"}
+#: Current image model wire ids observed in the Flow frontend. Unknown future
+#: ids are accepted by ``resolve_image_model`` instead of being silently
+#: replaced by the default, so callers can opt into a newly exposed model
+#: before Flow Kit itself ships another release.
+IMAGE_MODELS = {"GEM_PIX_2", "NARWHAL", "HARBOR_SEAL"}
 IMAGE_MODEL = "GEM_PIX_2"
 
-#: The nicknames models.json speaks, resolved to wire names.
-IMAGE_MODEL_BY_NICKNAME = {"NANO_BANANA_PRO": "GEM_PIX_2", "NANO_BANANA_2": "NARWHAL"}
+#: Friendly aliases. Exact Flow wire ids work too.
+IMAGE_MODEL_BY_NICKNAME = {
+    "NANO_BANANA_PRO": "GEM_PIX_2",
+    "NANO_BANANA_2": "NARWHAL",
+    "NANO_BANANA_2_LITE": "HARBOR_SEAL",
+    "NANO_BANANA_LITE": "HARBOR_SEAL",
+}
+IMAGE_MODEL_ID_RE = re.compile(r"^[A-Z][A-Z0-9_]{1,95}$")
 
 #: Image aspect ratios, measured by generating one of each and reading the
 #: JPEG header. This slot was mistaken for a variant count at first — 1 means
@@ -73,8 +80,16 @@
     "IMAGE_ASPECT_RATIO_SQUARE": ASPECT_SQUARE,
     "IMAGE_ASPECT_RATIO_PORTRAIT": ASPECT_PORTRAIT,
     "IMAGE_ASPECT_RATIO_LANDSCAPE": ASPECT_LANDSCAPE,
+    # Current spelling plus the old Flow Kit alias for compatibility.
+    "IMAGE_ASPECT_RATIO_PORTRAIT_THREE_FOUR": ASPECT_PORTRAIT_4_3,
     "IMAGE_ASPECT_RATIO_PORTRAIT_FOUR_THREE": ASPECT_PORTRAIT_4_3,
     "IMAGE_ASPECT_RATIO_LANDSCAPE_FOUR_THREE": ASPECT_LANDSCAPE_4_3,
+    # Friendly API aliases.
+    "1:1": ASPECT_SQUARE,
+    "9:16": ASPECT_PORTRAIT,
+    "16:9": ASPECT_LANDSCAPE,
+    "3:4": ASPECT_PORTRAIT_4_3,
+    "4:3": ASPECT_LANDSCAPE_4_3,
 }
 
 #: Video models this path accepts. The REST-era map was keyed by
@@ -117,10 +132,11 @@
 #: reframed by hand: a hair inside the edges, spanning 128/129 of the frame.
 FULL_FRAME_CROP = [None, 0.0038759689922481244, 1, 0.9961240310077519]
 
-#: A reference image, as the UI sends it: the media id FIRST and a type flag
-#: four slots later. Probing never found this — the id sat in the wrong
-#: position, the payload was accepted, and the picture quietly ignored it.
+#: Image inputs put the media id first and the input type four slots later.
+#: Type 1 is a reference; type 2 is the image being edited (BASE_IMAGE).
 REF_TYPE_IMAGE = 1
+BASE_TYPE_IMAGE = 2
+IMAGE_UPSCALE_RESOLUTIONS = {"2K": 1, "4K": 2}
 
 
 class RpcError(RuntimeError):
@@ -180,12 +196,18 @@ class MediaUrls:
 # ── model / aspect resolvers ─────────────────────────────────────────────────
 
 def resolve_image_model(key: Optional[str]) -> str:
-    """Nickname or wire name in, wire name out; anything unknown coerces."""
+    """Nickname or wire id in, wire id out.
+
+    Flow can add image models independently of Flow Kit releases. A
+    syntactically valid, previously unseen wire id therefore passes through
+    unchanged instead of being silently replaced by the default model.
+    """
     if isinstance(key, str):
-        if key in IMAGE_MODEL_BY_NICKNAME:
-            return IMAGE_MODEL_BY_NICKNAME[key]
-        if key in IMAGE_MODELS:
-            return key
+        normalized = key.strip().upper().replace("-", "_")
+        if normalized in IMAGE_MODEL_BY_NICKNAME:
+            return IMAGE_MODEL_BY_NICKNAME[normalized]
+        if IMAGE_MODEL_ID_RE.fullmatch(normalized):
+            return normalized
     return IMAGE_MODEL
 
 
@@ -211,8 +233,10 @@ def resolve_video_model(key: Optional[str]) -> str:
 
 
 def resolve_aspect(aspect: Any) -> int:
-    """Take either the wire value or the REST-era name."""
+    """Take a current/legacy wire name, friendly ratio, or integer 1-5."""
     if isinstance(aspect, int):
+        if aspect not in (1, 2, 3, 4, 5):
+            raise ValueError(f"image aspect must be 1-5, got {aspect}")
         return aspect
     try:
         return ASPECT_BY_NAME[aspect]
@@ -313,35 +337,64 @@ def _context(project_id: str) -> list:
             [CAPTCHA_SLOT, 1]]
 
 
+def _image_input(media_id: str, input_type: int) -> list:
+    return [media_id, None, None, None, input_type]
+
+
 def _reference(media_id: str) -> list:
-    return [media_id, None, None, None, REF_TYPE_IMAGE]
+    return _image_input(media_id, REF_TYPE_IMAGE)
+
+
+def _base_image(media_id: str) -> list:
+    return _image_input(media_id, BASE_TYPE_IMAGE)
 
 
 def image_request(prompt: str, project_id: str, count: int = 1,
                   aspect: Any = ASPECT_SQUARE, seed: Optional[int] = None,
                   prompts: Optional[list[str]] = None,
                   model: str = IMAGE_MODEL,
-                  ref_media_ids: Optional[list[str]] = None) -> str:
+                  ref_media_ids: Optional[list[str]] = None,
+                  base_media_id: Optional[str] = None) -> str:
     """One request item per variant, exactly as the REST payload did it.
 
     There is no "how many" field: Flow returns one image per item in the list,
     so `count` replicates the item under fresh seeds. `ref_media_ids` conditions
     the result on images already in the project — this is what keeps a character
     the same person from beat to beat.
     """
+    if not isinstance(count, int) or isinstance(count, bool) or not 1 <= count <= 4:
+        raise ValueError("image count must be an integer from 1 to 4")
     ratio = resolve_aspect(aspect)
+    resolved_model = resolve_image_model(model)
     base = seed if seed is not None else random.randint(1, 10**9)
     items = []
-    for index in range(max(1, count)):
+    for index in range(count):
         text = prompts[index] if prompts and index < len(prompts) else prompt
-        refs = [_reference(mid) for mid in (ref_media_ids or [])] or None
-        items.append([None, None, refs, base + index * 9973, ratio, model, None,
-                      _context(project_id), [[[text]]], None, None, None,
-                      _client_uuid(), _client_uuid()])
+        image_inputs = []
+        if base_media_id:
+            image_inputs.append(_base_image(base_media_id))
+        image_inputs.extend(
+            _reference(mid) for mid in (ref_media_ids or []) if mid != base_media_id
+        )
+        items.append([None, None, image_inputs or None, base + index * 9973, ratio,
+                      resolved_model, None, _context(project_id), [[[text]]],
+                      None, None, None, _client_uuid(), _client_uuid()])
     return build_envelope(RPC_GEN_IMAGE, [None, items, 1, _context(project_id),
                                           [_client_uuid()]])
 
 
+def image_upscale_request(media_id: str, resolution: str = "2K") -> str:
+    """Build the current FlowService.UpsampleImage request (RPC SPrCad)."""
+    key = str(resolution).strip().upper()
+    if key.startswith("UPSAMPLE_IMAGE_RESOLUTION_"):
+        key = key.removeprefix("UPSAMPLE_IMAGE_RESOLUTION_")
+    try:
+        code = IMAGE_UPSCALE_RESOLUTIONS[key]
+    except KeyError:
+        raise ValueError("image upscale resolution must be 2K or 4K") from None
+    return build_envelope(RPC_UPSCALE_IMAGE, [media_id, code, _context(None)])
+
+
 def video_request(prompt: str, project_id: str, source_media_id: str,
                   crop: Optional[list] = None,
                   aspect: Any = VIDEO_ASPECT_LANDSCAPE,
@@ -429,6 +482,14 @@ def read_uploaded_media_id(payload: Any) -> str:
     return media_id
 
 
+def read_upscaled_image(payload: Any) -> str:
+    """Return the base64 image body from FlowService.UpsampleImage."""
+    encoded = payload[1] if isinstance(payload, list) and len(payload) > 1 else None
+    if not isinstance(encoded, str) or len(encoded) < 100:
+        raise FlowBatchError("image upscale response carried no encoded image")
+    return encoded
+
+
 def read_operation(payload: Any) -> Operation:
     """`[null, 50, [[opId, projectId, sceneId, status, …]]]`.