2461 lines
82 KiB
Python
2461 lines
82 KiB
Python
"""AI generation helpers backed by configurable HTTP model endpoints."""
|
||
|
||
import base64
|
||
import ipaddress
|
||
from concurrent.futures import CancelledError, FIRST_COMPLETED, ThreadPoolExecutor, as_completed, wait
|
||
import copy
|
||
import json
|
||
import mimetypes
|
||
import os
|
||
import shutil
|
||
import socket
|
||
import subprocess
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
import urllib.error
|
||
import urllib.parse
|
||
import urllib.request
|
||
import uuid
|
||
|
||
import requests
|
||
from requests.adapters import HTTPAdapter
|
||
|
||
from . import appconfig, db, diagnostics, image_paths
|
||
from . import prompts as prompt_module
|
||
|
||
try:
|
||
from .version import APP_VERSION
|
||
except Exception:
|
||
APP_VERSION = "dev"
|
||
|
||
|
||
|
||
class AIError(RuntimeError):
|
||
"""Raised when AI generation cannot complete."""
|
||
|
||
class CMHubError(AIError):
|
||
"""Structured cmhub gateway error."""
|
||
|
||
def __init__(self, code, message, status=None, retryable=False, retry_after=None):
|
||
self.code = str(code or "unknown")
|
||
self.status = status
|
||
self.retryable = bool(retryable)
|
||
self.retry_after = retry_after
|
||
super().__init__("cmhub %s: %s" % (self.code, message))
|
||
|
||
|
||
CMHUB_IMAGE_MAX_BYTES = 20 * 1024 * 1024
|
||
CMHUB_TITLE_READ_TIMEOUT_SECONDS = 600
|
||
CMHUB_IMAGE_READ_TIMEOUT_SECONDS = 900
|
||
CMHUB_IMAGE_SUBMIT_READ_TIMEOUT_SECONDS = 36
|
||
CMHUB_IMAGE_POLL_READ_TIMEOUT_SECONDS = 15
|
||
CMHUB_IMAGE_POLL_DELAYS_SECONDS = (3.0, 5.0, 8.0, 10.0)
|
||
CMHUB_IMAGE_CONCURRENCY_LIMIT = 5
|
||
CMHUB_IMAGE_DOWNLOAD_ATTEMPTS = 3
|
||
CMHUB_IMAGE_SLOW_DOWNLOAD_SECONDS = 20.0
|
||
# 连接池要覆盖“生图并发 + 下载并发”（各上限 5）再留余量，避免连接不足排队或
|
||
# urllib3 "Connection pool is full" 警告。
|
||
CMHUB_HTTP_POOL_SIZE = 32
|
||
|
||
_CMHUB_SESSION = None
|
||
_CMHUB_SESSION_LOCK = threading.Lock()
|
||
|
||
|
||
def _cmhub_session():
|
||
"""返回 cmhub 专用的共享 `requests.Session`（连接复用 + 连接池）。
|
||
|
||
所有 cmhub HTTP（生成/下载/models/balance）都走同一个 Session，避免每次调用
|
||
新建 TCP+TLS 造成握手风暴与 connect_timeout；跨线程发请求安全，连接池足够大即可。
|
||
"""
|
||
|
||
global _CMHUB_SESSION
|
||
if _CMHUB_SESSION is None:
|
||
with _CMHUB_SESSION_LOCK:
|
||
if _CMHUB_SESSION is None:
|
||
session = requests.Session()
|
||
adapter = HTTPAdapter(
|
||
pool_connections=CMHUB_HTTP_POOL_SIZE,
|
||
pool_maxsize=CMHUB_HTTP_POOL_SIZE,
|
||
max_retries=0,
|
||
)
|
||
session.mount("http://", adapter)
|
||
session.mount("https://", adapter)
|
||
_CMHUB_SESSION = session
|
||
return _CMHUB_SESSION
|
||
|
||
|
||
def _apply_cmhub_proxy(use_system_proxy):
|
||
"""设置 cmhub Session 是否读取系统代理环境变量。
|
||
|
||
`use_system_proxy=False` → `trust_env=False`，忽略 `HTTP(S)_PROXY`/`ALL_PROXY` 直连；
|
||
公网 cmhub 网关默认直连，避免误走慢代理导致明文图片下载卡住。需要代理的环境可在
|
||
`config.json` 的 `ai.cmhub.use_system_proxy` 打开。
|
||
"""
|
||
|
||
session = _cmhub_session()
|
||
session.trust_env = bool(use_system_proxy)
|
||
return session
|
||
|
||
_RESOLUTION_SIZES = {
|
||
"512": (512, 512),
|
||
"1k": (1024, 1024),
|
||
"2k": (2048, 2048),
|
||
"4k": (4096, 4096),
|
||
}
|
||
|
||
|
||
def gen_title(
|
||
title_prompt,
|
||
old_title,
|
||
retry=None,
|
||
config=None,
|
||
models_path=appconfig.AI_MODELS_PATH,
|
||
on_step=None,
|
||
on_event=None,
|
||
cmhub_config_path=appconfig.CMHUB_CONFIG_PATH,
|
||
):
|
||
"""Generate a new product title from a prompt and the old title."""
|
||
|
||
cfg = appconfig.load_config() if config is None else config
|
||
ai_cfg = appconfig.ai_config(cfg)
|
||
if _ai_backend(ai_cfg) == "cmhub":
|
||
return _gen_title_cmhub(
|
||
title_prompt,
|
||
old_title,
|
||
retry=retry,
|
||
config=cfg,
|
||
cmhub_config_path=cmhub_config_path,
|
||
on_step=on_step,
|
||
on_event=on_event,
|
||
)
|
||
_notify_step(on_step, "load_text_model")
|
||
model = _role_model("text", ai_cfg.get("default_text_model"), models_path)
|
||
_notify_step(on_step, "title_build_request")
|
||
payload = _chat_payload(
|
||
model,
|
||
[
|
||
{"role": "user", "content": _compose_title_prompt(title_prompt, old_title)},
|
||
],
|
||
)
|
||
attempts = _attempt_count(ai_cfg, retry)
|
||
_notify_step(on_step, "title_request")
|
||
data = _call_with_retry(
|
||
model,
|
||
payload,
|
||
cfg,
|
||
attempts,
|
||
request_kind="json",
|
||
on_retry=lambda attempt, total_attempts, exc: _notify_retry(
|
||
on_step,
|
||
"title_request",
|
||
attempt,
|
||
total_attempts,
|
||
exc,
|
||
model,
|
||
),
|
||
)
|
||
_notify_step(on_step, "title_parse_response")
|
||
text = _extract_text(data).strip()
|
||
if not text:
|
||
raise AIError("AI 返回为空标题")
|
||
return text
|
||
|
||
|
||
def gen_cover(
|
||
cover_prompt,
|
||
old_cover_path,
|
||
out_path,
|
||
resolution=None,
|
||
jpg_quality=None,
|
||
retry=None,
|
||
config=None,
|
||
models_path=appconfig.AI_MODELS_PATH,
|
||
on_step=None,
|
||
on_event=None,
|
||
cmhub_config_path=appconfig.CMHUB_CONFIG_PATH,
|
||
):
|
||
"""Generate a new cover image and save it as a JPEG file."""
|
||
|
||
old_cover_path, out_path = _prepare_cover_input(old_cover_path, out_path, on_step)
|
||
|
||
cfg = appconfig.load_config() if config is None else config
|
||
ai_cfg = appconfig.ai_config(cfg)
|
||
if _ai_backend(ai_cfg) == "cmhub":
|
||
return _gen_cover_cmhub(
|
||
cover_prompt,
|
||
old_cover_path,
|
||
out_path,
|
||
resolution=resolution,
|
||
jpg_quality=jpg_quality,
|
||
retry=retry,
|
||
config=cfg,
|
||
cmhub_config_path=cmhub_config_path,
|
||
on_step=on_step,
|
||
on_event=on_event,
|
||
)
|
||
_notify_step(on_step, "load_image_model")
|
||
model = _role_model("image", ai_cfg.get("default_image_model"), models_path)
|
||
resolution = str(resolution or ai_cfg.get("resolution", "1k"))
|
||
quality = _jpg_quality(jpg_quality if jpg_quality is not None else ai_cfg.get("jpg_quality", 90))
|
||
attempts = _attempt_count(ai_cfg, retry)
|
||
|
||
api_type = model.get("api_type", "auto")
|
||
_notify_step(on_step, "cover_build_request")
|
||
if api_type == "images_edits":
|
||
body, content_type = _image_edit_body(model, cover_prompt, old_cover_path, resolution)
|
||
_notify_step(on_step, "cover_request")
|
||
data = _call_with_retry(
|
||
model,
|
||
body,
|
||
cfg,
|
||
attempts,
|
||
request_kind="multipart",
|
||
content_type=content_type,
|
||
on_retry=lambda attempt, total_attempts, exc: _notify_retry(
|
||
on_step,
|
||
"cover_request",
|
||
attempt,
|
||
total_attempts,
|
||
exc,
|
||
model,
|
||
),
|
||
)
|
||
else:
|
||
payload = _image_chat_payload(model, cover_prompt, old_cover_path, resolution)
|
||
_notify_step(on_step, "cover_request")
|
||
data = _call_with_retry(
|
||
model,
|
||
payload,
|
||
cfg,
|
||
attempts,
|
||
request_kind="json",
|
||
on_retry=lambda attempt, total_attempts, exc: _notify_retry(
|
||
on_step,
|
||
"cover_request",
|
||
attempt,
|
||
total_attempts,
|
||
exc,
|
||
model,
|
||
),
|
||
)
|
||
|
||
_notify_step(on_step, "cover_parse_response")
|
||
image_bytes = _extract_image_bytes(data, model, cfg)
|
||
_notify_step(on_step, "cover_save")
|
||
return _save_jpeg(image_bytes, out_path, resolution, quality)
|
||
|
||
|
||
def cmhub_image_concurrency_plan(ai_cfg):
|
||
"""Return protected cmhub image request/download concurrency."""
|
||
|
||
configured = _positive_int((ai_cfg or {}).get("image_concurrency", 1), 1)
|
||
actual = min(configured, CMHUB_IMAGE_CONCURRENCY_LIMIT)
|
||
return {
|
||
"configured_image_concurrency": configured,
|
||
"request_concurrency": actual,
|
||
"download_concurrency": actual,
|
||
"limit": CMHUB_IMAGE_CONCURRENCY_LIMIT,
|
||
}
|
||
|
||
|
||
def _positive_int(value, default=1):
|
||
try:
|
||
return max(1, int(value))
|
||
except (TypeError, ValueError):
|
||
return max(1, int(default or 1))
|
||
|
||
|
||
def _prepare_cover_input(old_cover_path, out_path, on_step=None):
|
||
_notify_step(on_step, "cover_validate_input")
|
||
old_cover_path = os.path.abspath(str(old_cover_path))
|
||
if not os.path.exists(old_cover_path):
|
||
raise FileNotFoundError("旧封面图片不存在: %s" % old_cover_path)
|
||
if not out_path:
|
||
raise AIError("缺少新封面输出路径")
|
||
return old_cover_path, out_path
|
||
|
||
|
||
def _notify_step(callback, step):
|
||
if callback is None:
|
||
return
|
||
try:
|
||
callback(step)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _notify_step_event(
|
||
callback,
|
||
step,
|
||
result="success",
|
||
detail=None,
|
||
level="info",
|
||
debug_only=False,
|
||
attempt=None,
|
||
attempts=None,
|
||
):
|
||
if callback is None:
|
||
return
|
||
payload = {
|
||
"step": step,
|
||
"result": result,
|
||
"level": level,
|
||
}
|
||
if detail is not None:
|
||
payload["detail"] = detail
|
||
if attempt is not None:
|
||
payload["attempt"] = attempt
|
||
if attempts is not None:
|
||
payload["attempts"] = attempts
|
||
if debug_only:
|
||
payload["debug_only"] = True
|
||
try:
|
||
callback(payload)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
|
||
def _notify_retry(callback, step, attempt, attempts, exc, model):
|
||
if callback is None:
|
||
return
|
||
try:
|
||
callback(
|
||
{
|
||
"step": step,
|
||
"result": "retry",
|
||
"attempt": attempt,
|
||
"attempts": attempts,
|
||
"detail": _redact(str(exc), model),
|
||
}
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _task_attempt_count(task, field):
|
||
try:
|
||
return int(getattr(task, field, 0) or 0)
|
||
except (TypeError, ValueError):
|
||
return 0
|
||
|
||
|
||
def _has_generated_title(task):
|
||
return bool(str(getattr(task, "new_title", "") or "").strip())
|
||
|
||
|
||
def _has_generated_cover(task):
|
||
return bool(str(getattr(task, "new_cover_path", "") or "").strip())
|
||
|
||
|
||
def cover_title_context(task):
|
||
"""Return the title used to describe a cover without changing task results."""
|
||
|
||
new_title = str(getattr(task, "new_title", "") or "").strip()
|
||
if new_title:
|
||
return new_title
|
||
return str(getattr(task, "old_title", "") or "").strip()
|
||
|
||
|
||
def generation_needs(task, generate_cover=False, generate_mode=None):
|
||
"""Return which AI components still need generation for this task."""
|
||
|
||
mode = appconfig.normalize_generate_mode(generate_mode, generate_cover=generate_cover)
|
||
stage = str(getattr(task, "stage", "") or "")
|
||
status = str(getattr(task, "status", "") or "")
|
||
if status in {"running", "skipped"}:
|
||
return {"title": False, "cover": False}
|
||
if stage not in {"collected", "generated"}:
|
||
return {"title": False, "cover": False}
|
||
if (
|
||
stage == "generated"
|
||
and status == "failed"
|
||
and _task_attempt_count(task, "apply_attempts") > 0
|
||
):
|
||
return {"title": False, "cover": False}
|
||
has_title = _has_generated_title(task)
|
||
needs_title = appconfig.generate_mode_includes_title(mode) and not has_title
|
||
needs_cover = False
|
||
if appconfig.generate_mode_includes_cover(mode) and not _has_generated_cover(task):
|
||
if mode == "cover":
|
||
needs_cover = bool(cover_title_context(task))
|
||
else:
|
||
needs_cover = has_title or needs_title
|
||
return {"title": needs_title, "cover": needs_cover}
|
||
|
||
|
||
def is_generatable_task(task, generate_cover=False, generate_mode=None):
|
||
"""判断任务是否能由② AI生成执行或重试。"""
|
||
|
||
needs = generation_needs(task, generate_cover=generate_cover, generate_mode=generate_mode)
|
||
return bool(needs["title"] or needs["cover"])
|
||
|
||
|
||
def generation_component_totals(tasks, generate_cover=False, generate_mode=None):
|
||
"""Count task and component gaps for a generation run."""
|
||
|
||
mode = appconfig.normalize_generate_mode(generate_mode, generate_cover=generate_cover)
|
||
eligible = [
|
||
task for task in list(tasks)
|
||
if is_generatable_task(task, generate_mode=mode)
|
||
]
|
||
title_total = 0
|
||
cover_total = 0
|
||
for task in eligible:
|
||
needs = generation_needs(task, generate_mode=mode)
|
||
if needs["title"]:
|
||
title_total += 1
|
||
if needs["cover"]:
|
||
cover_total += 1
|
||
return {
|
||
"total": len(eligible),
|
||
"title_total": title_total,
|
||
"cover_total": cover_total,
|
||
}
|
||
|
||
|
||
def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=None):
|
||
"""Generate titles first, then covers, and persist each successful task."""
|
||
|
||
runtime = dict(ai_cfg or {})
|
||
config = _runtime_config(runtime)
|
||
generation_cfg = appconfig.ai_config(config)
|
||
generation_cfg.update(
|
||
{
|
||
key: value
|
||
for key, value in runtime.items()
|
||
if key in {
|
||
"title_concurrency",
|
||
"image_concurrency",
|
||
"retry",
|
||
"jpg_quality",
|
||
"resolution",
|
||
"generate_cover",
|
||
"generate_mode",
|
||
}
|
||
}
|
||
)
|
||
db_path = runtime.get("db_path")
|
||
models_path = runtime.get("models_path", appconfig.AI_MODELS_PATH)
|
||
cmhub_config_path = runtime.get("cmhub_config_path", appconfig.CMHUB_CONFIG_PATH)
|
||
image_root = runtime.get("image_dir") or appconfig.image_dir(config)
|
||
account_by_alias = runtime.get("account_by_alias") or {}
|
||
on_task_update = runtime.get("on_task_update")
|
||
on_event = runtime.get("on_event")
|
||
on_error = runtime.get("on_error")
|
||
title_prompt = _prompt_value(prompts, "title")
|
||
cover_prompt = _prompt_value(prompts, "cover")
|
||
generate_mode = appconfig.normalize_generate_mode(
|
||
generation_cfg.get("generate_mode"),
|
||
generate_cover=generation_cfg.get("generate_cover", False),
|
||
)
|
||
generate_cover = appconfig.generate_mode_includes_cover(generate_mode)
|
||
should_stop = should_stop or (lambda: False)
|
||
eligible = [
|
||
task for task in list(tasks)
|
||
if is_generatable_task(task, generate_mode=generate_mode)
|
||
]
|
||
needs_by_task = {
|
||
getattr(task, "id", None): generation_needs(task, generate_mode=generate_mode)
|
||
for task in eligible
|
||
}
|
||
title_tasks = [
|
||
task for task in eligible
|
||
if needs_by_task.get(getattr(task, "id", None), {}).get("title")
|
||
]
|
||
cover_candidates = [
|
||
task for task in eligible
|
||
if needs_by_task.get(getattr(task, "id", None), {}).get("cover")
|
||
]
|
||
summary = {
|
||
"ok": True,
|
||
"total": len(eligible),
|
||
"title_total": len(title_tasks),
|
||
"title_done": 0,
|
||
"cover_done": 0,
|
||
"cover_total": len(cover_candidates) if generate_cover else 0,
|
||
"generated_done": 0,
|
||
"failed": 0,
|
||
"cancelled": False,
|
||
"generate_cover": generate_cover,
|
||
"generate_mode": generate_mode,
|
||
}
|
||
_emit_generation_progress(on_progress, summary)
|
||
title_results = {}
|
||
step_by_task = {}
|
||
step_lock = threading.Lock()
|
||
|
||
def set_step(task, step):
|
||
with step_lock:
|
||
step_by_task[getattr(task, "id", None)] = str(step)
|
||
|
||
def get_step(task, fallback):
|
||
with step_lock:
|
||
return step_by_task.get(getattr(task, "id", None), fallback)
|
||
|
||
def step_callback(task, phase):
|
||
def callback(event):
|
||
if isinstance(event, dict):
|
||
step = event.get("step") or "unknown"
|
||
result = event.get("result") or "start"
|
||
set_step(task, step)
|
||
_emit_generation_event(
|
||
on_event,
|
||
task,
|
||
phase,
|
||
step,
|
||
result,
|
||
detail=event.get("detail"),
|
||
level=event.get("level") or ("warning" if result == "retry" else "info"),
|
||
attempt=event.get("attempt"),
|
||
attempts=event.get("attempts"),
|
||
metadata=event.get("metadata"),
|
||
debug_only=event.get("debug_only"),
|
||
)
|
||
return
|
||
set_step(task, event)
|
||
_emit_generation_event(on_event, task, phase, event, "start")
|
||
return callback
|
||
|
||
with ThreadPoolExecutor(
|
||
max_workers=max(1, int(generation_cfg.get("title_concurrency", 1)))
|
||
) as executor:
|
||
futures = {}
|
||
for task in title_tasks:
|
||
if should_stop():
|
||
summary["cancelled"] = True
|
||
break
|
||
set_step(task, "title_submit")
|
||
_emit_generation_event(on_event, task, "title", "title_submit", "start")
|
||
futures[
|
||
executor.submit(
|
||
gen_title,
|
||
title_prompt,
|
||
getattr(task, "old_title", "") or "",
|
||
retry=generation_cfg.get("retry"),
|
||
config=config,
|
||
models_path=models_path,
|
||
on_step=step_callback(task, "title"),
|
||
on_event=step_callback(task, "title"),
|
||
cmhub_config_path=cmhub_config_path,
|
||
)
|
||
] = task
|
||
for future in as_completed(futures):
|
||
task = futures[future]
|
||
if should_stop():
|
||
summary["cancelled"] = True
|
||
_cancel_pending(futures)
|
||
try:
|
||
new_title = future.result()
|
||
title_results[task.id] = new_title
|
||
summary["title_done"] += 1
|
||
set_step(task, "title_done")
|
||
_emit_generation_event(on_event, task, "title", "title_done", "success")
|
||
set_step(task, "db_write")
|
||
_emit_generation_event(on_event, task, "title", "db_write", "start")
|
||
existing_cover_path = getattr(task, "new_cover_path", None)
|
||
db.set_generated(task.id, new_title, existing_cover_path, path=db_path)
|
||
needs_cover = needs_by_task.get(getattr(task, "id", None), {}).get("cover")
|
||
if not needs_cover:
|
||
summary["generated_done"] += 1
|
||
if on_task_update is not None:
|
||
on_task_update(
|
||
task.id,
|
||
{
|
||
"stage": "generated",
|
||
"status": "success",
|
||
"new_title": new_title,
|
||
"new_cover_path": existing_cover_path,
|
||
},
|
||
)
|
||
_emit_generation_event(
|
||
on_event,
|
||
task,
|
||
"title",
|
||
"db_write",
|
||
"success",
|
||
detail="标题已保存，等待封面" if needs_cover else "仅生成标题",
|
||
)
|
||
except CancelledError:
|
||
summary["cancelled"] = True
|
||
_emit_generation_event(on_event, task, "title", get_step(task, "title_request"), "cancelled", level="warning")
|
||
except Exception as exc:
|
||
summary["failed"] += 1
|
||
summary["ok"] = False
|
||
step = get_step(task, "title_request")
|
||
error = _mark_generate_failed(task, exc, db_path, on_task_update, step=step)
|
||
_emit_generation_event(on_event, task, "title", step, "failed", detail=error, level="error")
|
||
_emit_generation_error(on_error, task, "title", step, exc, error)
|
||
_emit_generation_progress(on_progress, summary)
|
||
|
||
if not generate_cover:
|
||
if summary["cancelled"]:
|
||
summary["ok"] = False
|
||
return summary
|
||
|
||
for task in eligible:
|
||
task_id = getattr(task, "id", None)
|
||
if task_id in title_results:
|
||
continue
|
||
if generate_mode == "cover":
|
||
title_context = cover_title_context(task)
|
||
elif _has_generated_title(task):
|
||
title_context = str(getattr(task, "new_title") or "")
|
||
else:
|
||
title_context = ""
|
||
if not title_context:
|
||
continue
|
||
title_results[task_id] = title_context
|
||
if _has_generated_title(task):
|
||
detail = "已有标题"
|
||
else:
|
||
detail = "没有新标题，本轮使用旧标题作为封面参考"
|
||
_emit_generation_event(
|
||
on_event,
|
||
task,
|
||
"title",
|
||
"title_submit",
|
||
"skipped",
|
||
detail=detail,
|
||
)
|
||
|
||
cover_tasks = [
|
||
task for task in cover_candidates
|
||
if getattr(task, "id", None) in title_results
|
||
]
|
||
|
||
def prepare_cover_task(task, new_title):
|
||
set_step(task, "cover_prompt_render")
|
||
_emit_generation_event(on_event, task, "cover", "cover_prompt_render", "start")
|
||
rendered_cover_prompt = prompt_module.render_prompt(
|
||
cover_prompt,
|
||
_prompt_context(task, new_title, account_by_alias),
|
||
)
|
||
_emit_generation_event(on_event, task, "cover", "cover_prompt_render", "success")
|
||
set_step(task, "cover_submit")
|
||
_emit_generation_event(on_event, task, "cover", "cover_submit", "start")
|
||
return rendered_cover_prompt, _new_cover_path(task, account_by_alias, image_root)
|
||
|
||
def record_cover_failure(task, exc, fallback_step="cover_request"):
|
||
summary["failed"] += 1
|
||
summary["ok"] = False
|
||
step = get_step(task, fallback_step)
|
||
error = _mark_generate_failed(task, exc, db_path, on_task_update, step=step)
|
||
_emit_generation_event(on_event, task, "cover", step, "failed", detail=error, level="error")
|
||
_emit_generation_error(on_error, task, "cover", step, exc, error)
|
||
|
||
def record_cover_cancelled(task, fallback_step="cover_request"):
|
||
summary["cancelled"] = True
|
||
_emit_generation_event(
|
||
on_event,
|
||
task,
|
||
"cover",
|
||
get_step(task, fallback_step),
|
||
"cancelled",
|
||
level="warning",
|
||
)
|
||
|
||
def persist_cover_success(task, title_context, new_cover_path):
|
||
set_step(task, "db_write")
|
||
_emit_generation_event(on_event, task, "cover", "db_write", "start")
|
||
db.set_generated_cover(task.id, new_cover_path, path=db_path)
|
||
summary["cover_done"] += 1
|
||
summary["generated_done"] += 1
|
||
needs_title = needs_by_task.get(getattr(task, "id", None), {}).get("title")
|
||
persisted_title = title_context if needs_title else getattr(task, "new_title", None)
|
||
if on_task_update is not None:
|
||
on_task_update(
|
||
task.id,
|
||
{
|
||
"stage": "generated",
|
||
"status": "success",
|
||
"new_title": persisted_title,
|
||
"new_cover_path": new_cover_path,
|
||
},
|
||
)
|
||
_emit_generation_event(on_event, task, "cover", "db_write", "success", detail=new_cover_path)
|
||
|
||
def run_direct_cover_tasks():
|
||
with ThreadPoolExecutor(
|
||
max_workers=_positive_int(generation_cfg.get("image_concurrency", 1), 1)
|
||
) as executor:
|
||
futures = {}
|
||
for task in cover_tasks:
|
||
if should_stop():
|
||
summary["cancelled"] = True
|
||
break
|
||
new_title = title_results[task.id]
|
||
try:
|
||
rendered_cover_prompt, new_cover_path = prepare_cover_task(task, new_title)
|
||
futures[
|
||
executor.submit(
|
||
gen_cover,
|
||
rendered_cover_prompt,
|
||
getattr(task, "old_cover_path", "") or "",
|
||
new_cover_path,
|
||
resolution=generation_cfg.get("resolution"),
|
||
jpg_quality=generation_cfg.get("jpg_quality"),
|
||
retry=generation_cfg.get("retry"),
|
||
config=config,
|
||
models_path=models_path,
|
||
on_step=step_callback(task, "cover"),
|
||
on_event=step_callback(task, "cover"),
|
||
cmhub_config_path=cmhub_config_path,
|
||
)
|
||
] = (task, new_title)
|
||
except Exception as exc:
|
||
record_cover_failure(task, exc, fallback_step="cover_prompt_render")
|
||
_emit_generation_progress(on_progress, summary)
|
||
for future in as_completed(futures):
|
||
task, new_title = futures[future]
|
||
if should_stop():
|
||
summary["cancelled"] = True
|
||
_cancel_pending(futures)
|
||
try:
|
||
persist_cover_success(task, new_title, future.result())
|
||
except CancelledError:
|
||
record_cover_cancelled(task)
|
||
except Exception as exc:
|
||
record_cover_failure(task, exc)
|
||
_emit_generation_progress(on_progress, summary)
|
||
|
||
def run_cmhub_cover_tasks():
|
||
plan = cmhub_image_concurrency_plan(generation_cfg)
|
||
request_futures = {}
|
||
download_futures = {}
|
||
next_index = 0
|
||
|
||
def submit_next_request(request_executor):
|
||
nonlocal next_index
|
||
while next_index < len(cover_tasks):
|
||
if should_stop():
|
||
summary["cancelled"] = True
|
||
return False
|
||
task = cover_tasks[next_index]
|
||
next_index += 1
|
||
new_title = title_results[task.id]
|
||
try:
|
||
rendered_cover_prompt, new_cover_path = prepare_cover_task(task, new_title)
|
||
future = request_executor.submit(
|
||
_request_cmhub_cover_image_async,
|
||
rendered_cover_prompt,
|
||
getattr(task, "old_cover_path", "") or "",
|
||
new_cover_path,
|
||
resolution=generation_cfg.get("resolution"),
|
||
jpg_quality=generation_cfg.get("jpg_quality"),
|
||
retry=generation_cfg.get("retry"),
|
||
config=config,
|
||
cmhub_config_path=cmhub_config_path,
|
||
on_step=step_callback(task, "cover"),
|
||
on_event=step_callback(task, "cover"),
|
||
task=task,
|
||
db_path=db_path,
|
||
should_stop=should_stop,
|
||
)
|
||
request_futures[future] = (task, new_title)
|
||
return True
|
||
except Exception as exc:
|
||
record_cover_failure(task, exc, fallback_step="cover_prompt_render")
|
||
_emit_generation_progress(on_progress, summary)
|
||
return False
|
||
|
||
with ThreadPoolExecutor(max_workers=plan["request_concurrency"]) as request_executor, \
|
||
ThreadPoolExecutor(max_workers=plan["download_concurrency"]) as download_executor:
|
||
for _ in range(plan["request_concurrency"]):
|
||
if not submit_next_request(request_executor):
|
||
break
|
||
|
||
while request_futures or download_futures:
|
||
if should_stop():
|
||
summary["cancelled"] = True
|
||
_cancel_pending(request_futures)
|
||
done, _ = wait(
|
||
set(request_futures.keys()) | set(download_futures.keys()),
|
||
return_when=FIRST_COMPLETED,
|
||
)
|
||
for future in done:
|
||
if future in request_futures:
|
||
task, new_title = request_futures.pop(future)
|
||
try:
|
||
request_result = future.result()
|
||
download_future = download_executor.submit(
|
||
_download_and_save_cmhub_cover,
|
||
request_result,
|
||
on_step=step_callback(task, "cover"),
|
||
)
|
||
download_futures[download_future] = (task, new_title)
|
||
except CancelledError:
|
||
record_cover_cancelled(task)
|
||
_emit_generation_progress(on_progress, summary)
|
||
except Exception as exc:
|
||
record_cover_failure(task, exc)
|
||
_emit_generation_progress(on_progress, summary)
|
||
else:
|
||
task, new_title = download_futures.pop(future)
|
||
try:
|
||
persist_cover_success(task, new_title, future.result())
|
||
except CancelledError:
|
||
record_cover_cancelled(task, fallback_step="cover_download")
|
||
except Exception as exc:
|
||
record_cover_failure(task, exc, fallback_step="cover_download")
|
||
_emit_generation_progress(on_progress, summary)
|
||
while (
|
||
not summary["cancelled"]
|
||
and not should_stop()
|
||
and len(request_futures) < plan["request_concurrency"]
|
||
and next_index < len(cover_tasks)
|
||
):
|
||
if not submit_next_request(request_executor):
|
||
break
|
||
|
||
if _ai_backend(generation_cfg) == "cmhub":
|
||
run_cmhub_cover_tasks()
|
||
else:
|
||
run_direct_cover_tasks()
|
||
|
||
if summary["cancelled"]:
|
||
summary["ok"] = False
|
||
return summary
|
||
|
||
|
||
def fetch_cmhub_models(base_url, api_key, connect_timeout=10, read_timeout=30, use_system_proxy=False):
|
||
"""Fetch cmhub model aliases for settings UI."""
|
||
|
||
base_url = str(base_url or "").strip()
|
||
api_key = str(api_key or "")
|
||
if not base_url:
|
||
raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub Base URL")
|
||
if not api_key:
|
||
raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub API Key")
|
||
_apply_cmhub_proxy(use_system_proxy)
|
||
data = _cmhub_call_with_retry(
|
||
"GET",
|
||
appconfig.cmhub_request_url(base_url, "/api/v1/models"),
|
||
api_key,
|
||
payload=None,
|
||
connect_timeout=max(1, int(connect_timeout or 10)),
|
||
read_timeout=max(1, int(read_timeout or 30)),
|
||
attempts=1,
|
||
on_retry=None,
|
||
)
|
||
models = data.get("models", [])
|
||
if not isinstance(models, list):
|
||
raise CMHubError("bad_response", "cmhub 模型列表格式错误")
|
||
return [copy.deepcopy(model) for model in models if isinstance(model, dict)]
|
||
|
||
|
||
def fetch_cmhub_balance(base_url, api_key, connect_timeout=10, read_timeout=30, use_system_proxy=False):
|
||
"""Fetch cmhub point balance for settings UI."""
|
||
|
||
base_url = str(base_url or "").strip()
|
||
api_key = str(api_key or "")
|
||
if not base_url:
|
||
raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub Base URL")
|
||
if not api_key:
|
||
raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub API Key")
|
||
_apply_cmhub_proxy(use_system_proxy)
|
||
data = _cmhub_call_with_retry(
|
||
"GET",
|
||
appconfig.cmhub_request_url(base_url, "/api/v1/balance"),
|
||
api_key,
|
||
payload=None,
|
||
connect_timeout=max(1, int(connect_timeout or 10)),
|
||
read_timeout=max(1, int(read_timeout or 30)),
|
||
attempts=1,
|
||
on_retry=None,
|
||
)
|
||
if "points_balance" not in data:
|
||
raise CMHubError("bad_response", "cmhub 余额返回格式错误")
|
||
return copy.deepcopy(data)
|
||
|
||
|
||
def _ai_backend(ai_cfg):
|
||
backend = str(ai_cfg.get("backend", "direct") or "direct").strip().lower()
|
||
if backend not in appconfig.AI_BACKENDS:
|
||
raise AIError("AI backend 必须是 direct 或 cmhub")
|
||
return backend
|
||
|
||
|
||
def _gen_title_cmhub(
|
||
title_prompt,
|
||
old_title,
|
||
retry,
|
||
config,
|
||
cmhub_config_path,
|
||
on_step=None,
|
||
on_event=None,
|
||
):
|
||
ai_cfg = appconfig.ai_config(config)
|
||
_notify_step(on_step, "load_text_model")
|
||
runtime = _cmhub_runtime(config, "title", cmhub_config_path)
|
||
resolution = _normalize_cmhub_resolution(ai_cfg.get("resolution", "1k"))
|
||
_notify_step(on_step, "title_build_request")
|
||
payload = {
|
||
"prompt": _compose_title_prompt(title_prompt, old_title),
|
||
"model": runtime["alias"],
|
||
"resolution": resolution,
|
||
}
|
||
attempts = _attempt_count(ai_cfg, retry)
|
||
_notify_step(on_step, "title_request")
|
||
data = _cmhub_call_with_retry(
|
||
"POST",
|
||
appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/title"),
|
||
runtime["api_key"],
|
||
payload=payload,
|
||
connect_timeout=runtime["connect_timeout"],
|
||
read_timeout=CMHUB_TITLE_READ_TIMEOUT_SECONDS,
|
||
attempts=attempts,
|
||
on_retry=lambda attempt, total_attempts, exc: _notify_cmhub_retry(
|
||
on_step,
|
||
"title_request",
|
||
attempt,
|
||
total_attempts,
|
||
exc,
|
||
),
|
||
)
|
||
_emit_cmhub_metadata(on_event, data, "title_request")
|
||
_notify_step(on_step, "title_parse_response")
|
||
titles = data.get("titles")
|
||
if not isinstance(titles, list) or not titles:
|
||
raise AIError("AI 返回为空标题")
|
||
text = str(titles[0] or "").strip()
|
||
if not text:
|
||
raise AIError("AI 返回为空标题")
|
||
return text
|
||
|
||
|
||
def _gen_cover_cmhub(
|
||
cover_prompt,
|
||
old_cover_path,
|
||
out_path,
|
||
resolution,
|
||
jpg_quality,
|
||
retry,
|
||
config,
|
||
cmhub_config_path,
|
||
on_step=None,
|
||
on_event=None,
|
||
):
|
||
request_result = _request_cmhub_cover_image_sync(
|
||
cover_prompt,
|
||
old_cover_path,
|
||
out_path,
|
||
resolution=resolution,
|
||
jpg_quality=jpg_quality,
|
||
retry=retry,
|
||
config=config,
|
||
cmhub_config_path=cmhub_config_path,
|
||
on_step=on_step,
|
||
on_event=on_event,
|
||
validate_input=False,
|
||
)
|
||
return _download_and_save_cmhub_cover(request_result, on_step=on_step)
|
||
|
||
|
||
def _request_cmhub_cover_image(
|
||
cover_prompt,
|
||
old_cover_path,
|
||
out_path,
|
||
resolution,
|
||
jpg_quality,
|
||
retry,
|
||
config,
|
||
cmhub_config_path,
|
||
on_step=None,
|
||
on_event=None,
|
||
validate_input=True,
|
||
):
|
||
return _request_cmhub_cover_image_sync(
|
||
cover_prompt,
|
||
old_cover_path,
|
||
out_path,
|
||
resolution=resolution,
|
||
jpg_quality=jpg_quality,
|
||
retry=retry,
|
||
config=config,
|
||
cmhub_config_path=cmhub_config_path,
|
||
on_step=on_step,
|
||
on_event=on_event,
|
||
validate_input=validate_input,
|
||
)
|
||
|
||
|
||
def _request_cmhub_cover_image_sync(
|
||
cover_prompt,
|
||
old_cover_path,
|
||
out_path,
|
||
resolution,
|
||
jpg_quality,
|
||
retry,
|
||
config,
|
||
cmhub_config_path,
|
||
on_step=None,
|
||
on_event=None,
|
||
validate_input=True,
|
||
):
|
||
if validate_input:
|
||
old_cover_path, out_path = _prepare_cover_input(old_cover_path, out_path, on_step)
|
||
ai_cfg = appconfig.ai_config(config)
|
||
_notify_step(on_step, "load_image_model")
|
||
runtime = _cmhub_runtime(config, "image", cmhub_config_path)
|
||
resolution = str(resolution or ai_cfg.get("resolution", "1k"))
|
||
quality = _jpg_quality(jpg_quality if jpg_quality is not None else ai_cfg.get("jpg_quality", 90))
|
||
_notify_step(on_step, "cover_build_request")
|
||
payload = {
|
||
"prompt": str(cover_prompt or ""),
|
||
"model": runtime["alias"],
|
||
"image_base64": _image_data_url(old_cover_path),
|
||
"resolution": _normalize_cmhub_resolution(resolution),
|
||
"aspect_ratio": "1:1",
|
||
}
|
||
attempts = _attempt_count(ai_cfg, retry)
|
||
read_timeout = _cmhub_read_timeout(config, resolution)
|
||
_notify_step(on_step, "cover_request")
|
||
request_started = time.perf_counter()
|
||
data = _cmhub_call_with_retry(
|
||
"POST",
|
||
appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/image"),
|
||
runtime["api_key"],
|
||
payload=payload,
|
||
connect_timeout=runtime["connect_timeout"],
|
||
read_timeout=read_timeout,
|
||
attempts=attempts,
|
||
on_retry=lambda attempt, total_attempts, exc: _notify_cmhub_retry(
|
||
on_step,
|
||
"cover_request",
|
||
attempt,
|
||
total_attempts,
|
||
exc,
|
||
),
|
||
)
|
||
request_elapsed = time.perf_counter() - request_started
|
||
_emit_cmhub_metadata(on_event, data, "cover_request")
|
||
_notify_step(on_step, "cover_parse_response")
|
||
image_url = _extract_cmhub_image_url(data, runtime["base_url"])
|
||
if not image_url:
|
||
raise AIError("AI 返回中没有图片数据")
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_request",
|
||
detail="cmhub 已返回 image_url，耗时 %s" % _format_seconds(request_elapsed),
|
||
)
|
||
if _debug_cmhub_image_url_enabled():
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_image_url",
|
||
result="debug",
|
||
detail="cmhub 图片 URL：%s" % image_url,
|
||
level="warning",
|
||
debug_only=True,
|
||
)
|
||
return {
|
||
"image_url": image_url,
|
||
"connect_timeout": runtime["connect_timeout"],
|
||
"read_timeout": read_timeout,
|
||
"out_path": out_path,
|
||
"resolution": resolution,
|
||
"quality": quality,
|
||
"use_system_proxy": runtime["use_system_proxy"],
|
||
"download_with_curl": runtime["download_with_curl"],
|
||
}
|
||
|
||
|
||
def _request_cmhub_cover_image_async(
|
||
cover_prompt,
|
||
old_cover_path,
|
||
out_path,
|
||
resolution,
|
||
jpg_quality,
|
||
retry,
|
||
config,
|
||
cmhub_config_path,
|
||
on_step=None,
|
||
on_event=None,
|
||
task=None,
|
||
db_path=None,
|
||
should_stop=None,
|
||
validate_input=True,
|
||
):
|
||
if validate_input:
|
||
old_cover_path, out_path = _prepare_cover_input(old_cover_path, out_path, on_step)
|
||
should_stop = should_stop or (lambda: False)
|
||
task_id = int(getattr(task, "id", 0) or 0)
|
||
if task_id <= 0:
|
||
raise AIError("缺少本地任务ID，无法提交 cmhub 生图任务")
|
||
ai_cfg = appconfig.ai_config(config)
|
||
_notify_step(on_step, "load_image_model")
|
||
runtime = _cmhub_runtime(config, "image", cmhub_config_path)
|
||
resolution = str(resolution or ai_cfg.get("resolution", "1k"))
|
||
quality = _jpg_quality(jpg_quality if jpg_quality is not None else ai_cfg.get("jpg_quality", 90))
|
||
_notify_step(on_step, "cover_build_request")
|
||
payload = {
|
||
"prompt": str(cover_prompt or ""),
|
||
"model": runtime["alias"],
|
||
"image_base64": _image_data_url(old_cover_path),
|
||
"resolution": _normalize_cmhub_resolution(resolution),
|
||
"aspect_ratio": "1:1",
|
||
}
|
||
request_result = {
|
||
"connect_timeout": runtime["connect_timeout"],
|
||
"read_timeout": _cmhub_read_timeout(config, resolution),
|
||
"out_path": out_path,
|
||
"resolution": resolution,
|
||
"quality": quality,
|
||
"use_system_proxy": runtime["use_system_proxy"],
|
||
"download_with_curl": runtime["download_with_curl"],
|
||
}
|
||
existing_image_task_id = str(getattr(task, "image_task_id", "") or "").strip()
|
||
if existing_image_task_id:
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_request",
|
||
detail="发现未完成的 cmhub 生图任务，继续查询结果",
|
||
)
|
||
return _poll_cmhub_cover_image_task(
|
||
existing_image_task_id,
|
||
runtime,
|
||
request_result,
|
||
task_id=task_id,
|
||
db_path=db_path,
|
||
should_stop=should_stop,
|
||
on_step=on_step,
|
||
)
|
||
|
||
image_task_key = str(getattr(task, "image_task_key", "") or "").strip()
|
||
if not image_task_key:
|
||
image_task_key = db.ensure_image_task_key(task_id, path=db_path)
|
||
_raise_if_cmhub_cover_cancelled(should_stop, on_step)
|
||
attempts = _attempt_count(ai_cfg, retry)
|
||
_notify_step(on_step, "cover_request")
|
||
request_started = time.perf_counter()
|
||
try:
|
||
data = _cmhub_call_with_retry(
|
||
"POST",
|
||
appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/image/tasks"),
|
||
runtime["api_key"],
|
||
payload=payload,
|
||
connect_timeout=runtime["connect_timeout"],
|
||
read_timeout=CMHUB_IMAGE_SUBMIT_READ_TIMEOUT_SECONDS,
|
||
attempts=attempts,
|
||
headers_extra={
|
||
"Idempotency-Key": image_task_key,
|
||
"X-Client-Version": str(APP_VERSION),
|
||
},
|
||
on_retry=lambda attempt, total_attempts, exc: _notify_cmhub_retry(
|
||
on_step,
|
||
"cover_request",
|
||
attempt,
|
||
total_attempts,
|
||
exc,
|
||
),
|
||
)
|
||
except CMHubError as exc:
|
||
if exc.code in {
|
||
"bad_request",
|
||
"content_blocked",
|
||
"insufficient_points",
|
||
"idempotency_conflict",
|
||
"unauthorized",
|
||
"account_disabled",
|
||
"model_not_allowed",
|
||
"no_pricing_rule",
|
||
"not_found",
|
||
}:
|
||
db.clear_image_task(task_id, path=db_path)
|
||
raise
|
||
request_elapsed = time.perf_counter() - request_started
|
||
_emit_cmhub_metadata(on_event, data, "cover_request")
|
||
image_task_id = str(data.get("task_id") or "").strip()
|
||
if not image_task_id:
|
||
raise CMHubError("bad_response", "cmhub 生图任务提交返回格式错误", retryable=False)
|
||
db.set_image_task_submitted(task_id, image_task_id, image_task_key, path=db_path)
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_request",
|
||
detail="cmhub 已提交生图任务 %s，耗时 %s"
|
||
% (image_task_id, _format_seconds(request_elapsed)),
|
||
)
|
||
return _poll_cmhub_cover_image_task(
|
||
image_task_id,
|
||
runtime,
|
||
request_result,
|
||
task_id=task_id,
|
||
db_path=db_path,
|
||
should_stop=should_stop,
|
||
on_step=on_step,
|
||
)
|
||
|
||
|
||
def _poll_cmhub_cover_image_task(
|
||
image_task_id,
|
||
runtime,
|
||
request_result,
|
||
task_id,
|
||
db_path,
|
||
should_stop,
|
||
on_step=None,
|
||
):
|
||
poll_url = appconfig.cmhub_request_url(
|
||
runtime["base_url"],
|
||
"/api/v1/generate/image/tasks/%s" % urllib.parse.quote(str(image_task_id), safe=""),
|
||
)
|
||
started = time.perf_counter()
|
||
deadline = started + max(1, int(CMHUB_IMAGE_READ_TIMEOUT_SECONDS))
|
||
poll_index = 0
|
||
while True:
|
||
_raise_if_cmhub_cover_cancelled(should_stop, on_step)
|
||
if time.perf_counter() >= deadline:
|
||
raise CMHubError(
|
||
"read_timeout",
|
||
"等待 cmhub 生图任务完成超时，下次可继续查询",
|
||
retryable=False,
|
||
)
|
||
_notify_step(on_step, "cover_poll")
|
||
try:
|
||
data = _cmhub_call_once(
|
||
"GET",
|
||
poll_url,
|
||
runtime["api_key"],
|
||
payload=None,
|
||
connect_timeout=runtime["connect_timeout"],
|
||
read_timeout=CMHUB_IMAGE_POLL_READ_TIMEOUT_SECONDS,
|
||
headers_extra={"X-Client-Version": str(APP_VERSION)},
|
||
)
|
||
except CMHubError as exc:
|
||
if exc.code in {"connect_timeout", "read_timeout", "network_error", "rate_limited"}:
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_poll",
|
||
result="retry",
|
||
detail=str(exc),
|
||
level="warning",
|
||
)
|
||
_sleep_cmhub_poll(poll_index, should_stop, on_step)
|
||
poll_index += 1
|
||
continue
|
||
db.clear_image_task(task_id, path=db_path)
|
||
raise
|
||
status = str(data.get("status") or "").strip().lower()
|
||
if status in {"queued", "running"}:
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_poll",
|
||
detail="cmhub 生图任务%s，继续等待" % ("排队中" if status == "queued" else "生成中"),
|
||
)
|
||
_sleep_cmhub_poll(poll_index, should_stop, on_step)
|
||
poll_index += 1
|
||
continue
|
||
if status == "succeeded":
|
||
image_url = _extract_cmhub_image_url(data, runtime["base_url"])
|
||
if not image_url:
|
||
raise CMHubError("bad_response", "cmhub 生图任务成功但没有图片地址", retryable=False)
|
||
_notify_step(on_step, "cover_parse_response")
|
||
elapsed = time.perf_counter() - started
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_request",
|
||
detail="cmhub 已返回 image_url，耗时 %s" % _format_seconds(elapsed),
|
||
)
|
||
if _debug_cmhub_image_url_enabled():
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_image_url",
|
||
result="debug",
|
||
detail="cmhub 图片 URL：%s" % image_url,
|
||
level="warning",
|
||
debug_only=True,
|
||
)
|
||
merged = dict(request_result)
|
||
merged["image_url"] = image_url
|
||
return merged
|
||
if status in {"failed", "expired"}:
|
||
error = data.get("error") if isinstance(data.get("error"), dict) else {}
|
||
code = _normalize_cmhub_error_code(error.get("code") or status)
|
||
message = _cmhub_user_message(code, str(error.get("message") or status))
|
||
db.clear_image_task(task_id, path=db_path)
|
||
raise CMHubError(code, message, retryable=_cmhub_retryable(code))
|
||
raise CMHubError("bad_response", "cmhub 生图任务状态返回格式错误", retryable=False)
|
||
|
||
|
||
def _sleep_cmhub_poll(poll_index, should_stop, on_step=None):
|
||
_raise_if_cmhub_cover_cancelled(should_stop, on_step)
|
||
delays = CMHUB_IMAGE_POLL_DELAYS_SECONDS
|
||
delay = delays[min(max(0, int(poll_index)), len(delays) - 1)]
|
||
time.sleep(max(0.0, float(delay)))
|
||
_raise_if_cmhub_cover_cancelled(should_stop, on_step)
|
||
|
||
|
||
def _raise_if_cmhub_cover_cancelled(should_stop, on_step=None):
|
||
try:
|
||
cancelled = bool(should_stop and should_stop())
|
||
except Exception:
|
||
cancelled = False
|
||
if not cancelled:
|
||
return
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_poll",
|
||
result="cancelled",
|
||
detail="已停止等待生图结果；服务端任务可能仍在完成，下次可继续查询",
|
||
level="warning",
|
||
)
|
||
raise CancelledError()
|
||
|
||
|
||
def _download_and_save_cmhub_cover(request_result, on_step=None):
|
||
image_url = request_result["image_url"]
|
||
connect_timeout = request_result["connect_timeout"]
|
||
read_timeout = request_result["read_timeout"]
|
||
out_path = request_result["out_path"]
|
||
resolution = request_result["resolution"]
|
||
quality = request_result["quality"]
|
||
_notify_step(on_step, "cover_download")
|
||
image_bytes, download_elapsed = _download_cmhub_image_with_retry(
|
||
image_url,
|
||
connect_timeout=connect_timeout,
|
||
read_timeout=read_timeout,
|
||
use_system_proxy=request_result.get("use_system_proxy", False),
|
||
download_with_curl=request_result.get("download_with_curl", "false"),
|
||
on_step=on_step,
|
||
)
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_download",
|
||
detail="下载完成，%s，耗时 %s"
|
||
% (_format_bytes(len(image_bytes)), _format_seconds(download_elapsed)),
|
||
)
|
||
_notify_step(on_step, "cover_save")
|
||
save_started = time.perf_counter()
|
||
saved_path = _save_jpeg(image_bytes, out_path, resolution, quality)
|
||
save_elapsed = time.perf_counter() - save_started
|
||
detail = "JPEG 已保存，耗时 %s" % _format_seconds(save_elapsed)
|
||
try:
|
||
detail += "，文件 %s" % _format_bytes(os.path.getsize(saved_path))
|
||
except OSError:
|
||
pass
|
||
_notify_step_event(on_step, "cover_save", detail=detail)
|
||
return saved_path
|
||
|
||
|
||
def _cmhub_runtime(config, operation, cmhub_config_path):
|
||
hub = appconfig.cmhub_config(config)
|
||
api_key = appconfig.get_cmhub_api_key(path=cmhub_config_path)
|
||
alias_key = "title_alias" if operation == "title" else "image_alias"
|
||
missing = []
|
||
if not hub.get("base_url"):
|
||
missing.append("Base URL")
|
||
if not api_key:
|
||
missing.append("API Key")
|
||
if not hub.get(alias_key):
|
||
missing.append("生文别名" if operation == "title" else "生图别名")
|
||
if missing:
|
||
raise CMHubError(
|
||
"cmhub_not_configured",
|
||
"请去⑤设置配置 cmhub：缺少 " + "、".join(missing),
|
||
retryable=False,
|
||
)
|
||
use_system_proxy = bool(hub.get("use_system_proxy", False))
|
||
_apply_cmhub_proxy(use_system_proxy)
|
||
return {
|
||
"base_url": hub["base_url"].rstrip("/"),
|
||
"api_key": api_key,
|
||
"alias": hub[alias_key],
|
||
"connect_timeout": max(1, int(hub.get("connect_timeout", 10) or 10)),
|
||
"use_system_proxy": use_system_proxy,
|
||
"download_with_curl": str(hub.get("download_with_curl", "auto") or "auto"),
|
||
}
|
||
|
||
|
||
def _compose_title_prompt(title_prompt, old_title):
|
||
prompt = str(title_prompt or "").strip()
|
||
old_title_text = str(old_title or "")
|
||
if "{旧标题}" in prompt:
|
||
body = prompt.replace("{旧标题}", old_title_text).strip()
|
||
else:
|
||
body = "%s\n\n旧标题：\n%s" % (prompt, old_title_text)
|
||
if body:
|
||
return "%s\n\n请只返回新标题，不要解释。" % body
|
||
return "请只返回新标题，不要解释。"
|
||
|
||
|
||
def _normalize_cmhub_resolution(resolution):
|
||
value = str(resolution or "1k").strip().lower()
|
||
mapping = {
|
||
"512": "512",
|
||
"512x512": "512",
|
||
"1k": "1K",
|
||
"1K": "1K",
|
||
"1024": "1K",
|
||
"2k": "2K",
|
||
"2K": "2K",
|
||
"2048": "2K",
|
||
"4k": "4K",
|
||
"4K": "4K",
|
||
"4096": "4K",
|
||
}
|
||
return mapping.get(value, str(resolution or "1K").upper())
|
||
|
||
|
||
def _cmhub_read_timeout(config, resolution):
|
||
return CMHUB_IMAGE_READ_TIMEOUT_SECONDS
|
||
|
||
|
||
def _debug_cmhub_image_url_enabled():
|
||
value = str(os.environ.get("CMSHOPEE_DEBUG_CMHUB_IMAGE_URL", "") or "")
|
||
return value.strip().lower() in {"1", "true", "yes", "on", "debug"}
|
||
|
||
|
||
def _extract_cmhub_image_url(data, base_url):
|
||
candidate = _find_image_ref(data)
|
||
if not candidate:
|
||
return ""
|
||
return _normalize_cmhub_image_url(candidate, base_url)
|
||
|
||
|
||
def _normalize_cmhub_image_url(value, base_url):
|
||
text = str(value or "").strip()
|
||
if not text:
|
||
return ""
|
||
parts = urllib.parse.urlsplit(text)
|
||
if parts.scheme in {"http", "https"}:
|
||
return text
|
||
if parts.scheme:
|
||
raise CMHubError(
|
||
"bad_response",
|
||
"cmhub 生图任务返回的图片地址格式错误",
|
||
retryable=False,
|
||
)
|
||
base = str(base_url or "").strip()
|
||
if not base:
|
||
raise CMHubError(
|
||
"bad_response",
|
||
"cmhub 生图任务返回了相对图片地址，但缺少 cmhub Base URL",
|
||
retryable=False,
|
||
)
|
||
base_parts = urllib.parse.urlsplit(base)
|
||
if text.startswith("//"):
|
||
scheme = base_parts.scheme or "https"
|
||
return f"{scheme}:{text}"
|
||
return urllib.parse.urljoin(base.rstrip("/") + "/", text)
|
||
|
||
|
||
def _download_cmhub_image_with_retry(
|
||
url,
|
||
connect_timeout,
|
||
read_timeout,
|
||
use_system_proxy=False,
|
||
download_with_curl="false",
|
||
on_step=None,
|
||
attempts=CMHUB_IMAGE_DOWNLOAD_ATTEMPTS,
|
||
slow_threshold=CMHUB_IMAGE_SLOW_DOWNLOAD_SECONDS,
|
||
):
|
||
total_attempts = max(1, int(attempts or 1))
|
||
total_started = time.perf_counter()
|
||
last_exc = None
|
||
for index in range(total_attempts):
|
||
try:
|
||
image_bytes = _download_cmhub_image(
|
||
url,
|
||
connect_timeout=connect_timeout,
|
||
read_timeout=read_timeout,
|
||
use_system_proxy=use_system_proxy,
|
||
download_with_curl=download_with_curl,
|
||
)
|
||
elapsed = time.perf_counter() - total_started
|
||
if elapsed >= float(slow_threshold or 0):
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_download",
|
||
result="warning",
|
||
detail="图片下载较慢，已用 %s，大小 %s"
|
||
% (_format_seconds(elapsed), _format_bytes(len(image_bytes))),
|
||
level="warning",
|
||
)
|
||
return image_bytes, elapsed
|
||
except Exception as exc:
|
||
last_exc = exc
|
||
if index + 1 >= total_attempts or not _cmhub_download_retryable(exc):
|
||
break
|
||
_notify_step_event(
|
||
on_step,
|
||
"cover_download",
|
||
result="retry",
|
||
detail=str(exc),
|
||
level="warning",
|
||
attempt=index + 1,
|
||
attempts=total_attempts,
|
||
)
|
||
time.sleep(min(2.0, 0.5 * (index + 1)))
|
||
if total_attempts > 1 and _cmhub_download_retryable(last_exc):
|
||
raise AIError(
|
||
"下载 cmhub 图片失败（已尝试 %s 次）: %s"
|
||
% (total_attempts, str(last_exc))
|
||
) from last_exc
|
||
raise last_exc
|
||
|
||
|
||
def _cmhub_download_retryable(exc):
|
||
if exc is None:
|
||
return False
|
||
message = str(exc or "")
|
||
if "图片超过大小上限" in message:
|
||
return False
|
||
if "HTTP 4" in message:
|
||
return False
|
||
return message.startswith("下载 cmhub 图片失败")
|
||
|
||
|
||
def _cmhub_call_with_retry(
|
||
method,
|
||
url,
|
||
api_key,
|
||
payload,
|
||
connect_timeout,
|
||
read_timeout,
|
||
attempts,
|
||
on_retry=None,
|
||
headers_extra=None,
|
||
):
|
||
attempts = max(1, int(attempts or 1))
|
||
last_exc = None
|
||
for index in range(attempts):
|
||
try:
|
||
return _cmhub_call_once(
|
||
method,
|
||
url,
|
||
api_key,
|
||
payload,
|
||
connect_timeout=connect_timeout,
|
||
read_timeout=read_timeout,
|
||
headers_extra=headers_extra,
|
||
)
|
||
except CMHubError as exc:
|
||
last_exc = exc
|
||
if not exc.retryable or index + 1 >= attempts:
|
||
raise
|
||
if on_retry is not None:
|
||
try:
|
||
on_retry(index + 1, attempts, exc)
|
||
except Exception:
|
||
pass
|
||
time.sleep(_cmhub_retry_delay(exc, index))
|
||
raise last_exc
|
||
|
||
|
||
def _cmhub_call_once(method, url, api_key, payload, connect_timeout, read_timeout, headers_extra=None):
|
||
headers = {
|
||
"Authorization": "Bearer " + str(api_key),
|
||
"Accept": "application/json",
|
||
}
|
||
if headers_extra:
|
||
headers.update({str(key): str(value) for key, value in dict(headers_extra).items()})
|
||
request_kwargs = {
|
||
"headers": headers,
|
||
"timeout": (max(1, int(connect_timeout)), max(1, int(read_timeout))),
|
||
}
|
||
if str(method).upper() != "GET":
|
||
request_kwargs["json"] = payload or {}
|
||
try:
|
||
response = _cmhub_session().request(str(method).upper(), url, **request_kwargs)
|
||
except requests.exceptions.ConnectTimeout as exc:
|
||
raise CMHubError("connect_timeout", "连接 cmhub 超时", retryable=True) from exc
|
||
except requests.exceptions.ReadTimeout as exc:
|
||
raise CMHubError("read_timeout", "等待 cmhub 返回超时", retryable=False) from exc
|
||
except requests.exceptions.RequestException as exc:
|
||
raise CMHubError(
|
||
"network_error",
|
||
_redact_cmhub(str(exc), api_key),
|
||
retryable=False,
|
||
) from exc
|
||
return _cmhub_response_json(response, api_key)
|
||
|
||
|
||
def _cmhub_response_json(response, api_key):
|
||
try:
|
||
data = response.json()
|
||
except ValueError:
|
||
data = {}
|
||
status = getattr(response, "status_code", None)
|
||
if status and status >= 400:
|
||
raise _cmhub_error_from_response(data, response, api_key)
|
||
if isinstance(data, dict) and isinstance(data.get("error"), dict):
|
||
raise _cmhub_error_from_response(data, response, api_key)
|
||
if not isinstance(data, dict):
|
||
raise CMHubError("bad_response", "cmhub 返回格式错误", status=status)
|
||
return data
|
||
|
||
|
||
def _cmhub_error_from_response(data, response, api_key):
|
||
status = getattr(response, "status_code", None)
|
||
error = data.get("error") if isinstance(data, dict) else None
|
||
if not isinstance(error, dict):
|
||
error = {}
|
||
code = _normalize_cmhub_error_code(error.get("code") or _cmhub_code_for_status(status), status)
|
||
raw_message = (error.get("message") or data.get("message")) if isinstance(data, dict) else ""
|
||
if not raw_message:
|
||
raw_message = getattr(response, "text", "")[:500]
|
||
message = _cmhub_user_message(code, _redact_cmhub(raw_message or code, api_key))
|
||
retry_after = _parse_retry_after(getattr(response, "headers", {}).get("Retry-After"))
|
||
return CMHubError(
|
||
code,
|
||
message,
|
||
status=status,
|
||
retryable=_cmhub_retryable(code),
|
||
retry_after=retry_after,
|
||
)
|
||
|
||
|
||
def _normalize_cmhub_error_code(code, status=None):
|
||
normalized = str(code or "unknown").strip().lower().replace("-", "_")
|
||
if status == 404 or normalized in {"notfound", "not_found"}:
|
||
return "not_found"
|
||
return normalized or "unknown"
|
||
|
||
|
||
def _cmhub_code_for_status(status):
|
||
return {
|
||
400: "bad_request",
|
||
401: "unauthorized",
|
||
402: "insufficient_points",
|
||
403: "account_disabled",
|
||
404: "not_found",
|
||
429: "rate_limited",
|
||
502: "upstream_error",
|
||
}.get(status, "unknown")
|
||
|
||
|
||
def _cmhub_retryable(code):
|
||
return str(code) in {
|
||
"upstream_error",
|
||
"upstream_timeout",
|
||
"task_timeout",
|
||
"rate_limited",
|
||
"connect_timeout",
|
||
}
|
||
|
||
|
||
def _cmhub_user_message(code, message):
|
||
defaults = {
|
||
"insufficient_points": "点数不足，请先充值",
|
||
"unauthorized": "cmhub API Key 无效，请去⑤设置重填",
|
||
"account_disabled": "cmhub 账号已禁用，请去网页端处理",
|
||
"bad_request": "cmhub 请求参数错误",
|
||
"model_not_allowed": "cmhub 模型别名无权限",
|
||
"no_pricing_rule": "cmhub 模型别名未配置价格",
|
||
"content_blocked": "cmhub 内容安全策略拒绝本次生成",
|
||
"upstream_error": "cmhub 上游生成失败，请稍后重试",
|
||
"upstream_timeout": "cmhub 上游生成超时，点数已退回，请稍后重试",
|
||
"task_timeout": "cmhub 生图任务超时，点数已退回，请稍后重试",
|
||
"idempotency_conflict": "cmhub 生图幂等键冲突，请重新生成",
|
||
"rate_limited": "cmhub 请求过于频繁，请稍后重试",
|
||
"not_found": "cmhub 接口不存在，请检查 Base URL 或该实例是否已部署 /api/v1/models",
|
||
}
|
||
default = defaults.get(str(code))
|
||
if str(code) == "not_found":
|
||
return default
|
||
if default and message and str(message) not in default:
|
||
return "%s：%s" % (default, message)
|
||
return default or str(message or code)
|
||
|
||
|
||
def _parse_retry_after(value):
|
||
try:
|
||
if value is None or value == "":
|
||
return None
|
||
return max(0.0, float(value))
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
|
||
def _cmhub_retry_delay(exc, index):
|
||
if exc.retry_after is not None:
|
||
return min(2.0, max(0.1, float(exc.retry_after)))
|
||
return min(2.0, 0.4 * (index + 1))
|
||
|
||
|
||
def _notify_cmhub_retry(callback, step, attempt, attempts, exc):
|
||
if callback is None:
|
||
return
|
||
try:
|
||
callback(
|
||
{
|
||
"step": step,
|
||
"result": "retry",
|
||
"attempt": attempt,
|
||
"attempts": attempts,
|
||
"detail": str(exc),
|
||
}
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _emit_cmhub_metadata(callback, data, step):
|
||
if callback is None or not isinstance(data, dict):
|
||
return
|
||
metadata = {
|
||
key: data.get(key)
|
||
for key in ("alias", "model_used", "points_cost", "points_balance", "call_id")
|
||
if data.get(key) is not None
|
||
}
|
||
if not metadata:
|
||
return
|
||
try:
|
||
callback(
|
||
{
|
||
"step": step,
|
||
"result": "meta",
|
||
"level": "info",
|
||
"metadata": metadata,
|
||
}
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _format_seconds(seconds):
|
||
value = max(0.0, float(seconds or 0.0))
|
||
return "%.1f秒" % value
|
||
|
||
|
||
def _format_bytes(size):
|
||
value = float(max(0, int(size or 0)))
|
||
units = ["B", "KB", "MB", "GB"]
|
||
unit = units[0]
|
||
for unit in units:
|
||
if value < 1024 or unit == units[-1]:
|
||
break
|
||
value /= 1024
|
||
if unit == "B":
|
||
return "%d%s" % (int(value), unit)
|
||
return "%.1f%s" % (value, unit)
|
||
|
||
|
||
def _download_cmhub_image(
|
||
url,
|
||
connect_timeout,
|
||
read_timeout,
|
||
max_bytes=CMHUB_IMAGE_MAX_BYTES,
|
||
use_system_proxy=False,
|
||
download_with_curl="false",
|
||
):
|
||
_assert_public_http_url(url)
|
||
if _should_use_curl_for_cmhub_download(download_with_curl):
|
||
try:
|
||
return _download_cmhub_image_with_curl(
|
||
url,
|
||
connect_timeout=connect_timeout,
|
||
read_timeout=read_timeout,
|
||
max_bytes=max_bytes,
|
||
use_system_proxy=use_system_proxy,
|
||
)
|
||
except AIError:
|
||
pass
|
||
return _download_cmhub_image_with_requests(
|
||
url,
|
||
connect_timeout=connect_timeout,
|
||
read_timeout=read_timeout,
|
||
max_bytes=max_bytes,
|
||
)
|
||
|
||
|
||
def _download_cmhub_image_with_requests(url, connect_timeout, read_timeout, max_bytes):
|
||
try:
|
||
response = _cmhub_session().get(
|
||
url,
|
||
stream=True,
|
||
timeout=(max(1, int(connect_timeout)), max(1, int(read_timeout))),
|
||
)
|
||
except requests.exceptions.RequestException as exc:
|
||
raise AIError("下载 cmhub 图片失败: %s" % exc) from exc
|
||
status = getattr(response, "status_code", 200)
|
||
if status >= 400:
|
||
raise AIError("下载 cmhub 图片失败: HTTP %s" % status)
|
||
chunks = []
|
||
total = 0
|
||
iterator = response.iter_content(chunk_size=65536) if hasattr(response, "iter_content") else [response.content]
|
||
for chunk in iterator:
|
||
if not chunk:
|
||
continue
|
||
total += len(chunk)
|
||
if total > max_bytes:
|
||
raise AIError("下载 cmhub 图片失败: 图片超过大小上限")
|
||
chunks.append(chunk)
|
||
return b"".join(chunks)
|
||
|
||
|
||
def _should_use_curl_for_cmhub_download(mode):
|
||
normalized = _normalize_curl_download_mode(mode)
|
||
if normalized == "false":
|
||
return False
|
||
if normalized == "auto" and os.name != "nt":
|
||
return False
|
||
return bool(_find_system_curl())
|
||
|
||
|
||
def _normalize_curl_download_mode(mode):
|
||
if isinstance(mode, bool):
|
||
return "true" if mode else "false"
|
||
normalized = str(mode or "auto").strip().lower()
|
||
if normalized in {"auto", "true", "false"}:
|
||
return normalized
|
||
return "auto"
|
||
|
||
|
||
def _find_system_curl():
|
||
candidates = []
|
||
if os.name == "nt":
|
||
system_root = os.environ.get("SystemRoot") or r"C:\Windows"
|
||
candidates.append(os.path.join(system_root, "System32", "curl.exe"))
|
||
discovered = shutil.which("curl")
|
||
if discovered:
|
||
candidates.append(discovered)
|
||
seen = set()
|
||
for candidate in candidates:
|
||
if not candidate:
|
||
continue
|
||
path = os.path.abspath(candidate)
|
||
lowered = path.lower()
|
||
if lowered in seen:
|
||
continue
|
||
seen.add(lowered)
|
||
if os.path.isfile(path):
|
||
return path
|
||
return ""
|
||
|
||
|
||
def _subprocess_hidden_window_kwargs():
|
||
if os.name != "nt":
|
||
return {}
|
||
creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0)
|
||
if creationflags:
|
||
return {"creationflags": creationflags}
|
||
startupinfo_cls = getattr(subprocess, "STARTUPINFO", None)
|
||
if startupinfo_cls is None:
|
||
return {}
|
||
startupinfo = startupinfo_cls()
|
||
startupinfo.dwFlags |= getattr(subprocess, "STARTF_USESHOWWINDOW", 1)
|
||
startupinfo.wShowWindow = getattr(subprocess, "SW_HIDE", 0)
|
||
return {"startupinfo": startupinfo}
|
||
|
||
|
||
def _download_cmhub_image_with_curl(
|
||
url,
|
||
connect_timeout,
|
||
read_timeout,
|
||
max_bytes,
|
||
use_system_proxy=False,
|
||
):
|
||
curl_path = _find_system_curl()
|
||
if not curl_path:
|
||
raise AIError("下载 cmhub 图片失败: 未找到系统 curl")
|
||
temp_config_path = None
|
||
temp_output_path = None
|
||
try:
|
||
with tempfile.NamedTemporaryFile(
|
||
"w",
|
||
encoding="utf-8",
|
||
suffix=".curlrc",
|
||
delete=False,
|
||
) as config_file:
|
||
temp_config_path = config_file.name
|
||
config_file.write("url = %s\n" % _curl_config_quote(url))
|
||
with tempfile.NamedTemporaryFile("wb", suffix=".img", delete=False) as output_file:
|
||
temp_output_path = output_file.name
|
||
args = [
|
||
curl_path,
|
||
"-K",
|
||
temp_config_path,
|
||
"--fail",
|
||
"--silent",
|
||
"--show-error",
|
||
"--connect-timeout",
|
||
str(max(1, int(connect_timeout))),
|
||
"--max-time",
|
||
str(max(1, int(read_timeout))),
|
||
"--max-filesize",
|
||
str(max(1, int(max_bytes))),
|
||
"--output",
|
||
temp_output_path,
|
||
]
|
||
if not bool(use_system_proxy):
|
||
args.extend(["--noproxy", "*"])
|
||
try:
|
||
completed = subprocess.run(
|
||
args,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE,
|
||
timeout=max(2, int(connect_timeout) + int(read_timeout) + 10),
|
||
check=False,
|
||
shell=False,
|
||
**_subprocess_hidden_window_kwargs(),
|
||
)
|
||
except (OSError, subprocess.TimeoutExpired) as exc:
|
||
raise AIError("下载 cmhub 图片失败: curl 执行失败") from exc
|
||
if completed.returncode != 0:
|
||
raise AIError("下载 cmhub 图片失败: curl 退出码 %s" % completed.returncode)
|
||
size = os.path.getsize(temp_output_path)
|
||
if size > max_bytes:
|
||
raise AIError("下载 cmhub 图片失败: 图片超过大小上限")
|
||
with open(temp_output_path, "rb") as fh:
|
||
return fh.read()
|
||
finally:
|
||
for path in (temp_config_path, temp_output_path):
|
||
if path:
|
||
try:
|
||
os.remove(path)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def _curl_config_quote(value):
|
||
text = str(value or "")
|
||
return '"' + text.replace("\\", "\\\\").replace('"', '\\"') + '"'
|
||
|
||
|
||
def _assert_public_http_url(url):
|
||
parts = urllib.parse.urlsplit(str(url or ""))
|
||
if parts.scheme not in {"http", "https"}:
|
||
raise AIError("cmhub 图片地址只允许 http/https")
|
||
host = parts.hostname
|
||
if not host:
|
||
raise AIError("cmhub 图片地址缺少域名")
|
||
if _is_local_hostname(host):
|
||
raise AIError("cmhub 图片地址不能指向本机或内网")
|
||
try:
|
||
_assert_public_ip(host)
|
||
return
|
||
except ValueError:
|
||
pass
|
||
try:
|
||
addresses = socket.getaddrinfo(
|
||
host,
|
||
parts.port or (443 if parts.scheme == "https" else 80),
|
||
type=socket.SOCK_STREAM,
|
||
)
|
||
except OSError as exc:
|
||
raise AIError("cmhub 图片地址无法解析: %s" % exc) from exc
|
||
if not addresses:
|
||
raise AIError("cmhub 图片地址无法解析")
|
||
for address in addresses:
|
||
ip_text = address[4][0]
|
||
_assert_public_ip(ip_text)
|
||
|
||
|
||
def _is_local_hostname(host):
|
||
lowered = str(host or "").strip().lower().rstrip(".")
|
||
return lowered in {"localhost"} or lowered.endswith(".localhost") or lowered.endswith(".local")
|
||
|
||
|
||
def _assert_public_ip(value):
|
||
ip = ipaddress.ip_address(value)
|
||
if (
|
||
ip.is_private
|
||
or ip.is_loopback
|
||
or ip.is_link_local
|
||
or ip.is_multicast
|
||
or ip.is_reserved
|
||
or ip.is_unspecified
|
||
):
|
||
raise AIError("cmhub 图片地址不能指向本机或内网")
|
||
|
||
|
||
def _redact_cmhub(text, api_key):
|
||
return appconfig.redact_secrets(text, [api_key])
|
||
|
||
def _role_model(category, name, models_path):
|
||
if not name:
|
||
raise AIError("未配置默认 %s 模型" % category)
|
||
model = appconfig.get_model(name, path=models_path)
|
||
if model.get("category") != category:
|
||
raise AIError("模型 %s 不是 %s 类别" % (name, category))
|
||
if not model.get("enabled", True):
|
||
raise AIError("模型已禁用: %s" % name)
|
||
missing = [
|
||
field
|
||
for field in ("url", "model", "api_key")
|
||
if not str(model.get(field, "")).strip()
|
||
]
|
||
if missing:
|
||
raise AIError("模型 %s 缺少字段: %s" % (name, ", ".join(missing)))
|
||
return model
|
||
|
||
|
||
def _runtime_config(runtime):
|
||
if runtime.get("config") is not None:
|
||
return runtime["config"]
|
||
config = appconfig.load_config()
|
||
ai_updates = {
|
||
key: value
|
||
for key, value in runtime.items()
|
||
if key in {
|
||
"default_text_model",
|
||
"default_image_model",
|
||
"title_concurrency",
|
||
"image_concurrency",
|
||
"retry",
|
||
"jpg_quality",
|
||
"resolution",
|
||
"generate_cover",
|
||
"resolution_timeouts",
|
||
"backend",
|
||
"cmhub",
|
||
}
|
||
}
|
||
if ai_updates:
|
||
config = copy.deepcopy(config)
|
||
config.setdefault("ai", {}).update(ai_updates)
|
||
return config
|
||
|
||
|
||
def _prompt_value(prompt_values, name):
|
||
if isinstance(prompt_values, dict):
|
||
return str(
|
||
prompt_values.get(name)
|
||
or prompt_values.get(f"{name}_prompt")
|
||
or ""
|
||
)
|
||
return ""
|
||
|
||
|
||
def _prompt_context(task, new_title, account_by_alias):
|
||
return {
|
||
"old_title": getattr(task, "old_title", ""),
|
||
"new_title": new_title,
|
||
"item_id": getattr(task, "item_id", ""),
|
||
"account_name": _account_name(task, account_by_alias),
|
||
"alias": getattr(task, "alias", ""),
|
||
}
|
||
|
||
|
||
def _account_name(task, account_by_alias):
|
||
alias = str(getattr(task, "alias", "") or "").strip()
|
||
account = account_by_alias.get(alias)
|
||
if account is not None:
|
||
return getattr(account, "account_name", "") or alias
|
||
return getattr(task, "account_name", "") or alias
|
||
|
||
|
||
def _new_cover_path(task, account_by_alias, image_root):
|
||
alias = str(getattr(task, "alias", "") or "").strip()
|
||
account = account_by_alias.get(alias)
|
||
return image_paths.task_image_path(image_root, task, account, "new")
|
||
|
||
|
||
def _mark_generate_failed(task, exc, db_path, on_task_update, step=None):
|
||
error = diagnostics.redact_log_text(str(exc) or exc.__class__.__name__)
|
||
error = db.format_failure_error(error, step)
|
||
db.mark_failed(task.id, "generate", error, path=db_path, step=step)
|
||
if on_task_update is not None:
|
||
on_task_update(task.id, {"status": "failed", "last_error": error})
|
||
return error
|
||
|
||
|
||
def _emit_generation_event(
|
||
callback,
|
||
task,
|
||
phase,
|
||
step,
|
||
result,
|
||
detail=None,
|
||
level="info",
|
||
attempt=None,
|
||
attempts=None,
|
||
metadata=None,
|
||
debug_only=None,
|
||
):
|
||
if callback is None:
|
||
return
|
||
payload = {
|
||
"task": task,
|
||
"phase": phase,
|
||
"step": str(step),
|
||
"result": result,
|
||
"level": level,
|
||
}
|
||
if detail is not None:
|
||
payload["detail"] = diagnostics.redact_log_text(detail)
|
||
if attempt is not None:
|
||
payload["attempt"] = attempt
|
||
if attempts is not None:
|
||
payload["attempts"] = attempts
|
||
if metadata is not None:
|
||
payload["metadata"] = appconfig.sanitize_for_log(metadata)
|
||
if debug_only:
|
||
payload["debug_only"] = True
|
||
try:
|
||
callback(payload)
|
||
except Exception:
|
||
return
|
||
|
||
|
||
def _emit_generation_error(callback, task, phase, step, exc, error):
|
||
if callback is None:
|
||
return
|
||
payload = {
|
||
"task": task,
|
||
"phase": phase,
|
||
"step": str(step),
|
||
"error": diagnostics.redact_log_text(error),
|
||
"exception": exc,
|
||
}
|
||
code = getattr(exc, "code", None)
|
||
status = getattr(exc, "status", None)
|
||
retryable = getattr(exc, "retryable", None)
|
||
if code is not None:
|
||
payload["code"] = str(code)
|
||
if status is not None:
|
||
payload["status"] = status
|
||
if retryable is not None:
|
||
payload["retryable"] = bool(retryable)
|
||
try:
|
||
callback(payload)
|
||
except Exception:
|
||
return
|
||
|
||
|
||
def _cancel_pending(futures):
|
||
for future in futures:
|
||
if not future.done():
|
||
future.cancel()
|
||
|
||
|
||
def _emit_generation_progress(on_progress, summary):
|
||
if on_progress is None:
|
||
return
|
||
payload = dict(summary)
|
||
try:
|
||
on_progress(payload)
|
||
except TypeError:
|
||
on_progress(
|
||
payload.get("title_done", 0),
|
||
payload.get("cover_done", 0),
|
||
payload.get("failed", 0),
|
||
)
|
||
|
||
|
||
def _attempt_count(ai_cfg, retry):
|
||
retry_count = ai_cfg.get("retry", 2) if retry is None else retry
|
||
return max(1, int(retry_count) + 1)
|
||
|
||
|
||
def _read_timeout(model, config):
|
||
return int(model.get("timeout_seconds") or appconfig.response_timeout(config))
|
||
|
||
|
||
def _connect_timeout(model):
|
||
return int(model.get("connect_timeout_seconds") or 30)
|
||
|
||
|
||
def _headers(model, content_type):
|
||
return {
|
||
"Authorization": "Bearer " + model["api_key"],
|
||
"Content-Type": content_type,
|
||
}
|
||
|
||
|
||
def _call_with_retry(
|
||
model,
|
||
body,
|
||
config,
|
||
attempts,
|
||
request_kind,
|
||
content_type=None,
|
||
on_retry=None,
|
||
):
|
||
last_exc = None
|
||
for index in range(attempts):
|
||
try:
|
||
return _call_once(model, body, config, request_kind, content_type=content_type)
|
||
except Exception as exc:
|
||
last_exc = exc
|
||
if index + 1 >= attempts:
|
||
break
|
||
if on_retry is not None:
|
||
try:
|
||
on_retry(index + 1, attempts, exc)
|
||
except Exception:
|
||
pass
|
||
time.sleep(min(2.0, 0.4 * (index + 1)))
|
||
raise AIError(
|
||
"AI 调用失败（已尝试 %s 次）: %s"
|
||
% (attempts, _redact(str(last_exc), model))
|
||
) from last_exc
|
||
|
||
|
||
def _call_once(model, body, config, request_kind, content_type=None):
|
||
if request_kind == "json":
|
||
data = json.dumps(body, ensure_ascii=False).encode("utf-8")
|
||
content_type = "application/json"
|
||
else:
|
||
data = body
|
||
request = urllib.request.Request(
|
||
appconfig.model_request_url(model),
|
||
data=data,
|
||
headers=_headers(model, content_type),
|
||
method="POST",
|
||
)
|
||
timeout = max(_connect_timeout(model), _read_timeout(model, config))
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||
raw = response.read()
|
||
except urllib.error.HTTPError as exc:
|
||
detail = _safe_http_error(exc, model)
|
||
raise AIError("HTTP %s: %s" % (exc.code, detail)) from exc
|
||
except urllib.error.URLError as exc:
|
||
raise AIError(_redact(str(exc.reason), model)) from exc
|
||
try:
|
||
return json.loads(raw.decode("utf-8"))
|
||
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
||
raise AIError("AI 返回不是有效 JSON") from exc
|
||
|
||
|
||
def _safe_http_error(exc, model):
|
||
try:
|
||
body = exc.read(2048).decode("utf-8", errors="replace")
|
||
except Exception:
|
||
body = ""
|
||
return _redact(body or str(exc), model)
|
||
|
||
|
||
def _chat_payload(model, messages):
|
||
payload = {
|
||
"model": model["model"],
|
||
"messages": messages,
|
||
}
|
||
payload.update(copy.deepcopy(model.get("extra_body", {})))
|
||
return payload
|
||
|
||
|
||
def _image_chat_payload(model, cover_prompt, old_cover_path, resolution):
|
||
prompt = "%s\n\n目标分辨率：%s。" % (str(cover_prompt or "").strip(), resolution)
|
||
content = [
|
||
{"type": "text", "text": prompt.strip()},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": _image_data_url(old_cover_path)},
|
||
},
|
||
]
|
||
return _chat_payload(model, [{"role": "user", "content": content}])
|
||
|
||
|
||
def _image_edit_body(model, cover_prompt, old_cover_path, resolution):
|
||
fields = {
|
||
"model": model["model"],
|
||
"prompt": str(cover_prompt or ""),
|
||
"size": _resolution_size_text(resolution),
|
||
}
|
||
fields.update(copy.deepcopy(model.get("extra_body", {})))
|
||
files = {
|
||
"image": (
|
||
os.path.basename(old_cover_path),
|
||
open(old_cover_path, "rb").read(),
|
||
mimetypes.guess_type(old_cover_path)[0] or "application/octet-stream",
|
||
)
|
||
}
|
||
return _multipart_body(fields, files)
|
||
|
||
|
||
def _multipart_body(fields, files):
|
||
boundary = "----cmshopee-%s" % uuid.uuid4().hex
|
||
chunks = []
|
||
for name, value in fields.items():
|
||
chunks.extend(
|
||
[
|
||
("--%s\r\n" % boundary).encode("utf-8"),
|
||
('Content-Disposition: form-data; name="%s"\r\n\r\n' % name).encode("utf-8"),
|
||
str(value).encode("utf-8"),
|
||
b"\r\n",
|
||
]
|
||
)
|
||
for name, file_info in files.items():
|
||
filename, data, content_type = file_info
|
||
chunks.extend(
|
||
[
|
||
("--%s\r\n" % boundary).encode("utf-8"),
|
||
(
|
||
'Content-Disposition: form-data; name="%s"; filename="%s"\r\n'
|
||
% (name, filename)
|
||
).encode("utf-8"),
|
||
("Content-Type: %s\r\n\r\n" % content_type).encode("utf-8"),
|
||
data,
|
||
b"\r\n",
|
||
]
|
||
)
|
||
chunks.append(("--%s--\r\n" % boundary).encode("utf-8"))
|
||
return b"".join(chunks), "multipart/form-data; boundary=%s" % boundary
|
||
|
||
|
||
def _image_data_url(path):
|
||
mime = mimetypes.guess_type(path)[0] or "image/jpeg"
|
||
with open(path, "rb") as fh:
|
||
encoded = base64.b64encode(fh.read()).decode("ascii")
|
||
return "data:%s;base64,%s" % (mime, encoded)
|
||
|
||
|
||
def _extract_text(data):
|
||
if isinstance(data, dict):
|
||
for key in ("output_text", "text", "content"):
|
||
value = data.get(key)
|
||
if isinstance(value, str):
|
||
return value
|
||
choices = data.get("choices")
|
||
if isinstance(choices, list) and choices:
|
||
first = choices[0]
|
||
if isinstance(first, dict):
|
||
if isinstance(first.get("text"), str):
|
||
return first["text"]
|
||
message = first.get("message") or {}
|
||
content = message.get("content")
|
||
return _content_text(content)
|
||
return ""
|
||
|
||
|
||
def _content_text(content):
|
||
if isinstance(content, str):
|
||
return content
|
||
if isinstance(content, list):
|
||
parts = []
|
||
for item in content:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
value = item.get("text") or item.get("content")
|
||
if isinstance(value, str):
|
||
parts.append(value)
|
||
return "".join(parts)
|
||
return ""
|
||
|
||
|
||
def _extract_image_bytes(data, model, config):
|
||
image_ref = _find_image_ref(data)
|
||
if not image_ref:
|
||
raise AIError("AI 返回中没有图片数据")
|
||
if image_ref.startswith("data:"):
|
||
return _decode_data_url(image_ref)
|
||
if _looks_base64(image_ref):
|
||
return base64.b64decode(image_ref)
|
||
return _download_image(image_ref, model, config)
|
||
|
||
|
||
def _find_image_ref(value):
|
||
if isinstance(value, dict):
|
||
for key in ("b64_json", "base64", "image_base64", "image", "url"):
|
||
candidate = value.get(key)
|
||
if isinstance(candidate, str) and candidate.strip():
|
||
return candidate.strip()
|
||
image_url = value.get("image_url")
|
||
if isinstance(image_url, str):
|
||
return image_url
|
||
if isinstance(image_url, dict):
|
||
candidate = image_url.get("url")
|
||
if isinstance(candidate, str):
|
||
return candidate
|
||
if isinstance(image_url, list):
|
||
candidate = _find_image_ref_from_list(image_url)
|
||
if candidate:
|
||
return candidate
|
||
for key in ("image_urls", "urls"):
|
||
candidate = value.get(key)
|
||
if isinstance(candidate, str) and candidate.strip():
|
||
return candidate.strip()
|
||
if isinstance(candidate, list):
|
||
found = _find_image_ref_from_list(candidate)
|
||
if found:
|
||
return found
|
||
for key in ("result", "data", "choices", "output", "content", "images", "files"):
|
||
candidate = _find_image_ref(value.get(key))
|
||
if candidate:
|
||
return candidate
|
||
message = value.get("message")
|
||
if message is not None:
|
||
candidate = _find_image_ref(message)
|
||
if candidate:
|
||
return candidate
|
||
elif isinstance(value, list):
|
||
return _find_image_ref_from_list(value)
|
||
return None
|
||
|
||
|
||
def _find_image_ref_from_list(values):
|
||
for item in values:
|
||
if isinstance(item, str) and item.strip():
|
||
return item.strip()
|
||
candidate = _find_image_ref(item)
|
||
if candidate:
|
||
return candidate
|
||
return None
|
||
|
||
|
||
def _decode_data_url(value):
|
||
if "," not in value:
|
||
raise AIError("图片 data URL 格式错误")
|
||
return base64.b64decode(value.split(",", 1)[1])
|
||
|
||
|
||
def _looks_base64(value):
|
||
compact = value.strip()
|
||
if compact.startswith(("http://", "https://")):
|
||
return False
|
||
if len(compact) < 32:
|
||
return False
|
||
allowed = set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=\n\r")
|
||
return all(char in allowed for char in compact)
|
||
|
||
|
||
def _download_image(url, model, config):
|
||
request = urllib.request.Request(url, method="GET")
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=_read_timeout(model, config)) as response:
|
||
return response.read()
|
||
except Exception as exc:
|
||
raise AIError("下载 AI 图片失败: %s" % _redact(str(exc), model)) from exc
|
||
|
||
|
||
def _save_jpeg(image_bytes, out_path, resolution, jpg_quality):
|
||
try:
|
||
from PIL import Image
|
||
except ImportError as exc:
|
||
raise AIError("缺少 Pillow，无法保存 AI 图片") from exc
|
||
|
||
import io
|
||
|
||
out_path = os.path.abspath(str(out_path))
|
||
directory = os.path.dirname(out_path)
|
||
if directory:
|
||
os.makedirs(directory, exist_ok=True)
|
||
size = _resolution_size(resolution)
|
||
try:
|
||
with Image.open(io.BytesIO(image_bytes)) as image:
|
||
image = image.convert("RGB")
|
||
if size:
|
||
image = image.resize(size, Image.LANCZOS)
|
||
image.save(out_path, "JPEG", quality=jpg_quality, optimize=True)
|
||
except Exception as exc:
|
||
raise AIError("AI 图片保存失败: %s" % exc) from exc
|
||
return out_path
|
||
|
||
|
||
def _resolution_size(resolution):
|
||
return _RESOLUTION_SIZES.get(str(resolution))
|
||
|
||
|
||
def _resolution_size_text(resolution):
|
||
size = _resolution_size(resolution)
|
||
if not size:
|
||
return str(resolution)
|
||
return "%sx%s" % size
|
||
|
||
|
||
def _jpg_quality(value):
|
||
value = int(value)
|
||
return min(100, max(1, value))
|
||
|
||
|
||
def _redact(text, model):
|
||
result = str(text)
|
||
for secret in (model.get("api_key"),):
|
||
if secret:
|
||
result = result.replace(str(secret), "***")
|
||
return result
|