Spaces:
Running
Running
| """AI 生图 API：生成结果存储 + 手表拉取端点 + 共享生图流程。 | |
| 生图走两条既有链路（与聊天一致，前端无需新协议）： | |
| 1. 阻塞式：POST /v1/xtc/chat —— 生图模型直接返回带 image_url 的 XTC 响应 | |
| （适合非手表端 / 生成快的模型） | |
| 2. 伪流式：POST /v1/xtc/chat/pseudo/start —— 后台生图，轮询到 done 后拉取 | |
| （手表默认路径，防 fetch.fetch 30s 超时） | |
| 本模块提供： | |
| - store_image() / get_image()：处理后图片的内存存储（TTL + 容量上限， | |
| 与 TTS 会话一致；手表生成后立即拉取，无需跨重启持久化） | |
| - run_image_generation()：完整生图流程（调上游 → 手表化处理 → 存储 → 组装 XTC 字段）， | |
| 供 chat / pseudo_stream 两个入口复用 | |
| - GET /v1/xtc/image/gen/{image_id}：手表拉取处理后的 JPEG（capability URL， | |
| 与 TTS audio 一致：image_id 为 128bit 随机 hex 不可猜测 + TTL，免 access_key， | |
| 因为手表 request.download 不便携带自定义 header） | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import re | |
| import time | |
| import uuid | |
| from typing import Optional | |
| from fastapi import APIRouter, Response | |
| from ..adapters import imagegen as imagegen_adapter | |
| from ..errors import HttpError | |
| from ..media.imagegen import DEFAULT_JPEG_QUALITY, DEFAULT_MAX_DIM, process_image_for_watch | |
| from ..providers.policy import is_image_model | |
| from ..services.config_store import load_config | |
| from ..services.prompt_optimizer import optimize_image_prompt | |
| from ._common import CORS_HEADERS, ok_with_cors, select_provider_and_key | |
| logger = logging.getLogger(__name__) | |
| router = APIRouter(prefix="/v1/xtc/image", tags=["imagegen"]) | |
| # 图片会话存活 30 分钟（生图比 TTS 慢，手表断网重试窗口也要留足） | |
| _IMAGE_TTL = 1800 | |
| # 内存上限：处理后图片典型 50~100KB，32 张约 2~3MB，HF Space 内存可控 | |
| _IMAGE_MAX = 32 | |
| _images: dict[str, dict] = {} | |
| def _cleanup_images() -> None: | |
| """清理过期图片 + 超限时淘汰最老的（与 TTS 缓存同策略）。""" | |
| now = time.time() | |
| expired = [k for k, v in _images.items() if now - v["created_at"] > _IMAGE_TTL] | |
| for k in expired: | |
| _images.pop(k, None) | |
| if len(_images) > _IMAGE_MAX: | |
| for k, _ in sorted(_images.items(), key=lambda kv: kv[1]["created_at"])[ | |
| : len(_images) - _IMAGE_MAX | |
| ]: | |
| _images.pop(k, None) | |
| def store_image(image_bytes: bytes, meta: Optional[dict] = None) -> str: | |
| """存入处理后图片，返回 image_id（随机 hex，用于 capability URL）。""" | |
| _cleanup_images() | |
| image_id = uuid.uuid4().hex | |
| _images[image_id] = { | |
| "bytes": image_bytes, | |
| "meta": meta or {}, | |
| "created_at": time.time(), | |
| } | |
| return image_id | |
| def get_image(image_id: str) -> Optional[dict]: | |
| entry = _images.get(image_id) | |
| if not entry: | |
| return None | |
| if time.time() - entry["created_at"] > _IMAGE_TTL: | |
| _images.pop(image_id, None) | |
| return None | |
| return entry | |
| def image_url_for(image_id: str) -> str: | |
| """手表端可拉取的相对 URL（前端拼 baseUrl）。""" | |
| return f"/v1/xtc/image/gen/{image_id}" | |
| async def download_generated_image(image_id: str) -> Response: | |
| """手表拉取处理后的生图 JPEG（capability URL，无需鉴权，见模块头注释）。""" | |
| entry = get_image(image_id) | |
| if not entry: | |
| raise HttpError("image not found or expired", status=404, code="not_found") | |
| meta = entry.get("meta") or {} | |
| return Response( | |
| content=entry["bytes"], | |
| media_type="image/jpeg", | |
| headers={ | |
| **CORS_HEADERS, | |
| "Cache-Control": "no-store", | |
| "x-xtc-image-width": str(meta.get("width") or ""), | |
| "x-xtc-image-height": str(meta.get("height") or ""), | |
| }, | |
| ) | |
| # ===== 共享生图流程（chat / pseudo 复用）===== | |
| # 附件块：normalize_chat_body 会把文本类附件追加成独立的一条 user 消息， | |
| # 形如 "[文件名]\n<内容>"（多个用空行连接）。生图 prompt 必须能把它和用户自己打的字区分开， | |
| # 否则「用户打了提示词 + 顺手带了个文件」时，生图/优化拿到的是文件内容而不是那句话。 | |
| _ATTACHMENT_HEAD_RE = re.compile(r"^\[[^\]\n]{1,120}\]\n") | |
| def _split_attachment(text: str) -> tuple[bool, str]: | |
| """返回 (是否附件块, 去掉文件名头的内容)。""" | |
| m = _ATTACHMENT_HEAD_RE.match(text or "") | |
| if not m: | |
| return False, text or "" | |
| return True, (text or "")[m.end():] | |
| def extract_image_prompt(messages: list) -> str: | |
| """从 messages 提取生图 prompt。 | |
| 取「最近一条用户自己打的文字」；如果还有附件块，把附件内容作为参考附在末尾 | |
| （截断到 600 字）。只有附件没有文字时，附件内容就是 prompt（此时它就是用户 | |
| 表达要画什么的方式）。 | |
| """ | |
| typed = "" | |
| attach = "" | |
| for msg in reversed(messages or []): | |
| if not isinstance(msg, dict) or msg.get("role") != "user": | |
| continue | |
| content = msg.get("content") | |
| text = "" | |
| if isinstance(content, str): | |
| text = content.strip() | |
| elif isinstance(content, list): | |
| parts = [ | |
| str(p.get("text") or "").strip() | |
| for p in content | |
| if isinstance(p, dict) and p.get("type") == "text" | |
| ] | |
| text = "\n".join(x for x in parts if x).strip() | |
| if not text: | |
| continue | |
| is_attach, body = _split_attachment(text) | |
| if is_attach: | |
| if not attach: | |
| attach = body.strip() | |
| continue | |
| if not typed: | |
| typed = text | |
| break | |
| if typed and attach: | |
| return (typed + "\n\n参考内容：\n" + attach[:600]).strip() | |
| if typed: | |
| return typed | |
| return attach[:2000] | |
| async def run_image_generation( | |
| *, | |
| provider, | |
| api_key: str, | |
| clean_model: str, | |
| messages: list, | |
| size: Optional[str] = None, | |
| access_key: Optional[str] = None, | |
| optimize: bool = False, | |
| optimizer_provider: Optional[str] = None, | |
| optimizer_model: Optional[str] = None, | |
| optimize_context: Optional[str] = None, | |
| ) -> dict: | |
| """完整生图流程：调上游 → 手表化处理 → 存储 → 返回 XTC 响应附加字段。 | |
| optimize=True 时先调用「优化模型」把用户提示词润色为专业生图提示词， | |
| 再送上游生成；优化模型缺失/调用失败自动回退原始 prompt（不阻断生图）， | |
| 并通过 optimize_warning 反馈前端。 | |
| 返回 dict 可直接并入 XTC 响应 / 伪流式会话结果： | |
| {image, image_id, image_url, image_width, image_height, text, usage, | |
| optimized, optimized_prompt?, original_prompt?, optimize_warning?} | |
| """ | |
| if provider.type == "gemini": | |
| raise HttpError( | |
| "Image generation is only supported on openai-compatible providers", | |
| status=400, | |
| code="bad_request", | |
| ) | |
| prompt = extract_image_prompt(messages) | |
| if not prompt: | |
| raise HttpError("prompt is required for image generation", status=400, code="bad_request") | |
| gen_prompt = prompt | |
| optimized_prompt: Optional[str] = None | |
| optimize_warning = "" | |
| if optimize: | |
| try: | |
| clean_opt = await _optimize_prompt( | |
| access_key=access_key, | |
| optimizer_provider=optimizer_provider, | |
| optimizer_model=optimizer_model, | |
| prompt_text=prompt, | |
| context_text=optimize_context or "", | |
| ) | |
| if clean_opt: | |
| optimized_prompt = clean_opt | |
| gen_prompt = clean_opt | |
| except HttpError as e: | |
| optimize_warning = "提示词优化失败：" + str(e.message or e.code or "unknown") | |
| logger.warning("[imagegen] optimize failed (fallback raw prompt) provider=%s model=%s: %s", | |
| provider.id, clean_model, e.message) | |
| except Exception as e: | |
| optimize_warning = "提示词优化异常：" + str(e) | |
| logger.warning("[imagegen] optimize error (fallback raw prompt) provider=%s model=%s", | |
| provider.id, clean_model, exc_info=True) | |
| gen = await imagegen_adapter.generate_image( | |
| base_url=provider.base_url or "", | |
| api_key=api_key, | |
| model=clean_model, | |
| prompt=gen_prompt, | |
| n=1, | |
| size=size, | |
| ) | |
| processed = process_image_for_watch(gen["image_bytes"]) | |
| image_id = store_image( | |
| processed["image_bytes"], | |
| meta={ | |
| "width": processed["width"], | |
| "height": processed["height"], | |
| "orig_width": processed["orig_width"], | |
| "orig_height": processed["orig_height"], | |
| }, | |
| ) | |
| revised = str(gen.get("revised_prompt") or "").strip() | |
| logger.info( | |
| "[imagegen] done provider=%s model=%s %dx%d->%dx%d %dbytes optimize=%s", | |
| provider.id, clean_model, | |
| processed["orig_width"], processed["orig_height"], | |
| processed["width"], processed["height"], len(processed["image_bytes"]), | |
| "on" if optimized_prompt else "off", | |
| ) | |
| out = { | |
| "image": True, | |
| "image_id": image_id, | |
| "image_url": image_url_for(image_id), | |
| "image_width": processed["width"], | |
| "image_height": processed["height"], | |
| # 手表消息列表需要一条占位文本；revised_prompt 有值时附带给前端展示 | |
| "revised_prompt": revised, | |
| "text": "图片已生成", | |
| "usage": gen.get("usage") or {}, | |
| } | |
| if optimize: | |
| out["optimized"] = bool(optimized_prompt) | |
| if optimized_prompt: | |
| out["optimized_prompt"] = optimized_prompt | |
| out["original_prompt"] = prompt | |
| if optimize_warning: | |
| out["optimize_warning"] = optimize_warning | |
| return out | |
| async def _optimize_prompt( | |
| *, | |
| access_key: Optional[str], | |
| optimizer_provider: Optional[str], | |
| optimizer_model: Optional[str], | |
| prompt_text: str, | |
| context_text: str, | |
| ) -> str: | |
| """把原始提示词交给用户指定的聊天模型（或最近用的聊天模型）润色。 | |
| 返回优化后的提示词；模型/厂商缺失时抛 HttpError 由调用方回退原始 prompt。 | |
| """ | |
| opt_model = str(optimizer_model or "").strip() | |
| if not opt_model: | |
| raise HttpError("optimizer model not specified", status=400, code="bad_request") | |
| config = await load_config() | |
| opt_provider, opt_key, clean_opt_model = select_provider_and_key( | |
| config=config, | |
| provider_id=(str(optimizer_provider or "").strip() or None), | |
| model=opt_model, | |
| fallback_model=None, | |
| ) | |
| return await optimize_image_prompt( | |
| provider=opt_provider, | |
| api_key=opt_key, | |
| clean_model=clean_opt_model, | |
| prompt_text=prompt_text, | |
| context_text=context_text, | |
| access_key=access_key, | |
| ) | |
| def provider_supports_image(provider, model: str) -> bool: | |
| """select_provider_and_key 之后调用：判断是否走生图流程。""" | |
| return is_image_model(provider, model) | |
| __all__ = [ | |
| "router", | |
| "DEFAULT_MAX_DIM", | |
| "DEFAULT_JPEG_QUALITY", | |
| "store_image", | |
| "get_image", | |
| "image_url_for", | |
| "extract_image_prompt", | |
| "run_image_generation", | |
| "provider_supports_image", | |
| ] | |