feat(ai-studio): resume hosted image jobs

This commit is contained in:
chengma
2026-07-11 14:29:58 +08:00
parent cddadf6364
commit dcaf68fd07
9 changed files with 240 additions and 7 deletions
+1
View File
@@ -28,6 +28,7 @@ if QT_IMPORT_ERROR is None:
ImageStudioExportWorker,
ImageStudioGenerateJobsWorker,
ImageStudioPullImagesWorker,
ImageStudioResumeJobsWorker,
WriteBackWorker,
)
from .tabs.accounts import AccountDialog, AccountsTab
+56 -1
View File
@@ -16,6 +16,7 @@ from ..workers import (
from ..workers import ImageStudioExportWorker as _RealImageStudioExportWorker
from ..workers import ImageStudioGenerateJobsWorker as _RealImageStudioGenerateJobsWorker
from ..workers import ImageStudioPullImagesWorker as _RealImageStudioPullImagesWorker
from ..workers import ImageStudioResumeJobsWorker as _RealImageStudioResumeJobsWorker
ASSET_MIME_TYPE = "application/x-cmshopee-image-studio-asset"
@@ -57,6 +58,15 @@ def ImageStudioExportWorker(*args, **kwargs):
)
def ImageStudioResumeJobsWorker(*args, **kwargs):
return _call_package_attr(
"ImageStudioResumeJobsWorker",
_RealImageStudioResumeJobsWorker,
*args,
**kwargs,
)
def _drop_event_position(event):
if hasattr(event, "position"):
return event.position().toPoint()
@@ -432,10 +442,13 @@ class ImageStudioTab(QWidget):
action_layout = QHBoxLayout()
self.start_button = QPushButton("开始生成")
self.start_button.setObjectName("imageStudioStartButton")
self.resume_button = QPushButton("继续查询任务")
self.resume_button.setObjectName("imageStudioResumeButton")
self.stop_button = QPushButton("停止")
self.stop_button.setObjectName("imageStudioStopButton")
self.stop_button.setEnabled(False)
action_layout.addWidget(self.start_button)
action_layout.addWidget(self.resume_button)
action_layout.addWidget(self.stop_button)
layout.addLayout(action_layout)
@@ -512,6 +525,7 @@ class ImageStudioTab(QWidget):
self.template_delete_button.clicked.connect(self.delete_template)
self.prompt_edit.textChanged.connect(self._save_project_prompt)
self.start_button.clicked.connect(self.start_generation)
self.resume_button.clicked.connect(self.resume_generation_jobs)
self.stop_button.clicked.connect(self.stop_generation)
self.export_button.clicked.connect(self.export_selections)
@@ -730,7 +744,7 @@ class ImageStudioTab(QWidget):
values = [
"任务",
"-",
self.JOB_STATUS_LABELS.get(obj.status, obj.status),
_job_status_text(obj, self.JOB_STATUS_LABELS),
f"源图 #{obj.source_asset_id or '-'}",
obj.error or "",
]
@@ -1055,6 +1069,35 @@ class ImageStudioTab(QWidget):
self._append_log("[AI工场] 已请求停止,正在等待安全边界")
self._status("AI工场生成已请求停止", "warning")
def resume_generation_jobs(self, checked=False):
if self.current_project is None:
self._message("项目未打开", "请先打开一个AI工场项目。")
return
resumable = image_studio.list_resumable_jobs(
path=self.db_path,
project_id=self.current_project.id,
include_failed_downloads=True,
)
if not resumable:
self._status("当前项目没有可继续查询的 cmhub 生图任务", "muted")
self._message("没有可继续查询任务", "当前项目没有已提交、生成中或下载失败的 cmhub 生图任务。")
return
self.progress_bar.setRange(0, len(resumable))
self.progress_bar.setValue(0)
self._append_log(f"[AI工场] 继续查询 {len(resumable)} 个已提交任务")
worker = ImageStudioResumeJobsWorker(
project_id=self.current_project.id,
aspect_ratio=self.aspect_combo.currentData(),
db_path=self.db_path,
config=self.config,
cmhub_config_path=self.cmhub_config_path,
)
worker.progress.connect(self._on_generate_progress)
worker.log.connect(self._append_log)
worker.finished.connect(self._on_generation_finished)
worker.failed.connect(self._on_worker_failed)
self._start_worker(worker, "AI工场继续查询")
def export_selections(self, checked=False):
if self.current_project is None:
self._message("项目未打开", "请先打开一个AI工场项目。")
@@ -1203,6 +1246,7 @@ class ImageStudioTab(QWidget):
self.count_spin.setEnabled(not running)
self.aspect_combo.setEnabled(not running)
self.start_button.setEnabled(not running)
self.resume_button.setEnabled(not running)
self.stop_button.setEnabled(running)
self.export_button.setEnabled(not running)
@@ -1380,3 +1424,14 @@ def _selection_tooltip(selection_type, asset):
f"{_selection_label(selection_type)}:{_asset_badge(getattr(asset, 'kind', ''))} "
f"#{getattr(asset, 'id', '')},比例 {getattr(asset, 'aspect_ratio', None) or '未知'}"
)
def _job_status_text(job, labels):
parts = [labels.get(job.status, job.status)]
if job.points_cost is not None:
parts.append(f"扣点{job.points_cost}")
if job.points_balance is not None:
parts.append(f"余额{job.points_balance}")
if job.call_id:
parts.append(f"call_id={job.call_id}")
return ",".join(parts)
+47
View File
@@ -192,6 +192,53 @@ class ImageStudioGenerateJobsWorker(BaseWorker):
return summary
class ImageStudioResumeJobsWorker(BaseWorker):
"""Resume submitted/running or failed-download AI studio jobs."""
def __init__(
self,
*,
project_id=None,
aspect_ratio="1:1",
db_path=None,
config=None,
cmhub_config_path=None,
):
super().__init__()
self.project_id = int(project_id) if project_id is not None else None
self.aspect_ratio = str(aspect_ratio or "1:1")
self.db_path = db_path
self.config = config
self.cmhub_config_path = cmhub_config_path
self._done = 0
self._failed = 0
self._lock = threading.Lock()
def execute(self):
def on_event(payload):
event = dict(payload or {})
self.log.emit(_format_image_studio_event(event))
if event.get("step") == "job_done":
with self._lock:
self._done += 1
if event.get("result") not in {"success"}:
self._failed += 1
progress = {"done": self._done, "failed": self._failed}
self.progress.emit(progress)
summary = image_studio_generation.resume_image_jobs(
project_id=self.project_id,
aspect_ratio=self.aspect_ratio,
config=self.config,
cmhub_config_path=self.cmhub_config_path,
path=self.db_path,
should_stop=self.should_cancel,
on_event=on_event,
)
summary["project_id"] = self.project_id
return summary
class ImageStudioExportWorker(BaseWorker):
"""Export AI studio final selections to local JPEG files."""
+8 -1
View File
@@ -738,7 +738,14 @@ def update_job_status(
return get_job(job_id, conn=database)
def list_resumable_jobs(path=None, conn=None, project_id=None):
def list_resumable_jobs(path=None, conn=None, project_id=None, include_failed_downloads=False):
if include_failed_downloads:
clauses = [
"task_id IS NOT NULL",
"(status IN (?, ?) OR (status = ? AND output_asset_id IS NULL))",
]
params = ["submitted", "running", "failed"]
else:
clauses = ["status IN (?, ?)", "task_id IS NOT NULL"]
params = ["submitted", "running"]
if project_id is not None:
+5 -1
View File
@@ -120,7 +120,11 @@ def resume_image_jobs(
should_stop=None,
on_event=None,
):
jobs = image_studio.list_resumable_jobs(path=path, project_id=project_id)
jobs = image_studio.list_resumable_jobs(
path=path,
project_id=project_id,
include_failed_downloads=True,
)
return run_jobs(
jobs,
aspect_ratio=aspect_ratio,
+1
View File
@@ -196,6 +196,7 @@
- 照片池展示原图、生成主图、生成详情图和在途/失败任务状态;单击可用图片设为源图,双击打开大图;右键移除只删除未被任务或终选引用的照片池记录,不删除本地图片文件。
- 右侧只有一个完整提示词框;模板目录固定为 `data/prompts/image_studio/`,与②标题/封面模板隔离。界面不显示“主提示词 / 每张动作词”。
- 生图固定走 cmhub 托管模型,使用⑤设置里的 cmhub Base URL/API Key/生图别名和图片并发;界面只显示扣点、余额、进度、失败,不展示自定义 Provider、API Key、生成来源选择或“导入本地图片”入口。
- 「继续查询任务」会恢复当前项目已提交、生成中或下载失败但已有 `task_id` 的 cmhub 生图任务;恢复时只 poll/download 原任务,不再次 submit,不重复扣点。照片池中的任务行显示排队/生成/失败/过期/停止状态,并附带扣点、余额和 `call_id`,便于运营和技术排障。
- 底部终选盘分为主图和详情图两列;照片池中已下载/已生成且本地文件可用的图片可拖入终选,落到已有位置时按插入顺延,同一类别内同一照片只能出现一次,主图和详情图之间允许复用同一照片。
- 终选列表内可拖动重排,Delete 或右键「移出终选」只移出终选,不删除照片池资产或本地文件;拖放/移出失败时刷新回 SQLite 中的持久化顺序。
- 主图推荐 1:1;比例不匹配只用黄色轻提示和 tooltip 提醒,不硬拦。文件缺失或尚未下载的照片不能拖入终选。
+8 -2
View File
@@ -3,7 +3,7 @@ id: T-594
title: AI工场 cmhub 托管生图故障恢复、计费一致性与运行回归
phase: 7
deps: [T-590, T-591]
status: TODO
status: DONE
created: 2026-07-11
---
@@ -36,4 +36,10 @@ AI工场会同时持有多张 cmhub 异步生图任务。基础编排和 GUI 接
## 执行记录
(完成后记录故障矩阵、状态决策与验证结果。)
- 2026-07-11:完成 AI工场 cmhub 托管生图恢复与状态收口。
- `image_studio.list_resumable_jobs(..., include_failed_downloads=True)` 支持把 failed 但已有 `task_id` 且无输出 asset 的任务纳入恢复队列,用于下载失败/本地保存失败后的继续查询。
- `image_studio_generation.resume_image_jobs()` 改用恢复队列:已有 `task_id` 的任务只 poll/download,不再 POST submit,不重复创建 job 或重复扣点。
- 新增 `ImageStudioResumeJobsWorker` 和 ⑥ AI工场「继续查询任务」按钮;无可续查任务时中文提示,有任务时复用现有进度/日志/刷新逻辑。
- 照片池任务行显示 submitted/running/failed/expired/cancelled 中文状态,并补充扣点、余额、`call_id`,便于计费和失败排障。
- 更新 `docs/routes.md`,明确继续查询不重复 submit、不重复扣点以及任务状态展示口径。
- 验证:主工作区 targeted 测试通过;因无关默认提示词脏文件仍会影响全量 unittest,已在干净 worktree 仅套用 T-594 diff 后运行 `python -m ruff check app tests main.py`、`py -3.10 -m compileall app main.py`、`py -3.10 -m unittest discover -s tests`(369 tests)和 `git diff --check`,全部通过。
+47
View File
@@ -602,6 +602,53 @@ class GuiTests(TempDirMixin, unittest.TestCase):
self.assert_removed(temp_dir)
def test_image_studio_tab_shows_resume_and_job_billing_status(self):
with self.make_temp_dir() as temp_dir:
cfg = self.make_config(temp_dir)
db.init_db(cfg["db_path"])
project = image_studio.create_or_get_project(
account_alias="alias-a",
account_slug="alias_a",
item_id="51100639510",
path=cfg["db_path"],
)
source = image_studio.add_asset(
project.id,
"original",
local_path=self.write_test_image(os.path.join(temp_dir, "source.jpg")),
path=cfg["db_path"],
)
job = image_studio.create_job(
project.id,
source_asset_id=source.id,
path=cfg["db_path"],
)
image_studio.set_job_submitted(
job.id,
"cmhub-task-1",
call_id="call-1",
points_cost=2,
points_balance=88,
path=cfg["db_path"],
)
tab = ImageStudioTab(config=cfg, db_path=cfg["db_path"])
self.addCleanup(tab.close)
tab._select_project(project.id)
self.assertEqual("继续查询任务", tab.resume_button.text())
statuses = [
tab.pool_table.item(row, 2).text()
for row in range(tab.pool_table.rowCount())
if tab.pool_table.item(row, 0).text() == "任务"
]
self.assertEqual(1, len(statuses))
self.assertIn("已提交", statuses[0])
self.assertIn("扣点2", statuses[0])
self.assertIn("余额88", statuses[0])
self.assertIn("call_id=call-1", statuses[0])
self.assert_removed(temp_dir)
def test_startup_update_gate_forced_blocks_and_opens_download(self):
boxes = []
+65
View File
@@ -231,6 +231,71 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
self.assert_removed(temp_dir)
def test_resume_failed_download_job_polls_without_new_submit(self):
with self.make_temp_dir() as temp_dir:
cfg, project, source = self._project_source(temp_dir)
job = image_studio.create_job(
project.id,
source_asset_id=source.id,
job_type="main",
prompt="恢复下载",
path=cfg["db_path"],
)
image_studio.set_job_submitted(
job.id,
"cmhub-task-download",
call_id="call-download",
points_cost=2,
points_balance=88,
path=cfg["db_path"],
)
image_studio.update_job_status(
job.id,
"failed",
error="下载新封面失败",
path=cfg["db_path"],
)
resumable = image_studio.list_resumable_jobs(
path=cfg["db_path"],
project_id=project.id,
include_failed_downloads=True,
)
self.assertEqual([job.id], [item.id for item in resumable])
def fake_poll(method, url, api_key, **kwargs):
self.assertEqual("GET", method)
return {
"task_id": "cmhub-task-download",
"status": "succeeded",
"result": {"image_url": "https://cdn.example.com/recovered.png"},
"points_balance": 88,
}
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
mock.patch("app.image_studio_generation.ai._cmhub_call_with_retry") as submit, \
mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \
mock.patch(
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
return_value=(self._png_bytes(), 0.1),
):
summary = image_studio_generation.resume_image_jobs(
project_id=project.id,
config=cfg,
path=cfg["db_path"],
)
self.assertEqual(1, summary["success"])
submit.assert_not_called()
updated = image_studio.get_job(job.id, path=cfg["db_path"])
self.assertEqual("succeeded", updated.status)
self.assertEqual("cmhub-task-download", updated.task_id)
self.assertEqual("call-download", updated.call_id)
assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"])
self.assertEqual(1, len(assets))
self.assert_removed(temp_dir)
def test_failed_cmhub_task_marks_only_that_job_failed(self):
with self.make_temp_dir() as temp_dir:
cfg, project, source = self._project_source(temp_dir)