feat: 完成T-303 AI批量生成
- 新增 generate_batch，先并发生成标题再并发生成封面，成功逐条 set_generated 落库 - Tab② 接入开始生成、停止、进度展示和双击新旧封面预览 - 新增 GenerateWorker，通过 worker signal 回传进度与行刷新 - 补充批量生成成功、失败、停止取消和 GUI worker 单元测试 - 同步任务看板、API、路由、当前状态与 progress 文档
This commit is contained in:
chengma
+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