diff --git a/backend/app/runners/comfy_runner.py b/backend/app/runners/comfy_runner.py index 31984fa..82c1198 100644 --- a/backend/app/runners/comfy_runner.py +++ b/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() diff --git a/backend/app/runners/creative_runner.py b/backend/app/runners/creative_runner.py index 6f2fb78..d350d63 100644 --- a/backend/app/runners/creative_runner.py +++ b/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, diff --git a/backend/tests/test_comfy_asset_transfer_contract.py b/backend/tests/test_comfy_asset_transfer_contract.py new file mode 100644 index 0000000..6833a73 --- /dev/null +++ b/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 diff --git a/backend/tests/test_comfy_upload.py b/backend/tests/test_comfy_upload.py new file mode 100644 index 0000000..7564b40 --- /dev/null +++ b/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)) diff --git a/backend/tests/test_creative_runner_upload.py b/backend/tests/test_creative_runner_upload.py new file mode 100644 index 0000000..3eba52b --- /dev/null +++ b/backend/tests/test_creative_runner_upload.py @@ -0,0 +1,381 @@ +"""Integration tests for creative_runner ComfyUI asset upload (issue #105).""" + +import pytest +from pathlib import Path +from unittest.mock import AsyncMock, Mock, patch, MagicMock +import tempfile + +PNG_BYTES = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01" + b"\x08\x06\x00\x00\x00\x1f\x15c4\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01\r\n-\xb4" + b"\x00\x00\x00\x00IEND\xaeB`\x82" +) + +from app.runners.creative_runner import CreativeRunner +from app.schemas.creative import CreativeActionRequest, CreativeActionType +from app.storage.asset_store import AssetRecord + + +@pytest.fixture +def temp_asset_files(): + """Create temporary image and mask files for testing.""" + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as img_file: + img_file.write(PNG_BYTES) + img_path = Path(img_file.name) + + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as mask_file: + mask_file.write(PNG_BYTES) + mask_path = Path(mask_file.name) + + yield img_path, mask_path + + # Cleanup + img_path.unlink(missing_ok=True) + mask_path.unlink(missing_ok=True) + + +@pytest.mark.asyncio +async def test_img2img_uploads_source_before_execution(temp_asset_files): + """IMG2IMG should upload source image to ComfyUI before submitting workflow.""" + img_path, _ = temp_asset_files + runner = CreativeRunner() + + req = CreativeActionRequest( + action=CreativeActionType.IMG2IMG, + prompt="test prompt", + model="test_model.safetensors", + engine_id="comfyui", + input_image_id="test_img_id", + width=512, + height=512, + seed=42, + ) + + # Mock asset store + mock_asset = AssetRecord( + id="test_img_id", + filename="test.png", + file_path="test.png", + media_type="image", + content_hash="abc123", + byte_size=1000, + created_at="2024-01-01T00:00:00Z", + ) + + with patch("app.runners.creative_runner.asset_store") as mock_store: + mock_store.get_asset = AsyncMock(return_value=mock_asset) + mock_store.get_absolute_path = Mock(return_value=img_path) + mock_store.save_image_from_url = AsyncMock(return_value=mock_asset) + + # Mock ComfyUI client + with patch("app.runners.creative_runner.comfy_client") as mock_comfy: + mock_comfy.upload_image = AsyncMock(return_value={ + "name": "uploaded_test.png", + "subfolder": "berry_assets", + "type": "input", + }) + mock_comfy.queue_prompt = AsyncMock(return_value={"prompt_id": "test_prompt_123"}) + mock_comfy.poll_history_outputs = AsyncMock(return_value=[{ + "filename": "output.png", + "subfolder": "", + "type": "output", + }]) + mock_comfy.base_url = "http://127.0.0.1:8188" + + # Mock cache + with patch("app.runners.creative_runner.cache_store") as mock_cache: + mock_cache.get_async = AsyncMock(return_value=None) + mock_cache.set_async = AsyncMock() + + result = await runner.execute(req) + + # Verify upload was called with correct path + mock_comfy.upload_image.assert_called_once() + upload_call_args = mock_comfy.upload_image.call_args + # Path may have different formats (forward/back slashes), just check it was called + assert upload_call_args is not None + + # Verify workflow was submitted after upload + mock_comfy.queue_prompt.assert_called_once() + + assert result.success is True + + +@pytest.mark.asyncio +async def test_inpaint_uploads_both_image_and_mask(temp_asset_files): + """INPAINT should upload both source image and mask before execution.""" + img_path, mask_path = temp_asset_files + runner = CreativeRunner() + + req = CreativeActionRequest( + action=CreativeActionType.INPAINT, + prompt="inpaint test", + model="test_model.safetensors", + engine_id="comfyui", + input_image_id="img_id", + mask_image_id="mask_id", + width=512, + height=512, + seed=42, + ) + + mock_img_asset = AssetRecord( + id="img_id", + filename="source.png", + file_path="source.png", + media_type="image", + content_hash="img_hash", + byte_size=1000, + created_at="2024-01-01T00:00:00Z", + ) + + mock_mask_asset = AssetRecord( + id="mask_id", + filename="mask.png", + file_path="mask.png", + media_type="image", + content_hash="mask_hash", + byte_size=500, + created_at="2024-01-01T00:00:00Z", + ) + + with patch("app.runners.creative_runner.asset_store") as mock_store: + def get_asset_side_effect(asset_id): + if asset_id == "img_id": + return mock_img_asset + elif asset_id == "mask_id": + return mock_mask_asset + return None + + def get_path_side_effect(asset_rec): + if asset_rec.id == "img_id": + return img_path + elif asset_rec.id == "mask_id": + return mask_path + return None + + mock_store.get_asset = AsyncMock(side_effect=get_asset_side_effect) + mock_store.get_absolute_path = Mock(side_effect=get_path_side_effect) + mock_store.save_image_from_url = AsyncMock(return_value=mock_img_asset) + + with patch("app.runners.creative_runner.comfy_client") as mock_comfy: + mock_comfy.upload_image = AsyncMock(return_value={ + "name": "uploaded_source.png", + "subfolder": "berry_assets", + "type": "input", + }) + mock_comfy.upload_mask = AsyncMock(return_value={ + "name": "uploaded_mask.png", + "subfolder": "berry_assets", + "type": "input", + }) + mock_comfy.queue_prompt = AsyncMock(return_value={"prompt_id": "inpaint_123"}) + mock_comfy.poll_history_outputs = AsyncMock(return_value=[{ + "filename": "inpainted.png", + "subfolder": "", + "type": "output", + }]) + mock_comfy.base_url = "http://127.0.0.1:8188" + + with patch("app.runners.creative_runner.cache_store") as mock_cache: + mock_cache.get_async = AsyncMock(return_value=None) + mock_cache.set_async = AsyncMock() + + result = await runner.execute(req) + + # Verify both uploads were called + mock_comfy.upload_image.assert_called_once() + mock_comfy.upload_mask.assert_called_once() + + # Verify workflow submission happened after uploads + mock_comfy.queue_prompt.assert_called_once() + + assert result.success is True + + +@pytest.mark.asyncio +async def test_upscale_uploads_source_image(temp_asset_files): + """UPSCALE should upload source image before execution.""" + img_path, _ = temp_asset_files + runner = CreativeRunner() + + req = CreativeActionRequest( + action=CreativeActionType.UPSCALE, + prompt="", + model="", + engine_id="comfyui", + input_image_id="upscale_img_id", + upscaler_name="RealESRGAN_x4plus.pth", + upscale_factor=4, + width=512, + height=512, + seed=42, + ) + + mock_asset = AssetRecord( + id="upscale_img_id", + filename="to_upscale.png", + file_path="to_upscale.png", + media_type="image", + content_hash="upscale_hash", + byte_size=1000, + created_at="2024-01-01T00:00:00Z", + ) + + with patch("app.runners.creative_runner.asset_store") as mock_store: + mock_store.get_asset = AsyncMock(return_value=mock_asset) + mock_store.get_absolute_path = Mock(return_value=img_path) + mock_store.save_image_from_url = AsyncMock(return_value=mock_asset) + + with patch("app.runners.creative_runner.comfy_client") as mock_comfy: + mock_comfy.upload_image = AsyncMock(return_value={ + "name": "uploaded_upscale.png", + "subfolder": "berry_assets", + "type": "input", + }) + mock_comfy.queue_prompt = AsyncMock(return_value={"prompt_id": "upscale_123"}) + mock_comfy.poll_history_outputs = AsyncMock(return_value=[{ + "filename": "upscaled.png", + "subfolder": "", + "type": "output", + }]) + mock_comfy.base_url = "http://127.0.0.1:8188" + + with patch("app.runners.creative_runner.cache_store") as mock_cache: + mock_cache.get_async = AsyncMock(return_value=None) + mock_cache.set_async = AsyncMock() + + result = await runner.execute(req) + + mock_comfy.upload_image.assert_called_once() + mock_comfy.queue_prompt.assert_called_once() + + assert result.success is True + + +@pytest.mark.asyncio +async def test_upload_failure_prevents_workflow_submission(): + """If upload fails, workflow should not be submitted and error should be reported.""" + runner = CreativeRunner() + + req = CreativeActionRequest( + action=CreativeActionType.IMG2IMG, + prompt="test", + model="test.safetensors", + engine_id="comfyui", + input_image_id="img_id", + width=512, + height=512, + seed=42, + ) + + mock_asset = AssetRecord( + id="img_id", + filename="test.png", + file_path="test.png", + media_type="image", + content_hash="hash123", + byte_size=1000, + created_at="2024-01-01T00:00:00Z", + ) + + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f: + f.write(b"fake_data") + temp_path = Path(f.name) + + try: + with patch("app.runners.creative_runner.asset_store") as mock_store: + mock_store.get_asset = AsyncMock(return_value=mock_asset) + mock_store.get_absolute_path = Mock(return_value=temp_path) + + with patch("app.runners.creative_runner.comfy_client") as mock_comfy: + # Upload fails + mock_comfy.upload_image = AsyncMock(side_effect=RuntimeError("Upload failed: Connection refused")) + mock_comfy.queue_prompt = AsyncMock() + + with patch("app.runners.creative_runner.cache_store") as mock_cache: + mock_cache.get_async = AsyncMock(return_value=None) + + result = await runner.execute(req) + + # Verify upload was attempted + mock_comfy.upload_image.assert_called_once() + + # Verify workflow was NOT submitted + mock_comfy.queue_prompt.assert_not_called() + + # Verify error result + assert result.success is False + assert "Failed to transfer source image to ComfyUI" in result.error_message + finally: + temp_path.unlink(missing_ok=True) + + +@pytest.mark.asyncio +async def test_uploaded_filename_used_in_workflow(): + """Verify that the uploaded filename from ComfyUI is used in the compiled workflow.""" + runner = CreativeRunner() + + req = CreativeActionRequest( + action=CreativeActionType.IMG2IMG, + prompt="test", + model="test.safetensors", + engine_id="comfyui", + input_image_id="img_id", + width=512, + height=512, + seed=42, + ) + + mock_asset = AssetRecord( + id="img_id", + filename="original.png", + file_path="original.png", + media_type="image", + content_hash="hash123", + byte_size=1000, + created_at="2024-01-01T00:00:00Z", + ) + + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f: + f.write(b"fake_data") + temp_path = Path(f.name) + + try: + with patch("app.runners.creative_runner.asset_store") as mock_store: + mock_store.get_asset = AsyncMock(return_value=mock_asset) + mock_store.get_absolute_path = Mock(return_value=temp_path) + mock_store.save_image_from_url = AsyncMock(return_value=mock_asset) + + with patch("app.runners.creative_runner.comfy_client") as mock_comfy: + # ComfyUI returns a different filename after upload + mock_comfy.upload_image = AsyncMock(return_value={ + "name": "comfy_renamed_12345.png", + "subfolder": "berry_assets", + "type": "input", + }) + mock_comfy.queue_prompt = AsyncMock(return_value={"prompt_id": "test_123"}) + mock_comfy.poll_history_outputs = AsyncMock(return_value=[{ + "filename": "output.png", + "subfolder": "", + "type": "output", + }]) + mock_comfy.base_url = "http://127.0.0.1:8188" + + with patch("app.runners.creative_runner.cache_store") as mock_cache: + mock_cache.get_async = AsyncMock(return_value=None) + mock_cache.set_async = AsyncMock() + + result = await runner.execute(req) + + # Get the workflow that was queued + queued_workflow = mock_comfy.queue_prompt.call_args[0][0] + + # The workflow should reference the ComfyUI-returned filename, not the original + workflow_str = str(queued_workflow) + assert "berry_assets/comfy_renamed_12345.png" in workflow_str + assert "original.png" not in workflow_str + + assert result.success is True + finally: + temp_path.unlink(missing_ok=True) diff --git a/backend/tests/test_m5_release.py b/backend/tests/test_m5_release.py index db19a4a..de96601 100644 --- a/backend/tests/test_m5_release.py +++ b/backend/tests/test_m5_release.py @@ -151,7 +151,13 @@ async def test_scenario_5_persistence_and_caching_across_restarts(): unique_prompt = f"tranquil forest {uuid.uuid4()}" b64_dummy = base64.b64encode(DUMMY_PNG).decode("utf-8") - with patch("httpx.AsyncClient.post", new_callable=AsyncMock) as mock_post: + with ( + patch("httpx.AsyncClient.post", new_callable=AsyncMock) as mock_post, + patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get, + ): + checkpoint_response = MagicMock() + checkpoint_response.json.return_value = [{"title": "test.safetensors", "model_name": "test"}] + mock_get.return_value = checkpoint_response mock_resp = MagicMock() mock_resp.status_code = 200 mock_resp.json.return_value = {"images": [b64_dummy]} @@ -161,10 +167,12 @@ async def test_scenario_5_persistence_and_caching_across_restarts(): action=CreativeActionType.TXT2IMG, prompt=unique_prompt, engine_id="managed_webui", + model="test.safetensors", seed=777, ) res1 = await runner.execute(req) + assert res1.success is True, res1.error_message assert res1.is_cached is False # Emulate restart: create a new CreativeRunner instance @@ -172,6 +180,7 @@ async def test_scenario_5_persistence_and_caching_across_restarts(): res2 = await runner_after_restart.execute(req) assert res2.is_cached is True assert res2.asset_id == res1.asset_id + mock_post.assert_awaited_once() def test_scenario_6_cancellation_semantics(): diff --git a/backend/tests/test_model_downloader.py b/backend/tests/test_model_downloader.py index f5f33c1..b3ed79a 100644 --- a/backend/tests/test_model_downloader.py +++ b/backend/tests/test_model_downloader.py @@ -25,7 +25,7 @@ def test_target_path_resolution(): # WebUI checkpoint path webui_path = downloader._resolve_target_path(model, target_engine="webui") assert "models" in str(webui_path) - assert "Stable-diffusion" in str(webui_path) + assert webui_path == comfy_path assert webui_path.name == "flux1-schnell-fp8.safetensors" # LoRA category path @@ -35,7 +35,7 @@ def test_target_path_resolution(): assert "loras" in str(comfy_lora_path) webui_lora_path = downloader._resolve_target_path(lora_model, target_engine="webui") - assert "Lora" in str(webui_lora_path) + assert webui_lora_path == comfy_lora_path def test_mirror_source_selection(): diff --git a/backend/tests/test_model_paths.py b/backend/tests/test_model_paths.py index 7da3fdf..ed67a32 100644 --- a/backend/tests/test_model_paths.py +++ b/backend/tests/test_model_paths.py @@ -1,181 +1,83 @@ -""" -Tests for model path consistency across downloader, supervisors, and engine search paths (Issue #107). +"""Regression tests for shared model paths and actual launch arguments (#107).""" -Verifies: -1. Model downloader uses engine_dir for all engines (ComfyUI and WebUI) -2. ComfyUI supervisor generates and passes extra-model-paths-config -3. WebUI supervisor passes correct model directory flags -4. Downloaded models appear in engine-accessible search paths -""" - -import tempfile from pathlib import Path +from unittest.mock import Mock, patch + import pytest +from app.runtime.installer import installer +from app.runtime.model_downloader import ModelDownloader from app.runtime.supervisor import ComfySupervisor from app.runtime.webui_supervisor import WebUISupervisor -from app.runtime.model_downloader import ModelDownloadManager -from app.schemas.model_hub import HubModelRecord, ModelSource - - -def test_comfy_supervisor_generates_extra_model_paths_config(): - """Verify ComfySupervisor generates extra_model_paths.yaml with shared model directories.""" - with tempfile.TemporaryDirectory() as tmpdir: - engine_dir = Path(tmpdir) / "engine" - supervisor = ComfySupervisor(engine_dir=engine_dir, port=8188) - supervisor.ensure_directories() - - # Generate config - config_path = supervisor._generate_extra_model_paths_config() - - assert config_path.exists() - assert config_path == engine_dir / "extra_model_paths.yaml" - - config_content = config_path.read_text(encoding="utf-8") - assert "berry_shared:" in config_content - assert f"base_path: {supervisor.models_dir}" in config_content - assert "checkpoints: checkpoints/" in config_content - assert "vae: vae/" in config_content - assert "loras: loras/" in config_content - assert "upscale_models: upscale_models/" in config_content - assert "controlnet: controlnet/" in config_content - - -def test_webui_supervisor_uses_engine_dir_models(): - """Verify WebUISupervisor model directories point to engine_dir/models/.""" - with tempfile.TemporaryDirectory() as tmpdir: - engine_dir = Path(tmpdir) / "engine" - supervisor = WebUISupervisor(engine_dir=engine_dir, port=7860) - - # Verify models_dir is under engine_dir - assert supervisor.models_dir == engine_dir / "models" - assert supervisor.models_dir.is_relative_to(engine_dir) - - -def test_model_downloader_resolves_consistent_paths(): - """Verify ModelDownloadManager resolves paths consistently for ComfyUI and WebUI.""" - with tempfile.TemporaryDirectory() as tmpdir: - from app.runtime.installer import installer - original_engine_dir = installer.engine_dir - installer.engine_dir = Path(tmpdir) / "engine" - - try: - manager = ModelDownloadManager() - - # Test ComfyUI checkpoint path - comfy_checkpoint = HubModelRecord( - id="test-comfy-ckpt", - title="Test ComfyUI Checkpoint", - category="checkpoint", - sources=[ModelSource(url="https://example.com/model.safetensors")], - ) - comfy_path = manager._resolve_target_path(comfy_checkpoint, "comfyui") - assert comfy_path.is_relative_to(installer.engine_dir / "models" / "checkpoints") - - # Test WebUI checkpoint path - should also use engine_dir, not home directory - webui_checkpoint = HubModelRecord( - id="test-webui-ckpt", - title="Test WebUI Checkpoint", - category="checkpoint", - sources=[ModelSource(url="https://example.com/model.safetensors")], - ) - webui_path = manager._resolve_target_path(webui_checkpoint, "webui") - assert webui_path.is_relative_to(installer.engine_dir / "models" / "checkpoints") - - # Verify ComfyUI and WebUI use the same shared model directory - assert comfy_path.parent == webui_path.parent - - # Test other categories - lora_model = HubModelRecord( - id="test-lora", - title="Test LoRA", - category="lora", - sources=[ModelSource(url="https://example.com/lora.safetensors")], - ) - lora_path = manager._resolve_target_path(lora_model, "comfyui") - assert lora_path.is_relative_to(installer.engine_dir / "models" / "loras") - - vae_model = HubModelRecord( - id="test-vae", - title="Test VAE", - category="vae", - sources=[ModelSource(url="https://example.com/vae.safetensors")], - ) - vae_path = manager._resolve_target_path(vae_model, "webui") - assert vae_path.is_relative_to(installer.engine_dir / "models" / "vae") - - finally: - installer.engine_dir = original_engine_dir - - -def test_downloaded_models_accessible_to_both_engines(): - """Verify downloaded models in shared paths are accessible to both ComfyUI and WebUI.""" - with tempfile.TemporaryDirectory() as tmpdir: - engine_dir = Path(tmpdir) / "engine" - - comfy_supervisor = ComfySupervisor(engine_dir=engine_dir) - webui_supervisor = WebUISupervisor(engine_dir=engine_dir) - - comfy_supervisor.ensure_directories() - webui_supervisor.ensure_directories() - - # Both should reference the same shared models directory - assert comfy_supervisor.models_dir == webui_supervisor.models_dir - assert comfy_supervisor.models_dir == engine_dir / "models" - - # Create a dummy checkpoint in the shared location - checkpoints_dir = engine_dir / "models" / "checkpoints" - checkpoints_dir.mkdir(parents=True, exist_ok=True) - dummy_checkpoint = checkpoints_dir / "test-model.safetensors" - dummy_checkpoint.write_bytes(b"dummy checkpoint data") - - # Verify it exists in the shared location accessible to both - assert dummy_checkpoint.exists() - assert dummy_checkpoint.is_relative_to(comfy_supervisor.models_dir) - assert dummy_checkpoint.is_relative_to(webui_supervisor.models_dir) - - -def test_comfy_start_command_includes_extra_model_paths(): - """Verify ComfySupervisor.start includes --extra-model-paths-config in launch command.""" - with tempfile.TemporaryDirectory() as tmpdir: - engine_dir = Path(tmpdir) / "engine" - supervisor = ComfySupervisor(engine_dir=engine_dir) - - # Create minimal installation structure - supervisor.ensure_directories() - supervisor.comfy_dir.mkdir(parents=True, exist_ok=True) - (supervisor.comfy_dir / "main.py").write_text("# dummy", encoding="utf-8") - supervisor.runtime_dir.mkdir(parents=True, exist_ok=True) - python_bin = supervisor.get_python_bin() - python_bin.parent.mkdir(parents=True, exist_ok=True) - python_bin.write_text("# dummy python", encoding="utf-8") - - # The start method will fail because python isn't real, but we can inspect - # that it attempts to generate the config - config_path = supervisor._generate_extra_model_paths_config() - assert config_path.exists() - assert "--extra-model-paths-config" in str(config_path) - - -def test_webui_start_command_includes_model_dir_flags(): - """Verify WebUISupervisor.start includes --ckpt-dir, --lora-dir, --vae-dir flags.""" - with tempfile.TemporaryDirectory() as tmpdir: - engine_dir = Path(tmpdir) / "engine" - supervisor = WebUISupervisor(engine_dir=engine_dir) - - # Verify the supervisor is configured to pass model directory flags - supervisor.ensure_directories() - - expected_ckpt_dir = supervisor.models_dir / "checkpoints" - expected_lora_dir = supervisor.models_dir / "loras" - expected_vae_dir = supervisor.models_dir / "vae" - - assert expected_ckpt_dir.exists() - assert expected_lora_dir.exists() - assert expected_vae_dir.exists() - - # These directories should be passed to WebUI launch command - # (we've verified in the source that start() passes these flags) - assert expected_ckpt_dir.is_relative_to(engine_dir) - assert expected_lora_dir.is_relative_to(engine_dir) - assert expected_vae_dir.is_relative_to(engine_dir) +from app.storage.hub_catalog import HubCatalog + + +def test_comfy_supervisor_generates_extra_model_paths_config(tmp_path: Path) -> None: + supervisor = ComfySupervisor(engine_dir=tmp_path / "engine") + supervisor.ensure_directories() + config_path = supervisor._generate_extra_model_paths_config() + assert config_path == supervisor.engine_dir / "extra_model_paths.yaml" + content = config_path.read_text(encoding="utf-8") + assert "berry_shared:" in content + assert f"base_path: {supervisor.models_dir}" in content + for subfolder in ("checkpoints", "vae", "loras", "upscale_models", "controlnet"): + assert f"{subfolder}: {subfolder}/" in content + + +def test_webui_supervisor_uses_engine_dir_models(tmp_path: Path) -> None: + supervisor = WebUISupervisor(engine_dir=tmp_path / "engine") + assert supervisor.models_dir == supervisor.engine_dir / "models" + + +def test_model_downloader_resolves_consistent_paths( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(installer, "engine_dir", tmp_path / "engine") + manager = ModelDownloader() + base = HubCatalog.get_model("flux-1-schnell-fp8") + assert base is not None + for category, subfolder in (("checkpoint", "checkpoints"), ("lora", "loras"), ("vae", "vae")): + model = base.model_copy(update={"category": category}) + comfy_path = manager._resolve_target_path(model, "comfyui") + webui_path = manager._resolve_target_path(model, "webui") + assert comfy_path == webui_path + assert comfy_path.parent == installer.engine_dir / "models" / subfolder + + +def test_downloaded_models_accessible_to_both_engines(tmp_path: Path) -> None: + engine_dir = tmp_path / "engine" + comfy = ComfySupervisor(engine_dir=engine_dir) + webui = WebUISupervisor(engine_dir=engine_dir) + comfy.ensure_directories() + webui.ensure_directories() + assert comfy.models_dir == webui.models_dir == engine_dir / "models" + checkpoint = comfy.models_dir / "checkpoints" / "test.safetensors" + checkpoint.write_bytes(b"dummy checkpoint data") + assert (webui.models_dir / "checkpoints" / checkpoint.name).read_bytes() == checkpoint.read_bytes() + + +def test_comfy_start_command_includes_extra_model_paths(tmp_path: Path) -> None: + supervisor = ComfySupervisor(engine_dir=tmp_path / "engine") + supervisor.ensure_directories() + supervisor.comfy_dir.mkdir(parents=True, exist_ok=True) + (supervisor.comfy_dir / "main.py").write_text("# dummy", encoding="utf-8") + python_bin = supervisor.get_python_bin() + python_bin.parent.mkdir(parents=True, exist_ok=True) + python_bin.write_text("# dummy", encoding="utf-8") + with ( + patch.object(supervisor, "is_running", return_value=False), + patch("app.runtime.hardware.get_hardware_launch_flags", return_value=[]), + patch("app.runtime.supervisor.subprocess.Popen", return_value=Mock(pid=12345)) as launch, + ): + assert supervisor.start()["success"] is True + cmd = launch.call_args.args[0] + index = cmd.index("--extra-model-paths-config") + assert cmd[index + 1] == str(supervisor.engine_dir / "extra_model_paths.yaml") + assert Path(cmd[index + 1]).is_file() + + +def test_webui_start_command_includes_model_dir_flags(tmp_path: Path) -> None: + supervisor = WebUISupervisor(engine_dir=tmp_path / "engine") + supervisor.ensure_directories() + for subfolder in ("checkpoints", "loras", "vae"): + assert (supervisor.models_dir / subfolder).is_dir()