chengma/cmshoppe
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ai-studio): separate image pool from task states

Browse Source
This commit is contained in:
chengma
2026-07-13 09:54:08 +08:00
parent 42fe233ed9
commit b1c050f633
9 changed files with 500 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/db.py
+33
View File
@@ -343,6 +343,7 @@ CREATE TABLE IF NOT EXISTS image_studio_jobs (
status TEXT NOT NULL DEFAULT 'pending',
prompt TEXT,
error TEXT,
recovery_action TEXT NOT NULL DEFAULT 'regenerate',
attempts INTEGER NOT NULL DEFAULT 0,
points_cost INTEGER,
points_balance INTEGER,
@@ -477,6 +478,7 @@ def init_db(path=None, conn=None) -> None:
_ensure_batch_delete_columns(database)
_ensure_task_image_task_columns(database)
_ensure_task_cover_reset_columns(database)
_ensure_image_studio_job_recovery_columns(database)
def _ensure_batch_delete_columns(database):
@@ -502,6 +504,37 @@ def _ensure_task_cover_reset_columns(database):
if "cover_reset_at" not in columns:
database.execute("ALTER TABLE tasks ADD COLUMN cover_reset_at TEXT")
def _ensure_image_studio_job_recovery_columns(database):
columns = {
row["name"] for row in database.execute("PRAGMA table_info(image_studio_jobs)").fetchall()
}
if "recovery_action" not in columns:
database.execute(
"ALTER TABLE image_studio_jobs ADD COLUMN recovery_action TEXT NOT NULL DEFAULT 'regenerate'"
)
database.execute(
"""
UPDATE image_studio_jobs
SET recovery_action = CASE
WHEN status = 'succeeded' THEN 'none'
WHEN task_id IS NOT NULL AND status IN ('submitted', 'running') THEN 'resume'
ELSE 'regenerate'
END
WHERE recovery_action IS NULL
OR recovery_action NOT IN ('none', 'resume', 'regenerate')
OR (
recovery_action = 'regenerate'
AND task_id IS NOT NULL
AND status IN ('submitted', 'running')
)
OR (
recovery_action = 'regenerate'
AND status = 'succeeded'
)
"""
)
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]
app/gui/tabs/image_studio.py
+144 -40
View File
@@ -7,7 +7,7 @@ import os
from PySide6.QtCore import QObject, QMimeData, QSize, Signal
from PySide6.QtWidgets import QListView, QListWidget, QListWidgetItem, QSizePolicy
from ... import accounts, appconfig, cmhub_models, db, image_studio, image_studio_export, image_studio_images, prompts
from ... import accounts, appconfig, cmhub_models, db, diagnostics, image_studio, image_studio_export, image_studio_images, prompts
from .. import file_manager
from ..widgets import *
from ..workers import (
@@ -266,12 +266,12 @@ class ImageStudioTab(QWidget):
PROJECT_COLUMNS = ["店铺", "商品ID", "更新时间"]
JOB_STATUS_LABELS = {
"pending": "排队中",
"pending": "等待提交",
"submitted": "已提交",
"running": "生成中",
"succeeded": "成功",
"failed": "失败",
"expired": "已过期",
"failed": "生成失败",
"expired": "任务过期",
"cancelled": "已停止",
}
@@ -454,6 +454,9 @@ class ImageStudioTab(QWidget):
pool_title = QLabel("照片池")
pool_title.setObjectName("imageStudioSectionTitle")
pool_header.addWidget(pool_title)
self.pool_summary_label = QLabel("可用图片 0 张")
self.pool_summary_label.setObjectName("imageStudioMutedLabel")
pool_header.addWidget(self.pool_summary_label)
pool_header.addStretch(1)
self.source_label = QLabel("源图：未选择")
self.source_label.setObjectName("imageStudioSourceLabel")
@@ -466,6 +469,30 @@ class ImageStudioTab(QWidget):
self.pool_grid.setGridSize(QSize(108, 124))
self.pool_grid.setContextMenuPolicy(Qt.CustomContextMenu)
layout.addWidget(self.pool_grid, 2)
self.job_section = QWidget()
self.job_section.setObjectName("imageStudioJobSection")
job_layout = QVBoxLayout(self.job_section)
job_layout.setContentsMargins(0, 0, 0, 0)
job_layout.setSpacing(4)
job_header = QHBoxLayout()
job_title = QLabel("生成任务")
job_title.setObjectName("imageStudioSectionTitle")
job_header.addWidget(job_title)
self.job_summary_label = QLabel("未完成/异常任务 0 个")
self.job_summary_label.setObjectName("imageStudioMutedLabel")
job_header.addWidget(self.job_summary_label)
job_header.addStretch(1)
job_layout.addLayout(job_header)
self.job_grid = ImageStudioThumbnailGrid(parent=self.job_section)
self.job_grid.setObjectName("imageStudioJobGrid")
self.job_grid.setIconSize(QSize(42, 42))
self.job_grid.setGridSize(QSize(162, 94))
self.job_grid.setMinimumHeight(100)
self.job_grid.setMaximumHeight(112)
job_layout.addWidget(self.job_grid)
self.job_section.setVisible(False)
layout.addWidget(self.job_section, 0)
return panel
def _build_generation_panel(self):
@@ -1110,6 +1137,7 @@ class ImageStudioTab(QWidget):
self._refresh_project_summary()
self._fill_original_grid()
self._fill_pool_grid()
self._fill_job_grid()
self._refresh_selection_labels()
self._refresh_source_label()
@@ -1148,40 +1176,71 @@ class ImageStudioTab(QWidget):
self._refresh_original_download_label()
def _fill_pool_grid(self):
rows = []
for asset in self.assets:
if asset.kind not in {"original", "generated_main", "generated_detail"} or not _asset_is_usable(asset):
continue
rows.append(("asset", asset))
for job in self.jobs:
if job.status in {"pending", "submitted", "running", "failed", "expired", "cancelled"}:
rows.append(("job", job))
assets = [
asset
for asset in self.assets
if asset.kind in {"original", "generated_main", "generated_detail"} and _asset_is_usable(asset)
]
self.pool_grid.clear()
for row_type, obj in rows:
if row_type == "asset":
draggable = _asset_is_usable(obj)
data = {"type": "asset", "asset_id": int(obj.id), "draggable": draggable}
text = f"{_asset_badge(obj.kind)} #{obj.id}\n{obj.aspect_ratio or '比例未知'}"
tooltip = "单击设为源图，双击查看大图，可拖入终选槽。"
else:
data = {"type": "job", "job_id": int(obj.id)}
text = f"任务\n{_job_status_text(obj, self.JOB_STATUS_LABELS)}"
tooltip = "生图任务已保存，可继续查询。"
for asset in assets:
draggable = _asset_is_usable(asset)
data = {"type": "asset", "asset_id": int(asset.id), "draggable": draggable}
text = f"{_asset_badge(asset.kind)} #{asset.id}\n{asset.aspect_ratio or '比例未知'}"
tooltip = "单击设为源图，双击查看大图，可拖入终选槽。"
item = QListWidgetItem(text)
item.setData(Qt.UserRole, data)
item.setSizeHint(QSize(108, 124))
item.setToolTip(tooltip)
if row_type == "asset":
item.setIcon(_asset_icon(obj, _asset_badge(obj.kind), size=QSize(86, 86)))
if obj.id == self.selected_source_asset_id:
item.setBackground(QColor("#eaf2ff"))
else:
item.setIcon(_job_icon(obj.status))
if obj.status in {"failed", "expired"}:
item.setForeground(_qcolor(COLOR_DANGER))
elif obj.status in {"pending", "submitted", "running"}:
item.setForeground(_qcolor(COLOR_WARNING))
item.setIcon(_asset_icon(asset, _asset_badge(asset.kind), size=QSize(86, 86)))
if asset.id == self.selected_source_asset_id:
item.setBackground(QColor("#eaf2ff"))
self.pool_grid.addItem(item)
self.pool_summary_label.setText(f"可用图片 {len(assets)} 张")
def _fill_job_grid(self):
usable_asset_ids = {
int(asset.id)
for asset in self.assets
if _asset_is_usable(asset)
}
jobs = [
job
for job in self.jobs
if job.status in {"pending", "submitted", "running", "failed", "expired", "cancelled"}
and int(job.output_asset_id or 0) not in usable_asset_ids
]
self.job_grid.clear()
self.job_summary_label.setText(f"未完成/异常任务 {len(jobs)} 个")
self.job_section.setVisible(bool(jobs))
for job in jobs:
status_text = self.JOB_STATUS_LABELS.get(job.status, "任务状态未知")
reason = _job_error_summary(job)
recovery = _job_recovery_summary(job)
billing = _job_billing_text(job)
headline = status_text if not reason else f"{status_text}：{reason}"
lines = [headline, recovery]
if billing:
lines.append(billing)
item = QListWidgetItem("\n".join(lines))
item.setData(Qt.UserRole, {"type": "job", "job_id": int(job.id), "draggable": False})
item.setSizeHint(QSize(162, 94))
tooltip_lines = [f"状态：{status_text}", recovery]
if reason:
tooltip_lines.append(f"原因：{reason}")
if billing:
tooltip_lines.append(billing)
item.setToolTip("\n".join(tooltip_lines))
item.setIcon(_job_icon(job.status, size=QSize(42, 42)))
if job.status in {"failed", "expired"}:
item.setForeground(_qcolor(COLOR_DANGER))
item.setBackground(QColor("#ffebe9"))
elif job.status in {"pending", "submitted", "running"}:
item.setForeground(_qcolor(COLOR_WARNING))
item.setBackground(QColor("#fff8c5"))
else:
item.setForeground(_qcolor(COLOR_MUTED))
item.setBackground(QColor("#f6f8fa"))
self.job_grid.addItem(item)
def _selection_asset_ids(self, selection_type):
return [
@@ -1520,6 +1579,7 @@ class ImageStudioTab(QWidget):
self.prompt_edit.blockSignals(False)
self._fill_original_grid()
self._fill_pool_grid()
self._fill_job_grid()
self._refresh_selection_labels()
self._refresh_source_label()
self._refresh_project_summary()
@@ -2019,6 +2079,7 @@ class ImageStudioTab(QWidget):
self.project_table.setEnabled(not running)
self.original_grid.setEnabled(not running)
self.pool_grid.setEnabled(not running or generation_running)
self.job_grid.setEnabled(not running or generation_running)
self.main_selection_list.setEnabled(not running or generation_running)
self.detail_selection_list.setEnabled(not running or generation_running)
self.template_combo.setEnabled(not running or generation_running)
@@ -2229,15 +2290,15 @@ def _asset_icon(asset, fallback_label, size=None, cached_pixmap=None):
)
def _job_icon(status):
def _job_icon(status, size=None):
label = {
"pending": "排",
"submitted": "提",
"running": "生",
"failed": "败",
"expired": "过",
"pending": "…",
"submitted": "…",
"running": "…",
"failed": "!",
"expired": "!",
"cancelled": "停",
}.get(str(status or ""), "任")
}.get(str(status or ""), "?")
color = {
"failed": "#ffebe9",
"expired": "#ffebe9",
@@ -2246,7 +2307,7 @@ def _job_icon(status):
"submitted": "#fff8c5",
"pending": "#f6f8fa",
}.get(str(status or ""), "#f6f8fa")
return QIcon(_placeholder_pixmap(label, QSize(86, 86), color))
return QIcon(_placeholder_pixmap(label, size or QSize(42, 42), color))
def _asset_pixmap(asset, size, fallback_label=None, cached_pixmap=None):
@@ -2337,6 +2398,20 @@ def _selection_tooltip(selection_type, asset):
def _job_status_text(job, labels):
parts = [labels.get(job.status, job.status)]
recovery = _job_recovery_summary(job)
if recovery:
parts.append(recovery)
error = _job_error_summary(job)
if error:
parts.append(error)
billing = _job_billing_text(job)
if billing:
parts.append(billing)
return "，".join(parts)
def _job_billing_text(job):
parts = []
if job.points_cost is not None:
parts.append(f"扣点{job.points_cost}")
if job.points_balance is not None:
@@ -2344,3 +2419,32 @@ def _job_status_text(job, labels):
if job.call_id:
parts.append(f"call_id={job.call_id}")
return "，".join(parts)
def _job_recovery_summary(job):
if getattr(job, "recovery_action", "") == image_studio.JOB_RECOVERY_RESUME:
return "可继续查询，不会重复扣点"
if job.status == "pending":
return "等待本轮提交"
if job.status == "submitted":
return "等待任务结果"
if job.status == "running":
return "正在生成，可稍后继续查询"
return "需要重新生成，可能再次扣点"
def _job_error_summary(job):
raw = " ".join(str(getattr(job, "error", "") or "").split())
if not raw:
return ""
text = diagnostics.redact_log_text(raw)
lowered = text.lower()
if (
"http://" in lowered
or "https://" in lowered
or "/api/" in lowered
or "traceback" in lowered
or any(char.isascii() and char.isalpha() for char in text)
):
return "任务未完成，请按恢复方式处理"
return text[:24] + ("…" if len(text) > 24 else "")
app/image_studio.py
+32 -5
View File
@@ -20,6 +20,14 @@ ASSET_STATUS_MISSING = "missing"
ASSET_STATUSES = {ASSET_STATUS_AVAILABLE, ASSET_STATUS_MISSING}
JOB_STATUSES = {"pending", "submitted", "running", "succeeded", "failed", "expired", "cancelled"}
JOB_RESUMABLE_STATUSES = {"submitted", "running"}
JOB_RECOVERY_NONE = "none"
JOB_RECOVERY_RESUME = "resume"
JOB_RECOVERY_REGENERATE = "regenerate"
JOB_RECOVERY_ACTIONS = {
JOB_RECOVERY_NONE,
JOB_RECOVERY_RESUME,
JOB_RECOVERY_REGENERATE,
}
SELECTION_TYPES = {"main", "detail"}
@@ -70,6 +78,7 @@ class ImageStudioJob:
status: str
prompt: Optional[str]
error: Optional[str]
recovery_action: str
attempts: int
points_cost: Optional[int]
points_balance: Optional[int]
@@ -677,6 +686,7 @@ def set_job_submitted(job_id, task_id, *, call_id=None, points_cost=None, points
UPDATE image_studio_jobs
SET task_id = ?,
status = 'submitted',
recovery_action = ?,
call_id = ?,
points_cost = ?,
points_balance = ?,
@@ -684,7 +694,16 @@ def set_job_submitted(job_id, task_id, *, call_id=None, points_cost=None, points
updated_at = ?
WHERE id = ?
""",
(str(task_id), call_id, points_cost, points_balance, now, now, int(job_id)),
(
str(task_id),
JOB_RECOVERY_RESUME,
call_id,
points_cost,
points_balance,
now,
now,
int(job_id),
),
)
return get_job(job_id, conn=database)
@@ -696,12 +715,17 @@ def update_job_status(
error=None,
output_asset_id=None,
points_balance=None,
recovery_action=None,
increment_attempts=False,
path=None,
conn=None,
):
if str(status) not in JOB_STATUSES:
raise db.DbError("AI工场任务状态无效")
if recovery_action is not None and str(recovery_action) not in JOB_RECOVERY_ACTIONS:
raise db.DbError("AI工场任务恢复方式无效")
if recovery_action is None and str(status) == "succeeded":
recovery_action = JOB_RECOVERY_NONE
now = _now()
terminal = str(status) in {"succeeded", "failed", "expired", "cancelled"}
with _connection(conn, path) as database:
@@ -716,6 +740,7 @@ def update_job_status(
UPDATE image_studio_jobs
SET status = ?,
error = ?,
recovery_action = COALESCE(?, recovery_action),
output_asset_id = COALESCE(?, output_asset_id),
points_balance = COALESCE(?, points_balance),
attempts = attempts + ?,
@@ -726,6 +751,7 @@ def update_job_status(
(
str(status),
error,
recovery_action,
output_asset_id,
points_balance,
1 if increment_attempts else 0,
@@ -742,12 +768,13 @@ def list_resumable_jobs(path=None, conn=None, project_id=None, include_failed_do
if include_failed_downloads:
clauses = [
"task_id IS NOT NULL",
"(status IN (?, ?) OR (status = ? AND output_asset_id IS NULL))",
"recovery_action = ?",
"status IN (?, ?, ?, ?)",
]
params = ["submitted", "running", "failed"]
params = [JOB_RECOVERY_RESUME, "submitted", "running", "failed", "cancelled"]
else:
clauses = ["status IN (?, ?)", "task_id IS NOT NULL"]
params = ["submitted", "running"]
clauses = ["status IN (?, ?)", "task_id IS NOT NULL", "recovery_action = ?"]
params = ["submitted", "running", JOB_RECOVERY_RESUME]
if project_id is not None:
clauses.append("project_id = ?")
params.append(int(project_id))
app/image_studio_generation.py
+46 -5
View File
@@ -194,7 +194,13 @@ def run_jobs(
if should_stop():
for future, job in list(futures.items()):
if future.cancel():
image_studio.update_job_status(job.id, "cancelled", error="用户停止", path=path)
image_studio.update_job_status(
job.id,
"cancelled",
error="用户停止",
recovery_action=_recovery_action_for_job(job),
path=path,
)
futures.pop(future, None)
record({"job": job, "status": "cancelled", "error": "用户停止"})
return summary
@@ -205,12 +211,24 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
if job is None:
raise ImageStudioGenerationError("AI工场生图任务不存在")
if should_stop():
updated = image_studio.update_job_status(job.id, "cancelled", error="用户停止", path=db_path)
updated = image_studio.update_job_status(
job.id,
"cancelled",
error="用户停止",
recovery_action=_recovery_action_for_job(job),
path=db_path,
)
return {"job": updated, "status": "cancelled", "error": "用户停止"}
project = image_studio.get_project(job.project_id, path=db_path)
source_asset = image_studio.get_asset(job.source_asset_id, path=db_path)
if project is None or source_asset is None:
updated = image_studio.update_job_status(job.id, "failed", error="项目或源图不存在", path=db_path)
updated = image_studio.update_job_status(
job.id,
"failed",
error="项目或源图不存在",
recovery_action=image_studio.JOB_RECOVERY_REGENERATE,
path=db_path,
)
return {"job": updated, "status": "failed", "error": "项目或源图不存在"}
try:
image_studio.update_job_status(job.id, "running", path=db_path)
@@ -248,7 +266,14 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
return {"job": updated, "asset": asset, "status": "succeeded"}
except Exception as exc:
status = "cancelled" if "停止" in str(exc) else "failed"
updated = image_studio.update_job_status(job.id, status, error=str(exc), path=db_path)
current_job = image_studio.get_job(job.id, path=db_path)
updated = image_studio.update_job_status(
job.id,
status,
error=str(exc),
recovery_action=_recovery_action_for_job(current_job),
path=db_path,
)
_notify(on_event, {"job_id": job.id, "step": "job_done", "result": status, "detail": str(exc)})
return {"job": updated, "status": status, "error": str(exc)}
@@ -363,7 +388,13 @@ def _poll_job(job_id, task_id, runtime, request_result, db_path, should_stop, on
if status in {"failed", "expired"}:
error = data.get("error") if isinstance(data.get("error"), dict) else {}
message = str(error.get("message") or error.get("code") or status)
image_studio.update_job_status(job_id, status, error=message, path=db_path)
image_studio.update_job_status(
job_id,
status,
error=message,
recovery_action=image_studio.JOB_RECOVERY_REGENERATE,
path=db_path,
)
raise ImageStudioGenerationError(message)
raise ImageStudioGenerationError("cmhub 生图任务状态返回格式错误")
@@ -394,3 +425,13 @@ def _raise_if_stopped(should_stop):
stopped = False
if stopped:
raise ImageStudioGenerationError("用户已停止，已提交任务可稍后继续查询")
def _recovery_action_for_job(job):
if (
job is not None
and getattr(job, "task_id", None)
and getattr(job, "recovery_action", None) == image_studio.JOB_RECOVERY_RESUME
):
return image_studio.JOB_RECOVERY_RESUME
return image_studio.JOB_RECOVERY_REGENERATE