feat: 完成T-303 AI批量生成
- 新增 generate_batch，先并发生成标题再并发生成封面，成功逐条 set_generated 落库 - Tab② 接入开始生成、停止、进度展示和双击新旧封面预览 - 新增 GenerateWorker，通过 worker signal 回传进度与行刷新 - 补充批量生成成功、失败、停止取消和 GUI worker 单元测试 - 同步任务看板、API、路由、当前状态与 progress 文档
This commit is contained in:
chengma
@@ -1,6 +1,7 @@
|
||||
"""AI generation helpers backed by configurable HTTP model endpoints."""
|
||||
|
||||
import base64
|
||||
from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed
|
||||
import copy
|
||||
import json
|
||||
import mimetypes
|
||||
@@ -10,7 +11,9 @@ import urllib.error
|
||||
import urllib.request
|
||||
import uuid
|
||||
|
||||
from . import appconfig
|
||||
from . import appconfig, db
|
||||
from . import prompts as prompt_module
|
||||
from .config import make_slug
|
||||
|
||||
|
||||
class AIError(RuntimeError):
|
||||
@@ -105,6 +108,144 @@ def gen_cover(
|
||||
return _save_jpeg(image_bytes, out_path, resolution, quality)
|
||||
|
||||
|
||||
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",
|
||||
}
|
||||
}
|
||||
)
|
||||
db_path = runtime.get("db_path")
|
||||
models_path = runtime.get("models_path", appconfig.AI_MODELS_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")
|
||||
title_prompt = _prompt_value(prompts, "title")
|
||||
cover_prompt = _prompt_value(prompts, "cover")
|
||||
should_stop = should_stop or (lambda: False)
|
||||
eligible = [
|
||||
task for task in list(tasks)
|
||||
if getattr(task, "stage", None) == "collected"
|
||||
]
|
||||
summary = {
|
||||
"ok": True,
|
||||
"total": len(eligible),
|
||||
"title_done": 0,
|
||||
"cover_done": 0,
|
||||
"failed": 0,
|
||||
"cancelled": False,
|
||||
}
|
||||
_emit_generation_progress(on_progress, summary)
|
||||
title_results = {}
|
||||
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=max(1, int(generation_cfg.get("title_concurrency", 1)))
|
||||
) as executor:
|
||||
futures = {}
|
||||
for task in eligible:
|
||||
if should_stop():
|
||||
summary["cancelled"] = True
|
||||
break
|
||||
futures[
|
||||
executor.submit(
|
||||
gen_title,
|
||||
title_prompt,
|
||||
getattr(task, "old_title", "") or "",
|
||||
retry=generation_cfg.get("retry"),
|
||||
config=config,
|
||||
models_path=models_path,
|
||||
)
|
||||
] = task
|
||||
for future in as_completed(futures):
|
||||
task = futures[future]
|
||||
if should_stop():
|
||||
summary["cancelled"] = True
|
||||
_cancel_pending(futures)
|
||||
try:
|
||||
title_results[task.id] = future.result()
|
||||
summary["title_done"] += 1
|
||||
except CancelledError:
|
||||
summary["cancelled"] = True
|
||||
except Exception as exc:
|
||||
summary["failed"] += 1
|
||||
summary["ok"] = False
|
||||
_mark_generate_failed(task, exc, db_path, on_task_update)
|
||||
_emit_generation_progress(on_progress, summary)
|
||||
|
||||
cover_tasks = [
|
||||
task for task in eligible
|
||||
if task.id in title_results
|
||||
]
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=max(1, int(generation_cfg.get("image_concurrency", 1)))
|
||||
) as executor:
|
||||
futures = {}
|
||||
for task in cover_tasks:
|
||||
if should_stop():
|
||||
summary["cancelled"] = True
|
||||
break
|
||||
new_title = title_results[task.id]
|
||||
rendered_cover_prompt = prompt_module.render_prompt(
|
||||
cover_prompt,
|
||||
_prompt_context(task, new_title, account_by_alias),
|
||||
)
|
||||
futures[
|
||||
executor.submit(
|
||||
gen_cover,
|
||||
rendered_cover_prompt,
|
||||
getattr(task, "old_cover_path", "") or "",
|
||||
_new_cover_path(task, account_by_alias, image_root),
|
||||
resolution=generation_cfg.get("resolution"),
|
||||
jpg_quality=generation_cfg.get("jpg_quality"),
|
||||
retry=generation_cfg.get("retry"),
|
||||
config=config,
|
||||
models_path=models_path,
|
||||
)
|
||||
] = (task, new_title)
|
||||
for future in as_completed(futures):
|
||||
task, new_title = futures[future]
|
||||
if should_stop():
|
||||
summary["cancelled"] = True
|
||||
_cancel_pending(futures)
|
||||
try:
|
||||
new_cover_path = future.result()
|
||||
db.set_generated(task.id, new_title, new_cover_path, path=db_path)
|
||||
summary["cover_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": new_cover_path,
|
||||
},
|
||||
)
|
||||
except CancelledError:
|
||||
summary["cancelled"] = True
|
||||
except Exception as exc:
|
||||
summary["failed"] += 1
|
||||
summary["ok"] = False
|
||||
_mark_generate_failed(task, exc, db_path, on_task_update)
|
||||
_emit_generation_progress(on_progress, summary)
|
||||
|
||||
if summary["cancelled"]:
|
||||
summary["ok"] = False
|
||||
return summary
|
||||
|
||||
|
||||
def _role_model(category, name, models_path):
|
||||
if not name:
|
||||
raise AIError("未配置默认 %s 模型" % category)
|
||||
@@ -123,6 +264,100 @@ def _role_model(category, name, models_path):
|
||||
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",
|
||||
"resolution_timeouts",
|
||||
}
|
||||
}
|
||||
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)
|
||||
slug = getattr(account, "slug", None) if account is not None else None
|
||||
if not slug:
|
||||
slug = make_slug(alias or getattr(task, "account_name", "") or "unknown")
|
||||
return os.path.abspath(
|
||||
os.path.join(
|
||||
image_root,
|
||||
slug,
|
||||
"%s_new.jpg" % getattr(task, "item_id", ""),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _mark_generate_failed(task, exc, db_path, on_task_update):
|
||||
error = str(exc) or exc.__class__.__name__
|
||||
db.mark_failed(task.id, "generate", error, path=db_path)
|
||||
if on_task_update is not None:
|
||||
on_task_update(task.id, {"status": "failed", "last_error": error})
|
||||
|
||||
|
||||
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)
|
||||
|
||||
+194
-1
@@ -7,6 +7,7 @@ import sys
|
||||
|
||||
try:
|
||||
from PySide6.QtCore import QAbstractTableModel, QModelIndex, Qt
|
||||
from PySide6.QtGui import QPixmap
|
||||
from PySide6.QtWidgets import (
|
||||
QAbstractItemView,
|
||||
QApplication,
|
||||
@@ -78,7 +79,7 @@ QTabBar::tab:hover:!selected {
|
||||
|
||||
|
||||
if QT_IMPORT_ERROR is None:
|
||||
from . import accounts, appconfig, chrome, db, editor, excel, prompts
|
||||
from . import accounts, ai, appconfig, chrome, db, editor, excel, prompts
|
||||
from . import config as account_config
|
||||
|
||||
|
||||
@@ -316,6 +317,8 @@ if QT_IMPORT_ERROR is None:
|
||||
self.title_prompt_path = title_prompt_path or prompts.TITLE_PROMPT_PATH
|
||||
self.cover_prompts_dir = cover_prompts_dir or prompts.COVER_PROMPTS_DIR
|
||||
self.current_cover_template = None
|
||||
self.generate_worker = None
|
||||
self.generate_thread = None
|
||||
|
||||
self.title_prompt_edit = QPlainTextEdit()
|
||||
self.title_prompt_edit.setObjectName("titlePromptEdit")
|
||||
@@ -336,6 +339,10 @@ if QT_IMPORT_ERROR is None:
|
||||
self.delete_cover_template_button = QPushButton("删除")
|
||||
self.insert_title_button = QPushButton("插入标题")
|
||||
self.preview_prompt_button = QPushButton("预览")
|
||||
self.generate_button = QPushButton("开始生成")
|
||||
self.stop_generate_button = QPushButton("停止")
|
||||
self.stop_generate_button.setEnabled(False)
|
||||
self.progress_label = QLabel("进度：标题0/0 · 封面0/0 · 失败0")
|
||||
|
||||
left_panel = QWidget()
|
||||
left_layout = QVBoxLayout(left_panel)
|
||||
@@ -401,9 +408,16 @@ if QT_IMPORT_ERROR is None:
|
||||
self.splitter.setStretchFactor(1, 3)
|
||||
self.splitter.setSizes([280, 860])
|
||||
|
||||
bottom_layout = QHBoxLayout()
|
||||
bottom_layout.addWidget(self.progress_label)
|
||||
bottom_layout.addStretch(1)
|
||||
bottom_layout.addWidget(self.generate_button)
|
||||
bottom_layout.addWidget(self.stop_generate_button)
|
||||
|
||||
layout = QVBoxLayout(self)
|
||||
layout.setContentsMargins(18, 18, 18, 18)
|
||||
layout.addWidget(self.splitter, 1)
|
||||
layout.addLayout(bottom_layout)
|
||||
|
||||
self.batch_filter.currentIndexChanged.connect(self.refresh_tasks)
|
||||
self.shop_filter.currentIndexChanged.connect(self.refresh_tasks)
|
||||
@@ -418,6 +432,9 @@ if QT_IMPORT_ERROR is None:
|
||||
self.delete_cover_template_button.clicked.connect(self.delete_cover_template)
|
||||
self.insert_title_button.clicked.connect(self.insert_title_placeholder)
|
||||
self.preview_prompt_button.clicked.connect(self.preview_cover_prompt)
|
||||
self.generate_button.clicked.connect(self.start_generate)
|
||||
self.stop_generate_button.clicked.connect(self.stop_generate)
|
||||
self.task_table.doubleClicked.connect(self.show_task_images)
|
||||
|
||||
self.refresh_cover_templates()
|
||||
self.refresh_tasks()
|
||||
@@ -566,6 +583,147 @@ if QT_IMPORT_ERROR is None:
|
||||
QMessageBox.information(self, "封面提示词预览", rendered)
|
||||
self._set_status("封面提示词预览已生成")
|
||||
|
||||
def start_generate(self, checked=False):
|
||||
if self.generate_thread is not None:
|
||||
self._set_status("AI 生成正在进行...")
|
||||
return
|
||||
tasks = [
|
||||
task for task in self.model.tasks
|
||||
if getattr(task, "stage", None) == "collected"
|
||||
]
|
||||
if not tasks:
|
||||
self._set_status("当前筛选结果没有可生成任务")
|
||||
return
|
||||
prompt_values = {
|
||||
"title": self.title_prompt_edit.toPlainText(),
|
||||
"cover": self.cover_prompt_edit.toPlainText(),
|
||||
}
|
||||
worker = GenerateWorker(
|
||||
tasks,
|
||||
prompt_values,
|
||||
db_path=self.db_path,
|
||||
config=self.config,
|
||||
)
|
||||
worker.progress.connect(self._on_generate_progress)
|
||||
worker.row_updated.connect(self._on_generate_row_updated)
|
||||
worker.log.connect(self._set_status)
|
||||
worker.failed.connect(self._on_generate_failed)
|
||||
worker.finished.connect(self._on_generate_finished)
|
||||
worker.cancelled.connect(self._on_generate_cancelled)
|
||||
thread = run_worker(worker, thread_name="GenerateWorker", start=False)
|
||||
thread.finished.connect(lambda: self._forget_generate_thread(thread))
|
||||
self.generate_worker = worker
|
||||
self.generate_thread = thread
|
||||
self._set_generate_running(True)
|
||||
self._update_generate_progress(
|
||||
{"total": len(tasks), "title_done": 0, "cover_done": 0, "failed": 0}
|
||||
)
|
||||
self._set_status(f"开始 AI 生成：{len(tasks)} 条")
|
||||
thread.start()
|
||||
|
||||
def stop_generate(self, checked=False):
|
||||
if self.generate_worker is not None:
|
||||
self.generate_worker.cancel()
|
||||
self._set_status("正在停止 AI 生成...")
|
||||
|
||||
def show_task_images(self, index):
|
||||
task = self.model.task_at(index.row()) if index.isValid() else self._selected_task()
|
||||
if task is None:
|
||||
self._set_status("没有可预览的任务")
|
||||
return
|
||||
dialog = QDialog(self)
|
||||
dialog.setWindowTitle(f"封面对照：{task.item_id}")
|
||||
layout = QVBoxLayout(dialog)
|
||||
images_layout = QHBoxLayout()
|
||||
images_layout.addWidget(self._image_panel("旧封面", task.old_cover_path))
|
||||
images_layout.addWidget(self._image_panel("新封面", task.new_cover_path))
|
||||
layout.addLayout(images_layout)
|
||||
buttons = QDialogButtonBox(QDialogButtonBox.Close)
|
||||
buttons.rejected.connect(dialog.reject)
|
||||
layout.addWidget(buttons)
|
||||
dialog.resize(720, 420)
|
||||
dialog.exec()
|
||||
|
||||
def _image_panel(self, title, path):
|
||||
panel = QWidget()
|
||||
layout = QVBoxLayout(panel)
|
||||
layout.addWidget(QLabel(title))
|
||||
image_label = QLabel()
|
||||
image_label.setAlignment(Qt.AlignCenter)
|
||||
image_label.setMinimumSize(260, 260)
|
||||
image_label.setWordWrap(True)
|
||||
if path and os.path.exists(str(path)):
|
||||
pixmap = QPixmap(str(path))
|
||||
if not pixmap.isNull():
|
||||
image_label.setPixmap(
|
||||
pixmap.scaled(
|
||||
260,
|
||||
260,
|
||||
Qt.KeepAspectRatio,
|
||||
Qt.SmoothTransformation,
|
||||
)
|
||||
)
|
||||
else:
|
||||
image_label.setText(f"图片无法读取\n{path}")
|
||||
else:
|
||||
image_label.setText(f"无图片\n{path or ''}".strip())
|
||||
layout.addWidget(image_label, 1)
|
||||
return panel
|
||||
|
||||
def _set_generate_running(self, running):
|
||||
self.generate_button.setEnabled(not running)
|
||||
self.stop_generate_button.setEnabled(running)
|
||||
self.refresh_button.setEnabled(not running)
|
||||
self.batch_filter.setEnabled(not running)
|
||||
self.shop_filter.setEnabled(not running)
|
||||
self.status_filter.setEnabled(not running)
|
||||
self.save_title_button.setEnabled(not running)
|
||||
self.save_cover_template_button.setEnabled(not running)
|
||||
self.save_cover_template_as_button.setEnabled(not running)
|
||||
self.rename_cover_template_button.setEnabled(not running)
|
||||
self.delete_cover_template_button.setEnabled(not running)
|
||||
|
||||
def _forget_generate_thread(self, thread):
|
||||
if self.generate_thread is thread:
|
||||
self.generate_thread = None
|
||||
self.generate_worker = None
|
||||
|
||||
def _on_generate_progress(self, payload):
|
||||
self._update_generate_progress(payload)
|
||||
self._set_status("生成进度：" + self._generate_progress_text(payload))
|
||||
|
||||
def _on_generate_row_updated(self, task_id, fields):
|
||||
self.refresh_tasks()
|
||||
|
||||
def _on_generate_failed(self, task_id, error):
|
||||
self._set_status(f"AI 生成失败：{error}")
|
||||
|
||||
def _on_generate_finished(self, payload):
|
||||
self._set_generate_running(False)
|
||||
self.refresh_tasks()
|
||||
self._update_generate_progress(payload)
|
||||
if payload.get("error"):
|
||||
self._set_status(f"AI 生成失败：{payload.get('error')}")
|
||||
return
|
||||
self._set_status("AI 生成完成：" + self._generate_progress_text(payload))
|
||||
|
||||
def _on_generate_cancelled(self, payload):
|
||||
self._set_generate_running(False)
|
||||
self.refresh_tasks()
|
||||
self._update_generate_progress(payload)
|
||||
self._set_status("AI 生成已停止：" + self._generate_progress_text(payload))
|
||||
|
||||
def _update_generate_progress(self, payload):
|
||||
self.progress_label.setText("进度：" + self._generate_progress_text(payload))
|
||||
|
||||
def _generate_progress_text(self, payload):
|
||||
return "标题{title}/{total} · 封面{cover}/{total} · 失败{failed}".format(
|
||||
title=payload.get("title_done", 0),
|
||||
cover=payload.get("cover_done", 0),
|
||||
total=payload.get("total", 0),
|
||||
failed=payload.get("failed", 0),
|
||||
)
|
||||
|
||||
def _selected_task(self):
|
||||
index = self.task_table.currentIndex()
|
||||
if index.isValid():
|
||||
@@ -1186,6 +1344,41 @@ if QT_IMPORT_ERROR is None:
|
||||
from .workers import BaseWorker, run_worker
|
||||
|
||||
|
||||
class GenerateWorker(BaseWorker):
|
||||
"""Generate titles and covers for collected tasks."""
|
||||
|
||||
def __init__(self, tasks, prompt_values, db_path=None, config=None):
|
||||
super().__init__()
|
||||
self.tasks = list(tasks)
|
||||
self.prompt_values = dict(prompt_values or {})
|
||||
self.db_path = db_path
|
||||
self.config = config
|
||||
|
||||
def execute(self):
|
||||
account_rows = accounts.list_accounts(path=self.db_path, config=self.config)
|
||||
account_by_alias = {
|
||||
str(account.alias).strip(): account
|
||||
for account in account_rows
|
||||
if str(account.alias).strip()
|
||||
}
|
||||
return ai.generate_batch(
|
||||
self.tasks,
|
||||
self.prompt_values,
|
||||
ai_cfg={
|
||||
"config": self.config,
|
||||
"db_path": self.db_path,
|
||||
"image_dir": appconfig.image_dir(self.config),
|
||||
"account_by_alias": account_by_alias,
|
||||
"on_task_update": self._emit_row_update,
|
||||
},
|
||||
on_progress=self.progress.emit,
|
||||
should_stop=self.should_cancel,
|
||||
)
|
||||
|
||||
def _emit_row_update(self, task_id, fields):
|
||||
self.row_updated.emit(int(task_id), dict(fields or {}))
|
||||
|
||||
|
||||
class CollectWorker(BaseWorker):
|
||||
"""Collect old title and cover for imported tasks."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user