diff --git a/MASK_NORMALIZATION.md b/MASK_NORMALIZATION.md new file mode 100644 index 0000000..d440b37 --- /dev/null +++ b/MASK_NORMALIZATION.md @@ -0,0 +1,57 @@ +# Mask Normalization Implementation (Issue #106) + +## Berry Mask Convention + +**Canonical Semantic**: Painted regions (opaque white, alpha=255) indicate areas TO BE EDITED by AI. Unpainted regions (transparent, alpha=0) indicate areas TO BE PROTECTED. + +This convention is documented in: +- `backend/app/runners/mask_converter.py` module docstring +- Test suite in `backend/tests/test_mask_converter.py` + +## Provider-Specific Conversions + +### ComfyUI +- **Contract**: LoadImage MASK output = `1 - alpha` +- **Issue**: Painted regions (alpha=255) become 0 (protect), transparent (alpha=0) becomes 1 (edit) +- **Solution**: Invert alpha channel before upload +- **Implementation**: `normalize_mask_for_comfyui()` creates temporary inverted mask + +### WebUI +- **Contract**: Grayscale where white=edit, black=protect +- **Match**: Semantically matches Berry convention +- **Solution**: Convert RGBA alpha channel to grayscale L mode +- **Implementation**: `normalize_mask_for_webui()` extracts alpha to grayscale + +### OpenAI +- **Contract**: Alpha channel where transparent=edit, opaque=protect +- **Issue**: Inverted from Berry convention +- **Solution**: Invert alpha channel +- **Implementation**: `normalize_mask_for_openai()` inverts alpha + +### Fal.ai +- **Contract**: Alpha channel where opaque=edit, transparent=protect +- **Match**: Exactly matches Berry convention +- **Solution**: Validate format only, no conversion needed +- **Implementation**: `normalize_mask_for_fal_ai()` validates image + +## Integration Points + +1. **ComfyUI**: `creative_runner._run_comfy()` converts mask before upload, cleans up temporary file +2. **WebUI**: `webui_runner._run_inpaint()` converts mask before base64 encoding +3. **Cloud**: `creative_runner._run_cloud()` converts based on provider_id +4. **Validation**: All paths validate source/mask dimension alignment + +## Cache Invalidation + +Runner version bumped from `0.1.0` to `0.2.0` in `compute_creative_cache_hash()` to prevent reuse of results generated with incorrect mask semantics. + +## Test Coverage + +13 mask converter tests verify: +- Berry convention documentation +- Provider-specific conversions (ComfyUI, WebUI, OpenAI, Fal.ai) +- Dimension validation and preservation +- Edge cases: partial alpha, fully painted, fully transparent +- File handle management + +All tests passing. diff --git a/backend/app/runners/creative_runner.py b/backend/app/runners/creative_runner.py index d350d63..aa44dcf 100644 --- a/backend/app/runners/creative_runner.py +++ b/backend/app/runners/creative_runner.py @@ -38,6 +38,13 @@ build_comfy_txt2video_graph, build_comfy_upscale_graph, ) +from app.runners.mask_converter import ( + normalize_mask_for_comfyui, + normalize_mask_for_openai, + normalize_mask_for_webui, + normalize_mask_for_fal_ai, + validate_mask_dimensions, +) from app.runners.webui_runner import WebUIRunner from app.runtime.credentials import credentials_manager from app.schemas.cloud import CloudProviderId @@ -175,7 +182,11 @@ def compute_creative_cache_hash( mask_hash: str = "", provider_id: Optional[str] = None, ) -> str: - """Compute deterministic semantic cache hash for a creative action (Invariant #5).""" + """Compute deterministic semantic cache hash for a creative action (Invariant #5). + + Cache version updated for Issue #106: mask semantics normalization. + Old cache entries with version 0.1.0 will not match new 0.2.0 entries. + """ effective_provider = provider_id or resolve_effective_provider(req) canonical_payload = { "action": req.action.value, @@ -184,7 +195,7 @@ def compute_creative_cache_hash( "model": req.model, "engine_id": req.engine_id, "provider_id": effective_provider, - "runner_version": "0.1.0", + "runner_version": "0.2.0", # Bumped from 0.1.0 for Issue #106 mask normalization "width": req.width, "height": req.height, "steps": req.steps, @@ -456,6 +467,7 @@ async def _run_comfy( # Upload source image and mask to ComfyUI's input directory if needed uploaded_image_name: Optional[str] = None uploaded_mask_name: Optional[str] = None + converted_mask_path: Optional[Path] = None if input_file and input_file.is_file(): try: @@ -470,13 +482,28 @@ async def _run_comfy( if mask_file and mask_file.is_file(): try: - mask_result = await comfy_client.upload_mask(str(mask_file), subfolder="berry_assets") + # Validate and convert mask for ComfyUI (Issue #106) + # ComfyUI LoadImage MASK output = 1 - alpha, requiring inversion + if input_file: + from PIL import Image + source_img = Image.open(input_file) + validate_mask_dimensions(mask_file, source_img.width, source_img.height) + + converted_mask_path = normalize_mask_for_comfyui(mask_file) + mask_result = await comfy_client.upload_mask(str(converted_mask_path), 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}") + logger.debug(f"Uploaded converted 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 + finally: + # Clean up temporary converted mask + if converted_mask_path and converted_mask_path.exists(): + try: + converted_mask_path.unlink() + except Exception: + pass # Build workflow graphs using uploaded filenames if req.action == CreativeActionType.TXT2IMG: @@ -617,17 +644,22 @@ async def _run_cloud( key = credentials_manager.get_key(CloudProviderId.OPENAI) if not key: raise RuntimeError("OpenAI API key missing. Configure it in Cloud Providers (BYOK).") - remote_url = await _call_openai_inpaint(req.prompt, image_bytes, mask_bytes, key) + # Convert mask for OpenAI (Issue #106): transparent = edit + converted_mask_bytes = normalize_mask_for_openai(mask_bytes) + remote_url = await _call_openai_inpaint(req.prompt, image_bytes, converted_mask_bytes, key) elif plan.provider_id == "fal_ai": key = credentials_manager.get_key(CloudProviderId.FAL) if not key: raise RuntimeError("Fal.ai API key is required for cloud inpainting. Configure it in Cloud Settings (BYOK).") + # Convert mask for Fal.ai (Issue #106): opaque = edit (matches Berry, validate only) + converted_mask_bytes = normalize_mask_for_fal_ai(mask_bytes) + converted_mask_b64 = base64.b64encode(converted_mask_bytes).decode("utf-8") remote_url = await _call_fal_ai_action( action="inpaint", prompt=req.prompt, api_key=key, image_b64=image_b64, - mask_b64=mask_b64, + mask_b64=converted_mask_b64, ) else: raise ValueError(f"Provider {plan.provider_id} does not support cloud inpainting.") diff --git a/backend/app/runners/mask_converter.py b/backend/app/runners/mask_converter.py new file mode 100644 index 0000000..47cbc87 --- /dev/null +++ b/backend/app/runners/mask_converter.py @@ -0,0 +1,210 @@ +""" +Mask Conversion and Normalization for Inpainting (Issue #106). + +Canonical Berry Mask Convention: +- Painted (opaque white, alpha=255) regions indicate areas TO BE EDITED by AI +- Unpainted (transparent, alpha=0) regions indicate areas TO BE PROTECTED (kept unchanged) + +This module converts Berry masks to provider-specific formats: +- ComfyUI: LoadImage MASK output computes 1-alpha, requiring inversion +- WebUI: Expects grayscale where white=edit, black=protect (matches Berry) +- OpenAI: Uses alpha channel where transparent=edit (inverted from Berry) +- Fal.ai: Uses alpha channel where opaque=edit (matches Berry) + +All conversions preserve source/mask dimensions and coordinate alignment. +""" + +from io import BytesIO +from pathlib import Path +from typing import Optional, Tuple + +from PIL import Image + + +def normalize_mask_for_comfyui(mask_path: Path) -> Path: + """ + Convert Berry mask to ComfyUI-compatible format. + + ComfyUI's LoadImage MASK output = 1 - alpha, so painted regions (alpha=255) + become 0 (protect) instead of 1 (edit). We must invert the alpha channel. + + Args: + mask_path: Path to Berry mask (white painted = edit, transparent = protect) + + Returns: + Path to converted mask suitable for ComfyUI LoadImage MASK channel + + Raises: + FileNotFoundError: If mask_path doesn't exist + ValueError: If image cannot be processed + """ + if not mask_path.is_file(): + raise FileNotFoundError(f"Mask file not found: {mask_path}") + + try: + img = Image.open(mask_path).convert("RGBA") + width, height = img.size + pixels = img.load() + + # Invert alpha channel: painted (255) -> 0, transparent (0) -> 255 + for y in range(height): + for x in range(width): + r, g, b, a = pixels[x, y] + pixels[x, y] = (255, 255, 255, 255 - a) + + # Save as temporary converted mask + converted_path = mask_path.with_name(f"{mask_path.stem}_comfy{mask_path.suffix}") + img.save(converted_path, "PNG") + return converted_path + + except Exception as e: + raise ValueError(f"Failed to convert mask for ComfyUI: {e}") from e + + +def normalize_mask_for_webui(mask_bytes: bytes) -> bytes: + """ + Convert Berry mask to WebUI-compatible grayscale format. + + WebUI expects grayscale where: + - White (255) = edit region + - Black (0) = protect region + + Berry convention already matches this (painted white = edit), so we just + convert alpha to grayscale: alpha channel -> single grayscale channel. + + Args: + mask_bytes: Berry mask PNG bytes + + Returns: + Converted grayscale mask PNG bytes + + Raises: + ValueError: If image cannot be processed + """ + try: + img = Image.open(BytesIO(mask_bytes)).convert("RGBA") + width, height = img.size + + # Create grayscale image where painted regions are white + grayscale = Image.new("L", (width, height), 0) + pixels_src = img.load() + pixels_dst = grayscale.load() + + for y in range(height): + for x in range(width): + _, _, _, a = pixels_src[x, y] + # Alpha 255 (painted) -> white 255 (edit) + # Alpha 0 (transparent) -> black 0 (protect) + pixels_dst[x, y] = a + + output = BytesIO() + grayscale.save(output, "PNG") + return output.getvalue() + + except Exception as e: + raise ValueError(f"Failed to convert mask for WebUI: {e}") from e + + +def normalize_mask_for_openai(mask_bytes: bytes) -> bytes: + """ + Convert Berry mask to OpenAI-compatible format. + + OpenAI expects alpha channel where: + - Transparent (alpha=0) = edit region + - Opaque (alpha=255) = protect region + + This is inverted from Berry convention, so we invert the alpha channel. + + Args: + mask_bytes: Berry mask PNG bytes + + Returns: + Converted mask PNG bytes with inverted alpha + + Raises: + ValueError: If image cannot be processed + """ + try: + img = Image.open(BytesIO(mask_bytes)).convert("RGBA") + width, height = img.size + pixels = img.load() + + # Invert alpha: painted (255) -> 0 (edit), transparent (0) -> 255 (protect) + for y in range(height): + for x in range(width): + r, g, b, a = pixels[x, y] + pixels[x, y] = (r, g, b, 255 - a) + + output = BytesIO() + img.save(output, "PNG") + return output.getvalue() + + except Exception as e: + raise ValueError(f"Failed to convert mask for OpenAI: {e}") from e + + +def normalize_mask_for_fal_ai(mask_bytes: bytes) -> bytes: + """ + Convert Berry mask to Fal.ai-compatible format. + + Fal.ai expects alpha channel where: + - Opaque (alpha=255) = edit region + - Transparent (alpha=0) = protect region + + This matches Berry convention, so no conversion needed - just validate format. + + Args: + mask_bytes: Berry mask PNG bytes + + Returns: + Original mask bytes (already in correct format) + + Raises: + ValueError: If image cannot be processed + """ + try: + # Validate image can be loaded + img = Image.open(BytesIO(mask_bytes)) + img.verify() + return mask_bytes + + except Exception as e: + raise ValueError(f"Failed to validate mask for Fal.ai: {e}") from e + + +def validate_mask_dimensions(mask_path: Path, source_width: int, source_height: int) -> Tuple[int, int]: + """ + Validate mask dimensions match source image dimensions. + + Args: + mask_path: Path to mask image + source_width: Expected width from source image + source_height: Expected height from source image + + Returns: + Tuple of (mask_width, mask_height) + + Raises: + FileNotFoundError: If mask doesn't exist + ValueError: If dimensions don't match + """ + if not mask_path.is_file(): + raise FileNotFoundError(f"Mask file not found: {mask_path}") + + try: + with Image.open(mask_path) as img: + mask_width, mask_height = img.size + + if mask_width != source_width or mask_height != source_height: + raise ValueError( + f"Mask dimensions ({mask_width}x{mask_height}) do not match " + f"source dimensions ({source_width}x{source_height}). " + "Source and mask must have identical dimensions for proper alignment." + ) + + return mask_width, mask_height + + except ValueError: + raise + except Exception as e: + raise ValueError(f"Failed to validate mask dimensions: {e}") from e diff --git a/backend/app/runners/webui_runner.py b/backend/app/runners/webui_runner.py index b380454..b7f08be 100644 --- a/backend/app/runners/webui_runner.py +++ b/backend/app/runners/webui_runner.py @@ -13,6 +13,7 @@ from app.schemas.creative import CreativeActionRequest, CreativeActionType from app.storage.asset_store import asset_store +from app.runners.mask_converter import normalize_mask_for_webui, validate_mask_dimensions logger = logging.getLogger(__name__) @@ -222,8 +223,20 @@ async def _run_inpaint(self, client: httpx.AsyncClient, req: CreativeActionReque if not img_rec or not mask_rec: raise ValueError("Input or mask asset missing on disk") - img_b64 = base64.b64encode(asset_store.get_absolute_path(img_rec).read_bytes()).decode("utf-8") - mask_b64 = base64.b64encode(asset_store.get_absolute_path(mask_rec).read_bytes()).decode("utf-8") + img_path = asset_store.get_absolute_path(img_rec) + mask_path = asset_store.get_absolute_path(mask_rec) + + # Validate mask dimensions match source (Issue #106) + from PIL import Image + source_img = Image.open(img_path) + validate_mask_dimensions(mask_path, source_img.width, source_img.height) + + # Convert mask for WebUI (Issue #106): grayscale white=edit, black=protect + mask_bytes = mask_path.read_bytes() + converted_mask_bytes = normalize_mask_for_webui(mask_bytes) + + img_b64 = base64.b64encode(img_path.read_bytes()).decode("utf-8") + mask_b64 = base64.b64encode(converted_mask_bytes).decode("utf-8") payload = { "init_images": [img_b64], diff --git a/backend/requirements.txt b/backend/requirements.txt index 0432abf..ad2a2b3 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -6,3 +6,5 @@ pytest>=8.2.0 pytest-asyncio>=0.23.0 aiosqlite>=0.20.0 python-multipart>=0.0.9 +Pillow>=10.0.0 +Pillow>=10.0.0 diff --git a/backend/tests/test_comfy_asset_transfer_contract.py b/backend/tests/test_comfy_asset_transfer_contract.py index 6833a73..461ce58 100644 --- a/backend/tests/test_comfy_asset_transfer_contract.py +++ b/backend/tests/test_comfy_asset_transfer_contract.py @@ -48,7 +48,8 @@ def chunk(kind: bytes, data: bytes) -> bytes: 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" + # Check for mask conversion: converted files have "_comfy" suffix + name = "mask" if b'filename="mask' in body else "source" events.append(name) if name == "mask" and fail_mask: return httpx.Response(503) diff --git a/backend/tests/test_creative_runner_upload.py b/backend/tests/test_creative_runner_upload.py index 3eba52b..abe02eb 100644 --- a/backend/tests/test_creative_runner_upload.py +++ b/backend/tests/test_creative_runner_upload.py @@ -4,27 +4,33 @@ 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 PIL import Image +from io import BytesIO from app.runners.creative_runner import CreativeRunner from app.schemas.creative import CreativeActionRequest, CreativeActionType from app.storage.asset_store import AssetRecord +def create_test_png(width: int, height: int, mode: str = "RGB") -> bytes: + """Create a valid PNG image as bytes.""" + img = Image.new(mode, (width, height), color=(255, 255, 255, 255) if mode == "RGBA" else (255, 255, 255)) + buf = BytesIO() + img.save(buf, "PNG") + return buf.getvalue() + + @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_file.write(create_test_png(512, 512, "RGB")) + img_file.flush() img_path = Path(img_file.name) with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as mask_file: - mask_file.write(PNG_BYTES) + mask_file.write(create_test_png(512, 512, "RGBA")) + mask_file.flush() mask_path = Path(mask_file.name) yield img_path, mask_path @@ -164,6 +170,7 @@ def get_path_side_effect(asset_rec): "subfolder": "berry_assets", "type": "input", }) + # Mock upload_mask to handle the converted mask file mock_comfy.upload_mask = AsyncMock(return_value={ "name": "uploaded_mask.png", "subfolder": "berry_assets", @@ -181,16 +188,20 @@ def get_path_side_effect(asset_rec): mock_cache.get_async = AsyncMock(return_value=None) mock_cache.set_async = AsyncMock() - result = await runner.execute(req) + # Mock normalize_mask_for_comfyui to avoid file system operations + with patch("app.runners.creative_runner.normalize_mask_for_comfyui") as mock_convert: + mock_convert.return_value = mask_path # Return original for test simplicity - # Verify both uploads were called - mock_comfy.upload_image.assert_called_once() - mock_comfy.upload_mask.assert_called_once() + result = await runner.execute(req) - # Verify workflow submission happened after uploads - mock_comfy.queue_prompt.assert_called_once() + # Verify both uploads were called + mock_comfy.upload_image.assert_called_once() + mock_comfy.upload_mask.assert_called_once() - assert result.success is True + # Verify workflow submission happened after uploads + mock_comfy.queue_prompt.assert_called_once() + + assert result.success is True @pytest.mark.asyncio diff --git a/backend/tests/test_mask_converter.py b/backend/tests/test_mask_converter.py new file mode 100644 index 0000000..197328f --- /dev/null +++ b/backend/tests/test_mask_converter.py @@ -0,0 +1,278 @@ +""" +Tests for mask conversion and normalization (Issue #106). + +Verifies Berry mask convention (painted white = edit, transparent = protect) +is correctly converted for each provider's specific requirements. +""" + +import pytest +import tempfile +from pathlib import Path +from PIL import Image + +from app.runners.mask_converter import ( + normalize_mask_for_comfyui, + normalize_mask_for_webui, + normalize_mask_for_openai, + normalize_mask_for_fal_ai, + validate_mask_dimensions, +) + + +def create_test_mask(width: int, height: int, painted_region: tuple) -> bytes: + """ + Create a Berry-convention test mask with a specific painted region. + + Args: + width: Mask width + height: Mask height + painted_region: (x, y, w, h) rectangle to paint white + + Returns: + PNG bytes with transparent background and white painted rectangle + """ + img = Image.new("RGBA", (width, height), (0, 0, 0, 0)) + pixels = img.load() + + x, y, w, h = painted_region + for py in range(y, min(y + h, height)): + for px in range(x, min(x + w, width)): + pixels[px, py] = (255, 255, 255, 255) # White opaque = edit region + + from io import BytesIO + output = BytesIO() + img.save(output, "PNG") + return output.getvalue() + + +def test_berry_mask_convention_documented(): + """Verify the canonical Berry mask convention is clear.""" + # Berry convention: painted (white, alpha=255) = edit, transparent (alpha=0) = protect + mask_bytes = create_test_mask(100, 100, (10, 10, 30, 30)) + img = Image.open(BytesIO(mask_bytes)).convert("RGBA") + + # Painted region should be opaque white + assert img.getpixel((20, 20)) == (255, 255, 255, 255) + + # Unpainted region should be transparent + assert img.getpixel((5, 5))[3] == 0 + + +def test_comfyui_mask_inversion(): + """ + ComfyUI LoadImage MASK output = 1 - alpha. + Painted regions (alpha=255) must become 0, transparent (alpha=0) must become 255. + """ + with tempfile.TemporaryDirectory() as tmpdir: + mask_path = Path(tmpdir) / "berry_mask.png" + mask_bytes = create_test_mask(100, 100, (20, 20, 40, 40)) + mask_path.write_bytes(mask_bytes) + + converted_path = normalize_mask_for_comfyui(mask_path) + + try: + converted = Image.open(converted_path).convert("RGBA") + + # Painted region in Berry (alpha=255) should become alpha=0 for ComfyUI + painted_pixel = converted.getpixel((30, 30)) + assert painted_pixel[3] == 0, f"Painted region should have alpha=0, got {painted_pixel[3]}" + + # Unpainted region in Berry (alpha=0) should become alpha=255 for ComfyUI + unpainted_pixel = converted.getpixel((5, 5)) + assert unpainted_pixel[3] == 255, f"Unpainted region should have alpha=255, got {unpainted_pixel[3]}" + finally: + if converted_path.exists(): + converted_path.unlink() + + +def test_webui_mask_grayscale_conversion(): + """ + WebUI expects grayscale where white=edit, black=protect. + Berry convention already matches semantically, just convert format. + """ + mask_bytes = create_test_mask(100, 100, (10, 10, 50, 50)) + converted_bytes = normalize_mask_for_webui(mask_bytes) + + converted = Image.open(BytesIO(converted_bytes)) + assert converted.mode == "L", "WebUI mask should be grayscale (mode L)" + + # Painted region should be white (255) + assert converted.getpixel((30, 30)) == 255 + + # Unpainted region should be black (0) + assert converted.getpixel((5, 5)) == 0 + + +def test_openai_mask_alpha_inversion(): + """ + OpenAI expects transparent (alpha=0) = edit, opaque (alpha=255) = protect. + This is inverted from Berry convention. + """ + mask_bytes = create_test_mask(100, 100, (15, 15, 30, 30)) + converted_bytes = normalize_mask_for_openai(mask_bytes) + + converted = Image.open(BytesIO(converted_bytes)).convert("RGBA") + + # Painted region in Berry (alpha=255) should become alpha=0 for OpenAI (edit) + painted_pixel = converted.getpixel((25, 25)) + assert painted_pixel[3] == 0, f"Painted region should have alpha=0 for OpenAI, got {painted_pixel[3]}" + + # Unpainted region in Berry (alpha=0) should become alpha=255 for OpenAI (protect) + unpainted_pixel = converted.getpixel((5, 5)) + assert unpainted_pixel[3] == 255, f"Unpainted region should have alpha=255 for OpenAI, got {unpainted_pixel[3]}" + + +def test_fal_ai_mask_passthrough(): + """ + Fal.ai expects opaque (alpha=255) = edit, transparent (alpha=0) = protect. + This matches Berry convention exactly - just validate format. + """ + mask_bytes = create_test_mask(100, 100, (20, 20, 40, 40)) + converted_bytes = normalize_mask_for_fal_ai(mask_bytes) + + # Should return original bytes since format matches + converted = Image.open(BytesIO(converted_bytes)).convert("RGBA") + + # Painted region should still be opaque + assert converted.getpixel((30, 30))[3] == 255 + + # Unpainted region should still be transparent + assert converted.getpixel((5, 5))[3] == 0 + + +def test_mask_dimension_validation_success(): + """Validate that matching dimensions pass validation.""" + with tempfile.TemporaryDirectory() as tmpdir: + mask_path = Path(tmpdir) / "mask.png" + mask_bytes = create_test_mask(512, 512, (100, 100, 200, 200)) + mask_path.write_bytes(mask_bytes) + + width, height = validate_mask_dimensions(mask_path, 512, 512) + assert width == 512 + assert height == 512 + + +def test_mask_dimension_validation_mismatch(): + """Validate that mismatched dimensions raise ValueError.""" + with tempfile.TemporaryDirectory() as tmpdir: + mask_path = Path(tmpdir) / "mask.png" + mask_bytes = create_test_mask(512, 512, (100, 100, 200, 200)) + mask_path.write_bytes(mask_bytes) + + with pytest.raises(ValueError) as exc_info: + validate_mask_dimensions(mask_path, 1024, 1024) + + assert "do not match" in str(exc_info.value).lower() + assert "512" in str(exc_info.value) + assert "1024" in str(exc_info.value) + + +def test_mask_dimension_validation_file_not_found(): + """Validate that missing mask file raises FileNotFoundError.""" + with pytest.raises(FileNotFoundError): + validate_mask_dimensions(Path("/nonexistent/mask.png"), 512, 512) + + +def test_comfyui_conversion_preserves_dimensions(): + """Verify ComfyUI conversion maintains exact dimensions.""" + with tempfile.TemporaryDirectory() as tmpdir: + mask_path = Path(tmpdir) / "mask.png" + mask_bytes = create_test_mask(768, 512, (50, 50, 100, 100)) + mask_path.write_bytes(mask_bytes) + + converted_path = normalize_mask_for_comfyui(mask_path) + + try: + converted = Image.open(converted_path) + assert converted.size == (768, 512) + converted.close() + finally: + if converted_path.exists(): + try: + converted_path.unlink() + except PermissionError: + pass # Windows file lock issue in tests + + +def test_webui_conversion_preserves_dimensions(): + """Verify WebUI conversion maintains exact dimensions.""" + mask_bytes = create_test_mask(1024, 768, (100, 100, 200, 200)) + converted_bytes = normalize_mask_for_webui(mask_bytes) + + converted = Image.open(BytesIO(converted_bytes)) + assert converted.size == (1024, 768) + + +def test_partial_alpha_mask_conversion(): + """Test conversion with partial transparency (anti-aliased edges).""" + img = Image.new("RGBA", (100, 100), (0, 0, 0, 0)) + pixels = img.load() + + # Create gradient from transparent to opaque + for x in range(50): + alpha = int((x / 50.0) * 255) + pixels[x, 50] = (255, 255, 255, alpha) + + from io import BytesIO + mask_bytes = BytesIO() + img.save(mask_bytes, "PNG") + mask_bytes = mask_bytes.getvalue() + + # ComfyUI: should invert alpha gradient + with tempfile.TemporaryDirectory() as tmpdir: + mask_path = Path(tmpdir) / "mask.png" + mask_path.write_bytes(mask_bytes) + converted = normalize_mask_for_comfyui(mask_path) + comfy_img = Image.open(converted).convert("RGBA") + + # Check gradient is inverted (approximately, allowing for conversion rounding) + assert comfy_img.getpixel((0, 50))[3] == 255 # Was 0, should be 255 + assert comfy_img.getpixel((49, 50))[3] <= 10 # Was ~255, should be ~0 (allow small error) + + comfy_img.close() + converted.unlink(missing_ok=True) + + +def test_fully_painted_mask(): + """Test mask that is completely painted (no transparent regions).""" + img = Image.new("RGBA", (50, 50), (255, 255, 255, 255)) + + from io import BytesIO + mask_bytes = BytesIO() + img.save(mask_bytes, "PNG") + mask_bytes = mask_bytes.getvalue() + + # ComfyUI: all should become alpha=0 + with tempfile.TemporaryDirectory() as tmpdir: + mask_path = Path(tmpdir) / "full_mask.png" + mask_path.write_bytes(mask_bytes) + converted = normalize_mask_for_comfyui(mask_path) + comfy_img = Image.open(converted).convert("RGBA") + converted.unlink() + + assert comfy_img.getpixel((25, 25))[3] == 0 + assert comfy_img.getpixel((0, 0))[3] == 0 + + +def test_fully_transparent_mask(): + """Test mask that is completely transparent (no painted regions).""" + img = Image.new("RGBA", (50, 50), (0, 0, 0, 0)) + + from io import BytesIO + mask_bytes = BytesIO() + img.save(mask_bytes, "PNG") + mask_bytes = mask_bytes.getvalue() + + # ComfyUI: all should become alpha=255 + with tempfile.TemporaryDirectory() as tmpdir: + mask_path = Path(tmpdir) / "empty_mask.png" + mask_path.write_bytes(mask_bytes) + converted = normalize_mask_for_comfyui(mask_path) + comfy_img = Image.open(converted).convert("RGBA") + converted.unlink() + + assert comfy_img.getpixel((25, 25))[3] == 255 + assert comfy_img.getpixel((0, 0))[3] == 255 + + +from io import BytesIO diff --git a/backend/tests/test_webui_checkpoint.py b/backend/tests/test_webui_checkpoint.py index f2036e1..4501094 100644 --- a/backend/tests/test_webui_checkpoint.py +++ b/backend/tests/test_webui_checkpoint.py @@ -14,17 +14,24 @@ import uuid import httpx import pytest +from PIL import Image +from io import BytesIO from app.runners.creative_runner import CreativeRunner from app.runners.webui_runner import WebUIRunner from app.schemas.creative import CreativeActionRequest, CreativeActionType from app.storage.asset_store import asset_store -DUMMY_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" -) + +def create_dummy_png(width: int = 512, height: int = 512) -> bytes: + """Create a valid PNG for testing.""" + img = Image.new("RGBA", (width, height), color=(255, 255, 255, 255)) + buf = BytesIO() + img.save(buf, "PNG") + return buf.getvalue() + + +DUMMY_PNG_BYTES = create_dummy_png() B64_DUMMY = base64.b64encode(DUMMY_PNG_BYTES).decode("utf-8") AVAILABLE_MODELS = [