fix(runners): transfer source images and masks to ComfyUI before execution (fixes #105)

backend/app/runners/comfy_runner.py

@@ -9,7 +9,9 @@
 import asyncio
 import json
 import logging
+import mimetypes
 import uuid
+from pathlib import Path
 from typing import Any, AsyncGenerator, Dict, List, Optional
 import httpx
 
@@ -205,6 +207,58 @@ async def interrupt(self) -> bool:
             logger.debug("ComfyUI interrupt failed or not reachable: %s", err)
             return False
 
+    async def upload_image(self, image_path: str, subfolder: str = "", overwrite: bool = False) -> Dict[str, Any]:
+        """
+        Upload an image to ComfyUI's input directory via POST /upload/image.
+
+        Returns the engine's filename and subfolder references for use in LoadImage nodes.
+        Raises RuntimeError if upload fails.
+        """
+        file_path = Path(image_path)
+        if not await asyncio.to_thread(file_path.is_file):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+
+        try:
+            client = self._get_client()
+
+            content = await asyncio.to_thread(file_path.read_bytes)
+            mime_type = mimetypes.guess_type(file_path.name)[0] or "application/octet-stream"
+            files = {"image": (file_path.name, content, mime_type)}
+            data = {"subfolder": subfolder, "overwrite": str(overwrite).lower(), "type": "input"}
+            resp = await client.post(f"{self.base_url}/upload/image", files=files, data=data)
+            resp.raise_for_status()
+
+            result = resp.json()
+            if (
+                not isinstance(result, dict)
+                or not isinstance(result.get("name"), str)
+                or not result["name"].strip()
+                or not isinstance(result.get("subfolder", subfolder), str)
+                or result.get("type", "input") != "input"
+            ):
+                raise RuntimeError(f"ComfyUI upload returned unexpected response: {result}")
+
+            return {
+                "name": result["name"],
+                "subfolder": result.get("subfolder", subfolder),
+                "type": result.get("type", "input"),
+            }
+        except httpx.HTTPStatusError as http_err:
+            raise RuntimeError(f"ComfyUI upload failed with status {http_err.response.status_code}: {http_err}") from http_err
+        except Exception as err:
+            raise RuntimeError(f"Failed to upload image to ComfyUI: {err}") from err
+
+    async def upload_mask(self, mask_path: str, subfolder: str = "", overwrite: bool = False) -> Dict[str, Any]:
+        """
+        Upload an independent mask file without changing its pixel/alpha data.
+
+        ComfyUI's /upload/mask edits the alpha channel of an existing image and
+        requires original_ref. Our graphs load a separate mask asset, so both
+        source images and masks use /upload/image. Mask semantics are handled
+        separately by the inpainting adapter.
+        """
+        return await self.upload_image(mask_path, subfolder=subfolder, overwrite=overwrite)
+
 
 # Global default ComfyUI client instance
 comfy_client = ComfyUIClient()

backend/app/runners/creative_runner.py

@@ -452,7 +452,33 @@ async def _run_comfy(
         input_file: Optional[Path],
         mask_file: Optional[Path],
     ) -> Dict[str, Any]:
-        """Compile and submit ComfyUI macro graph."""
+        """Compile and submit ComfyUI macro graph, uploading required assets first."""
+        # Upload source image and mask to ComfyUI's input directory if needed
+        uploaded_image_name: Optional[str] = None
+        uploaded_mask_name: Optional[str] = None
+
+        if input_file and input_file.is_file():
+            try:
+                upload_result = await comfy_client.upload_image(str(input_file), subfolder="berry_assets")
+                uploaded_image_name = upload_result["name"]
+                # Include subfolder in the filename if ComfyUI expects it
+                if upload_result.get("subfolder"):
+                    uploaded_image_name = f"{upload_result['subfolder']}/{uploaded_image_name}"
+                logger.debug(f"Uploaded source image to ComfyUI: {uploaded_image_name}")
+            except Exception as upload_err:
+                raise RuntimeError(f"Failed to transfer source image to ComfyUI: {upload_err}") from upload_err
+
+        if mask_file and mask_file.is_file():
+            try:
+                mask_result = await comfy_client.upload_mask(str(mask_file), subfolder="berry_assets")
+                uploaded_mask_name = mask_result["name"]
+                if mask_result.get("subfolder"):
+                    uploaded_mask_name = f"{mask_result['subfolder']}/{uploaded_mask_name}"
+                logger.debug(f"Uploaded mask to ComfyUI: {uploaded_mask_name}")
+            except Exception as upload_err:
+                raise RuntimeError(f"Failed to transfer mask to ComfyUI: {upload_err}") from upload_err
+
+        # Build workflow graphs using uploaded filenames
         if req.action == CreativeActionType.TXT2IMG:
             prompt_graph = build_comfy_txt2img_graph(
                 prompt=req.prompt,
@@ -464,11 +490,11 @@ async def _run_comfy(
                 seed=req.seed,
             )
         elif req.action == CreativeActionType.IMG2IMG:
-            if not input_file:
-                raise ValueError("Source image required for img2img")
+            if not uploaded_image_name:
+                raise ValueError("Source image upload failed for img2img")
             prompt_graph = build_comfy_img2img_graph(
                 prompt=req.prompt,
-                image_filename=input_file.name,
+                image_filename=uploaded_image_name,
                 negative_prompt=req.negative_prompt,
                 checkpoint=req.model,
                 steps=req.steps,
@@ -477,12 +503,12 @@ async def _run_comfy(
                 seed=req.seed,
             )
         elif req.action == CreativeActionType.INPAINT:
-            if not input_file or not mask_file:
-                raise ValueError("Source image and mask required for inpaint")
+            if not uploaded_image_name or not uploaded_mask_name:
+                raise ValueError("Source image and mask upload required for inpaint")
             prompt_graph = build_comfy_inpaint_graph(
                 prompt=req.prompt,
-                image_filename=input_file.name,
-                mask_filename=mask_file.name,
+                image_filename=uploaded_image_name,
+                mask_filename=uploaded_mask_name,
                 negative_prompt=req.negative_prompt,
                 checkpoint=req.model,
                 steps=req.steps,
@@ -491,17 +517,17 @@ async def _run_comfy(
                 seed=req.seed,
             )
         elif req.action == CreativeActionType.UPSCALE:
-            if not input_file:
-                raise ValueError("Source image required for upscale")
+            if not uploaded_image_name:
+                raise ValueError("Source image upload failed for upscale")
             prompt_graph = build_comfy_upscale_graph(
-                image_filename=input_file.name,
+                image_filename=uploaded_image_name,
                 upscaler_model=req.upscaler_name or "RealESRGAN_x4plus.pth",
             )
         elif req.action == CreativeActionType.IMG2VIDEO:
-            if not input_file:
-                raise ValueError("Source image required for ComfyUI img2video")
+            if not uploaded_image_name:
+                raise ValueError("Source image upload failed for ComfyUI img2video")
             prompt_graph = build_comfy_img2video_graph(
-                image_filename=input_file.name,
+                image_filename=uploaded_image_name,
                 checkpoint=req.model if "svd" in req.model.lower() else "svd_xt.safetensors",
                 width=req.width,
                 height=req.height,

backend/tests/test_comfy_asset_transfer_contract.py

@@ -0,0 +1,108 @@
+"""Verify upload ordering, engine references, failure handling, and cache reuse."""
+
+import json
+import struct
+import zlib
+from pathlib import Path
+from typing import Any
+from unittest.mock import AsyncMock, Mock, patch
+
+import httpx
+import pytest
+
+from app.runners.comfy_runner import ComfyUIClient
+from app.runners.creative_runner import CreativeRunner
+from app.schemas.creative import CreativeActionRequest, CreativeActionType
+from app.storage.asset_store import AssetRecord
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("action,fail_mask", [
+    ("img2img", False), ("inpaint", False), ("upscale", False),
+    ("img2video", False), ("inpaint", True),
+])
+async def test_upload_then_queue_and_cache_reuse(
+    tmp_path: Path, action: str, fail_mask: bool,
+) -> None:
+    originals: dict[Path, bytes] = {}
+    records: dict[str, AssetRecord] = {}
+    for name in ("source", "mask"):
+        path = tmp_path / f"{name}.png"
+        def chunk(kind: bytes, data: bytes) -> bytes:
+            return struct.pack(">I", len(data)) + kind + data + struct.pack(">I", zlib.crc32(kind + data))
+        path.write_bytes(
+            b"\x89PNG\r\n\x1a\n"
+            + chunk(b"IHDR", struct.pack(">IIBBBBB", 8, 8, 8, 6, 0, 0, 0))
+            + chunk(b"IDAT", zlib.compress((b"\x00" + bytes((255, 255, 255, 128)) * 8) * 8))
+            + chunk(b"IEND", b"")
+        )
+        originals[path] = path.read_bytes()
+        records[name] = AssetRecord(
+            id=name, filename=path.name, file_path=path.name, media_type="image",
+            content_hash=f"{name}-hash", byte_size=path.stat().st_size,
+            created_at="2026-10-06T00:00:00Z",
+        )
+    events: list[str] = []
+    cached: dict[str, Any] = {}
+
+    async def handle(request: httpx.Request) -> httpx.Response:
+        if request.url.path == "/upload/image":
+            body = await request.aread()
+            name = "mask" if b'filename="mask.png"' in body else "source"
+            events.append(name)
+            if name == "mask" and fail_mask:
+                return httpx.Response(503)
+            return httpx.Response(200, json={
+                "name": f"{name} (1).png", "subfolder": "engine/folder", "type": "input",
+            })
+        assert request.url.path == "/prompt"
+        events.append("queue")
+        graph = json.loads(await request.aread())["prompt"]
+        loaded = {
+            node["inputs"]["image"] for node in graph.values()
+            if node["class_type"] == "LoadImage"
+        }
+        expected = {"engine/folder/source (1).png"}
+        if action == "inpaint":
+            expected.add("engine/folder/mask (1).png")
+        assert loaded == expected
+        assert events == (["source", "mask", "queue"] if action == "inpaint" else ["source", "queue"])
+        return httpx.Response(200, json={"prompt_id": "test-prompt"})
+
+    async def save_cache(key: str, value: dict[str, Any]) -> None:
+        cached[key] = value
+
+    client = ComfyUIClient()
+    client.poll_history_outputs = AsyncMock(return_value=[{"filename": "output.png", "type": "output"}])
+    req = CreativeActionRequest(
+        action=CreativeActionType(action), engine_id="comfyui", input_image_id="source",
+        mask_image_id="mask" if action == "inpaint" else None, seed=42,
+    )
+    with (
+        patch("app.runners.creative_runner.comfy_client", client),
+        patch("app.runners.creative_runner.asset_store") as store,
+        patch("app.runners.creative_runner.cache_store") as cache,
+    ):
+        store.get_asset = AsyncMock(side_effect=records.get)
+        store.get_absolute_path = Mock(side_effect=lambda rec: tmp_path / rec.filename)
+        store.save_image_from_url = AsyncMock(return_value=Mock(id="output"))
+        store.save_media_from_url = AsyncMock(return_value=Mock(id="output"))
+        cache.get_async = AsyncMock(side_effect=cached.get)
+        cache.set_async = AsyncMock(side_effect=save_cache)
+        async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as transport:
+            client._client = transport
+            first = await CreativeRunner().execute(req)
+            if fail_mask:
+                assert not first.success
+                assert "Failed to transfer mask" in first.error_message
+                assert events == ["source", "mask"]
+                cache.set_async.assert_not_awaited()
+                client.poll_history_outputs.assert_not_awaited()
+            else:
+                assert first.success, first.error_message
+                second = await CreativeRunner().execute(req)
+                assert second.success and second.is_cached
+                assert second.asset_id == first.asset_id
+                assert events.count("source") == events.count("queue") == 1
+                client.poll_history_outputs.assert_awaited_once()
+    assert {path: path.read_bytes() for path in originals} == originals

backend/tests/test_comfy_upload.py

@@ -0,0 +1,120 @@
+"""Exercise real multipart serialization for ComfyUI asset transfers (#105)."""
+
+from email import policy
+from email.parser import BytesParser
+from pathlib import Path
+from typing import Any
+
+import httpx
+import pytest
+
+from app.runners.comfy_runner import ComfyUIClient
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("host,port", [("127.0.0.1", 8188), ("engine.example", 9000)])
+@pytest.mark.parametrize("method,suffix,mime", [
+    ("upload_image", ".jpg", "image/jpeg"),
+    ("upload_mask", ".png", "image/png"),
+])
+async def test_upload_preserves_bytes_and_uses_engine_reference(
+    tmp_path: Path, host: str, port: int, method: str, suffix: str, mime: str,
+) -> None:
+    source = tmp_path / f"original{suffix}"
+    content = b"\x00\xff\x80source-or-mask-bytes\r\n"
+    source.write_bytes(content)
+    requests: list[httpx.Request] = []
+
+    async def handle(request: httpx.Request) -> httpx.Response:
+        requests.append(request)
+        assert str(request.url) == f"http://{host}:{port}/upload/image"
+        body = await request.aread()
+        message = BytesParser(policy=policy.default).parsebytes(
+            f"Content-Type: {request.headers['content-type']}\r\n\r\n".encode() + body
+        )
+        parts = {
+            part.get_param("name", header="content-disposition"): part
+            for part in message.iter_parts()
+        }
+        assert parts["image"].get_filename() == source.name
+        assert parts["image"].get_content_type() == mime
+        assert parts["image"].get_payload(decode=True) == content
+        assert parts["subfolder"].get_payload(decode=True) == b"berry_assets"
+        assert parts["type"].get_payload(decode=True) == b"input"
+        assert parts["overwrite"].get_payload(decode=True) == b"false"
+        return httpx.Response(200, json={
+            "name": "renamed (1).png", "subfolder": "engine/subfolder", "type": "input",
+        })
+
+    client = ComfyUIClient(host=host, port=port)
+    async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as transport:
+        client._client = transport
+        result = await getattr(client, method)(str(source), subfolder="berry_assets")
+    assert len(requests) == 1
+    assert result == {"name": "renamed (1).png", "subfolder": "engine/subfolder", "type": "input"}
+    assert source.read_bytes() == content
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("response", [
+    {}, {"name": ""}, {"name": None}, {"name": 12},
+    {"name": "x.png", "subfolder": None},
+    {"name": "x.png", "type": "output"}, [],
+])
+async def test_invalid_upload_reference_is_rejected(tmp_path: Path, response: Any) -> None:
+    source = tmp_path / "source.png"
+    source.write_bytes(b"data")
+    client = ComfyUIClient()
+    async with httpx.AsyncClient(
+        transport=httpx.MockTransport(lambda request: httpx.Response(200, json=response))
+    ) as transport:
+        client._client = transport
+        with pytest.raises(RuntimeError, match="unexpected response"):
+            await client.upload_image(str(source))
+
+
+@pytest.mark.asyncio
+async def test_upload_defaults_and_explicit_overwrite(tmp_path: Path) -> None:
+    source = tmp_path / "source.png"
+    source.write_bytes(b"data")
+
+    async def handle(request: httpx.Request) -> httpx.Response:
+        assert b'\r\n\r\ntrue\r\n' in await request.aread()
+        return httpx.Response(200, json={"name": "renamed.png"})
+
+    client = ComfyUIClient()
+    async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as transport:
+        client._client = transport
+        result = await client.upload_image(str(source), subfolder="berry_assets", overwrite=True)
+    assert result == {"name": "renamed.png", "subfolder": "berry_assets", "type": "input"}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["upload_image", "upload_mask"])
+async def test_missing_file_does_not_upload(tmp_path: Path, method: str) -> None:
+    def handle(request: httpx.Request) -> httpx.Response:
+        pytest.fail("Missing files must not send HTTP requests")
+
+    client = ComfyUIClient()
+    async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as transport:
+        client._client = transport
+        with pytest.raises(FileNotFoundError, match="not found"):
+            await getattr(client, method)(str(tmp_path / "missing.png"))
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failure", ["http", "connection"])
+async def test_upload_reports_failure(tmp_path: Path, failure: str) -> None:
+    source = tmp_path / "source.png"
+    source.write_bytes(b"data")
+
+    def handle(request: httpx.Request) -> httpx.Response:
+        if failure == "connection":
+            raise httpx.ConnectError("Connection refused", request=request)
+        return httpx.Response(500)
+
+    client = ComfyUIClient()
+    async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as transport:
+        client._client = transport
+        with pytest.raises(RuntimeError, match="500|Connection refused"):
+            await client.upload_mask(str(source))