feat(muse2api): send entity reference images with scene images

README.md

@@ -405,9 +405,14 @@ gateway may run on another host), and downloads every result into
 so `/fk-refresh-urls` is never needed for it.
 
 Supports image and i2v; no edit, upscale or TTS, and no end frame or r2v unless
-`MUSE2API_ALLOW_DEGRADED=1`. Entity reference images are not sent (muse.ai takes
-none), so character consistency across scenes is weaker than on Flow.
-Troubleshooting lives in `/fk-provider` and `/fk-doctor`.
+`MUSE2API_ALLOW_DEGRADED=1`. Entity reference images are sent with each scene
+image (up to four, inline as data URLs; extra refs are dropped and logged), and
+a ref that cannot be read fails the job rather than rendering the scene without
+its character.
+
+If the gateway runs with billing enabled, every call is charged to the wallet of
+the key in `MUSE2API_KEY`. Use the admin key or a key marked `unlimited` for
+Flow Kit. Troubleshooting lives in `/fk-provider` and `/fk-doctor`.
 
 ### Post-production (provider-neutral)
 

agent/sdk/services/muse2api_provider.py

@@ -16,8 +16,8 @@
 
 What muse.ai can and cannot do shapes the capabilities:
 
-* image         — text-to-image. Reference images are not an input the
-                  gateway accepts, so entity refs are dropped (logged).
+* image         — text-to-image, with up to four entity refs sent inline as
+                  data URLs (extra refs are dropped, logged).
 * video (i2v)   — first frame + prompt.
 * chained video — no end-frame input: fails unless MUSE2API_ALLOW_DEGRADED=1,
                   which drops the end frame and renders plain i2v.
@@ -34,6 +34,8 @@
 from pathlib import Path
 from urllib.parse import unquote, urlparse
 
+import httpx
+
 from agent import config
 from agent.sdk.services.assistant_provider import _image_result, _video_result
 from agent.sdk.services.provider_base import (
@@ -44,7 +46,12 @@
     ProviderCapabilities,
     ProviderJob,
 )
-from agent.services.muse2api_client import Muse2APIClient, Muse2APIError, to_data_url
+from agent.services.muse2api_client import (
+    MAX_IMAGE_REFS,
+    Muse2APIClient,
+    Muse2APIError,
+    to_data_url,
+)
 from agent.utils.paths import file_url_to_path
 
 logger = logging.getLogger(__name__)
@@ -113,8 +120,8 @@ async def run(self, job: ProviderJob) -> dict:
         except Muse2APIError as e:
             logger.error("muse2api: job %s (%s) failed: %s", job.job_id[:12], job.kind, e)
             return {"error": str(e)}
-        except (OSError, ValueError) as e:
-            # Unreadable input image, full disk, bad base64 …
+        except (OSError, ValueError, httpx.HTTPError) as e:
+            # Unreadable or unreachable input image, full disk, bad base64 …
             logger.error("muse2api: job %s (%s) failed: %s", job.job_id[:12], job.kind, e)
             return {"error": f"muse2api: {e}"}
         return {"error": f"muse2api provider does not support job kind '{job.kind}'"}
@@ -124,11 +131,17 @@ async def _run_image(self, job: ProviderJob) -> dict:
         composition = (job.extra or {}).get("composition")
         if composition:
             prompt = f"{prompt}\n\nComposition: {composition}"
-        if job.reference_urls:
-            logger.info("muse2api: job %s — %d reference image(s) dropped, the "
-                        "gateway takes no image refs", job.job_id[:12], len(job.reference_urls))
+        refs = list(dict.fromkeys(job.reference_urls or []))
+        if len(refs) > MAX_IMAGE_REFS:
+            logger.warning("muse2api: job %s — %d reference images, the gateway takes "
+                           "%d; dropping the rest", job.job_id[:12], len(refs), MAX_IMAGE_REFS)
+            refs = refs[:MAX_IMAGE_REFS]
+        # An unreadable ref fails the job (OSError/ValueError in run()) rather
+        # than quietly rendering a scene without its character.
+        images = [await to_data_url(ref) for ref in refs]
         data = await self.client.generate_image(
-            prompt, model=config.MUSE2API_IMAGE_MODEL, size=_ASPECT.get(job.orientation))
+            prompt, model=config.MUSE2API_IMAGE_MODEL, size=_ASPECT.get(job.orientation),
+            images=images or None)
         url = await self._save(job.job_id, data, _image_ext(data))
         mid = str(uuid.uuid4())
         logger.info("muse2api: image %s ready (media_id=%s)", job.job_id[:12], mid[:8])

agent/services/muse2api_client.py

@@ -3,7 +3,7 @@
 muse2api exposes the muse.ai web app as an OpenAI-compatible API. Flow Kit
 only needs the media half of it:
 
-    POST /v1/images/generations   prompt -> image(s)   (synchronous)
+    POST /v1/images/generations   prompt (+ up to 4 refs) -> image   (synchronous)
     POST /v1/videos               prompt (+first frame) -> task
     GET  /v1/videos/{task_id}     task status; result.url when succeeded
     GET  /v1/media/{name}         download generated media (public)
