Spaces:
Running on Zero
Running on Zero
| """Shared, CPU-only validation and image transport for the private Dev service.""" | |
| import base64 | |
| import binascii | |
| import io | |
| import warnings | |
| from typing import Literal | |
| from PIL import Image | |
| from pydantic import BaseModel, ConfigDict, Field, model_validator | |
| MAX_IMAGE_BYTES = 20 * 1024 * 1024 | |
| MAX_BODY_BYTES = 28 * 1024 * 1024 | |
| MAX_PIXELS = 20_000_000 | |
| SIZES = { | |
| (1024, 1024), | |
| (1152, 896), | |
| (896, 1152), | |
| (1280, 768), | |
| (768, 1280), | |
| (1344, 768), | |
| (768, 1344), | |
| (512, 512), | |
| (576, 448), | |
| (448, 576), | |
| (640, 384), | |
| (384, 640), | |
| } | |
| def validate_image(image): | |
| if image.width * image.height > MAX_PIXELS or min(image.size) < 16: | |
| raise ValueError("Use an image between 16 pixels per side and 20 megapixels.") | |
| if max(image.size) / min(image.size) > 8: | |
| raise ValueError("Image aspect ratio must be at most 8:1.") | |
| def decode_image(encoded): | |
| try: | |
| raw = base64.b64decode(encoded, validate=True) | |
| if len(raw) > MAX_IMAGE_BYTES: | |
| raise ValueError("Image exceeds 20 MB.") | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("error", Image.DecompressionBombWarning) | |
| with Image.open(io.BytesIO(raw)) as image: | |
| validate_image(image) | |
| if image.format not in {"PNG", "JPEG", "WEBP"} or getattr(image, "n_frames", 1) != 1: | |
| raise ValueError("Use a single PNG, JPEG or WebP image.") | |
| image.load() | |
| return image.copy() | |
| except (OSError, binascii.Error, Image.DecompressionBombError, Image.DecompressionBombWarning) as error: | |
| raise ValueError("Invalid or oversized image.") from error | |
| def encode_image(image): | |
| validate_image(image) | |
| buffer = io.BytesIO() | |
| image.save(buffer, format="PNG") | |
| if buffer.tell() > MAX_IMAGE_BYTES: | |
| raise ValueError("Image exceeds 20 MB; resize it first.") | |
| return base64.b64encode(buffer.getvalue()).decode("ascii") | |
| class JobRequest(BaseModel): | |
| model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False) | |
| task: Literal["image_to_depth", "text_to_image", "text_to_rgbd", "depth_to_image", "reimagine"] | |
| prompt: str = Field(default="", max_length=4000) | |
| image: str | None = Field(default=None, max_length=28 * 1024 * 1024) | |
| width: int = 1024 | |
| height: int = 1024 | |
| resolution: Literal[512, 1024] = 512 | |
| steps: int = Field(default=50, ge=1, le=100) | |
| guidance: float = Field(default=3.5, ge=1, le=10) | |
| seed: int = Field(default=0, ge=0, lt=2**31) | |
| depth_format: Literal["depth", "disparity", "colours", "encoded"] = "colours" | |
| def validate_task(self): | |
| if (self.width, self.height) not in SIZES: | |
| raise ValueError("Unsupported output dimensions.") | |
| needs_image = self.task in {"image_to_depth", "depth_to_image", "reimagine"} | |
| if needs_image != (self.image is not None): | |
| raise ValueError("This task requires exactly its supported inputs.") | |
| if self.task == "image_to_depth" and self.steps > 8: | |
| raise ValueError("Depth estimation supports 1–8 steps.") | |
| return self | |