fix: sync inpainting.py/model_loader.py with main repo - deploy real StableDiffusionInpainter

The Space had been running a stale LamaModelInpainter this whole time (the
crop-back-to-original-size fix for SimpleLama's multiple-of-8 padding was
never pushed here), which crashed with 'images do not match' on any crop
whose dimensions weren't already divisible by 8 - i.e. almost every real
photo, only missed by tiny square 256x256 test images. Also brings over
StableDiffusionInpainter, RemoteInpainter, proportional crop padding, and
CORS middleware from the main repo.

server/app.py

@@ -1,6 +1,7 @@
 import io
 
 from fastapi import FastAPI, File, HTTPException, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
 from pydantic import BaseModel
 
@@ -38,6 +39,15 @@ def create_fastapi_app(
     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=["*"],
+    )
 
     @app.get("/health")
     def health():
@@ -72,7 +82,10 @@ def create_fastapi_app(
         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)
-        result = composite_inpaint(pil_image, pil_mask, inpainter)
+        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

server/inpainting.py

@@ -6,27 +6,93 @@ from PIL import Image, ImageFilter
 class Inpainter(ABC):
     """Fills the masked region of an image."""
 
+    # Set True on an Inpainter whose inpaint() already does its own
+    # crop/composite around the mask (e.g. RemoteInpainter, which forwards
+    # to another server's full /inpaint pipeline) - callers should pass it
+    # the full image/mask directly instead of routing through
+    # composite_inpaint(), which would crop before calling it and cause a
+    # redundant, inconsistently-sized second crop/composite downstream.
+    does_own_compositing = False
+
     @abstractmethod
     def inpaint(self, image: Image.Image, mask: Image.Image) -> Image.Image:
         """image and mask are the same size; mask 255 = area to fill. Returns a filled RGB image."""
 
 
+class RemoteInpainter(Inpainter):
+    """Forwards inpainting to a remote server's /inpaint endpoint (e.g. the
+    HF Space) instead of running a local model.
+
+    Used for a hybrid self-host setup: SAM2 segmentation runs locally (light
+    enough for this machine's GPU), but Stable Diffusion inpainting - which
+    repeatedly crashed this machine's GPU driver even with reduced steps and
+    inter-step pacing - runs on the already-verified-stable cloud Space
+    instead.
+    """
+
+    # The remote /inpaint endpoint already crops around the mask, inpaints,
+    # and feather-composites the result back into the full image itself
+    # (see composite_inpaint below) - composite_inpaint() must not wrap this
+    # inpainter too, or the crop it hands to inpaint() gets composited a
+    # second time server-side, producing a differently-sized result that
+    # fails Image.composite()'s "images do not match" size check.
+    does_own_compositing = True
+
+    def __init__(self, base_url: str, timeout_seconds: float = 300.0):
+        self._base_url = base_url.rstrip("/")
+        self._timeout_seconds = timeout_seconds
+
+    def inpaint(self, image: Image.Image, mask: Image.Image) -> Image.Image:
+        import io
+
+        import httpx
+
+        image_bytes = io.BytesIO()
+        image.convert("RGB").save(image_bytes, format="PNG")
+        mask_bytes = io.BytesIO()
+        mask.convert("L").save(mask_bytes, format="PNG")
+
+        response = httpx.post(
+            f"{self._base_url}/inpaint",
+            files={
+                "image": ("image.png", image_bytes.getvalue(), "image/png"),
+                "mask": ("mask.png", mask_bytes.getvalue(), "image/png"),
+            },
+            timeout=self._timeout_seconds,
+        )
+        response.raise_for_status()
+        return Image.open(io.BytesIO(response.content)).convert("RGB")
+
+
 def composite_inpaint(
     image: Image.Image,
     mask: Image.Image,
     inpainter: Inpainter,
     padding: int = 64,
 ) -> Image.Image:
-    """Crop around the mask with padding, inpaint the crop, feather-blend it back into the full image."""
+    """Crop around the mask with padding, inpaint the crop, feather-blend it back into the full image.
+
+    `padding` is a floor, not the actual amount used: the crop grows with
+    the masked object's own size (half its bounding-box width/height) so
+    LaMa always gets a ring of real background roughly as wide as the
+    object itself, not just a fixed 64px margin - a fixed-size margin
+    starves inpainting of surrounding context for anything bigger than a
+    small object, which shows up as a blurry/smudged fill instead of a
+    clean reconstruction.
+    """
     bbox = mask.getbbox()
     if bbox is None:
         return image.copy()
 
     left, top, right, bottom = bbox
