update to use the official Finegrain API directly

README.md

@@ -4,7 +4,8 @@ emoji: 🧽
 colorFrom: gray
 colorTo: blue
 sdk: gradio
-sdk_version: 4.42.0
+sdk_version: 4.44.1
+python_version: 3.12
 app_file: src/app.py
 pinned: false
 license: other

pyproject.toml

@@ -6,13 +6,14 @@ authors = [
     { name = "Pierre Chapuis", email = "[email protected]" }
 ]
 dependencies = [
-    "gradio>=4.41.0",
+    "gradio>=4.41.0,<5",   # gradio-imageslider requires <5
     "environs>=11.0.0",
-    "gradio-image-annotation @ https://huggingface.co/spaces/finegrain/finegrain-object-eraser/resolve/main/gradio_image_annotation-0.2.3-py3-none-any.whl",
+    "gradio-image-annotation>=0.2.5",
     "httpx>=0.27.0",
     "pillow>=10.4.0",
     "gradio-imageslider>=0.0.20",
     "pillow-heif>=0.18.0",
+    "httpx-sse>=0.4.0",
 ]
 readme = "README.md"
 requires-python = ">= 3.12, <3.13"

requirements.lock

@@ -14,93 +14,95 @@ aiofiles==23.2.1
     # via gradio
 annotated-types==0.7.0
     # via pydantic
-anyio==4.4.0
+anyio==4.8.0
     # via gradio
     # via httpx
     # via starlette
-certifi==2024.8.30
+certifi==2024.12.14
     # via httpcore
     # via httpx
     # via requests
-charset-normalizer==3.3.2
+charset-normalizer==3.4.1
     # via requests
-click==8.1.7
+click==8.1.8
     # via typer
     # via uvicorn
-contourpy==1.3.0
+contourpy==1.3.1
     # via matplotlib
 cycler==0.12.1
     # via matplotlib
-environs==11.0.0
+environs==14.1.0
     # via eraser
-fastapi==0.112.2
+fastapi==0.115.6
     # via gradio
-ffmpy==0.4.0
+ffmpy==0.5.0
     # via gradio
-filelock==3.15.4
+filelock==3.16.1
     # via huggingface-hub
-fonttools==4.53.1
+fonttools==4.55.3
     # via matplotlib
-fsspec==2024.6.1
+fsspec==2024.12.0
     # via gradio-client
     # via huggingface-hub
-gradio==4.42.0
+gradio==4.44.1
     # via eraser
     # via gradio-image-annotation
     # via gradio-imageslider
 gradio-client==1.3.0
     # via gradio
-gradio-image-annotation @ https://huggingface.co/spaces/finegrain/finegrain-object-eraser/resolve/main/gradio_image_annotation-0.2.3-py3-none-any.whl
+gradio-image-annotation==0.2.5
     # via eraser
 gradio-imageslider==0.0.20
     # via eraser
 h11==0.14.0
     # via httpcore
     # via uvicorn
-httpcore==1.0.5
+httpcore==1.0.7
     # via httpx
-httpx==0.27.2
+httpx==0.28.1
     # via eraser
     # via gradio
     # via gradio-client
-huggingface-hub==0.24.6
+httpx-sse==0.4.0
+    # via eraser
+huggingface-hub==0.27.1
     # via gradio
     # via gradio-client
-idna==3.8
+idna==3.10
     # via anyio
     # via httpx
     # via requests
-importlib-resources==6.4.4
+importlib-resources==6.5.2
     # via gradio
-jinja2==3.1.4
+jinja2==3.1.5
     # via gradio
-kiwisolver==1.4.5
+kiwisolver==1.4.8
     # via matplotlib
 markdown-it-py==3.0.0
     # via rich
 markupsafe==2.1.5
     # via gradio
     # via jinja2
-marshmallow==3.22.0
+marshmallow==3.25.1
     # via environs
-matplotlib==3.9.2
+matplotlib==3.10.0
     # via gradio
 mdurl==0.1.2
     # via markdown-it-py
-numpy==2.1.0
+numpy==2.2.2
     # via contourpy
     # via gradio
     # via matplotlib
     # via pandas
-orjson==3.10.7
+orjson==3.10.15
     # via gradio
-packaging==24.1
+packaging==24.2
     # via gradio
     # via gradio-client
     # via huggingface-hub
     # via marshmallow
     # via matplotlib
-pandas==2.2.2
+pandas==2.2.3
     # via gradio
 pillow==10.4.0
     # via eraser
@@ -108,55 +110,55 @@ pillow==10.4.0
     # via gradio-imageslider
     # via matplotlib
     # via pillow-heif