@@ -38,6 +38,8 @@
 logger = logging.getLogger(__name__)
 
 TERMINAL_OK = "succeeded"
+#: The gateway's ``image`` field on /v1/images/generations takes up to four refs.
+MAX_IMAGE_REFS = 4
 TERMINAL_FAILED = ("failed", "cancelled")
 
 
@@ -142,11 +144,18 @@ async def readyz(self) -> dict:
 
     # ── Images ──
 
-    async def generate_image(self, prompt: str, *, model: str, size: str | None) -> bytes:
-        """Generate one image and return its bytes (b64_json: no second download)."""
+    async def generate_image(self, prompt: str, *, model: str, size: str | None,
+                             images: list[str] | None = None) -> bytes:
+        """Generate one image and return its bytes (b64_json: no second download).
+
+        ``images`` are reference images as data URLs; the gateway attaches
+        them to the muse.ai chat and takes at most :data:`MAX_IMAGE_REFS`.
+        """
         body = {"prompt": prompt, "model": model, "n": 1, "response_format": "b64_json"}
         if size:
             body["size"] = size
+        if images:
+            body["image"] = images[:MAX_IMAGE_REFS]
         data = await self._request("POST", "/v1/images/generations",
                                    json=body, timeout=self.timeout_s)
         items = data.get("data") or []

skills/fk-provider.md

@@ -52,7 +52,8 @@ audio — it is rejected with 400.)
 - `assistant`: image, edit_image, i2v, r2v, **audio (TTS)**. No upscale.
 - `muse2api`: image, i2v. No edit_image, upscale or audio. No end frame and no
   r2v unless `MUSE2API_ALLOW_DEGRADED=1` (then both render as plain i2v).
-  Entity reference images are not sent — muse.ai takes none.
+  Entity reference images are sent with scene images (up to four; an
+  unreadable ref fails the job).
 
 Asking a provider for a kind it doesn't support returns 400 with a message
 naming the right backend to use instead.

tests/unit/test_muse2api_provider.py

@@ -88,9 +88,9 @@ def start_image(tmp_path):
     return "file://" + str(p)
 
 
-async def test_image_is_saved_locally_with_minted_uuid(provider, gateway):
+async def test_image_is_saved_locally_with_minted_uuid(provider, gateway, start_image):
     job = ProviderJob(job_id="img1", kind=KIND_IMAGE, prompt="a fox",
-                      orientation="HORIZONTAL", reference_urls=("file:///x.png",),
+                      orientation="HORIZONTAL", reference_urls=(start_image,),
                       extra={"composition": "wide establishing shot"})
     result = await provider.run(job)
 
@@ -106,6 +106,48 @@ async def test_image_is_saved_locally_with_minted_uuid(provider, gateway):
     assert body["model"] == "muse-image"
     assert body["response_format"] == "b64_json"
     assert "wide establishing shot" in body["prompt"]
+    assert body["image"] == ["data:image/png;base64," + base64.b64encode(PNG).decode()]
+
+
+async def test_image_without_refs_sends_no_image_field(provider, gateway):
+    job = ProviderJob(job_id="img2", kind=KIND_IMAGE, prompt="a fox", orientation="VERTICAL")
+    assert parse_result(await provider.run(job), "GENERATE_IMAGE").success
+    assert "image" not in gateway.calls[0][2]
+
+
+async def test_image_refs_are_deduped_and_capped_at_four(provider, gateway, tmp_path):
+    refs = []
+    for i in range(6):
+        p = tmp_path / f"ref{i}.png"
+        p.write_bytes(PNG + bytes([i]))
+        refs.append("file://" + str(p))
+    job = ProviderJob(job_id="img3", kind=KIND_IMAGE, prompt="a fox",
+                      orientation="VERTICAL", reference_urls=(refs[0], refs[0], *refs[1:]))
+    assert parse_result(await provider.run(job), "GENERATE_IMAGE").success
+
+    sent = gateway.calls[0][2]["image"]
+    assert len(sent) == 4
+    assert sent == ["data:image/png;base64," + base64.b64encode(PNG + bytes([i])).decode()
+                    for i in range(4)]
+
+
+async def test_unreadable_ref_fails_the_job_instead_of_dropping_the_character(provider, gateway):
+    job = ProviderJob(job_id="img4", kind=KIND_IMAGE, prompt="a fox",
+                      orientation="VERTICAL", reference_urls=("file:///no/such/ref.png",))
+    result = await provider.run(job)
+    assert "error" in result and "ref.png" in result["error"]
+    assert gateway.calls == []
+
+
+async def test_unreachable_http_ref_fails_the_job(provider, gateway, monkeypatch):
+    async def boom(ref, **_):
+        raise httpx.ConnectError("no route", request=httpx.Request("GET", ref))
+    monkeypatch.setattr(mp, "to_data_url", boom)
+    job = ProviderJob(job_id="img5", kind=KIND_IMAGE, prompt="a fox",
+                      orientation="VERTICAL", reference_urls=("https://cdn.test/ref.png",))
+    result = await provider.run(job)
+    assert result["error"].startswith("muse2api: ")
+    assert gateway.calls == []
 
 
 async def test_video_sends_first_frame_inline_and_polls(provider, gateway, start_image):