T-564 async cmhub image tasks

This commit is contained in:
chengma
2026-07-09 00:07:09 +08:00
parent 487408ba0f
commit c435456977
10 changed files with 761 additions and 38 deletions
+299 -4
View File
@@ -24,6 +24,11 @@ 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):
@@ -43,6 +48,9 @@ class CMHubError(AIError):
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 = 30
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
@@ -685,7 +693,7 @@ def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=No
try:
rendered_cover_prompt, new_cover_path = prepare_cover_task(task, new_title)
future = request_executor.submit(
_request_cmhub_cover_image,
_request_cmhub_cover_image_async,
rendered_cover_prompt,
getattr(task, "old_cover_path", "") or "",
new_cover_path,
@@ -696,6 +704,9 @@ def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=No
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
@@ -881,7 +892,7 @@ def _gen_cover_cmhub(
on_step=None,
on_event=None,
):
request_result = _request_cmhub_cover_image(
request_result = _request_cmhub_cover_image_sync(
cover_prompt,
old_cover_path,
out_path,
@@ -909,6 +920,34 @@ def _request_cmhub_cover_image(
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)
@@ -977,6 +1016,249 @@ def _request_cmhub_cover_image(
}
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":
result = data.get("result") if isinstance(data.get("result"), dict) else {}
image_url = str(result.get("image_url") or data.get("image_url") or "").strip()
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"]
@@ -1155,6 +1437,7 @@ def _cmhub_call_with_retry(
read_timeout,
attempts,
on_retry=None,
headers_extra=None,
):
attempts = max(1, int(attempts or 1))
last_exc = None
@@ -1167,6 +1450,7 @@ def _cmhub_call_with_retry(
payload,
connect_timeout=connect_timeout,
read_timeout=read_timeout,
headers_extra=headers_extra,
)
except CMHubError as exc:
last_exc = exc
@@ -1181,11 +1465,13 @@ def _cmhub_call_with_retry(
raise last_exc
def _cmhub_call_once(method, url, api_key, payload, connect_timeout, read_timeout):
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))),
@@ -1262,7 +1548,13 @@ def _cmhub_code_for_status(status):
def _cmhub_retryable(code):
return str(code) in {"upstream_error", "rate_limited", "connect_timeout"}
return str(code) in {
"upstream_error",
"upstream_timeout",
"task_timeout",
"rate_limited",
"connect_timeout",
}
def _cmhub_user_message(code, message):
@@ -1275,6 +1567,9 @@ def _cmhub_user_message(code, message):
"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",
}
+92 -1
View File
@@ -46,6 +46,7 @@ FAILURE_STEP_LABELS = {
"cover_submit": "提交封面任务",
"cover_build_request": "构建封面请求",
"cover_request": "请求生成封面",
"cover_poll": "等待封面结果",
"cover_download": "下载新封面",
"cover_save": "保存新封面",
"cover_parse_response": "解析封面响应",
@@ -138,6 +139,8 @@ class Task:
old_cover_path: Optional[str]
new_title: Optional[str]
new_cover_path: Optional[str]
image_task_id: Optional[str]
image_task_key: Optional[str]
committed: int
stage: str
status: str
@@ -231,6 +234,8 @@ CREATE TABLE IF NOT EXISTS tasks (
old_cover_path TEXT,
new_title TEXT,
new_cover_path TEXT,
image_task_id TEXT,
image_task_key TEXT,
committed INTEGER NOT NULL DEFAULT 0,
stage TEXT NOT NULL DEFAULT 'imported',
status TEXT NOT NULL DEFAULT 'pending',
@@ -355,6 +360,7 @@ def init_db(path=None, conn=None) -> None:
with database:
database.executescript(SCHEMA_SQL)
_ensure_batch_delete_columns(database)
_ensure_task_image_task_columns(database)
def _ensure_batch_delete_columns(database):
@@ -364,6 +370,14 @@ def _ensure_batch_delete_columns(database):
if "deleted_reason" not in columns:
database.execute("ALTER TABLE batches ADD COLUMN deleted_reason TEXT")
def _ensure_task_image_task_columns(database):
columns = {row["name"] for row in database.execute("PRAGMA table_info(tasks)").fetchall()}
if "image_task_id" not in columns:
database.execute("ALTER TABLE tasks ADD COLUMN image_task_id TEXT")
if "image_task_key" not in columns:
database.execute("ALTER TABLE tasks ADD COLUMN image_task_key TEXT")
def create_batch(file_paths: Iterable[str], note=None, path=None, conn=None) -> str:
batch_id = datetime.now().strftime("%Y%m%d_%H%M%S_") + uuid.uuid4().hex[:8]
files = [os.path.abspath(file_path) for file_path in file_paths]
@@ -688,6 +702,79 @@ def mark_skipped(task_id, reason, path=None, conn=None) -> None:
)
def ensure_image_task_key(task_id, path=None, conn=None) -> str:
"""Return a stable cmhub image idempotency key for the current cover attempt."""
task_id = int(task_id)
with _connection(conn, path) as database:
with database:
row = database.execute(
"SELECT image_task_key FROM tasks WHERE id = ?",
(task_id,),
).fetchone()
if row is None:
raise DbError(f"任务不存在: {task_id}")
key = str(row["image_task_key"] or "").strip()
if not key:
key = f"cmshopee-task-{task_id}-{uuid.uuid4().hex}"
database.execute(
"""
UPDATE tasks
SET image_task_key = ?, updated_at = ?
WHERE id = ?
""",
(key, _now(), task_id),
)
return key
def set_image_task_submitted(task_id, image_task_id, image_task_key=None, path=None, conn=None) -> None:
"""Persist the cmhub image task id immediately after submit succeeds."""
task_id = int(task_id)
image_task_id = str(image_task_id or "").strip()
if not image_task_id:
raise DbError("缺少 cmhub 生图任务ID")
with _connection(conn, path) as database:
with database:
row = database.execute(
"SELECT image_task_key FROM tasks WHERE id = ?",
(task_id,),
).fetchone()
if row is None:
raise DbError(f"任务不存在: {task_id}")
key = str(image_task_key or row["image_task_key"] or "").strip()
if not key:
key = f"cmshopee-task-{task_id}-{uuid.uuid4().hex}"
database.execute(
"""
UPDATE tasks
SET image_task_id = ?,
image_task_key = ?,
updated_at = ?
WHERE id = ?
""",
(image_task_id, key, _now(), task_id),
)
def clear_image_task(task_id, path=None, conn=None) -> None:
"""Clear persisted cmhub image task state so the next cover run can submit anew."""
with _connection(conn, path) as database:
with database:
database.execute(
"""
UPDATE tasks
SET image_task_id = NULL,
image_task_key = NULL,
updated_at = ?
WHERE id = ?
""",
(_now(), int(task_id)),
)
def set_collected(task_id, old_title, old_cover_path, path=None, conn=None) -> None:
now = _now()
with _connection(conn, path) as database:
@@ -821,19 +908,23 @@ def reset_generated(
now = _now()
new_title = None if reset_title else before.new_title
new_cover = None if reset_cover else before.new_cover_path
image_task_id = None if reset_cover else before.image_task_id
image_task_key = None if reset_cover else before.image_task_key
with database:
database.execute(
"""
UPDATE tasks
SET new_title = ?,
new_cover_path = ?,
image_task_id = ?,
image_task_key = ?,
stage = 'generated',
status = 'success',
last_error = NULL,
updated_at = ?
WHERE id = ?
""",
(new_title, new_cover, now, int(task_id)),
(new_title, new_cover, image_task_id, image_task_key, now, int(task_id)),
)
after = get_task(task_id, conn=database)
return {