feat(invoke): "Use as Ref Image" button + resilient auth fallback

photomap/backend/metadata_modules/invoke_formatter.py

@@ -61,19 +61,37 @@
     "</svg>"
 )
 
+# Photo-frame icon — a rectangle with a small sun and a mountain inside,
+# drawn in the same stroked style as the remix icon so the three buttons
+# share a consistent visual language.
+_USE_REF_SVG = (
+    '<svg width="18" height="18" viewBox="0 0 24 24" fill="none" '
+    'stroke="currentColor" stroke-width="2.2" stroke-linecap="round" '
+    'stroke-linejoin="round" aria-hidden="true">'
+    '<rect x="3" y="3" width="18" height="18" rx="2" ry="2"/>'
+    '<circle cx="8.5" cy="8.5" r="1.5"/>'
+    '<polyline points="21 15 16 10 5 21"/>'
+    "</svg>"
+)
+
 
 def _recall_buttons_html() -> str:
-    """Render the recall / remix button group shown at the bottom of the drawer."""
+    """Render the recall / remix / use-ref button group shown at the bottom of the drawer."""
     return (
         '<div class="invoke-recall-controls" data-invoke-recall="1">'
+        '<button type="button" class="invoke-recall-btn" data-recall-mode="remix" '
+        'title="Remix (recall parameters without the seed) to InvokeAI">'
+        f'{_REMIX_SVG}<span class="invoke-recall-label">Remix</span>'
+        '<span class="invoke-recall-status" aria-live="polite"></span>'
+        "</button>"
         '<button type="button" class="invoke-recall-btn" data-recall-mode="recall" '
         'title="Recall parameters (including seed) to InvokeAI">'
         f'{_RECALL_SVG}<span class="invoke-recall-label">Recall</span>'
         '<span class="invoke-recall-status" aria-live="polite"></span>'
         "</button>"
-        '<button type="button" class="invoke-recall-btn" data-recall-mode="remix" '
-        'title="Remix (recall parameters without the seed) to InvokeAI">'
-        f'{_REMIX_SVG}<span class="invoke-recall-label">Remix</span>'
+        '<button type="button" class="invoke-recall-btn" data-recall-mode="use_ref" '
+        'title="Upload this image to InvokeAI and use it as a reference image">'
+        f'{_USE_REF_SVG}<span class="invoke-recall-label">Use as Ref Image</span>'
         '<span class="invoke-recall-status" aria-live="polite"></span>'
         "</button>"
         "</div>"

photomap/backend/routers/invoke.py

@@ -8,17 +8,27 @@
   tuple, load the image metadata, build a recall payload from it, and proxy
   it to the configured InvokeAI backend's ``/api/v1/recall/{queue_id}``
   endpoint.
+* ``POST /invokeai/use_ref_image`` — upload the selected image to InvokeAI
+  and then call the same recall endpoint with the uploaded image as a
+  reference image parameter, so the next generation uses it for visual
+  guidance.
 
 When the configured InvokeAI backend runs in multi-user mode, the
 ``username`` / ``password`` fields are used to obtain a JWT bearer token
 via ``/api/v1/auth/login``.  The token is cached in-process and
-automatically refreshed on 401.
+automatically refreshed on 401.  If the backend has since been
+reconfigured into single-user mode it will reject a token-bearing
+request with a 403 — that causes the cached token to be discarded and
+the call to be retried anonymously without requiring a restart.
 """
 
 from __future__ import annotations
 
 import logging
+import mimetypes
 import time
+from collections.abc import Awaitable, Callable
+from pathlib import Path
 
 import httpx
 from fastapi import APIRouter, HTTPException
@@ -46,31 +56,31 @@
 _token_username: str | None = None
 
 
-async def _get_auth_headers(base_url: str, username: str | None, password: str | None) -> dict[str, str]:
-    """Return an ``Authorization`` header dict, or empty dict for anonymous access.
+def _cached_auth_headers(base_url: str, username: str | None) -> dict[str, str]:
+    """Return ``{"Authorization": "Bearer ..."}`` if we still hold a valid
+    cached token for this ``(base_url, username)`` pair, else ``{}``.
 
-    If a valid cached token exists it is reused.  If no token is cached the
-    caller should first try the request without credentials — the InvokeAI
-    backend running in single-user mode will accept it.  When the caller
-    receives a 401 it should call ``_invalidate_token_cache`` and call this
-    function again; the second call will perform the login.
+    This never talks to the network.  Deliberate: the first attempt at any
+    request always uses whatever auth we already have (or none), so that a
+    backend that has since been reconfigured into single-user mode is given
+    a chance to accept the call anonymously.
     """
-    global _cached_token, _token_expires_at, _token_base_url, _token_username  # noqa: PLW0603
-
-    # Reuse cached token if still valid for this backend+user
     if (
         _cached_token
         and time.monotonic() < _token_expires_at
         and _token_base_url == base_url
         and _token_username == username
     ):
         return {"Authorization": f"Bearer {_cached_token}"}
+    return {}
+
 
-    # No credentials configured — anonymous access
-    if not username or not password:
-        return {}
+async def _login(base_url: str, username: str, password: str) -> dict[str, str]:
+    """Exchange ``username``/``password`` for a JWT via the InvokeAI auth
+    endpoint, cache the token, and return the ``Authorization`` header.
+    """
+    global _cached_token, _token_expires_at, _token_base_url, _token_username  # noqa: PLW0603
 
-    # No cached token yet — try to obtain one
     login_url = f"{base_url.rstrip('/')}/api/v1/auth/login"
     try:
         async with httpx.AsyncClient(timeout=_HTTP_TIMEOUT) as client:
@@ -99,9 +109,50 @@ async def _get_auth_headers(base_url: str, username: str | None, password: str |
 
 def _invalidate_token_cache() -> None:
     """Clear the cached token so the next request re-authenticates."""
-    global _cached_token, _token_expires_at  # noqa: PLW0603
+    global _cached_token, _token_expires_at, _token_base_url, _token_username  # noqa: PLW0603
     _cached_token = None
     _token_expires_at = 0.0
+    _token_base_url = None
+    _token_username = None
+
+
+async def _post_with_auth_fallback(
+    base_url: str,
+    username: str | None,
+    password: str | None,
+    request_fn: Callable[[dict[str, str]], Awaitable[httpx.Response]],
+) -> httpx.Response:
+    """Perform an InvokeAI request with graceful handling of auth transitions.
+
+    ``request_fn`` is an async callable that takes a headers dict and
+    performs the HTTP call — using a factory lets the caller re-open file
+    streams (needed for multipart uploads) on a retry.
+
+    Three-step flow:
+
+    1. First attempt uses whatever token we have cached (or no auth at all).
+       A freshly-restarted single-user backend then accepts the call even
+       if credentials are stored in PhotoMap.
+    2. If the first attempt returns **401**, the backend demands
+       authentication: if credentials are configured we log in, cache a
+       fresh token, and retry.
+    3. If the first attempt was made *with* a token and returns **403**
+       (most commonly "Multiuser mode is disabled. Authentication is not
+       required…"), the backend was reconfigured to single-user mode — we
+       invalidate the cached token and retry anonymously.
+    """
+    auth_headers = _cached_auth_headers(base_url, username)
+    response = await request_fn(auth_headers)
+
+    if response.status_code == 401 and username and password:
+        _invalidate_token_cache()
+        auth_headers = await _login(base_url, username, password)
+        response = await request_fn(auth_headers)
+    elif response.status_code == 403 and auth_headers:
+        _invalidate_token_cache()
+        response = await request_fn({})
+
+    return response
 
 
 class InvokeAISettings(BaseModel):
@@ -124,6 +175,14 @@ class RecallRequest(BaseModel):
     queue_id: str = Field("default", description="InvokeAI queue id to target")
 
 
+class UseRefImageRequest(BaseModel):
+    """Payload posted by the drawer's "Use Ref Image" button."""
+
+    album_key: str = Field(..., description="Album containing the image")
+    index: int = Field(..., ge=0, description="Image index within the album")
+    queue_id: str = Field("default", description="InvokeAI queue id to target")
+
+
 @invoke_router.get("/config")
 async def get_invokeai_config() -> dict:
     """Return the persisted InvokeAI connection settings.
@@ -175,6 +234,17 @@ def _load_raw_metadata(album_key: str, index: int) -> dict:
     return entry if isinstance(entry, dict) else {}
 
 
+def _load_image_path(album_key: str, index: int) -> Path:
+    embeddings = get_embeddings_for_album(album_key)
+    if not embeddings:
+        raise HTTPException(status_code=404, detail="Album not found")
+    indexes = embeddings.indexes
+    filenames = indexes["sorted_filenames"]
+    if index < 0 or index >= len(filenames):
+        raise HTTPException(status_code=404, detail="Index out of range")
+    return Path(str(filenames[index]))
+
+
 def _build_recall_payload(raw_metadata: dict, include_seed: bool) -> dict:
     if not raw_metadata:
         raise HTTPException(
@@ -222,18 +292,17 @@ async def recall_parameters(request: RecallRequest) -> dict:
     username = settings["username"]
     password = settings["password"]
 
-    # Try without credentials first (single-user mode). If the backend
-    # requires authentication (401), obtain a token and retry.
-    auth_headers = await _get_auth_headers(base_url, username, password)
-
     try:
         async with httpx.AsyncClient(timeout=_HTTP_TIMEOUT) as client:
-            response = await client.post(url, json=payload, params={"strict": "true"}, headers=auth_headers)
 
-            if response.status_code == 401 and username and password:
-                _invalidate_token_cache()
-                auth_headers = await _get_auth_headers(base_url, username, password)
-                response = await client.post(url, json=payload, params={"strict": "true"}, headers=auth_headers)
+            async def _do(headers: dict[str, str]) -> httpx.Response:
+                return await client.post(
+                    url, json=payload, params={"strict": "true"}, headers=headers
+                )
+
+            response = await _post_with_auth_fallback(
+                base_url, username, password, _do
+            )
     except httpx.RequestError as exc:
         logger.warning("InvokeAI recall request failed: %s", exc)
         raise HTTPException(
@@ -262,3 +331,138 @@ async def recall_parameters(request: RecallRequest) -> dict:
         "sent": payload,
         "response": remote,
     }
+
+
+async def _upload_image_to_invokeai(
+    client: httpx.AsyncClient,
+    base_url: str,
+    image_path: Path,
+    username: str | None,
+    password: str | None,
+) -> str:
+    """Upload ``image_path`` to InvokeAI and return the assigned ``image_name``.
+
+    Auth transitions (anonymous ↔ token) are handled transparently by
+    ``_post_with_auth_fallback``; the multipart stream is re-opened on each
+    retry since the previous one will have been consumed.
+    """
+    upload_url = f"{base_url.rstrip('/')}/api/v1/images/upload"
+    mime_type = mimetypes.guess_type(image_path.name)[0] or "image/png"
+    params = {"image_category": "user", "is_intermediate": "false"}
+
+    async def _do(headers: dict[str, str]) -> httpx.Response:
+        with image_path.open("rb") as fh:
+            files = {"file": (image_path.name, fh, mime_type)}
+            return await client.post(
+                upload_url, files=files, params=params, headers=headers
+            )
+
+    upload_resp = await _post_with_auth_fallback(base_url, username, password, _do)
+
+    if upload_resp.status_code >= 400:
+        raise HTTPException(
+            status_code=502,
+            detail=(
+                f"InvokeAI image upload returned {upload_resp.status_code}: "
+                f"{upload_resp.text[:200]}"
+            ),
+        )
+
+    try:
+        body = upload_resp.json()
+    except ValueError as exc:
+        raise HTTPException(
+            status_code=502,
+            detail="InvokeAI image upload returned a non-JSON response",
+        ) from exc
+
+    image_name = body.get("image_name")
+    if not image_name:
+        raise HTTPException(
+            status_code=502,
+            detail="InvokeAI image upload response did not include image_name",
+        )
+    return image_name
+
+
+@invoke_router.post("/use_ref_image")
+async def use_ref_image(request: UseRefImageRequest) -> dict:
+    """Upload the selected image to InvokeAI and set it as a reference image.
+
+    Implements the two-step flow: first ``POST /api/v1/images/upload`` so
+    InvokeAI knows the file, then ``POST /api/v1/recall/{queue_id}`` with the
+    returned ``image_name`` in ``reference_images`` so the next generation
+    picks it up as visual guidance.
+    """
+    settings = config_manager.get_invokeai_settings()
+    base_url = settings["url"]
+    if not base_url:
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                "InvokeAI backend URL is not configured. Set it in the "
+                "PhotoMap settings panel."
+            ),
+        )
+
+    image_path = _load_image_path(request.album_key, request.index)
+    if not image_path.is_file():
+        raise HTTPException(
+            status_code=404, detail=f"Image file not found on disk: {image_path.name}"
+        )
+
+    username = settings["username"]
+    password = settings["password"]
+
+    recall_url = f"{base_url.rstrip('/')}/api/v1/recall/{request.queue_id}"
+
+    try:
+        async with httpx.AsyncClient(timeout=_HTTP_TIMEOUT) as client:
+            image_name = await _upload_image_to_invokeai(
+                client, base_url, image_path, username, password
+            )
+
+            # Deliberately omit ``strict=true`` so that the recall only
+            # *adds* the reference image to whatever the user already has
+            # set up in InvokeAI rather than resetting every other
+            # parameter back to defaults.
+            payload = {"reference_images": [{"image_name": image_name}]}
+
+            async def _do_recall(headers: dict[str, str]) -> httpx.Response:
+                return await client.post(recall_url, json=payload, headers=headers)
+
+            response = await _post_with_auth_fallback(
+                base_url, username, password, _do_recall
+            )
+    except httpx.RequestError as exc:
+        logger.warning("InvokeAI use_ref_image request failed: %s", exc)
+        raise HTTPException(
+            status_code=502,
+            detail=f"Could not reach InvokeAI backend at {base_url}: {exc}",
+        ) from exc
+
+    if response.status_code >= 400:
+        logger.warning(
+            "InvokeAI use_ref_image recall returned %s: %s",
+            response.status_code,
+            response.text,
+        )
+        raise HTTPException(
+            status_code=502,
+            detail=(
+                f"InvokeAI backend returned {response.status_code}: "
+                f"{response.text[:200]}"
+            ),
+        )
+
+    try:
+        remote = response.json()
+    except ValueError:
+        remote = {"raw": response.text}
+
+    return {
+        "success": True,
+        "sent": payload,
+        "uploaded_image_name": image_name,
+        "response": remote,
+    }