feat: add ai image service
This commit is contained in:
@@ -0,0 +1,383 @@
|
||||
import base64
|
||||
import binascii
|
||||
import logging
|
||||
import mimetypes
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable, Optional, Tuple
|
||||
from urllib.parse import urljoin, urlparse
|
||||
|
||||
import requests
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
API_AUTO = "auto"
|
||||
API_CHAT = "chat"
|
||||
API_GEMINI = "gemini"
|
||||
API_IMAGES = "images"
|
||||
API_IMAGES_EDITS = "images_edits"
|
||||
|
||||
SUPPORTED_API_TYPES = {
|
||||
API_AUTO,
|
||||
API_CHAT,
|
||||
API_GEMINI,
|
||||
API_IMAGES,
|
||||
API_IMAGES_EDITS,
|
||||
}
|
||||
|
||||
_BASE64_KEYS = {"image_base64", "base64", "b64_json", "data"}
|
||||
_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".gif"}
|
||||
|
||||
|
||||
class AiImageServiceError(RuntimeError):
|
||||
"""Base error for AI image generation service failures."""
|
||||
|
||||
|
||||
class AiConfigError(ValueError):
|
||||
"""Raised when model configuration is missing or invalid."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class AiModelConfig:
|
||||
"""AI image model configuration loaded from ai_models.json."""
|
||||
|
||||
url: str
|
||||
model: str
|
||||
api_key: str
|
||||
api_type: str = API_AUTO
|
||||
timeout_seconds: int = 240
|
||||
connect_timeout_seconds: int = 30
|
||||
extra_body: Any = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
if not isinstance(data, dict):
|
||||
raise AiConfigError("AI 模型配置必须是对象")
|
||||
return cls(
|
||||
url=str(data.get("url", "") or ""),
|
||||
model=str(data.get("model", "") or ""),
|
||||
api_key=str(data.get("api_key", "") or ""),
|
||||
api_type=str(data.get("api_type", API_AUTO) or API_AUTO),
|
||||
timeout_seconds=_to_int(data.get("timeout_seconds", 240), 240),
|
||||
connect_timeout_seconds=_to_int(data.get("connect_timeout_seconds", 30), 30),
|
||||
extra_body=data.get("extra_body") or {},
|
||||
)
|
||||
|
||||
|
||||
def api_config_errors(config):
|
||||
"""Return human-readable configuration errors without raising."""
|
||||
try:
|
||||
cfg = _coerce_config(config)
|
||||
except AiConfigError as exc:
|
||||
return [str(exc)]
|
||||
errors = []
|
||||
|
||||
if not cfg.url.strip():
|
||||
errors.append("缺少 url")
|
||||
else:
|
||||
scheme = urlparse(cfg.url).scheme.lower()
|
||||
if scheme not in ("http", "https"):
|
||||
errors.append("url 必须是 http(s) 地址")
|
||||
if not cfg.model.strip():
|
||||
errors.append("缺少 model")
|
||||
if not cfg.api_key.strip():
|
||||
errors.append("缺少 api_key")
|
||||
if cfg.api_type not in SUPPORTED_API_TYPES:
|
||||
errors.append("api_type 不支持: {}".format(cfg.api_type))
|
||||
if cfg.timeout_seconds <= 0:
|
||||
errors.append("timeout_seconds 必须大于 0")
|
||||
if cfg.connect_timeout_seconds <= 0:
|
||||
errors.append("connect_timeout_seconds 必须大于 0")
|
||||
if not isinstance(cfg.extra_body, dict):
|
||||
errors.append("extra_body 必须是对象")
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
def validate_api_config(config):
|
||||
errors = api_config_errors(config)
|
||||
if errors:
|
||||
raise AiConfigError("; ".join(errors))
|
||||
|
||||
|
||||
def detect_api_type(url, api_type=API_AUTO):
|
||||
"""Resolve auto/chat/gemini/images/images_edits for a model endpoint."""
|
||||
if api_type and api_type != API_AUTO:
|
||||
if api_type not in SUPPORTED_API_TYPES:
|
||||
raise AiConfigError("api_type 不支持: {}".format(api_type))
|
||||
return api_type
|
||||
|
||||
lower_url = str(url).lower()
|
||||
if "generatecontent" in lower_url or "gemini" in lower_url:
|
||||
return API_GEMINI
|
||||
if "/images/edits" in lower_url:
|
||||
return API_IMAGES_EDITS
|
||||
if "/images" in lower_url:
|
||||
return API_IMAGES
|
||||
return API_CHAT
|
||||
|
||||
|
||||
def normalize_api_url(url, api_type):
|
||||
"""Return an endpoint URL for the selected API type.
|
||||
|
||||
Full endpoint URLs are kept. Base URLs are completed using OpenAI-compatible
|
||||
defaults or Gemini generateContent where appropriate.
|
||||
"""
|
||||
raw = str(url).strip()
|
||||
if not raw:
|
||||
return raw
|
||||
|
||||
endpoint = {
|
||||
API_CHAT: "chat/completions",
|
||||
API_IMAGES: "images/generations",
|
||||
API_IMAGES_EDITS: "images/edits",
|
||||
}.get(api_type)
|
||||
|
||||
lower_path = urlparse(raw).path.lower().rstrip("/")
|
||||
if api_type == API_GEMINI:
|
||||
if "generatecontent" in lower_path:
|
||||
return raw
|
||||
return _join_url(raw, "v1beta/models/{model}:generateContent")
|
||||
|
||||
if endpoint is None:
|
||||
return raw
|
||||
if lower_path.endswith("/" + endpoint):
|
||||
return raw
|
||||
if lower_path.endswith("/v1"):
|
||||
return _join_url(raw, endpoint)
|
||||
return _join_url(raw, "v1/" + endpoint)
|
||||
|
||||
|
||||
def image_to_data_url(image_path):
|
||||
"""Read a local image file and return data:<mime>;base64,..."""
|
||||
path = Path(image_path)
|
||||
mime_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
|
||||
with open(str(path), "rb") as f:
|
||||
encoded = base64.b64encode(f.read()).decode("ascii")
|
||||
return "data:{};base64,{}".format(mime_type, encoded)
|
||||
|
||||
|
||||
def decode_image_data_url(data_url):
|
||||
marker = ";base64,"
|
||||
if marker not in data_url:
|
||||
raise AiImageServiceError("Unsupported data URL image format")
|
||||
return base64.b64decode(data_url.split(marker, 1)[1])
|
||||
|
||||
|
||||
def build_payload(config, prompt, image_data_url, resolution="1K", aspect_ratio="1:1"):
|
||||
"""Build JSON body for chat/gemini/images request types."""
|
||||
cfg = _coerce_config(config)
|
||||
api_type = detect_api_type(cfg.url, cfg.api_type)
|
||||
|
||||
if api_type == API_CHAT:
|
||||
payload = {
|
||||
"model": cfg.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": image_data_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
"stream": False,
|
||||
}
|
||||
elif api_type == API_GEMINI:
|
||||
mime_type, data = _split_data_url(image_data_url)
|
||||
payload = {
|
||||
"contents": [
|
||||
{
|
||||
"parts": [
|
||||
{"text": prompt},
|
||||
{"inlineData": {"mimeType": mime_type, "data": data}},
|
||||
]
|
||||
}
|
||||
],
|
||||
"generationConfig": {"responseModalities": ["TEXT", "IMAGE"]},
|
||||
}
|
||||
elif api_type == API_IMAGES:
|
||||
payload = {
|
||||
"model": cfg.model,
|
||||
"prompt": prompt,
|
||||
"image_urls": [image_data_url],
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"resolution": resolution,
|
||||
"n": 1,
|
||||
}
|
||||
else:
|
||||
raise AiConfigError("build_payload 不支持 {}".format(api_type))
|
||||
|
||||
payload.update(cfg.extra_body)
|
||||
return payload
|
||||
|
||||
|
||||
def build_multipart_fields(config, prompt, image_path, resolution="1K"):
|
||||
"""Build data/files for images_edits requests.
|
||||
|
||||
Caller owns and must close the returned file object.
|
||||
"""
|
||||
cfg = _coerce_config(config)
|
||||
handle = open(str(image_path), "rb")
|
||||
mime_type = mimetypes.guess_type(str(image_path))[0] or "application/octet-stream"
|
||||
data = {
|
||||
"model": cfg.model,
|
||||
"prompt": prompt,
|
||||
"n": "1",
|
||||
"size": _resolution_to_size(resolution),
|
||||
}
|
||||
data.update(cfg.extra_body)
|
||||
files = {"image": (Path(image_path).name, handle, mime_type)}
|
||||
return data, files
|
||||
|
||||
|
||||
def extract_image_from_response(data, session=None, timeout=60):
|
||||
"""Return image bytes found in an AI JSON response, or None.
|
||||
|
||||
Searches data-url/base64 fields first, then image URLs.
|
||||
"""
|
||||
for key, value in _walk_json_items(data):
|
||||
if not isinstance(value, str):
|
||||
continue
|
||||
text = value.strip()
|
||||
if text.startswith("data:image/"):
|
||||
return decode_image_data_url(text)
|
||||
if key and key.lower() in _BASE64_KEYS and _looks_like_base64(text):
|
||||
try:
|
||||
return base64.b64decode(text)
|
||||
except (TypeError, ValueError, binascii.Error):
|
||||
pass
|
||||
|
||||
for _key, value in _walk_json_items(data):
|
||||
if isinstance(value, str) and _is_image_url(value):
|
||||
return download_image_url(value, session=session, timeout=timeout)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def download_image_url(url, session=None, timeout=60):
|
||||
client = session or requests.Session()
|
||||
if hasattr(client, "trust_env"):
|
||||
client.trust_env = False
|
||||
response = client.get(url, timeout=timeout)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
|
||||
|
||||
class ImageApiClient:
|
||||
"""Small HTTP client for AI image APIs.
|
||||
|
||||
GUI code should not use requests directly; it should call this service from
|
||||
worker threads and send progress back through Qt signals.
|
||||
"""
|
||||
|
||||
def __init__(self, config, session=None):
|
||||
self.config = _coerce_config(config)
|
||||
self.session = session or requests.Session()
|
||||
if hasattr(self.session, "trust_env"):
|
||||
self.session.trust_env = False
|
||||
|
||||
def generate(self, prompt, image_path, resolution="1K", aspect_ratio="1:1"):
|
||||
validate_api_config(self.config)
|
||||
api_type = detect_api_type(self.config.url, self.config.api_type)
|
||||
url = normalize_api_url(self.config.url, api_type)
|
||||
if api_type == API_GEMINI:
|
||||
url = url.replace("{model}", self.config.model)
|
||||
|
||||
headers = {"Authorization": "Bearer {}".format(self.config.api_key)}
|
||||
timeout = (self.config.connect_timeout_seconds, self.config.timeout_seconds)
|
||||
|
||||
if api_type == API_IMAGES_EDITS:
|
||||
data, files = build_multipart_fields(self.config, prompt, image_path, resolution)
|
||||
try:
|
||||
response = self.session.post(
|
||||
url,
|
||||
headers=headers,
|
||||
data=data,
|
||||
files=files,
|
||||
timeout=timeout,
|
||||
)
|
||||
finally:
|
||||
files["image"][1].close()
|
||||
else:
|
||||
data_url = image_to_data_url(image_path)
|
||||
payload = build_payload(self.config, prompt, data_url, resolution, aspect_ratio)
|
||||
headers["Content-Type"] = "application/json"
|
||||
response = self.session.post(url, headers=headers, json=payload, timeout=timeout)
|
||||
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
image_bytes = extract_image_from_response(
|
||||
payload,
|
||||
session=self.session,
|
||||
timeout=self.config.timeout_seconds,
|
||||
)
|
||||
if not image_bytes:
|
||||
raise AiImageServiceError("AI 响应中未找到图片")
|
||||
return image_bytes
|
||||
|
||||
|
||||
def _coerce_config(config):
|
||||
if isinstance(config, AiModelConfig):
|
||||
return config
|
||||
return AiModelConfig.from_dict(config)
|
||||
|
||||
|
||||
def _to_int(value, default):
|
||||
try:
|
||||
return int(value or default)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def _join_url(base_url, suffix):
|
||||
base = str(base_url).rstrip("/") + "/"
|
||||
return urljoin(base, suffix)
|
||||
|
||||
|
||||
def _split_data_url(data_url):
|
||||
prefix, encoded = data_url.split(",", 1)
|
||||
mime_type = prefix[len("data:"):].split(";", 1)[0]
|
||||
return mime_type, encoded
|
||||
|
||||
|
||||
def _resolution_to_size(resolution):
|
||||
mapping = {
|
||||
"512": "512x512",
|
||||
"512px": "512x512",
|
||||
"1K": "1024x1024",
|
||||
"2K": "2048x2048",
|
||||
"4K": "4096x4096",
|
||||
}
|
||||
return mapping.get(str(resolution), str(resolution))
|
||||
|
||||
|
||||
def _walk_json_items(value, key=None):
|
||||
# type: (Any, Optional[str]) -> Iterable[Tuple[Optional[str], Any]]
|
||||
yield key, value
|
||||
if isinstance(value, dict):
|
||||
for child_key, child_value in value.items():
|
||||
for item in _walk_json_items(child_value, str(child_key)):
|
||||
yield item
|
||||
elif isinstance(value, list):
|
||||
for child_value in value:
|
||||
for item in _walk_json_items(child_value, key):
|
||||
yield item
|
||||
|
||||
|
||||
def _looks_like_base64(text):
|
||||
if len(text) < 8:
|
||||
return False
|
||||
try:
|
||||
base64.b64decode(text, validate=True)
|
||||
return True
|
||||
except (TypeError, ValueError, binascii.Error):
|
||||
return False
|
||||
|
||||
|
||||
def _is_image_url(text):
|
||||
parsed = urlparse(text.strip())
|
||||
if parsed.scheme.lower() not in ("http", "https"):
|
||||
return False
|
||||
suffix = Path(parsed.path).suffix.lower()
|
||||
return suffix in _IMAGE_EXTENSIONS
|
||||
Reference in New Issue
Block a user