From fe1dc7778e70b122af908c882bec058a224ff2ac Mon Sep 17 00:00:00 2001 From: ila Date: Thu, 18 Jun 2026 17:40:56 +0800 Subject: [PATCH] feat: add ai image service --- src/services/ai_image_service.py | 383 +++++++++++++++++++++++++++++++ tasks.md | 2 +- tests/test_ai_image_service.py | 255 ++++++++++++++++++++ 3 files changed, 639 insertions(+), 1 deletion(-) create mode 100644 src/services/ai_image_service.py create mode 100644 tests/test_ai_image_service.py diff --git a/src/services/ai_image_service.py b/src/services/ai_image_service.py new file mode 100644 index 0000000..afacfb9 --- /dev/null +++ b/src/services/ai_image_service.py @@ -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:;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 diff --git a/tasks.md b/tasks.md index c0dd735..aa8b4ef 100644 --- a/tasks.md +++ b/tasks.md @@ -1048,7 +1048,7 @@ - [x] 前置依赖：`requirements.txt` 增加 Python 3.7 兼容的 `requests` / `urllib3` / `openpyxl` 锁定版本，并同步 `docs/03-technical-stack.md` - [x] `core/models.py` 新增 `OutfitTask` / `OutfitResult`（纯 dataclass，Python 3.7 兼容，不依赖 PySide6） - [x] `services/excel_service.py`：读行 → `List[OutfitTask]`、写回 D/E/F、占用检测、跳过「完成」/按设置重试「失败」/空字段安全跳过、每行即存 -- [ ] `services/ai_image_service.py`：移植旧项目 `ImageApiClient`（多模型、多请求格式 `chat/gemini/images/images_edits`、传图 data-url、递归取图、URL 归一化、字段校验）；注意 PEP585 类型注解改 Python 3.7 写法 +- [x] `services/ai_image_service.py`：移植旧项目 `ImageApiClient`（多模型、多请求格式 `chat/gemini/images/images_edits`、传图 data-url、递归取图、URL 归一化、字段校验）；注意 PEP585 类型注解改 Python 3.7 写法 - [ ] `core/ai_outfit.py`：单行生成纯逻辑编排（提示词渲染 + 调用 + 保存 JPG + 产出 `OutfitResult`） - [ ] 单测：Excel 读写、提示词渲染、取图、命名去重（API 用 mock）；可在 Python 3.7 运行、不依赖 GUI diff --git a/tests/test_ai_image_service.py b/tests/test_ai_image_service.py new file mode 100644 index 0000000..b1969b1 --- /dev/null +++ b/tests/test_ai_image_service.py @@ -0,0 +1,255 @@ +"""Tests for AI image service helpers and client request construction.""" +import base64 +import shutil +import sys +import tempfile +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + + +class TestAiImageServiceHelpers(unittest.TestCase): + def setUp(self): + self.tmp = Path(tempfile.mkdtemp()) + self.image_path = self.tmp / "sample.png" + self.image_bytes = b"\x89PNG\r\n\x1a\nfake-png" + self.image_path.write_bytes(self.image_bytes) + + def tearDown(self): + shutil.rmtree(str(self.tmp), ignore_errors=True) + + def test_api_config_errors_require_core_fields(self): + from services.ai_image_service import api_config_errors + + errors = api_config_errors({"url": "", "model": "", "api_key": ""}) + + self.assertIn("缺少 url", errors) + self.assertIn("缺少 model", errors) + self.assertIn("缺少 api_key", errors) + + def test_api_config_errors_reject_non_object_config(self): + from services.ai_image_service import api_config_errors + + self.assertEqual(api_config_errors(["bad"]), ["AI 模型配置必须是对象"]) + + def test_api_config_errors_reject_invalid_extra_body_and_timeout(self): + from services.ai_image_service import api_config_errors + + errors = api_config_errors({ + "url": "https://api.example.test", + "model": "m", + "api_key": "k", + "timeout_seconds": "bad", + "extra_body": ["bad"], + }) + + self.assertIn("timeout_seconds 必须大于 0", errors) + self.assertIn("extra_body 必须是对象", errors) + + def test_detect_api_type_from_url(self): + from services.ai_image_service import ( + API_CHAT, + API_GEMINI, + API_IMAGES_EDITS, + detect_api_type, + ) + + self.assertEqual(detect_api_type("https://x/v1/chat/completions"), API_CHAT) + self.assertEqual(detect_api_type("https://x/v1/images/edits"), API_IMAGES_EDITS) + self.assertEqual(detect_api_type("https://x/v1beta/models/m:generateContent"), API_GEMINI) + + def test_normalize_openai_urls(self): + from services.ai_image_service import API_CHAT, API_IMAGES_EDITS, normalize_api_url + + self.assertEqual( + normalize_api_url("https://api.example.test", API_CHAT), + "https://api.example.test/v1/chat/completions", + ) + self.assertEqual( + normalize_api_url("https://api.example.test/v1", API_IMAGES_EDITS), + "https://api.example.test/v1/images/edits", + ) + + def test_image_to_data_url(self): + from services.ai_image_service import image_to_data_url + + data_url = image_to_data_url(self.image_path) + + self.assertTrue(data_url.startswith("data:image/png;base64,")) + encoded = data_url.split(",", 1)[1] + self.assertEqual(base64.b64decode(encoded), self.image_bytes) + + def test_build_chat_payload(self): + from services.ai_image_service import build_payload + + payload = build_payload( + { + "url": "https://api.example.test/v1/chat/completions", + "model": "m1", + "api_key": "k", + "api_type": "chat", + "extra_body": {"temperature": 0}, + }, + "生成穿搭", + "data:image/png;base64,abc", + ) + + self.assertEqual(payload["model"], "m1") + self.assertEqual(payload["messages"][0]["content"][0]["text"], "生成穿搭") + self.assertEqual( + payload["messages"][0]["content"][1]["image_url"]["url"], + "data:image/png;base64,abc", + ) + self.assertEqual(payload["temperature"], 0) + + def test_build_gemini_payload(self): + from services.ai_image_service import build_payload + + payload = build_payload( + { + "url": "https://generativelanguage.googleapis.com/v1beta/models/m:generateContent", + "model": "gemini", + "api_key": "k", + "api_type": "gemini", + }, + "prompt", + "data:image/jpeg;base64,aW1n", + ) + + inline = payload["contents"][0]["parts"][1]["inlineData"] + self.assertEqual(inline["mimeType"], "image/jpeg") + self.assertEqual(inline["data"], "aW1n") + + def test_build_images_payload(self): + from services.ai_image_service import build_payload + + payload = build_payload( + { + "url": "https://api.example.test/v1/images/generations", + "model": "img", + "api_key": "k", + "api_type": "images", + }, + "prompt", + "data:image/png;base64,abc", + resolution="1K", + ) + + self.assertEqual(payload["image_urls"], ["data:image/png;base64,abc"]) + self.assertEqual(payload["aspect_ratio"], "1:1") + self.assertEqual(payload["resolution"], "1K") + + def test_build_multipart_fields(self): + from services.ai_image_service import build_multipart_fields + + data, files = build_multipart_fields( + { + "url": "https://api.example.test/v1/images/edits", + "model": "edit", + "api_key": "k", + "api_type": "images_edits", + }, + "prompt", + self.image_path, + resolution="1K", + ) + try: + self.assertEqual(data["model"], "edit") + self.assertEqual(data["size"], "1024x1024") + self.assertEqual(files["image"][0], "sample.png") + self.assertEqual(files["image"][2], "image/png") + finally: + files["image"][1].close() + + def test_extract_data_url_from_response(self): + from services.ai_image_service import extract_image_from_response + + encoded = base64.b64encode(b"image-bytes").decode("ascii") + payload = {"choices": [{"message": {"content": "x", "image": "data:image/png;base64," + encoded}}]} + + self.assertEqual(extract_image_from_response(payload), b"image-bytes") + + def test_extract_base64_from_response(self): + from services.ai_image_service import extract_image_from_response + + encoded = base64.b64encode(b"image-bytes").decode("ascii") + payload = {"data": [{"b64_json": encoded}]} + + self.assertEqual(extract_image_from_response(payload), b"image-bytes") + + def test_extract_image_url_downloads(self): + from services.ai_image_service import extract_image_from_response + + session = _FakeSession() + payload = {"data": [{"url": "https://cdn.example.test/out.jpg"}]} + + self.assertEqual(extract_image_from_response(payload, session=session), b"downloaded") + self.assertEqual(session.last_url, "https://cdn.example.test/out.jpg") + + +class TestImageApiClient(unittest.TestCase): + def setUp(self): + self.tmp = Path(tempfile.mkdtemp()) + self.image_path = self.tmp / "sample.png" + self.image_path.write_bytes(b"\x89PNG\r\n\x1a\nfake-png") + + def tearDown(self): + shutil.rmtree(str(self.tmp), ignore_errors=True) + + def test_generate_posts_chat_payload_and_extracts_image(self): + from services.ai_image_service import ImageApiClient + + session = _FakeSession() + config = { + "url": "https://api.example.test", + "model": "model-x", + "api_key": "secret", + "api_type": "chat", + } + + result = ImageApiClient(config, session=session).generate("prompt", self.image_path) + + self.assertEqual(result, b"image-bytes") + self.assertEqual(session.last_post_url, "https://api.example.test/v1/chat/completions") + self.assertEqual(session.last_headers["Authorization"], "Bearer secret") + self.assertEqual(session.last_json["model"], "model-x") + + +class _FakeResponse: + def __init__(self, payload=None, content=b"downloaded"): + self._payload = payload or { + "data": [ + {"b64_json": base64.b64encode(b"image-bytes").decode("ascii")} + ] + } + self.content = content + + def raise_for_status(self): + return None + + def json(self): + return self._payload + + +class _FakeSession: + def __init__(self): + self.trust_env = True + self.last_url = None + self.last_post_url = None + self.last_headers = None + self.last_json = None + + def get(self, url, timeout=60): + self.last_url = url + return _FakeResponse(content=b"downloaded") + + def post(self, url, headers=None, json=None, data=None, files=None, timeout=None): + self.last_post_url = url + self.last_headers = headers + self.last_json = json + return _FakeResponse() + + +if __name__ == "__main__": + unittest.main()