tw2/qingjian_cs/services/service_gemini_stable.py

#!/usr/bin/env python3
from __future__ import annotations

import asyncio
import base64
import logging
import os
from typing import Any

import aiohttp

logger = logging.getLogger(__name__)

GEMINI_API_KEY = "sk-8i7uYE0RtnQwDImV8a5f7014DcAb46F6BcEb72Df92218aC8"
GEMINI_API_URL = "https://api.laozhang.ai/v1beta/models"
GEMINI_MODEL = "gemini-3.1-flash-image-preview"
GEMINI_IMAGE_SIZE = "4K"
GEMINI_MAX_RETRIES = 3

DEFAULT_PROMPT = (
    "提取印花图案，去褶皱并补齐缺失区域，生成完整清晰的平面图。"
    "严格保持原图元素位置、颜色和细节，不要改风格。"
)


class GeminiExtractStableService:
    def image_to_base64(self, image_path: str) -> str | None:
        try:
            if not os.path.exists(image_path):
                logger.error("文件不存在: %s", image_path)
                return None
            with open(image_path, "rb") as f:
                return base64.b64encode(f.read()).decode("utf-8")
        except Exception as e:
            logger.error("Base64转换失败: %s", e)
            return None

    async def extract_pattern(
        self,
        input_path: str,
        output_path: str,
        custom_prompt: str | None = None,
        aspect_ratio: str = "1:1",
    ) -> tuple[bool, str, dict[str, Any]]:
        img64 = self.image_to_base64(input_path)
        if not img64:
            return False, "图片编码失败", {}

        prompt = custom_prompt or DEFAULT_PROMPT
        api_url = f"{GEMINI_API_URL}/{GEMINI_MODEL}:generateContent"
        valid_ratios = {"1:1", "9:16", "16:9", "3:4", "4:3", "3:2", "2:3", "5:4", "4:5"}
        image_config: dict[str, Any] = {"imageSize": GEMINI_IMAGE_SIZE}
        if aspect_ratio in valid_ratios:
            image_config["aspectRatio"] = aspect_ratio

        headers = {
            "Authorization": f"Bearer {GEMINI_API_KEY}",
            "Content-Type": "application/json",
        }
        payload = {
            "contents": [
                {
                    "parts": [
                        {"text": prompt},
                    ]
                }
            ],
            "generationConfig": {
                "responseModalities": ["IMAGE"],
                "imageConfig": image_config,
            },
        }

        logger.info("[稳定作图] model=%s ratio=%s size=%s", GEMINI_MODEL, image_config.get("aspectRatio", "1:1"), GEMINI_IMAGE_SIZE)
        timeout = aiohttp.ClientTimeout(total=180, connect=30)
        for i in range(1, GEMINI_MAX_RETRIES + 1):
            try:
                logger.info("[稳定作图] 请求中 第%s/%s次", i, GEMINI_MAX_RETRIES)
                async with aiohttp.ClientSession(timeout=timeout) as session:
                    async with session.post(api_url, headers=headers, json=payload) as resp:
                        txt = await resp.text()
                        if resp.status != 200:
                            logger.error("[稳定作图] 请求失败 status=%s body=%s", resp.status, txt[:240])
                            if i < GEMINI_MAX_RETRIES:
                                await asyncio.sleep(i)
                                continue
                            return False, f"API失败:{resp.status}", {}
                        result = await resp.json()

                try:
                    parts = result["candidates"][0]["content"]["parts"]
                    b64 = ""
                    for p in parts:
                        if "inlineData" in p and isinstance(p["inlineData"], dict):
                            b64 = str(p["inlineData"].get("data", "") or "")
                            if b64:
                                break
                    if not b64:
                        logger.error("[稳定作图] 未找到 inlineData")
                        return False, "未找到图片数据", {}
                    image_data = base64.b64decode(b64)
                except Exception as e:
                    logger.error("[稳定作图] 解析返回失败: %s", e)
                    return False, f"解析失败:{e}", {}

                out_dir = os.path.dirname(output_path)
                if out_dir:
                    os.makedirs(out_dir, exist_ok=True)
                with open(output_path, "wb") as f:
                    f.write(image_data)
                logger.info("[稳定作图] 输出成功: %s", output_path)
                return True, "ok", {"output_path": output_path, "file_size": len(image_data), "api_used": "gemini_stable"}
            except Exception as e:
                logger.error("[稳定作图] 异常: %s", e)
                if i < GEMINI_MAX_RETRIES:
                    await asyncio.sleep(i)
                    continue
                return False, f"异常:{e}", {}
        return False, "未知失败", {}

    async def cleanup(self) -> None:
        return None