Merge pull request #6819 from dxrpl/fix/6818-image-routes-bad-input

fix(gallery): answer 400 for a missing or invalid image on the image-edit routes (#6818)

routes/gallery/gallery_routes.py

@@ -2433,11 +2433,10 @@ def setup_gallery_routes() -> APIRouter:
         image_b64 = body.get("image")
         amount = body.get("amount", 50) / 100.0
 
-        from PIL import Image, ImageFilter
+        from PIL import ImageFilter
         import base64, io
 
-        img_bytes = _b64_bytes(image_b64)
+        img = _b64_to_pil_image(image_b64, mode="RGB")
-        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
 
         # Unsharp mask: radius=2, percent=amount*200, threshold=3
         sharpened = img.filter(ImageFilter.UnsharpMask(radius=2, percent=int(amount * 200), threshold=3))
@@ -2472,8 +2471,7 @@ def setup_gallery_routes() -> APIRouter:
         except ImportError:
             raise HTTPException(500, "Server missing a required dependency")
         # Decode source image (RGB; Real-ESRGAN doesn't preserve alpha).
-        img_bytes = _b64_bytes(image_b64)
+        src = _b64_to_pil_image(image_b64, mode="RGB")
-        src = Image.open(io.BytesIO(img_bytes)).convert("RGB")
         try:
             patch_realesrgan_torchvision_compat()
             from realesrgan import RealESRGANer
@@ -2525,8 +2523,7 @@ def setup_gallery_routes() -> APIRouter:
             import numpy as np
         except ImportError:
             raise HTTPException(500, "Server missing a required dependency")
-        img_bytes = _b64_bytes(image_b64)
+        src = _b64_to_pil_image(image_b64, mode="RGB")
-        src = Image.open(io.BytesIO(img_bytes)).convert("RGB")
         try:
             patch_realesrgan_torchvision_compat()
             from basicsr.archs.rrdbnet_arch import RRDBNet
@@ -2778,8 +2775,7 @@ def setup_gallery_routes() -> APIRouter:
         from PIL import Image
         import base64, io
 
-        img_bytes = _b64_bytes(image_b64)
+        img = _b64_to_pil_image(image_b64)
-        img = Image.open(io.BytesIO(img_bytes)).convert("RGBA")
         W, H = img.size
 
         hint = None
@@ -2868,8 +2864,7 @@ def setup_gallery_routes() -> APIRouter:
         from PIL import Image, ImageFilter, ImageEnhance
         import numpy as np
 
-        img_bytes = _b64_bytes(image_b64)
+        img = _b64_to_pil_image(image_b64, mode="RGB")
-        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
 
         # Try GFPGAN first (AI face restoration)
         try:

tests/test_gallery_image_bad_input.py

@@ -0,0 +1,46 @@
+import asyncio
+import base64
+import io
+
+import pytest
+from fastapi import HTTPException
+from PIL import Image
+
+import routes.gallery_routes as gallery
+
+ROUTES = ['sharpen', 'remove-bg', 'denoise', 'upscale-local', 'enhance-face']
+
+
+def route(name):
+    return next(r for r in gallery.setup_gallery_routes().routes if r.path == f'/api/image/{name}')
+
+
+def call(name, monkeypatch, body):
+    monkeypatch.setattr(gallery, 'require_privilege', lambda *args: 'alice')
+    monkeypatch.setattr(gallery, '_local_image_models_enabled', lambda: True)
+    monkeypatch.setattr(gallery, '_remote_image_ops_base', lambda: '')
+
+    class Request:
+        async def json(self):
+            return body
+
+    return asyncio.run(route(name).endpoint(Request()))
+
+
+@pytest.mark.parametrize('name', ROUTES)
+@pytest.mark.parametrize('body', [
+    {},
+    {'image': '!!notb64'},
+    {'image': base64.b64encode(b'hello').decode()},
+])
+def test_bad_image_is_a_client_error(name, body, monkeypatch):
+    with pytest.raises(HTTPException) as exc:
+        call(name, monkeypatch, body)
+    assert exc.value.status_code == 400
+
+
+def test_sharpen_still_accepts_a_valid_image(monkeypatch):
+    output = io.BytesIO()
+    Image.new('RGB', (4, 4), 'red').save(output, format='PNG')
+    result = call('sharpen', monkeypatch, {'image': base64.b64encode(output.getvalue()).decode()})
+    assert Image.open(io.BytesIO(base64.b64decode(result['image']))).size == (4, 4)