add size checks to other options (#612)

Author: wirthual

libs/infinity_emb/infinity_emb/transformer/vision/utils.py

@@ -21,6 +21,7 @@
 
 def resolve_from_img_obj(img_obj: "ImageClassType") -> ImageSingle:
     """Resolve an image from a ImageClassType Object."""
+    assert_image_has_valid_size(img_obj)
     return ImageSingle(image=img_obj)
 
 
@@ -34,22 +35,27 @@ async def resolve_from_img_url(img_url: str, session: "aiohttp.ClientSession") -
 
     try:
         img = Image.open(io.BytesIO(downloaded_img))
-        if img.size[0] < 3 or img.size[1] < 3:
-            # https://upload.wikimedia.org/wikipedia/commons/c/ca/1x1.png
-            raise ImageCorruption(
-                f"An image in your request is too small for processing {img.size}"
-            )
+        assert_image_has_valid_size(img)
         return ImageSingle(image=img)
     except Exception as e:
         raise ImageCorruption(
             f"error opening the payload from an image in your request from url: {e}"
         )
 
 
+def assert_image_has_valid_size(img: "ImageClassType") -> None:
+    if img.size[0] < 3 or img.size[1] < 3:
+        # https://upload.wikimedia.org/wikipedia/commons/c/ca/1x1.png
+        raise ImageCorruption(
+            f"An image in your request is too small for processing {img.size}. Minimum required size is 3x3."
+        )
+
+
 def resolve_from_img_bytes(bytes_img: bytes) -> ImageSingle:
     """Resolve an image from a Data URI"""
     try:
         img = Image.open(io.BytesIO(bytes_img))
+        assert_image_has_valid_size(img)
         return ImageSingle(image=img)
     except Exception as e:
         raise ImageCorruption(f"error decoding data URI: {e}")

libs/infinity_emb/tests/end_to_end/test_torch_vision.py

@@ -220,3 +220,24 @@ async def test_vision_fail(client):
         json={"model": MODEL, "input": ["test"]},
     )
     assert response_unsupported.status_code == status.HTTP_400_BAD_REQUEST
+
+
+@pytest.mark.anyio
+async def test_vision_base64_too_small(client, image_sample):
+    response = await client.post(
+        f"{PREFIX}/embeddings_image",
+        json={
+            "model": MODEL,
+            "input": [
+                "data:image/jpeg;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAusB9sXsdrgAAAAASUVORK5CYII=",
+                pytest.IMAGE_SAMPLE_URL,
+            ],
+        },
+    )
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    rdata = response.json()
+    message = rdata["error"]["message"]
+    assert (
+        message
+        == "<class 'infinity_emb.primitives.ImageCorruption'> -> error decoding data URI: An image in your request is too small for processing (1, 1). Minimum required size is 3x3."
+    )

libs/infinity_emb/tests/unit_test/test_engine.py

@@ -14,6 +14,7 @@
     EmbeddingDtype,
     InferenceEngine,
     ModelNotDeployedError,
+    ImageCorruption,
 )
 
 # Only compile on Linux 3.9-3.11 with torch
@@ -238,6 +239,15 @@ async def test_torch_clip_embed():
     for i in range(1, len(sentences)):
         assert np.dot(emb_text_np[0], emb_image_np[0]) > np.dot(emb_text_np[i], emb_image_np[0])
 
+    async with engine:
+        with pytest.raises(
+            ImageCorruption,
+            match=r"An image in your request is too small for processing \(1, 1\)\. Minimum required size is 3x3\.",
+        ):
+            small_image = Image.new("L", (1, 1))
+            t_small_image = asyncio.create_task(engine.image_embed(images=[small_image]))
+            await t_small_image
+
 
 @pytest.mark.anyio
 async def test_clap_like_model(audio_sample):