Spaces:
Sleeping
Sleeping
| """Default real provider — Replicate (SDXL inpainting + IP-Adapter). | |
| Sends the room image, the inpaint mask, the catalog reference image, optional | |
| depth conditioning, and a prompt to a Replicate model that supports image- | |
| conditioned inpainting, then returns the generated image. | |
| Credentials and the model id come from the environment only (see config.py). | |
| Input field names are configurable so different compatible models can be used | |
| without code changes. | |
| """ | |
| from __future__ import annotations | |
| import io | |
| import requests | |
| from PIL import Image | |
| from ..config import settings | |
| from .base import ImageGenProvider | |
| def _to_png_file(img: Image.Image, name: str) -> io.BytesIO: | |
| buf = io.BytesIO() | |
| img.save(buf, format="PNG") | |
| buf.seek(0) | |
| buf.name = name # helps the client infer the content type | |
| return buf | |
| def _output_to_image(output) -> Image.Image: | |
| """Replicate may return a URL, a list of URLs, or FileOutput objects.""" | |
| item = output[0] if isinstance(output, (list, tuple)) else output | |
| if hasattr(item, "read"): # FileOutput (newer client) | |
| return Image.open(io.BytesIO(item.read())).convert("RGB") | |
| resp = requests.get(str(item), timeout=180) | |
| resp.raise_for_status() | |
| return Image.open(io.BytesIO(resp.content)).convert("RGB") | |
| class ReplicateImageGenProvider(ImageGenProvider): | |
| name = "replicate" | |
| def __init__(self) -> None: | |
| if not settings.REPLICATE_API_TOKEN: | |
| raise RuntimeError("REPLICATE_API_TOKEN is not set") | |
| if not settings.REPLICATE_MODEL: | |
| raise RuntimeError("REPLICATE_MODEL is not set") | |
| import replicate # imported lazily so mock mode needs no creds | |
| self._client = replicate.Client(api_token=settings.REPLICATE_API_TOKEN) | |
| def generate( | |
| self, | |
| room_image: Image.Image, | |
| mask: Image.Image, | |
| reference_image: Image.Image, | |
| depth_map: Image.Image | None = None, | |
| prompt: str = "", | |
| **kwargs, | |
| ) -> Image.Image: | |
| inputs = { | |
| settings.REPLICATE_INPUT_IMAGE_KEY: _to_png_file(room_image.convert("RGB"), "room.png"), | |
| settings.REPLICATE_INPUT_MASK_KEY: _to_png_file(mask.convert("L"), "mask.png"), | |
| settings.REPLICATE_INPUT_REFERENCE_KEY: _to_png_file( | |
| reference_image.convert("RGB"), "reference.png" | |
| ), | |
| settings.REPLICATE_INPUT_PROMPT_KEY: prompt, | |
| } | |
| if depth_map is not None and kwargs.get("send_depth", True): | |
| inputs.setdefault("depth_image", _to_png_file(depth_map.convert("RGB"), "depth.png")) | |
| # Allow callers/models to inject any additional inputs verbatim. | |
| inputs.update(kwargs.get("extra_inputs") or {}) | |
| output = self._client.run(settings.REPLICATE_MODEL, input=inputs) | |
| return _output_to_image(output) | |