-pillow-heif==0.18.0
+pillow-heif==0.21.0
     # via eraser
-pydantic==2.8.2
+pydantic==2.10.5
     # via fastapi
     # via gradio
-pydantic-core==2.20.1
+pydantic-core==2.27.2
     # via pydantic
 pydub==0.25.1
     # via gradio
-pygments==2.18.0
+pygments==2.19.1
     # via rich
-pyparsing==3.1.4
+pyparsing==3.2.1
     # via matplotlib
 python-dateutil==2.9.0.post0
     # via matplotlib
     # via pandas
 python-dotenv==1.0.1
     # via environs
-python-multipart==0.0.9
+python-multipart==0.0.20
     # via gradio
-pytz==2024.1
+pytz==2024.2
     # via pandas
 pyyaml==6.0.2
     # via gradio
     # via huggingface-hub
 requests==2.32.3
     # via huggingface-hub
-rich==13.8.0
+rich==13.9.4
     # via typer
-ruff==0.6.3
+ruff==0.9.2
     # via gradio
 semantic-version==2.10.0
     # via gradio
 shellingham==1.5.4
     # via typer
-six==1.16.0
+six==1.17.0
     # via python-dateutil
 sniffio==1.3.1
     # via anyio
-    # via httpx
-starlette==0.38.4
+starlette==0.41.3
     # via fastapi
 tomlkit==0.12.0
     # via gradio
-tqdm==4.66.5
+tqdm==4.67.1
     # via huggingface-hub
-typer==0.12.5
+typer==0.15.1
     # via gradio
 typing-extensions==4.12.2
+    # via anyio
     # via fastapi
     # via gradio
     # via gradio-client
@@ -164,12 +166,12 @@ typing-extensions==4.12.2
     # via pydantic
     # via pydantic-core
     # via typer
-tzdata==2024.1
+tzdata==2024.2
     # via pandas
-urllib3==2.2.2
+urllib3==2.3.0
     # via gradio
     # via requests
-uvicorn==0.30.6
+uvicorn==0.34.0
     # via gradio
 websockets==12.0
     # via gradio-client

requirements.txt

@@ -1,6 +1,7 @@
-https://huggingface.co/spaces/finegrain/finegrain-object-eraser/resolve/main/gradio_image_annotation-0.2.3-py3-none-any.whl
+gradio_image_annotation>=0.2.5
 gradio_imageslider>=0.0.20
 environs>=11.0.0
 httpx>=0.27.0
+httpx-sse>=0.4.0
 pillow>=10.4.0
 pillow-heif>=0.18.0

src/app.py

@@ -1,14 +1,16 @@
+import dataclasses as dc
 import io
 from typing import Any
 
 import gradio as gr
-import httpx
 import pillow_heif
 from environs import Env
 from gradio_image_annotation import image_annotator
 from gradio_imageslider import ImageSlider
 from PIL import Image
 
+from fg import EditorAPIContext
+
 pillow_heif.register_heif_opener()
 pillow_heif.register_avif_opener()
 
@@ -16,11 +18,16 @@ env = Env()
 env.read_env()
 
 with env.prefixed("ERASER_"):
-    API_URL: str = str(env.str("API_URL", "https://spaces.finegrain.ai/eraser"))
-    API_KEY: str | None = env.str("API_KEY", None)
+    API_URL: str = str(env.str("API_URL", "https://api.finegrain.ai/editor"))
+    API_USER: str | None = env.str("API_USER")
+    API_PASSWORD: str | None = env.str("API_PASSWORD")
     CA_BUNDLE: str | None = env.str("CA_BUNDLE", None)
 
-auth = None if API_KEY is None else httpx.BasicAuth("hf", API_KEY)
+assert API_USER is not None
+assert API_PASSWORD is not None
+CTX = EditorAPIContext(uri=API_URL, user=API_USER, password=API_PASSWORD)
+if CA_BUNDLE:
+    CTX.verify = CA_BUNDLE
 
 
 def resize(image: Image.Image, shortest_side: int = 768) -> Image.Image:
@@ -31,38 +38,97 @@ def resize(image: Image.Image, shortest_side: int = 768) -> Image.Image:
     return image.resize(size=(int(shortest_side * image.width / image.height), shortest_side))
 
 
