Spaces:
Running on Zero
Running on Zero
nmdbbb
fix: sync inpainting.py/model_loader.py with main repo - deploy real StableDiffusionInpainter
0fb4a38 | import io | |
| from fastapi import FastAPI, File, HTTPException, UploadFile | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from PIL import Image | |
| from pydantic import BaseModel | |
| from server.inpainting import Inpainter, composite_inpaint | |
| from server.segmentation import Segmenter, resize_for_upload | |
| from server.session_store import SessionStore | |
| class Point(BaseModel): | |
| x: int | |
| y: int | |
| positive: bool | |
| class SegmentPointRequest(BaseModel): | |
| session_id: str | |
| points: list[Point] | |
| class InpaintRequest(BaseModel): | |
| session_id: str | |
| def _png_response(image: Image.Image): | |
| from fastapi.responses import Response | |
| buffer = io.BytesIO() | |
| image.save(buffer, format="PNG") | |
| return Response(content=buffer.getvalue(), media_type="image/png") | |
| def create_fastapi_app( | |
| segmenter: Segmenter, | |
| inpainter: Inpainter, | |
| session_store: SessionStore, | |
| ) -> FastAPI: | |
| app = FastAPI() | |
| # Lets the Flutter web build (served from a browser origin) call this | |
| # API directly - browsers block cross-origin requests by default, and | |
| # this server has no cookies/session-cookie auth to protect. | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| def health(): | |
| return {"status": "ok"} | |
| async def segment_init(image: UploadFile = File(...)): | |
| contents = await image.read() | |
| pil_image = Image.open(io.BytesIO(contents)).convert("RGB") | |
| pil_image = resize_for_upload(pil_image) | |
| segmenter.encode(pil_image) | |
| session_id = session_store.create({"image": pil_image}) | |
| return {"session_id": session_id, "width": pil_image.width, "height": pil_image.height} | |
| def segment_point(request: SegmentPointRequest): | |
| session = session_store.get(request.session_id) | |
| if session is None: | |
| raise HTTPException(status_code=404, detail="session not found") | |
| points = [(p.x, p.y, p.positive) for p in request.points] | |
| mask = segmenter.decode(session["image"], points) | |
| return _png_response(mask) | |
| async def inpaint( | |
| image: UploadFile = File(...), | |
| mask: UploadFile = File(...), | |
| ): | |
| image_bytes = await image.read() | |
| mask_bytes = await mask.read() | |
| pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") | |
| pil_mask = Image.open(io.BytesIO(mask_bytes)).convert("L") | |
| if pil_mask.size != pil_image.size: | |
| pil_mask = pil_mask.resize(pil_image.size) | |
| if inpainter.does_own_compositing: | |
| result = inpainter.inpaint(pil_image, pil_mask) | |
| else: | |
| result = composite_inpaint(pil_image, pil_mask, inpainter) | |
| return _png_response(result) | |
| return app | |
| # ZeroGPU requires @spaces.GPU to be a real decorator on a module-level | |
| # function - assigning `obj.method = spaces.GPU(obj.method)` at runtime | |
| # (the previous approach) is not detected during Space startup ("No | |
| # @spaces.GPU function detected during startup"). `spaces` isn't installed | |
| # in the local dev/test environment (only on the real HF Space runtime), so | |
| # this whole block is guarded - the test suite only ever uses | |
| # create_fastapi_app() with fakes and never imports this branch's names. | |
| try: | |
| import spaces | |
| def _gpu_encode(image): | |
| from server.model_loader import get_segmenter | |
| return get_segmenter().encode(image) | |
| def _gpu_decode(image, points): | |
| from server.model_loader import get_segmenter | |
| return get_segmenter().decode(image, points) | |
| def _gpu_inpaint(image, mask): | |
| from server.model_loader import get_inpainter | |
| return get_inpainter().inpaint(image, mask) | |
| except ImportError: | |
| spaces = None | |
| class _GpuSegmenter(Segmenter): | |
| """Delegates to the module-level @spaces.GPU-decorated functions above.""" | |
| def encode(self, image): | |
| return _gpu_encode(image) | |
| def decode(self, image, points): | |
| return _gpu_decode(image, points) | |
| class _GpuInpainter(Inpainter): | |
| """Delegates to the module-level @spaces.GPU-decorated function above.""" | |
| def inpaint(self, image, mask): | |
| return _gpu_inpaint(image, mask) | |
| def build_production_app(): | |
| """Entrypoint used by the Hugging Face Space runtime (`python server/app.py`).""" | |
| import gradio as gr | |
| from server.session_store import SessionStore | |
| session_store = SessionStore(ttl_seconds=600) | |
| fastapi_app = create_fastapi_app(_GpuSegmenter(), _GpuInpainter(), session_store) | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| "# Object Removal API\n" | |
| "This Space exposes `/health`, `/segment/init`, `/segment/point`, " | |
| "`/inpaint` for the Object Removal mobile app. No UI is served here." | |
| ) | |
| # ZeroGPU's startup handshake (spaces/zero/gradio.py's one_launch) | |
| # monkey-patches gr.Blocks.launch itself: the very first call to | |
| # .launch() on any Blocks instance runs a startup() function that | |
| # reports the @spaces.GPU-decorated functions to HF's ZeroGPU infra. | |
| # gr.mount_gradio_app() + a manually-run uvicorn server (the previous | |
| # approach here) never calls .launch() at all, so that handshake | |
| # never fires and the infra times out waiting for it ("No | |
| # @spaces.GPU function detected during startup") - the message is | |
| # misleading, it actually means "no startup report was ever sent". | |
| # _gpu_inpaint is wired directly as a Gradio fn= here (Image+Image -> | |
| # Image maps cleanly onto gr.Image components) so the decorator scan | |
| # that runs during .launch() finds it; the real image data still | |
| # flows through the /segment/init, /segment/point, /inpaint FastAPI | |
| # routes above, which call the same underlying _gpu_encode/ | |
| # _gpu_decode/_gpu_inpaint functions directly. | |
| with gr.Row(visible=False): | |
| gr.Interface( | |
| fn=_gpu_inpaint, | |
| inputs=[gr.Image(type="pil"), gr.Image(type="pil", image_mode="L")], | |
| outputs=gr.Image(type="pil"), | |
| ) | |
| # demo.launch() (not gr.mount_gradio_app) is required so ZeroGPU's | |
| # startup handshake fires. prevent_thread_lock=True lets it start its | |
| # own uvicorn server in a background thread and return immediately, so | |
| # our custom FastAPI routes can be merged onto the same app afterwards. | |
| # ssr_mode=False is required too: HF Spaces sets GRADIO_SSR_MODE=true by | |
| # default, which installs a middleware that proxies every non-internal | |
| # path (including our /health etc. and even Gradio's own /docs) to a | |
| # Node SSR renderer before Starlette routing ever runs - every path | |
| # silently returned the same SPA shell HTML until this was disabled. | |
| served_app, _local_url, _share_url = demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| prevent_thread_lock=True, | |
| show_error=True, | |
| quiet=True, | |
| ssr_mode=False, | |
| ) | |
| served_app.include_router(fastapi_app.router) | |
| return demo | |
| if __name__ == "__main__": | |
| demo = build_production_app() | |
| demo.block_thread() | |