Spaces:
Running on Zero
Running on Zero
File size: 3,170 Bytes
7423487 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """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"
@model_validator(mode="after")
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
|