-def process_bbox(
-    prompts: dict[str, Any],
-    request: gr.Request | None,
-) -> tuple[Image.Image, Image.Image]:
+@dc.dataclass(kw_only=True)
+class ProcessParams:
+    image: Image.Image
+    prompt: str | None = None
+    bbox: tuple[int, int, int, int] | None = None
+
+
+async def _process(ctx: EditorAPIContext, params: ProcessParams) -> Image.Image:
+    with io.BytesIO() as f:
+        params.image.save(f, format="JPEG")
+        async with ctx as client:
+            response = await client.post(
+                f"{ctx.uri}/state/upload",
+                files={"file": f},
+                headers=ctx.auth_headers,
+            )
+    response.raise_for_status()
+    st_input = response.json()["state"]
+
+    if params.bbox:
+        segment_input_st = st_input
+        segment_params = {"bbox": list(params.bbox)}
+    else:
+        assert params.prompt
+        async with ctx as client:
+            response = await client.post(
+                f"{ctx.uri}/skills/infer-bbox/{st_input}",
+                json={"product_name": params.prompt},
+                headers=ctx.auth_headers,
+            )
+        response.raise_for_status()
+        st_bbox = response.json()["state"]
+        await ctx.sse_await(st_bbox)
+        segment_input_st = st_bbox
+        segment_params = {}
+
+    async with ctx as client:
+        response = await client.post(
+            f"{ctx.uri}/skills/segment/{segment_input_st}",
+            json=segment_params,
+            headers=ctx.auth_headers,
+        )
+    response.raise_for_status()
+    st_mask = response.json()["state"]
+    await ctx.sse_await(st_mask)
+
+    erase_params: dict[str, str | bool] = {
+        "mode": "free",  # new API
+        "restore_original_resolution": False,  # legacy API
+    }
+    async with ctx as client:
+        response = await client.post(
+            f"{ctx.uri}/skills/erase/{st_input}/{st_mask}",
+            json=erase_params,
+            headers=ctx.auth_headers,
+        )
+    response.raise_for_status()
+    st_erased = response.json()["state"]
+    await ctx.sse_await(st_erased)
+
+    async with ctx as client:
+        response = await client.get(
+            f"{ctx.uri}/state/image/{st_erased}",
+            params={"format": "JPEG", "resolution": "DISPLAY"},
+            headers=ctx.auth_headers,
+        )
+    response.raise_for_status()
+    f = io.BytesIO()
+    f.write(response.content)
+    f.seek(0)
+    return Image.open(f)
+
+
+def process_bbox(prompts: dict[str, Any]) -> tuple[Image.Image, Image.Image]:
     assert isinstance(img := prompts["image"], Image.Image)
     assert isinstance(boxes := prompts["boxes"], list)
     assert len(boxes) == 1
     assert isinstance(box := boxes[0], dict)
-    headers = {}
-    if request:  # avoid DOS - can be None despite type hint!
-        client_ip = request.headers.get("x-forwarded-for") or request.client.host
-        headers = {"X-HF-Client-IP": client_ip}
 
     resized_img = resize(img)
     bbox = [box[k] for k in ["xmin", "ymin", "xmax", "ymax"]]
     if resized_img.width != img.width:
         bbox = [int(v * resized_img.width / img.width) for v in bbox]
 
-    with io.BytesIO() as f:
-        resized_img.save(f, format="JPEG")
-        r = httpx.post(
-            API_URL,
-            data={"bbox": ",".join([str(v) for v in bbox])},
-            files={"file": f},
-            verify=CA_BUNDLE or True,
-            timeout=30.0,
-            auth=auth,
-            headers=headers,
-        )
-    r.raise_for_status()
-
-    output_image = Image.open(io.BytesIO(r.content))
+    output_image = CTX.run_one_sync(
+        _process,
+        ProcessParams(
+            image=resized_img,
+            bbox=(bbox[0], bbox[1], bbox[2], bbox[3]),
+        ),
+    )
     return (img, output_image)
 
 
@@ -70,32 +136,12 @@ def on_change_bbox(prompts: dict[str, Any] | None):
     return gr.update(interactive=prompts is not None and len(prompts["boxes"]) > 0)
 
 
-def process_prompt(
-    img: Image.Image,
-    prompt: str,
-    request: gr.Request | None,
-) -> tuple[Image.Image, Image.Image]:
-    headers = {}
-    if request:  # avoid DOS - can be None despite type hint!
-        client_ip = request.headers.get("x-forwarded-for") or request.client.host
-        headers = {"X-HF-Client-IP": client_ip}
-
+def process_prompt(img: Image.Image, prompt: str) -> tuple[Image.Image, Image.Image]:
     resized_img = resize(img)
-
-    with io.BytesIO() as f:
-        resized_img.save(f, format="JPEG")
-        r = httpx.post(
-            API_URL,
-            data={"prompt": prompt},
-            files={"file": f},
-            verify=CA_BUNDLE or True,
-            timeout=30.0,
-            auth=auth,
-            headers=headers,
-        )
-    r.raise_for_status()
-
-    output_image = Image.open(io.BytesIO(r.content))
+    output_image = CTX.run_one_sync(
+        _process,
+        ProcessParams(image=resized_img, prompt=prompt),
+    )
     return (img, output_image)
 
 