-    left = max(0, left - padding)
-    top = max(0, top - padding)
-    right = min(image.width, right + padding)
-    bottom = min(image.height, bottom + padding)
+    object_width = right - left
+    object_height = bottom - top
+    horizontal_padding = max(padding, object_width // 2)
+    vertical_padding = max(padding, object_height // 2)
+    left = max(0, left - horizontal_padding)
+    top = max(0, top - vertical_padding)
+    right = min(image.width, right + horizontal_padding)
+    bottom = min(image.height, bottom + vertical_padding)
     crop_box = (left, top, right, bottom)
 
     image_crop = image.crop(crop_box)
@@ -66,5 +132,95 @@ class LamaModelInpainter(Inpainter):
 
     def inpaint(self, image: Image.Image, mask: Image.Image) -> Image.Image:
         self._ensure_loaded()
+        original_size = image.size
         result = self._model(image.convert("RGB"), mask.convert("L"))
+        # SimpleLama pads the input up to a multiple of 8 internally and
+        # returns the padded size without cropping back down, so the result
+        # can be larger than the input - composite_inpaint() requires the
+        # returned image to match the original crop's size exactly.
+        if result.size != original_size:
+            result = result.crop((0, 0, original_size[0], original_size[1]))
+        return result.convert("RGB")
+
+
+class StableDiffusionInpainter(Inpainter):
+    """Generative inpainting via Stable Diffusion 1.5's inpainting checkpoint.
+
+    Produces noticeably better results than LaMa on cluttered/detailed
+    backgrounds (patterned wallpaper, text, shelves of objects) - LaMa is a
+    single discriminative pass that falls back to a blurry average fill
+    when it can't confidently reconstruct fine detail, while SD generatively
+    hallucinates plausible texture. Trades speed for quality: expect several
+    seconds per removal on a modern GPU, and it needs a real GPU (~4-5GB
+    VRAM) to run at a usable speed at all.
+    """
+
+    def __init__(
+        self,
+        device: str = "cuda",
+        model_id: str = "runwayml/stable-diffusion-inpainting",
+        working_size: tuple[int, int] = (512, 512),
+        num_inference_steps: int = 30,
+        step_pause_seconds: float = 0.0,
+    ):
+        self._device = device
+        self._model_id = model_id
+        self._working_size = working_size
+        self._num_inference_steps = num_inference_steps
+        # A gaming GPU's own workload (frame rendering) naturally pulses -
+        # short bursts of full load separated by idle gaps - which is why it
+        # can sustain high clocks without the rapid thermal spikes seen
+        # under uninterrupted ML compute (e.g. 73C -> 83C in ~2s observed on
+        # this machine's laptop GPU during SAM2/SD inference). Sleeping
+        # briefly between denoising steps recreates that pulsed pattern in
+        # software when a real fix (proper power/thermal limiting) isn't
+        # available - set via SD_STEP_PAUSE_SECONDS for self-hosted local
+        # runs; leave at 0 (no pause) for ZeroGPU, where calls are billed by
+        # wall-clock time and the cloud GPU doesn't have this problem.
+        self._step_pause_seconds = step_pause_seconds
+        self._pipe = None
+
+    def _ensure_loaded(self) -> None:
+        if self._pipe is not None:
+            return
+        import torch
+        from diffusers import AutoPipelineForInpainting
+
+        dtype = torch.float16 if self._device.startswith("cuda") else torch.float32
+        pipe = AutoPipelineForInpainting.from_pretrained(self._model_id, torch_dtype=dtype)
+        pipe = pipe.to(self._device)
+        pipe.set_progress_bar_config(disable=True)
+        self._pipe = pipe
+
+    def inpaint(self, image: Image.Image, mask: Image.Image) -> Image.Image:
+        self._ensure_loaded()
+        original_size = image.size
+
+        # SD1.5 was trained at 512x512; run inference at that scale
+        # regardless of crop size and resize the result back afterwards -
+        # composite_inpaint() requires the returned image to match the
+        # original crop's exact size.
+        image_resized = image.convert("RGB").resize(self._working_size, Image.LANCZOS)
+        mask_resized = mask.convert("L").resize(self._working_size, Image.NEAREST)
+
+        callback = None
+        if self._step_pause_seconds > 0:
+            import time
+
+            def callback(pipe, step_index, timestep, callback_kwargs):
+                time.sleep(self._step_pause_seconds)
+                return callback_kwargs
+
+        result = self._pipe(
+            prompt="",
+            image=image_resized,
+            mask_image=mask_resized,
+            height=self._working_size[1],
+            width=self._working_size[0],
+            num_inference_steps=self._num_inference_steps,
+            callback_on_step_end=callback,
+        ).images[0]
+
+        if result.size != original_size:
+            result = result.resize(original_size, Image.LANCZOS)
         return result.convert("RGB")

server/model_loader.py

@@ -19,9 +19,27 @@ def get_segmenter():
 def get_inpainter():
     global _inpainter
     if _inpainter is None:
-        from server.inpainting import LamaModelInpainter
+        device = os.environ.get("INFERENCE_DEVICE", "cuda")
+        backend = os.environ.get("INPAINT_BACKEND", "sd")  # "sd" or "lama"
 
-        _inpainter = LamaModelInpainter(
-            device=os.environ.get("INFERENCE_DEVICE", "cuda"),
-        )
+        if backend == "lama":
+            from server.inpainting import LamaModelInpainter
+
+            _inpainter = LamaModelInpainter(device=device)
+        elif backend == "remote":
+            from server.inpainting import RemoteInpainter
+
+            _inpainter = RemoteInpainter(
+                base_url=os.environ["REMOTE_INPAINT_URL"],
+                timeout_seconds=float(os.environ.get("REMOTE_INPAINT_TIMEOUT_SECONDS", "300")),
+            )
+        else:
+            from server.inpainting import StableDiffusionInpainter
+
+            _inpainter = StableDiffusionInpainter(
+                device=device,
+                model_id=os.environ.get("SD_INPAINT_MODEL_ID", "runwayml/stable-diffusion-inpainting"),
+                step_pause_seconds=float(os.environ.get("SD_STEP_PAUSE_SECONDS", "0")),
+                num_inference_steps=int(os.environ.get("SD_NUM_INFERENCE_STEPS", "30")),
+            )
     return _inpainter