Spaces:
Running
Running
| from __future__ import annotations | |
| import hashlib | |
| import ipaddress | |
| import io | |
| import json | |
| import logging | |
| import os | |
| import secrets | |
| import socket | |
| import threading | |
| import time | |
| import warnings | |
| from contextlib import asynccontextmanager | |
| from dataclasses import dataclass, field | |
| from typing import Any, Callable, Literal | |
| from urllib.parse import SplitResult, urlsplit | |
| import httpx | |
| from fastapi import APIRouter, HTTPException, Request | |
| from fastapi import status as http_status | |
| from fastapi.responses import JSONResponse | |
| from PIL import Image, ImageOps, UnidentifiedImageError | |
| from pydantic import BaseModel, ConfigDict, Field, field_validator | |
| LOGGER = logging.getLogger(__name__) | |
| DNS_RESOLVER_SLOTS = threading.BoundedSemaphore(4) | |
| IMAGE_FETCH_SLOTS = threading.BoundedSemaphore(4) | |
| MIN_IMAGE_DIMENSION = 32 | |
| DEFAULT_MAX_IMAGE_BYTES = 8 * 1024 * 1024 | |
| MAX_CONFIGURABLE_IMAGE_BYTES = 20 * 1024 * 1024 | |
| DEFAULT_MAX_IMAGE_PIXELS = 4_194_304 | |
| MAX_CONFIGURABLE_IMAGE_PIXELS = 20_000_000 | |
| DEFAULT_MAX_IMAGE_DIMENSION = 4096 | |
| MAX_CONFIGURABLE_IMAGE_DIMENSION = 8192 | |
| DEFAULT_MAX_RECORDS = 256 | |
| MIN_MAX_RECORDS = 8 | |
| MAX_MAX_RECORDS = 4096 | |
| JobStatus = Literal["queued", "running", "succeeded", "failed"] | |
| SummarySeparator = Literal["comma", "newline", "space"] | |
| AnalysisExecutor = Callable[ | |
| ["ImageInversionJobRequest", Image.Image, str], | |
| dict[str, Any], | |
| ] | |
| ImageFetcher = Callable[[str, "ImageInversionJobSettings"], Image.Image] | |
| def _read_int_env(name: str, default: int) -> int: | |
| """读取整数环境变量。 | |
| Args: | |
| name: 环境变量名称。 | |
| default: 变量为空时采用的默认值。 | |
| Returns: | |
| 解析完成的整数。 | |
| """ | |
| raw_value = os.getenv(name) | |
| if raw_value is None or not raw_value.strip(): | |
| return default | |
| try: | |
| return int(raw_value) | |
| except ValueError as exc: | |
| raise RuntimeError(f"{name} must be an integer") from exc | |
| def _read_float_env(name: str, default: float) -> float: | |
| """读取有限浮点环境变量。 | |
| Args: | |
| name: 环境变量名称。 | |
| default: 变量为空时采用的默认值。 | |
| Returns: | |
| 解析完成的有限浮点数。 | |
| """ | |
| raw_value = os.getenv(name) | |
| if raw_value is None or not raw_value.strip(): | |
| return default | |
| try: | |
| parsed_value = float(raw_value) | |
| except ValueError as exc: | |
| raise RuntimeError(f"{name} must be a number") from exc | |
| if not float("-inf") < parsed_value < float("inf"): | |
| raise RuntimeError(f"{name} must be finite") | |
| return parsed_value | |
| def normalize_allowed_hosts(raw_hosts: str) -> frozenset[str]: | |
| """规范化远程图片精确主机白名单。 | |
| Args: | |
| raw_hosts: 以逗号分隔的主机名,不允许协议、端口、路径或通配符。 | |
| Returns: | |
| 完成小写与 IDNA 转换的不可变主机集合。 | |
| """ | |
| normalized_hosts: set[str] = set() | |
| for raw_host in raw_hosts.split(","): | |
| host = raw_host.strip().rstrip(".") | |
| if not host: | |
| continue | |
| if any(marker in host for marker in ("://", "/", "?", "#", "*", "@", ":")): | |
| raise RuntimeError( | |
| "JOB_IMAGE_ALLOWED_HOSTS must contain exact hostnames only" | |
| ) | |
| try: | |
| normalized_hosts.add(host.encode("idna").decode("ascii").lower()) | |
| except UnicodeError as exc: | |
| raise RuntimeError( | |
| "JOB_IMAGE_ALLOWED_HOSTS contains an invalid hostname" | |
| ) from exc | |
| if not normalized_hosts: | |
| raise RuntimeError("JOB_IMAGE_ALLOWED_HOSTS must contain at least one hostname") | |
| return frozenset(normalized_hosts) | |
| class ImageInversionJobSettings: | |
| """保存 Image Inversion 自定义任务 API 的部署边界与资源限制。""" | |
| api_key: str | |
| allowed_hosts: frozenset[str] | |
| result_ttl_seconds: int = 1800 | |
| poll_after_seconds: int = 2 | |
| fetch_timeout_seconds: float = 15.0 | |
| max_image_bytes: int = DEFAULT_MAX_IMAGE_BYTES | |
| max_image_pixels: int = DEFAULT_MAX_IMAGE_PIXELS | |
| max_image_dimension: int = DEFAULT_MAX_IMAGE_DIMENSION | |
| max_records: int = DEFAULT_MAX_RECORDS | |
| space_host: str = "" | |
| def from_env(cls) -> "ImageInversionJobSettings": | |
| """从 Space Secret 与 Variables 读取并校验任务 API 配置。 | |
| Args: | |
| 此方法不接收参数,配置统一从当前进程环境读取。 | |
| Returns: | |
| 已完成密钥、网络及图片资源边界校验的设置。 | |
| """ | |
| api_key = os.getenv("JOB_API_KEY", "").strip() | |
| if len(api_key) < 32: | |
| raise RuntimeError( | |
| "JOB_API_KEY must be configured as a Space Secret with at least 32 characters" | |
| ) | |
| result_ttl_seconds = _read_int_env("JOB_RESULT_TTL_SECONDS", 1800) | |
| poll_after_seconds = _read_int_env("JOB_POLL_AFTER_SECONDS", 2) | |
| fetch_timeout_seconds = _read_float_env( | |
| "JOB_IMAGE_FETCH_TIMEOUT_SECONDS", | |
| 15.0, | |
| ) | |
| max_image_bytes = _read_int_env( | |
| "JOB_IMAGE_MAX_BYTES", | |
| DEFAULT_MAX_IMAGE_BYTES, | |
| ) | |
| max_image_pixels = _read_int_env( | |
| "JOB_IMAGE_MAX_PIXELS", | |
| DEFAULT_MAX_IMAGE_PIXELS, | |
| ) | |
| max_image_dimension = _read_int_env( | |
| "JOB_IMAGE_MAX_DIMENSION", | |
| DEFAULT_MAX_IMAGE_DIMENSION, | |
| ) | |
| max_records = _read_int_env("JOB_MAX_RECORDS", DEFAULT_MAX_RECORDS) | |
| if not 60 <= result_ttl_seconds <= 86400: | |
| raise RuntimeError("JOB_RESULT_TTL_SECONDS must be between 60 and 86400") | |
| if not 1 <= poll_after_seconds <= 30: | |
| raise RuntimeError("JOB_POLL_AFTER_SECONDS must be between 1 and 30") | |
| if not 1.0 <= fetch_timeout_seconds <= 15.0: | |
| raise RuntimeError( | |
| "JOB_IMAGE_FETCH_TIMEOUT_SECONDS must be between 1 and 15" | |
| ) | |
| if not 1024 <= max_image_bytes <= MAX_CONFIGURABLE_IMAGE_BYTES: | |
| raise RuntimeError("JOB_IMAGE_MAX_BYTES must be between 1 KiB and 20 MiB") | |
| if not 1 <= max_image_pixels <= MAX_CONFIGURABLE_IMAGE_PIXELS: | |
| raise RuntimeError( | |
| "JOB_IMAGE_MAX_PIXELS must be between 1 and 20000000" | |
| ) | |
| if not MIN_IMAGE_DIMENSION <= max_image_dimension <= MAX_CONFIGURABLE_IMAGE_DIMENSION: | |
| raise RuntimeError( | |
| "JOB_IMAGE_MAX_DIMENSION must be between 32 and 8192" | |
| ) | |
| if not MIN_MAX_RECORDS <= max_records <= MAX_MAX_RECORDS: | |
| raise RuntimeError("JOB_MAX_RECORDS must be between 8 and 4096") | |
| return cls( | |
| api_key=api_key, | |
| allowed_hosts=normalize_allowed_hosts( | |
| os.getenv("JOB_IMAGE_ALLOWED_HOSTS", "") | |
| ), | |
| result_ttl_seconds=result_ttl_seconds, | |
| poll_after_seconds=poll_after_seconds, | |
| fetch_timeout_seconds=fetch_timeout_seconds, | |
| max_image_bytes=max_image_bytes, | |
| max_image_pixels=max_image_pixels, | |
| max_image_dimension=max_image_dimension, | |
| max_records=max_records, | |
| space_host=os.getenv("SPACE_HOST", "").strip().rstrip("/"), | |
| ) | |
| class ImageInversionJobRequest(BaseModel): | |
| """定义 Image Inversion 标签分析任务的公开请求参数。""" | |
| model_config = ConfigDict( | |
| extra="forbid", | |
| str_strip_whitespace=True, | |
| allow_inf_nan=False, | |
| ) | |
| input_image_url: str = Field(min_length=1, max_length=4096) | |
| general_threshold: float = Field(default=0.35, ge=0.0, le=1.0) | |
| character_threshold: float = Field(default=0.85, ge=0.0, le=1.0) | |
| show_confidence: bool = True | |
| show_general: bool = True | |
| show_character: bool = True | |
| show_ip: bool = True | |
| separator: SummarySeparator = "comma" | |
| show_chinese: bool = True | |
| def require_https_image_url(cls, value: str) -> str: | |
| """在 Schema 阶段拒绝非 HTTPS、凭据、片段和非标准端口。 | |
| Args: | |
| value: 调用方提交的远程图片 URL。 | |
| Returns: | |
| 通过基础 HTTPS 语法检查的原始 URL。 | |
| """ | |
| try: | |
| parsed = urlsplit(value) | |
| port = parsed.port | |
| except ValueError as exc: | |
| raise ValueError("image URL is invalid") from exc | |
| if ( | |
| parsed.scheme.lower() != "https" | |
| or not parsed.hostname | |
| or parsed.username is not None | |
| or parsed.password is not None | |
| or parsed.fragment | |
| or port not in {None, 443} | |
| ): | |
| raise ValueError( | |
| "image URL must use credential-free HTTPS on the default port" | |
| ) | |
| return value | |
| class ValidatedImageURL: | |
| """保存通过协议、主机与公网地址检查的图片 URL。""" | |
| value: str | |
| host: str | |
| class ImageSourceError(RuntimeError): | |
| """表示图片来源不合法或暂时无法抓取。""" | |
| def __init__(self, status_code: int, public_message: str) -> None: | |
| """创建携带脱敏公开信息的图片来源错误。 | |
| Args: | |
| status_code: API 应返回的 HTTP 状态码。 | |
| public_message: 不包含 URL、凭据或底层网络细节的公开文案。 | |
| Returns: | |
| 此初始化方法不返回数据。 | |
| """ | |
| super().__init__(public_message) | |
| self.status_code = status_code | |
| self.public_message = public_message | |
| def _resolve_host_addresses(host: str, timeout_seconds: float) -> list[Any]: | |
| """在硬期限内解析白名单主机的地址。 | |
| Args: | |
| host: 已通过精确白名单校验的规范化主机名。 | |
| timeout_seconds: DNS 解析允许占用的最长墙钟秒数。 | |
| Returns: | |
| socket.getaddrinfo 返回的非空地址记录列表。 | |
| """ | |
| completed = threading.Event() | |
| result: dict[str, Any] = {} | |
| if not DNS_RESOLVER_SLOTS.acquire(blocking=False): | |
| raise ImageSourceError( | |
| http_status.HTTP_504_GATEWAY_TIMEOUT, | |
| "The image host resolution is temporarily busy.", | |
| ) | |
| def resolve() -> None: | |
| try: | |
| result["records"] = socket.getaddrinfo( | |
| host, | |
| 443, | |
| type=socket.SOCK_STREAM, | |
| ) | |
| except Exception as exc: | |
| result["error"] = exc | |
| finally: | |
| DNS_RESOLVER_SLOTS.release() | |
| completed.set() | |
| resolver = threading.Thread( | |
| target=resolve, | |
| name="image-inversion-dns", | |
| daemon=True, | |
| ) | |
| try: | |
| resolver.start() | |
| except Exception: | |
| DNS_RESOLVER_SLOTS.release() | |
| raise | |
| if not completed.wait(timeout=max(0.001, timeout_seconds)): | |
| raise ImageSourceError( | |
| http_status.HTTP_504_GATEWAY_TIMEOUT, | |
| "The image host resolution timed out.", | |
| ) | |
| resolution_error = result.get("error") | |
| if resolution_error is not None: | |
| raise ImageSourceError( | |
| http_status.HTTP_502_BAD_GATEWAY, | |
| "The image host could not be resolved.", | |
| ) from resolution_error | |
| address_records = result.get("records") | |
| if not address_records: | |
| raise ImageSourceError( | |
| http_status.HTTP_502_BAD_GATEWAY, | |
| "The image host did not return an address.", | |
| ) | |
| return address_records | |
| def validate_image_url( | |
| url: str, | |
| settings: ImageInversionJobSettings, | |
| ) -> ValidatedImageURL: | |
| """校验远程图片 URL 只指向白名单内的公网 HTTPS 主机。 | |
| Args: | |
| url: 调用方提交的完整图片 URL。 | |
| settings: 包含白名单与网络资源限制的任务 API 设置。 | |
| Returns: | |
| 已规范化主机名并可安全用于抓取的 URL 描述。 | |
| """ | |
| try: | |
| parsed: SplitResult = urlsplit(url) | |
| port = parsed.port | |
| except ValueError as exc: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The image URL is invalid.", | |
| ) from exc | |
| if parsed.scheme.lower() != "https" or not parsed.hostname: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "Image URLs must use HTTPS.", | |
| ) | |
| if parsed.username is not None or parsed.password is not None or parsed.fragment: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The image URL contains unsupported credentials or fragments.", | |
| ) | |
| if port not in {None, 443}: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The image URL must use the default HTTPS port.", | |
| ) | |
| try: | |
| normalized_host = ( | |
| parsed.hostname.rstrip(".").encode("idna").decode("ascii").lower() | |
| ) | |
| except UnicodeError as exc: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The image URL hostname is invalid.", | |
| ) from exc | |
| if normalized_host not in settings.allowed_hosts: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The image URL host is not allowed.", | |
| ) | |
| address_records = _resolve_host_addresses( | |
| normalized_host, | |
| settings.fetch_timeout_seconds, | |
| ) | |
| # 白名单限定预期域名;逐个拒绝非公网解析结果可阻断错误配置和 DNS 绕过。 | |
| for address_record in address_records: | |
| raw_address = address_record[4][0].split("%", 1)[0] | |
| try: | |
| resolved_ip = ipaddress.ip_address(raw_address) | |
| except ValueError as exc: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The image host resolved to an invalid address.", | |
| ) from exc | |
| if not resolved_ip.is_global: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The image host must resolve only to public addresses.", | |
| ) | |
| return ValidatedImageURL(value=url, host=normalized_host) | |
| def _validate_connected_peer(response: httpx.Response) -> None: | |
| """确认 HTTP 客户端实际连接的对端仍是公网 IP。 | |
| Args: | |
| response: 已建立 TLS 连接并收到响应头的 httpx 响应。 | |
| Returns: | |
| 对端地址可验证且属于公网时不返回数据。 | |
| """ | |
| network_stream = response.extensions.get("network_stream") | |
| if network_stream is None: | |
| raise ImageSourceError( | |
| http_status.HTTP_502_BAD_GATEWAY, | |
| "The image connection peer could not be verified.", | |
| ) | |
| try: | |
| peer_address = network_stream.get_extra_info("server_addr") | |
| raw_address = peer_address[0].split("%", 1)[0] | |
| peer_ip = ipaddress.ip_address(raw_address) | |
| except (AttributeError, IndexError, TypeError, ValueError) as exc: | |
| raise ImageSourceError( | |
| http_status.HTTP_502_BAD_GATEWAY, | |
| "The image connection peer could not be verified.", | |
| ) from exc | |
| if not peer_ip.is_global: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The image connection peer must be a public address.", | |
| ) | |
| def decode_image_bytes( | |
| image_bytes: bytes, | |
| content_type: str, | |
| settings: ImageInversionJobSettings, | |
| ) -> Image.Image: | |
| """校验远程内容并解码为脱离底层流的 RGB 图片。 | |
| Args: | |
| image_bytes: 在字节上限内完整读取的图片内容。 | |
| content_type: 远程响应声明的 MIME 类型,仅为接口兼容保留。 | |
| settings: 图片字节、像素及单边尺寸限制。 | |
| Returns: | |
| 完成 EXIF 方向修正并复制到内存的 RGB PIL 图片。 | |
| """ | |
| if not image_bytes or len(image_bytes) > settings.max_image_bytes: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The remote image is empty or exceeds the configured size limit.", | |
| ) | |
| # 远端可合法返回 application/octet-stream;只信任 Pillow 识别的真实格式。 | |
| _ = content_type | |
| try: | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("error", Image.DecompressionBombWarning) | |
| with Image.open(io.BytesIO(image_bytes)) as source_image: | |
| if source_image.format not in {"JPEG", "PNG", "WEBP"}: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The remote image format is not supported.", | |
| ) | |
| width, height = source_image.size | |
| if width < MIN_IMAGE_DIMENSION or height < MIN_IMAGE_DIMENSION: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The remote image dimensions are not supported.", | |
| ) | |
| if ( | |
| width > settings.max_image_dimension | |
| or height > settings.max_image_dimension | |
| or width * height > settings.max_image_pixels | |
| ): | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The remote image exceeds the configured dimension limit.", | |
| ) | |
| source_image.load() | |
| return ImageOps.exif_transpose(source_image).convert("RGB").copy() | |
| except ImageSourceError: | |
| raise | |
| except ( | |
| UnidentifiedImageError, | |
| OSError, | |
| ValueError, | |
| Image.DecompressionBombError, | |
| Image.DecompressionBombWarning, | |
| ) as exc: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The remote image could not be decoded safely.", | |
| ) from exc | |
| def _map_image_request_error( | |
| exc: Exception, | |
| deadline_reached: threading.Event, | |
| ) -> ImageSourceError: | |
| """把底层 HTTP 错误映射为不含远程 URL 的公开错误。 | |
| Args: | |
| exc: httpx 或客户端关闭路径产生的底层异常。 | |
| deadline_reached: 硬期限看门狗是否已经触发。 | |
| Returns: | |
| 可由 API 层安全公开的图片来源错误。 | |
| """ | |
| if deadline_reached.is_set() or isinstance(exc, httpx.TimeoutException): | |
| return ImageSourceError( | |
| http_status.HTTP_504_GATEWAY_TIMEOUT, | |
| "The image download timed out.", | |
| ) | |
| return ImageSourceError( | |
| http_status.HTTP_502_BAD_GATEWAY, | |
| "The image host is temporarily unavailable.", | |
| ) | |
| def _fetch_remote_image_inner( | |
| url: str, | |
| settings: ImageInversionJobSettings, | |
| ) -> Image.Image: | |
| """从白名单公网主机流式抓取并安全解码图片。 | |
| Args: | |
| url: 调用方提交的图片 HTTPS URL。 | |
| settings: 图片白名单、超时、体积、像素与尺寸边界。 | |
| Returns: | |
| 可直接交给标签模型的 RGB PIL 图片。 | |
| """ | |
| started_at = time.monotonic() | |
| validated_url = validate_image_url(url, settings) | |
| remaining_seconds = settings.fetch_timeout_seconds - ( | |
| time.monotonic() - started_at | |
| ) | |
| if remaining_seconds <= 0: | |
| raise ImageSourceError( | |
| http_status.HTTP_504_GATEWAY_TIMEOUT, | |
| "The image download timed out.", | |
| ) | |
| timeout = httpx.Timeout( | |
| remaining_seconds, | |
| connect=min(5.0, remaining_seconds), | |
| ) | |
| deadline_reached = threading.Event() | |
| client = httpx.Client( | |
| timeout=timeout, | |
| follow_redirects=False, | |
| trust_env=False, | |
| ) | |
| def stop_at_deadline() -> None: | |
| deadline_reached.set() | |
| try: | |
| client.close() | |
| except Exception: | |
| pass | |
| deadline_timer = threading.Timer(remaining_seconds, stop_at_deadline) | |
| deadline_timer.daemon = True | |
| try: | |
| with client: | |
| deadline_timer.start() | |
| with client.stream( | |
| "GET", | |
| validated_url.value, | |
| headers={"Accept": "image/png,image/jpeg,image/webp"}, | |
| ) as response: | |
| if time.monotonic() - started_at > settings.fetch_timeout_seconds: | |
| raise ImageSourceError( | |
| http_status.HTTP_504_GATEWAY_TIMEOUT, | |
| "The image download timed out.", | |
| ) | |
| _validate_connected_peer(response) | |
| if 300 <= response.status_code < 400: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "Image URL redirects are not allowed.", | |
| ) | |
| if response.status_code == http_status.HTTP_408_REQUEST_TIMEOUT: | |
| raise ImageSourceError( | |
| http_status.HTTP_504_GATEWAY_TIMEOUT, | |
| "The image download timed out.", | |
| ) | |
| if response.status_code == 429 or response.status_code >= 500: | |
| raise ImageSourceError( | |
| http_status.HTTP_502_BAD_GATEWAY, | |
| "The image host is temporarily unavailable.", | |
| ) | |
| if response.status_code != http_status.HTTP_200_OK: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The image URL did not return a readable resource.", | |
| ) | |
| declared_length = response.headers.get("content-length") | |
| if declared_length: | |
| try: | |
| parsed_length = int(declared_length) | |
| if parsed_length < 0: | |
| raise ValueError("negative content length") | |
| if parsed_length > settings.max_image_bytes: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The remote image exceeds the configured size limit.", | |
| ) | |
| except ValueError as exc: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The image host returned an invalid content length.", | |
| ) from exc | |
| image_buffer = bytearray() | |
| for chunk in response.iter_bytes(): | |
| if time.monotonic() - started_at > settings.fetch_timeout_seconds: | |
| raise ImageSourceError( | |
| http_status.HTTP_504_GATEWAY_TIMEOUT, | |
| "The image download timed out.", | |
| ) | |
| image_buffer.extend(chunk) | |
| if len(image_buffer) > settings.max_image_bytes: | |
| raise ImageSourceError( | |
| http_status.HTTP_400_BAD_REQUEST, | |
| "The remote image exceeds the configured size limit.", | |
| ) | |
| decoded_image = decode_image_bytes( | |
| bytes(image_buffer), | |
| response.headers.get("content-type", ""), | |
| settings, | |
| ) | |
| deadline_timer.cancel() | |
| if ( | |
| deadline_reached.is_set() | |
| or time.monotonic() - started_at > settings.fetch_timeout_seconds | |
| ): | |
| decoded_image.close() | |
| raise ImageSourceError( | |
| http_status.HTTP_504_GATEWAY_TIMEOUT, | |
| "The image download timed out.", | |
| ) | |
| return decoded_image | |
| except ImageSourceError: | |
| raise | |
| except (httpx.RequestError, RuntimeError) as exc: | |
| raise _map_image_request_error(exc, deadline_reached) from exc | |
| finally: | |
| deadline_timer.cancel() | |
| def fetch_remote_image( | |
| url: str, | |
| settings: ImageInversionJobSettings, | |
| ) -> Image.Image: | |
| """以有界后台抓取器执行下载并强制墙钟总期限。 | |
| Args: | |
| url: 调用方提交的图片 HTTPS URL。 | |
| settings: 白名单、硬超时、字节、像素与尺寸边界。 | |
| Returns: | |
| 在总期限内完成验证、EXIF 转正和 RGB 转换的 PIL 图片。 | |
| """ | |
| if not IMAGE_FETCH_SLOTS.acquire(blocking=False): | |
| raise ImageSourceError( | |
| http_status.HTTP_502_BAD_GATEWAY, | |
| "The image fetch service is temporarily busy.", | |
| ) | |
| completed = threading.Event() | |
| state_lock = threading.Lock() | |
| state: dict[str, Any] = {"abandoned": False} | |
| def fetch() -> None: | |
| fetched_image: Image.Image | None = None | |
| try: | |
| fetched_image = _fetch_remote_image_inner(url, settings) | |
| should_close = False | |
| with state_lock: | |
| if state["abandoned"]: | |
| should_close = True | |
| else: | |
| state["image"] = fetched_image | |
| fetched_image = None | |
| if should_close and fetched_image is not None: | |
| fetched_image.close() | |
| fetched_image = None | |
| except Exception as exc: | |
| with state_lock: | |
| if not state["abandoned"]: | |
| state["error"] = exc | |
| finally: | |
| if fetched_image is not None: | |
| try: | |
| fetched_image.close() | |
| except Exception as exc: | |
| LOGGER.warning( | |
| "Failed to close an abandoned fetched image: %s", | |
| type(exc).__name__, | |
| ) | |
| IMAGE_FETCH_SLOTS.release() | |
| completed.set() | |
| fetch_thread = threading.Thread( | |
| target=fetch, | |
| name="image-inversion-fetch", | |
| daemon=True, | |
| ) | |
| try: | |
| fetch_thread.start() | |
| except Exception as exc: | |
| IMAGE_FETCH_SLOTS.release() | |
| raise ImageSourceError( | |
| http_status.HTTP_502_BAD_GATEWAY, | |
| "The image fetch service could not start.", | |
| ) from exc | |
| if not completed.wait(timeout=settings.fetch_timeout_seconds): | |
| abandoned_image: Image.Image | None = None | |
| with state_lock: | |
| state["abandoned"] = True | |
| abandoned_image = state.pop("image", None) | |
| if abandoned_image is not None: | |
| try: | |
| abandoned_image.close() | |
| except Exception as exc: | |
| LOGGER.warning( | |
| "Failed to close a timed-out fetched image: %s", | |
| type(exc).__name__, | |
| ) | |
| raise ImageSourceError( | |
| http_status.HTTP_504_GATEWAY_TIMEOUT, | |
| "The image download timed out.", | |
| ) | |
| with state_lock: | |
| fetch_error = state.get("error") | |
| fetched_image = state.get("image") | |
| if fetch_error is not None: | |
| if isinstance(fetch_error, ImageSourceError): | |
| raise fetch_error | |
| raise ImageSourceError( | |
| http_status.HTTP_502_BAD_GATEWAY, | |
| "The image host is temporarily unavailable.", | |
| ) from fetch_error | |
| if fetched_image is None: | |
| raise ImageSourceError( | |
| http_status.HTTP_502_BAD_GATEWAY, | |
| "The image response could not be completed.", | |
| ) | |
| return fetched_image | |
| def _constant_time_text_equal(supplied: str, expected: str) -> bool: | |
| """以字节形式比较不可信文本。 | |
| Args: | |
| supplied: 请求携带的待校验文本。 | |
| expected: 服务端保存的期望文本。 | |
| Returns: | |
| 两个 UTF-8 字节序列完全一致时返回 True。 | |
| """ | |
| return secrets.compare_digest( | |
| supplied.encode("utf-8", errors="surrogatepass"), | |
| expected.encode("utf-8", errors="surrogatepass"), | |
| ) | |
| def fingerprint_request(payload: ImageInversionJobRequest) -> str: | |
| """生成与 JSON 字段顺序无关的任务请求指纹。 | |
| Args: | |
| payload: 已通过 Pydantic 校验的标签分析参数。 | |
| Returns: | |
| 用于幂等键复用校验的 SHA-256 摘要。 | |
| """ | |
| canonical_payload = json.dumps( | |
| payload.model_dump(mode="json"), | |
| ensure_ascii=False, | |
| separators=(",", ":"), | |
| sort_keys=True, | |
| ) | |
| return hashlib.sha256(canonical_payload.encode("utf-8")).hexdigest() | |
| class ImageInversionJobRecord: | |
| """保存单个标签分析任务在当前 Space 进程内的生命周期数据。""" | |
| job_id: str | |
| payload: ImageInversionJobRequest | |
| request_fingerprint: str | |
| input_image: Image.Image | None | |
| idempotency_key: str | None = None | |
| status: JobStatus = "queued" | |
| created_at: float = field(default_factory=time.time) | |
| started_at: float | None = None | |
| completed_at: float | None = None | |
| result: dict[str, Any] | None = None | |
| error_type: str | None = None | |
| class ImageInversionJobAPI: | |
| """实现 Image Inversion 异步 JSON 任务 API 状态机。""" | |
| def __init__( | |
| self, | |
| settings: ImageInversionJobSettings, | |
| executor: AnalysisExecutor, | |
| inference_slot: threading.Lock, | |
| image_fetcher: ImageFetcher = fetch_remote_image, | |
| ) -> None: | |
| """创建任务服务并注册 Image Inversion 自定义路由。 | |
| Args: | |
| settings: 已校验的部署配置和资源限制。 | |
| executor: 复用现有标签分析链路的执行适配函数。 | |
| inference_slot: 与 UI 共用的非阻塞单推理门闩。 | |
| image_fetcher: 可替换的安全远程图片抓取器。 | |
| Returns: | |
| 此初始化方法不返回数据。 | |
| """ | |
| self.settings = settings | |
| self.executor = executor | |
| self.inference_slot = inference_slot | |
| self.image_fetcher = image_fetcher | |
| self.jobs: dict[str, ImageInversionJobRecord] = {} | |
| self.idempotency_jobs: dict[str, str] = {} | |
| self.jobs_lock = threading.RLock() | |
| self.active_job_id: str | None = None | |
| self.router = APIRouter(tags=["jobs"]) | |
| self.router.add_api_route( | |
| "/api/jobs", | |
| self.create_job, | |
| methods=["POST"], | |
| name="image_inversion_create_job", | |
| ) | |
| self.router.add_api_route( | |
| "/api/jobs/{job_id}", | |
| self.get_job_status, | |
| methods=["GET"], | |
| name="image_inversion_get_job_status", | |
| ) | |
| def install_on_app(self, app: Any) -> None: | |
| """把自定义任务路由安装到 Gradio FastAPI 应用前部。 | |
| Args: | |
| app: Gradio 在 launch 阶段创建的 FastAPI 应用。 | |
| Returns: | |
| 路由已存在或安装完成后不返回数据。 | |
| """ | |
| if getattr(app.state, "image_inversion_job_api_registered", False): | |
| return | |
| original_route_count = len(app.router.routes) | |
| app.include_router(self.router) | |
| added_routes = app.router.routes[original_route_count:] | |
| original_routes = app.router.routes[:original_route_count] | |
| # Gradio 含宽泛路由,自定义 API 必须排在其前面才能稳定命中。 | |
| app.router.routes[:] = [*added_routes, *original_routes] | |
| app.state.image_inversion_job_api_registered = True | |
| def _require_api_key(self, request: Request) -> None: | |
| supplied_key = request.headers.get("x-api-key", "") | |
| if not supplied_key or not _constant_time_text_equal( | |
| supplied_key, | |
| self.settings.api_key, | |
| ): | |
| raise HTTPException( | |
| status_code=http_status.HTTP_401_UNAUTHORIZED, | |
| detail="A valid X-API-Key header is required.", | |
| ) | |
| def _public_route_url( | |
| self, | |
| request: Request, | |
| route_name: str, | |
| **path_params: str, | |
| ) -> str: | |
| path = str(request.app.url_path_for(route_name, **path_params)) | |
| if self.settings.space_host: | |
| space_host = self.settings.space_host | |
| if not space_host.startswith(("http://", "https://")): | |
| space_host = f"https://{space_host}" | |
| return f"{space_host}{path}" | |
| return str(request.url_for(route_name, **path_params)) | |
| def _is_expired( | |
| self, | |
| record: ImageInversionJobRecord, | |
| now: float | None = None, | |
| ) -> bool: | |
| current_time = time.time() if now is None else now | |
| if record.completed_at is None: | |
| return False | |
| return current_time - record.completed_at > self.settings.result_ttl_seconds | |
| def _close_images(*images: Image.Image | None) -> None: | |
| """逐一关闭已知图片并隔离单个清理异常。 | |
| Args: | |
| images: 要释放的零个或多个明确图片对象。 | |
| Returns: | |
| 清理完成后不返回数据;关闭错误仅写入脱敏日志。 | |
| """ | |
| for image in images: | |
| if image is not None: | |
| try: | |
| image.close() | |
| except Exception as exc: | |
| LOGGER.warning( | |
| "Failed to close an Image Inversion job image: %s", | |
| type(exc).__name__, | |
| ) | |
| def _close_job_image(record: ImageInversionJobRecord | None) -> None: | |
| if record is None: | |
| return | |
| ImageInversionJobAPI._close_images(record.input_image) | |
| record.input_image = None | |
| def _remove_record_locked(self, record: ImageInversionJobRecord) -> None: | |
| self.jobs.pop(record.job_id, None) | |
| if record.idempotency_key is not None: | |
| if self.idempotency_jobs.get(record.idempotency_key) == record.job_id: | |
| self.idempotency_jobs.pop(record.idempotency_key, None) | |
| def remove_expired_jobs(self) -> None: | |
| """清理当前所有过期终态任务及其输入图片引用。 | |
| Args: | |
| 此方法不接收参数。 | |
| Returns: | |
| 清理完成后不返回数据。 | |
| """ | |
| expired_records: list[ImageInversionJobRecord] = [] | |
| now = time.time() | |
| with self.jobs_lock: | |
| for record in tuple(self.jobs.values()): | |
| if self._is_expired(record, now): | |
| self._remove_record_locked(record) | |
| expired_records.append(record) | |
| for expired_record in expired_records: | |
| self._close_job_image(expired_record) | |
| def _ensure_capacity_for_new_record(self) -> None: | |
| """为一个新任务腾出有限记录容量。 | |
| Args: | |
| 此方法不接收参数。 | |
| Returns: | |
| 容量已经可用时不返回数据;无法安全腾出容量时抛出 HTTP 503。 | |
| """ | |
| evicted_records: list[ImageInversionJobRecord] = [] | |
| with self.jobs_lock: | |
| now = time.time() | |
| for record in tuple(self.jobs.values()): | |
| if self._is_expired(record, now): | |
| self._remove_record_locked(record) | |
| evicted_records.append(record) | |
| while len(self.jobs) >= self.settings.max_records: | |
| terminal_records = [ | |
| record | |
| for record in self.jobs.values() | |
| if record.status in {"succeeded", "failed"} | |
| and record.completed_at is not None | |
| ] | |
| if not terminal_records: | |
| raise HTTPException( | |
| status_code=http_status.HTTP_503_SERVICE_UNAVAILABLE, | |
| detail="The prediction job store is full.", | |
| headers={"Retry-After": "5"}, | |
| ) | |
| # 仅淘汰最老终态;绝不删除 queued/running 任务。 | |
| oldest_terminal = min( | |
| terminal_records, | |
| key=lambda record: record.completed_at or record.created_at, | |
| ) | |
| self._remove_record_locked(oldest_terminal) | |
| evicted_records.append(oldest_terminal) | |
| for evicted_record in evicted_records: | |
| self._close_job_image(evicted_record) | |
| def _job_or_404(self, job_id: str) -> ImageInversionJobRecord: | |
| expired_record: ImageInversionJobRecord | None = None | |
| with self.jobs_lock: | |
| record = self.jobs.get(job_id) | |
| if record is not None and self._is_expired(record): | |
| expired_record = record | |
| self._remove_record_locked(record) | |
| record = None | |
| self._close_job_image(expired_record) | |
| if record is None: | |
| raise HTTPException( | |
| status_code=http_status.HTTP_404_NOT_FOUND, | |
| detail="Job not found or result expired.", | |
| ) | |
| return record | |
| def _status_response( | |
| self, | |
| request: Request, | |
| record: ImageInversionJobRecord, | |
| ) -> JSONResponse: | |
| status_url = self._public_route_url( | |
| request, | |
| "image_inversion_get_job_status", | |
| job_id=record.job_id, | |
| ) | |
| return JSONResponse( | |
| status_code=http_status.HTTP_202_ACCEPTED, | |
| headers={ | |
| "Location": status_url, | |
| "Cache-Control": "no-store", | |
| }, | |
| content={ | |
| "job_id": record.job_id, | |
| "status_url": status_url, | |
| "poll_after_seconds": self.settings.poll_after_seconds, | |
| }, | |
| ) | |
| def _lookup_idempotent_job( | |
| self, | |
| idempotency_key: str | None, | |
| request_fingerprint: str, | |
| ) -> ImageInversionJobRecord | None: | |
| if idempotency_key is None: | |
| return None | |
| with self.jobs_lock: | |
| existing_job_id = self.idempotency_jobs.get(idempotency_key) | |
| record = self.jobs.get(existing_job_id or "") | |
| if record is not None and self._is_expired(record): | |
| self._remove_record_locked(record) | |
| record = None | |
| if record is None: | |
| if existing_job_id is not None: | |
| self.idempotency_jobs.pop(idempotency_key, None) | |
| return None | |
| if not secrets.compare_digest( | |
| record.request_fingerprint, | |
| request_fingerprint, | |
| ): | |
| raise HTTPException( | |
| status_code=http_status.HTTP_409_CONFLICT, | |
| detail=( | |
| "Idempotency-Key was already used with a different " | |
| "request body." | |
| ), | |
| ) | |
| return record | |
| def create_job( | |
| self, | |
| payload: ImageInversionJobRequest, | |
| request: Request, | |
| ) -> JSONResponse: | |
| """验证远程图片并创建立即执行的异步标签分析任务。 | |
| Args: | |
| payload: 已通过 Pydantic 校验的命名分析参数。 | |
| request: 包含共享密钥的 FastAPI 请求。 | |
| Returns: | |
| HTTP 202 任务标识、状态 URL 与建议轮询间隔。 | |
| """ | |
| self._require_api_key(request) | |
| self.remove_expired_jobs() | |
| idempotency_key = request.headers.get("idempotency-key", "").strip() or None | |
| if idempotency_key is not None and len(idempotency_key) > 200: | |
| raise HTTPException( | |
| status_code=http_status.HTTP_400_BAD_REQUEST, | |
| detail="Idempotency-Key must be 200 characters or fewer.", | |
| ) | |
| request_fingerprint = fingerprint_request(payload) | |
| existing_record = self._lookup_idempotent_job( | |
| idempotency_key, | |
| request_fingerprint, | |
| ) | |
| if existing_record is not None: | |
| return self._status_response(request, existing_record) | |
| self._ensure_capacity_for_new_record() | |
| with self.jobs_lock: | |
| active_record = self.jobs.get(self.active_job_id or "") | |
| api_job_active = active_record is not None and active_record.status in { | |
| "queued", | |
| "running", | |
| } | |
| # 抓图前先做快速忙检查;抓图后仍在状态锁内做权威二次判定。 | |
| if api_job_active or self.inference_slot.locked(): | |
| raise HTTPException( | |
| status_code=http_status.HTTP_503_SERVICE_UNAVAILABLE, | |
| detail="The prediction service is busy.", | |
| headers={"Retry-After": "5"}, | |
| ) | |
| slot_owned_by_request = False | |
| record: ImageInversionJobRecord | None = None | |
| replay_record: ImageInversionJobRecord | None = None | |
| input_image: Image.Image | None = None | |
| try: | |
| try: | |
| input_image = self.image_fetcher( | |
| payload.input_image_url, | |
| self.settings, | |
| ) | |
| except ImageSourceError as exc: | |
| raise HTTPException( | |
| status_code=exc.status_code, | |
| detail=exc.public_message, | |
| ) from exc | |
| with self.jobs_lock: | |
| # 相同幂等键可并发抓图;只有二次判定有权创建唯一任务。 | |
| replay_record = self._lookup_idempotent_job( | |
| idempotency_key, | |
| request_fingerprint, | |
| ) | |
| if replay_record is None: | |
| self._ensure_capacity_for_new_record() | |
| active_record = self.jobs.get(self.active_job_id or "") | |
| if active_record is not None and active_record.status in { | |
| "queued", | |
| "running", | |
| }: | |
| raise HTTPException( | |
| status_code=http_status.HTTP_503_SERVICE_UNAVAILABLE, | |
| detail="The prediction service is busy.", | |
| headers={"Retry-After": "5"}, | |
| ) | |
| if not self.inference_slot.acquire(blocking=False): | |
| raise HTTPException( | |
| status_code=http_status.HTTP_503_SERVICE_UNAVAILABLE, | |
| detail="The prediction service is busy.", | |
| headers={"Retry-After": "5"}, | |
| ) | |
| slot_owned_by_request = True | |
| job_id = secrets.token_urlsafe(24) | |
| record = ImageInversionJobRecord( | |
| job_id=job_id, | |
| payload=payload, | |
| request_fingerprint=request_fingerprint, | |
| input_image=input_image, | |
| idempotency_key=idempotency_key, | |
| ) | |
| self.jobs[job_id] = record | |
| self.active_job_id = job_id | |
| if idempotency_key is not None: | |
| self.idempotency_jobs[idempotency_key] = job_id | |
| if replay_record is not None: | |
| self._close_images(input_image) | |
| input_image = None | |
| return self._status_response(request, replay_record) | |
| if record is None: | |
| raise RuntimeError("Job initialization did not produce a record.") | |
| response = self._status_response(request, record) | |
| worker = threading.Thread( | |
| target=self._execute_job, | |
| args=(record,), | |
| name=f"image-inversion-job-{record.job_id[:8]}", | |
| daemon=True, | |
| ) | |
| worker.start() | |
| # 线程启动成功后,槽位所有权转交给 worker 最外层 finally。 | |
| slot_owned_by_request = False | |
| return response | |
| except Exception: | |
| with self.jobs_lock: | |
| if record is not None: | |
| self._remove_record_locked(record) | |
| if record is not None and self.active_job_id == record.job_id: | |
| self.active_job_id = None | |
| self._close_job_image(record) | |
| if record is None: | |
| self._close_images(input_image) | |
| if slot_owned_by_request: | |
| self.inference_slot.release() | |
| raise | |
| def _execute_job(self, record: ImageInversionJobRecord) -> None: | |
| try: | |
| with self.jobs_lock: | |
| current_record = self.jobs.get(record.job_id) | |
| if current_record is not record or record.status != "queued": | |
| return | |
| record.status = "running" | |
| record.started_at = time.time() | |
| if record.input_image is None: | |
| raise RuntimeError("The job input image is unavailable.") | |
| result = self.executor( | |
| record.payload, | |
| record.input_image, | |
| record.job_id, | |
| ) | |
| if not isinstance(result, dict): | |
| raise RuntimeError("The prediction executor did not return an object.") | |
| # 任务结果必须在后台线程内完成 JSON 校验,避免轮询阶段才暴露序列化错误。 | |
| normalized_result = json.loads( | |
| json.dumps(result, ensure_ascii=False, allow_nan=False) | |
| ) | |
| with self.jobs_lock: | |
| current_record = self.jobs.get(record.job_id) | |
| if current_record is record and record.status == "running": | |
| record.status = "succeeded" | |
| record.result = normalized_result | |
| record.completed_at = time.time() | |
| except Exception as exc: | |
| LOGGER.error( | |
| "Image Inversion job %s failed with %s", | |
| record.job_id[:8], | |
| type(exc).__name__, | |
| ) | |
| with self.jobs_lock: | |
| current_record = self.jobs.get(record.job_id) | |
| if current_record is record and record.status in {"queued", "running"}: | |
| record.status = "failed" | |
| record.error_type = type(exc).__name__ | |
| record.completed_at = time.time() | |
| finally: | |
| with self.jobs_lock: | |
| if self.active_job_id == record.job_id: | |
| self.active_job_id = None | |
| self._close_job_image(record) | |
| self.inference_slot.release() | |
| self.remove_expired_jobs() | |
| def get_job_status(self, job_id: str, request: Request) -> JSONResponse: | |
| """读取标签分析任务状态并在成功后返回 JSON 结果。 | |
| Args: | |
| job_id: POST 创建任务时返回的自定义任务标识。 | |
| request: 用于鉴权的 FastAPI 请求。 | |
| Returns: | |
| 运行中返回 202,成功返回 result,失败返回脱敏 500。 | |
| """ | |
| self._require_api_key(request) | |
| record = self._job_or_404(job_id) | |
| with self.jobs_lock: | |
| status_value = record.status | |
| result_value = record.result | |
| if status_value in {"queued", "running"}: | |
| return JSONResponse( | |
| status_code=http_status.HTTP_202_ACCEPTED, | |
| headers={ | |
| "Retry-After": str(self.settings.poll_after_seconds), | |
| "Cache-Control": "no-store", | |
| }, | |
| content={"status": status_value}, | |
| ) | |
| if status_value == "failed": | |
| return JSONResponse( | |
| status_code=http_status.HTTP_500_INTERNAL_SERVER_ERROR, | |
| headers={"Cache-Control": "no-store"}, | |
| content={"error": {"code": "PREDICTION_FAILED"}}, | |
| ) | |
| if result_value is None: | |
| raise RuntimeError("A succeeded prediction job has no result.") | |
| return JSONResponse( | |
| content=result_value, | |
| headers={"Cache-Control": "no-store"}, | |
| ) | |
| def create_job_api_lifespan(job_api: ImageInversionJobAPI): | |
| """创建供 Gradio launch 组合使用的自定义路由 lifespan。 | |
| Args: | |
| job_api: 已配置执行器和共享推理锁的标签分析任务服务。 | |
| Returns: | |
| 可传给 Gradio app_kwargs 的异步 lifespan 上下文管理器。 | |
| """ | |
| async def job_api_lifespan(app: Any): | |
| """在 Gradio 应用接收请求前安装自定义任务路由。 | |
| Args: | |
| app: Gradio 创建并传入 lifespan 的 FastAPI 应用。 | |
| Returns: | |
| lifespan 启动与关闭阶段不返回业务数据。 | |
| """ | |
| job_api.install_on_app(app) | |
| yield | |
| return job_api_lifespan | |