@@ -112,8 +158,9 @@ TITLE = """
     padding: 0.5rem 1rem;
     font-size: 1.25rem;
   ">
-    🚀 For an optimized version of this space, try out the 
-    <a href="https://finegrain.ai/editor?utm_source=hf&utm_campaign=object-eraser" target="_blank">Finegrain Editor</a>! You'll find there all our AI tools made available in a nice UI. 🚀
+    🚀 For an optimized version of this space, try out the
+    <a href="https://finegrain.ai/editor?utm_source=hf&utm_campaign=object-eraser" target="_blank">Finegrain Editor</a>!
+    You'll find there all our AI tools made available in a nice UI. 🚀
   </div>
 
   <h1 style="font-size: 1.5rem; margin-bottom: 0.5rem;">

src/fg.py

@@ -0,0 +1,117 @@
+import asyncio
+import dataclasses as dc
+import json
+from collections import defaultdict
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+import httpx
+import httpx_sse
+
+
+def _new_future() -> asyncio.Future[Any]:
+    return asyncio.get_running_loop().create_future()
+
+
+@dc.dataclass(kw_only=True)
+class EditorAPIContext:
+    uri: str
+    user: str
+    password: str
+    token: str | None = None
+    verify: bool | str = True
+    _client: httpx.AsyncClient | None = None
+
+    sse_futures: dict[str, asyncio.Future[dict[str, Any]]] = dc.field(default_factory=lambda: defaultdict(_new_future))
+
+    async def __aenter__(self) -> httpx.AsyncClient:
+        if self._client:
+            return self._client
+        self._client = httpx.AsyncClient(verify=self.verify)
+        return self._client
+
+    async def __aexit__(self, *args: Any) -> None:
+        if self._client:
+            await self._client.__aexit__(*args)
+            self._client = None
+
+    @property
+    def auth_headers(self) -> dict[str, str]:
+        assert self.token
+        return {"Authorization": f"Bearer {self.token}"}
+
+    async def login(self) -> None:
+        async with self as client:
+            response = await client.post(
+                f"{self.uri}/auth/login",
+                json={"username": self.user, "password": self.password},
+            )
+            response.raise_for_status()
+            self.token = response.json()["token"]
+
+    async def sse_loop(self) -> None:
+        async with self as client:
+            response = await client.post(f"{self.uri}/sub-auth", headers=self.auth_headers)
+        response.raise_for_status()
+        sub_token = response.json()["token"]
+        url = f"{self.uri}/sub/{sub_token}"
+        async with (
+            httpx.AsyncClient(timeout=None, verify=self.verify) as c,
+            httpx_sse.aconnect_sse(c, "GET", url) as es,
+        ):
+            future = self.sse_futures["_sse_loop"]
+            future.set_result({"status": "ok"})
+            async for sse in es.aiter_sse():
+                jdata = json.loads(sse.data)
+                future = self.sse_futures[jdata["state"]]
+                future.set_result(jdata)
+
+    async def sse_await(self, state_id: str) -> None:
+        future = self.sse_futures[state_id]
+        jdata = await future
+        if jdata["status"] != "ok":
+            print("ERROR", jdata)
+        assert jdata["status"] == "ok"
+        del self.sse_futures[state_id]
+
+    async def get_meta(self, state_id: str) -> dict[str, Any]:
+        async with self as client:
+            response = await client.get(
+                f"{self.uri}/state/meta/{state_id}",
+                headers=self.auth_headers,
+            )
+            response.raise_for_status()
+            return response.json()
+
+    async def run_one[Tin, Tout](
+        self,
+        co: Callable[["EditorAPIContext", Tin], Awaitable[Tout]],
+        params: Tin,
+    ) -> Tout:
+        await self.login()
+        async with asyncio.TaskGroup() as tg:
+            sse_task = tg.create_task(self.sse_loop())
+
+            async def outer_co(params: Tin) -> Tout:
+                # _sse_loop is a fake event to wait until the SSE loop is properly setup.
+                await self.sse_await("_sse_loop")
+                r = await co(self, params)
+                sse_task.cancel()
+                return r
+
+            r = tg.create_task(outer_co(params))
+
+        return r.result()
+
+    def run_one_sync[Tin, Tout](
+        self,
+        co: Callable[["EditorAPIContext", Tin], Awaitable[Tout]],
+        params: Tin,
+    ) -> Tout:
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        return loop.run_until_complete(self.run_one(co, params))