985 lines
41 KiB
Python
985 lines
41 KiB
Python
"""Tab 6: AI image studio UI."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
|
||
from ... import accounts, appconfig, db, image_studio, prompts
|
||
from .. import file_manager
|
||
from ..widgets import *
|
||
from ..workers import (
|
||
ImageStudioDownloadOriginalWorker as _RealImageStudioDownloadOriginalWorker,
|
||
)
|
||
from ..workers import ImageStudioGenerateJobsWorker as _RealImageStudioGenerateJobsWorker
|
||
from ..workers import ImageStudioPullImagesWorker as _RealImageStudioPullImagesWorker
|
||
|
||
|
||
def ImageStudioPullImagesWorker(*args, **kwargs):
|
||
return _call_package_attr(
|
||
"ImageStudioPullImagesWorker",
|
||
_RealImageStudioPullImagesWorker,
|
||
*args,
|
||
**kwargs,
|
||
)
|
||
|
||
|
||
def ImageStudioDownloadOriginalWorker(*args, **kwargs):
|
||
return _call_package_attr(
|
||
"ImageStudioDownloadOriginalWorker",
|
||
_RealImageStudioDownloadOriginalWorker,
|
||
*args,
|
||
**kwargs,
|
||
)
|
||
|
||
|
||
def ImageStudioGenerateJobsWorker(*args, **kwargs):
|
||
return _call_package_attr(
|
||
"ImageStudioGenerateJobsWorker",
|
||
_RealImageStudioGenerateJobsWorker,
|
||
*args,
|
||
**kwargs,
|
||
)
|
||
|
||
|
||
class ImageStudioPreviewDialog(QDialog):
|
||
"""Simple large image preview used by original and pool tables."""
|
||
|
||
def __init__(self, asset, parent=None):
|
||
super().__init__(parent)
|
||
self.asset = asset
|
||
self.setWindowTitle(self._title_for_asset(asset))
|
||
layout = QVBoxLayout(self)
|
||
scroll = QScrollArea()
|
||
scroll.setWidgetResizable(False)
|
||
image_label = QLabel()
|
||
image_label.setAlignment(Qt.AlignCenter)
|
||
path = str(getattr(asset, "local_path", "") or "")
|
||
image = QImage(path) if path and os.path.isfile(path) else QImage()
|
||
if image.isNull():
|
||
image_label.setText("图片尚未下载或读取失败")
|
||
image_label.setMinimumSize(420, 260)
|
||
else:
|
||
image_label.setPixmap(QPixmap.fromImage(image))
|
||
image_label.resize(image.size())
|
||
self.setWindowTitle(
|
||
f"{self._title_for_asset(asset)} · {image.width()}x{image.height()}"
|
||
)
|
||
scroll.setWidget(image_label)
|
||
layout.addWidget(scroll, 1)
|
||
buttons = QHBoxLayout()
|
||
buttons.addStretch(1)
|
||
close_button = QPushButton("关闭")
|
||
close_button.clicked.connect(self.reject)
|
||
buttons.addWidget(close_button)
|
||
layout.addLayout(buttons)
|
||
self.resize(720, 520)
|
||
|
||
def _title_for_asset(self, asset):
|
||
badge = _asset_badge(getattr(asset, "kind", ""))
|
||
asset_id = getattr(asset, "id", "")
|
||
return f"AI工场图片预览：{badge} #{asset_id}"
|
||
|
||
|
||
class ImageStudioTab(QWidget):
|
||
"""Sixth tab: project-based AI image studio."""
|
||
|
||
PROJECT_COLUMNS = ["项目", "账号", "商品ID", "更新时间"]
|
||
ORIGINAL_COLUMNS = ["序号", "状态", "远程地址"]
|
||
POOL_COLUMNS = ["类型", "比例", "状态", "来源", "本地文件"]
|
||
JOB_STATUS_LABELS = {
|
||
"pending": "排队中",
|
||
"submitted": "已提交",
|
||
"running": "生成中",
|
||
"succeeded": "成功",
|
||
"failed": "失败",
|
||
"expired": "已过期",
|
||
"cancelled": "已停止",
|
||
}
|
||
|
||
def __init__(
|
||
self,
|
||
parent=None,
|
||
db_path=None,
|
||
config=None,
|
||
config_path=None,
|
||
status_callback=None,
|
||
prompts_dir=None,
|
||
):
|
||
super().__init__(parent)
|
||
self.setObjectName("imageStudioTab")
|
||
self.config = appconfig.load_config(config_path or appconfig.CONFIG_PATH) if config is None else config
|
||
self.db_path = db_path or appconfig.db_path(self.config)
|
||
self.config_path = config_path or self.config.get("config_path") or appconfig.CONFIG_PATH
|
||
self.prompts_dir = prompts_dir or appconfig.image_studio_prompts_dir(self.config)
|
||
self.cmhub_config_path = self.config.get("cmhub_config_path") or appconfig.cmhub_config_file_path(self.config)
|
||
self.status_callback = status_callback
|
||
self.projects = []
|
||
self.accounts = []
|
||
self.current_project = None
|
||
self.assets = []
|
||
self.jobs = []
|
||
self.selections = []
|
||
self.selected_source_asset_id = None
|
||
self._running_worker = None
|
||
self._running_thread = None
|
||
self._download_open_after = {}
|
||
|
||
self._build_ui()
|
||
self._connect_signals()
|
||
self.refresh_accounts()
|
||
self.refresh_templates()
|
||
self.refresh_projects()
|
||
|
||
def _build_ui(self):
|
||
root = QVBoxLayout(self)
|
||
root.setContentsMargins(10, 8, 10, 8)
|
||
root.setSpacing(8)
|
||
|
||
splitter = QSplitter(Qt.Horizontal)
|
||
splitter.setObjectName("imageStudioMainSplitter")
|
||
splitter.addWidget(self._build_project_panel())
|
||
splitter.addWidget(self._build_pool_panel())
|
||
splitter.addWidget(self._build_generation_panel())
|
||
splitter.setStretchFactor(0, 1)
|
||
splitter.setStretchFactor(1, 3)
|
||
splitter.setStretchFactor(2, 2)
|
||
root.addWidget(splitter, 1)
|
||
root.addWidget(self._build_final_panel(), 0)
|
||
|
||
def _build_project_panel(self):
|
||
panel = QWidget()
|
||
panel.setObjectName("imageStudioProjectPanel")
|
||
layout = QVBoxLayout(panel)
|
||
layout.setContentsMargins(0, 0, 0, 0)
|
||
layout.setSpacing(8)
|
||
|
||
project_form = QFormLayout()
|
||
project_form.setLabelAlignment(Qt.AlignRight)
|
||
self.account_combo = QComboBox()
|
||
self.account_combo.setObjectName("imageStudioAccountCombo")
|
||
self.item_id_edit = QLineEdit()
|
||
self.item_id_edit.setObjectName("imageStudioItemIdEdit")
|
||
self.item_id_edit.setPlaceholderText("商品ID")
|
||
project_form.addRow("账号", self.account_combo)
|
||
project_form.addRow("商品ID", self.item_id_edit)
|
||
layout.addLayout(project_form)
|
||
|
||
toolbar = QHBoxLayout()
|
||
self.open_project_button = QPushButton("打开项目")
|
||
self.open_project_button.setObjectName("imageStudioOpenProjectButton")
|
||
self.pull_images_button = QPushButton("拉取主图")
|
||
self.pull_images_button.setObjectName("imageStudioPullImagesButton")
|
||
self.open_folder_button = QPushButton("打开项目文件夹")
|
||
self.open_folder_button.setObjectName("imageStudioOpenFolderButton")
|
||
toolbar.addWidget(self.open_project_button)
|
||
toolbar.addWidget(self.pull_images_button)
|
||
toolbar.addWidget(self.open_folder_button)
|
||
layout.addLayout(toolbar)
|
||
|
||
self.project_table = QTableWidget(0, len(self.PROJECT_COLUMNS))
|
||
self.project_table.setObjectName("imageStudioProjectList")
|
||
self.project_table.setHorizontalHeaderLabels(self.PROJECT_COLUMNS)
|
||
self.project_table.setSelectionBehavior(QAbstractItemView.SelectRows)
|
||
self.project_table.setSelectionMode(QAbstractItemView.SingleSelection)
|
||
self.project_table.setEditTriggers(QAbstractItemView.NoEditTriggers)
|
||
self.project_table.horizontalHeader().setStretchLastSection(True)
|
||
self.project_table.verticalHeader().setVisible(False)
|
||
layout.addWidget(self.project_table, 1)
|
||
return panel
|
||
|
||
def _build_pool_panel(self):
|
||
panel = QWidget()
|
||
panel.setObjectName("imageStudioPoolPanel")
|
||
layout = QVBoxLayout(panel)
|
||
layout.setContentsMargins(0, 0, 0, 0)
|
||
layout.setSpacing(8)
|
||
|
||
original_header = QHBoxLayout()
|
||
original_header.addWidget(QLabel("蝦皮原主图"))
|
||
original_header.addStretch(1)
|
||
self.original_hint_label = QLabel("单击下载并加入照片池，双击查看大图")
|
||
self.original_hint_label.setObjectName("imageStudioOriginalHintLabel")
|
||
original_header.addWidget(self.original_hint_label)
|
||
layout.addLayout(original_header)
|
||
|
||
self.original_table = QTableWidget(0, len(self.ORIGINAL_COLUMNS))
|
||
self.original_table.setObjectName("imageStudioOriginalTable")
|
||
self.original_table.setHorizontalHeaderLabels(self.ORIGINAL_COLUMNS)
|
||
self.original_table.setSelectionBehavior(QAbstractItemView.SelectRows)
|
||
self.original_table.setSelectionMode(QAbstractItemView.SingleSelection)
|
||
self.original_table.setEditTriggers(QAbstractItemView.NoEditTriggers)
|
||
self.original_table.horizontalHeader().setStretchLastSection(True)
|
||
self.original_table.verticalHeader().setVisible(False)
|
||
layout.addWidget(self.original_table, 1)
|
||
|
||
pool_header = QHBoxLayout()
|
||
pool_header.addWidget(QLabel("照片池"))
|
||
pool_header.addStretch(1)
|
||
self.source_label = QLabel("源图：未选择")
|
||
self.source_label.setObjectName("imageStudioSourceLabel")
|
||
pool_header.addWidget(self.source_label)
|
||
layout.addLayout(pool_header)
|
||
|
||
self.pool_table = QTableWidget(0, len(self.POOL_COLUMNS))
|
||
self.pool_table.setObjectName("imageStudioPoolTable")
|
||
self.pool_table.setHorizontalHeaderLabels(self.POOL_COLUMNS)
|
||
self.pool_table.setSelectionBehavior(QAbstractItemView.SelectRows)
|
||
self.pool_table.setSelectionMode(QAbstractItemView.SingleSelection)
|
||
self.pool_table.setEditTriggers(QAbstractItemView.NoEditTriggers)
|
||
self.pool_table.setContextMenuPolicy(Qt.CustomContextMenu)
|
||
self.pool_table.horizontalHeader().setStretchLastSection(True)
|
||
self.pool_table.verticalHeader().setVisible(False)
|
||
layout.addWidget(self.pool_table, 2)
|
||
return panel
|
||
|
||
def _build_generation_panel(self):
|
||
panel = QWidget()
|
||
panel.setObjectName("imageStudioGenerationPanel")
|
||
layout = QVBoxLayout(panel)
|
||
layout.setContentsMargins(0, 0, 0, 0)
|
||
layout.setSpacing(8)
|
||
|
||
template_layout = QGridLayout()
|
||
self.template_combo = QComboBox()
|
||
self.template_combo.setObjectName("imageStudioTemplateCombo")
|
||
self.template_new_button = QPushButton("新建")
|
||
self.template_new_button.setObjectName("imageStudioTemplateNewButton")
|
||
self.template_rename_button = QPushButton("重命名")
|
||
self.template_rename_button.setObjectName("imageStudioTemplateRenameButton")
|
||
self.template_save_button = QPushButton("保存")
|
||
self.template_save_button.setObjectName("imageStudioTemplateSaveButton")
|
||
self.template_delete_button = QPushButton("删除")
|
||
self.template_delete_button.setObjectName("imageStudioTemplateDeleteButton")
|
||
template_layout.addWidget(self.template_combo, 0, 0, 1, 2)
|
||
template_layout.addWidget(self.template_new_button, 0, 2)
|
||
template_layout.addWidget(self.template_rename_button, 1, 0)
|
||
template_layout.addWidget(self.template_save_button, 1, 1)
|
||
template_layout.addWidget(self.template_delete_button, 1, 2)
|
||
layout.addLayout(template_layout)
|
||
|
||
self.prompt_edit = QPlainTextEdit()
|
||
self.prompt_edit.setObjectName("imageStudioPromptEdit")
|
||
self.prompt_edit.setPlaceholderText("输入完整图片生成提示词")
|
||
self.prompt_edit.setMinimumHeight(160)
|
||
layout.addWidget(self.prompt_edit, 2)
|
||
|
||
form = QFormLayout()
|
||
self.job_type_combo = QComboBox()
|
||
self.job_type_combo.setObjectName("imageStudioJobTypeCombo")
|
||
self.job_type_combo.addItem("主图", "main")
|
||
self.job_type_combo.addItem("详情图", "detail")
|
||
self.count_spin = QSpinBox()
|
||
self.count_spin.setObjectName("imageStudioCountSpin")
|
||
self.count_spin.setRange(1, 12)
|
||
self.count_spin.setValue(4)
|
||
self.aspect_combo = QComboBox()
|
||
self.aspect_combo.setObjectName("imageStudioAspectCombo")
|
||
for value in ("1:1", "3:4", "4:3", "9:16", "16:9"):
|
||
self.aspect_combo.addItem(value, value)
|
||
form.addRow("类型", self.job_type_combo)
|
||
form.addRow("数量", self.count_spin)
|
||
form.addRow("比例", self.aspect_combo)
|
||
layout.addLayout(form)
|
||
|
||
self.billing_label = QLabel("cmhub 托管模型：扣点以返回结果为准")
|
||
self.billing_label.setObjectName("imageStudioBillingLabel")
|
||
self.billing_label.setWordWrap(True)
|
||
layout.addWidget(self.billing_label)
|
||
|
||
action_layout = QHBoxLayout()
|
||
self.start_button = QPushButton("开始生成")
|
||
self.start_button.setObjectName("imageStudioStartButton")
|
||
self.stop_button = QPushButton("停止")
|
||
self.stop_button.setObjectName("imageStudioStopButton")
|
||
self.stop_button.setEnabled(False)
|
||
action_layout.addWidget(self.start_button)
|
||
action_layout.addWidget(self.stop_button)
|
||
layout.addLayout(action_layout)
|
||
|
||
self.progress_bar = QProgressBar()
|
||
self.progress_bar.setObjectName("imageStudioProgressBar")
|
||
self.progress_bar.setRange(0, 1)
|
||
self.progress_bar.setValue(0)
|
||
layout.addWidget(self.progress_bar)
|
||
|
||
self.log_view = QPlainTextEdit()
|
||
self.log_view.setObjectName("imageStudioLogView")
|
||
self.log_view.setReadOnly(True)
|
||
self.log_view.setPlaceholderText("运行日志会在开始后显示")
|
||
layout.addWidget(self.log_view, 1)
|
||
return panel
|
||
|
||
def _build_final_panel(self):
|
||
panel = QWidget()
|
||
panel.setObjectName("imageStudioFinalPanel")
|
||
layout = QHBoxLayout(panel)
|
||
layout.setContentsMargins(0, 0, 0, 0)
|
||
self.main_selection_label = QLabel("主图终选 0/9（拖放排序将在 T-592 接入）")
|
||
self.main_selection_label.setObjectName("imageStudioMainSelectionLabel")
|
||
self.detail_selection_label = QLabel("详情图终选 0/12（拖放排序将在 T-592 接入）")
|
||
self.detail_selection_label.setObjectName("imageStudioDetailSelectionLabel")
|
||
layout.addWidget(self.main_selection_label)
|
||
layout.addWidget(self.detail_selection_label)
|
||
layout.addStretch(1)
|
||
return panel
|
||
|
||
def _connect_signals(self):
|
||
self.open_project_button.clicked.connect(self.open_project)
|
||
self.pull_images_button.clicked.connect(self.pull_main_images)
|
||
self.open_folder_button.clicked.connect(self.open_project_folder)
|
||
self.project_table.itemSelectionChanged.connect(self._on_project_selection_changed)
|
||
self.original_table.cellClicked.connect(self._on_original_clicked)
|
||
self.original_table.cellDoubleClicked.connect(self._on_original_double_clicked)
|
||
self.pool_table.cellClicked.connect(self._on_pool_clicked)
|
||
self.pool_table.cellDoubleClicked.connect(self._on_pool_double_clicked)
|
||
self.pool_table.customContextMenuRequested.connect(self._show_pool_context_menu)
|
||
self.template_combo.currentIndexChanged.connect(self.load_selected_template)
|
||
self.template_new_button.clicked.connect(self.create_template)
|
||
self.template_rename_button.clicked.connect(self.rename_template)
|
||
self.template_save_button.clicked.connect(self.save_template)
|
||
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.stop_button.clicked.connect(self.stop_generation)
|
||
|
||
def refresh_accounts(self):
|
||
self.account_combo.clear()
|
||
try:
|
||
self.accounts = accounts.list_accounts(path=self.db_path, config=self.config)
|
||
except Exception as exc:
|
||
self.accounts = []
|
||
self._status(f"账号读取失败：{exc}", "danger")
|
||
for account in self.accounts:
|
||
self.account_combo.addItem(
|
||
f"{account.account_name}（{account.alias}）",
|
||
account.alias,
|
||
)
|
||
if not self.accounts:
|
||
self.account_combo.addItem("暂无账号，请先到④账号管理添加", "")
|
||
|
||
def refresh_projects(self):
|
||
try:
|
||
db.init_db(self.db_path)
|
||
self.projects = image_studio.list_projects(path=self.db_path)
|
||
except Exception as exc:
|
||
self.projects = []
|
||
self._status(f"AI工场项目读取失败：{exc}", "danger")
|
||
self._fill_project_table()
|
||
if self.current_project is None and self.projects:
|
||
self._select_project(self.projects[0].id)
|
||
elif self.current_project is not None:
|
||
self._select_project(self.current_project.id, quiet=True)
|
||
|
||
def _fill_project_table(self):
|
||
self.project_table.setRowCount(len(self.projects))
|
||
for row, project in enumerate(self.projects):
|
||
values = [
|
||
f"{project.account_alias} / {project.item_id}",
|
||
project.account_name or project.account_alias,
|
||
project.item_id,
|
||
project.updated_at,
|
||
]
|
||
for column, value in enumerate(values):
|
||
item = QTableWidgetItem(str(value or ""))
|
||
item.setData(Qt.UserRole, int(project.id))
|
||
self.project_table.setItem(row, column, item)
|
||
self.project_table.resizeColumnsToContents()
|
||
|
||
def open_project(self, checked=False):
|
||
alias = str(self.account_combo.currentData() or "").strip()
|
||
item_id = self.item_id_edit.text().strip()
|
||
if not alias:
|
||
self._message("账号未选择", "请先在④账号管理添加并选择账号。")
|
||
return
|
||
if not item_id:
|
||
self._message("商品ID不能为空", "请输入要打开的蝦皮商品ID。")
|
||
return
|
||
account = accounts.get_account(alias, path=self.db_path, config=self.config)
|
||
try:
|
||
project = image_studio.create_or_get_project(
|
||
account,
|
||
item_id=item_id,
|
||
path=self.db_path,
|
||
)
|
||
except Exception as exc:
|
||
self._message("打开项目失败", str(exc))
|
||
self._status(f"打开AI工场项目失败：{exc}", "danger")
|
||
return
|
||
self.current_project = project
|
||
self.item_id_edit.setText(project.item_id)
|
||
self._set_account_combo(project.account_alias)
|
||
self.refresh_projects()
|
||
self._select_project(project.id)
|
||
self._status("AI工场项目已打开", "success")
|
||
|
||
def pull_main_images(self, checked=False):
|
||
alias = str(self.account_combo.currentData() or "").strip()
|
||
item_id = self.item_id_edit.text().strip()
|
||
if self.current_project is not None:
|
||
alias = self.current_project.account_alias
|
||
item_id = self.current_project.item_id
|
||
if not alias or not item_id:
|
||
self._message("项目未打开", "请先选择账号和商品ID并打开项目。")
|
||
return
|
||
worker = ImageStudioPullImagesWorker(
|
||
alias,
|
||
item_id,
|
||
db_path=self.db_path,
|
||
config=self.config,
|
||
)
|
||
worker.log.connect(self._append_log)
|
||
worker.finished.connect(self._on_pull_finished)
|
||
worker.failed.connect(self._on_worker_failed)
|
||
self._start_worker(worker, "AI工场拉取主图")
|
||
self._append_log("[AI工场] 拉取主图开始")
|
||
|
||
def _on_pull_finished(self, summary):
|
||
if summary.get("ok") is False:
|
||
if self._running_worker is not None:
|
||
self._on_worker_failed(-1, summary.get("error") or "拉取主图失败")
|
||
return
|
||
project = summary.get("project")
|
||
if project is not None:
|
||
self.current_project = project
|
||
self._set_account_combo(project.account_alias)
|
||
self.item_id_edit.setText(project.item_id)
|
||
self._finish_worker()
|
||
self.refresh_projects()
|
||
if project is not None:
|
||
self._select_project(project.id)
|
||
self._status(f"已拉取 {summary.get('count', 0)} 张蝦皮原主图", "success")
|
||
|
||
def open_project_folder(self, checked=False):
|
||
if self.current_project is None:
|
||
self._message("项目未打开", "请先打开一个AI工场项目。")
|
||
return
|
||
dirs = image_studio.default_project_image_dirs(self.current_project, self.config)
|
||
try:
|
||
os.makedirs(dirs["root"], exist_ok=True)
|
||
opened = file_manager.open_in_file_manager(dirs["root"])
|
||
except Exception as exc:
|
||
self._message("打开项目文件夹失败", str(exc))
|
||
self._status(f"打开项目文件夹失败：{exc}", "warning")
|
||
return
|
||
self._status(f"已打开项目文件夹：{opened}", "success")
|
||
|
||
def _on_project_selection_changed(self):
|
||
items = self.project_table.selectedItems()
|
||
if not items:
|
||
return
|
||
project_id = items[0].data(Qt.UserRole)
|
||
if project_id is not None:
|
||
self._select_project(project_id)
|
||
|
||
def _select_project(self, project_id, quiet=False):
|
||
try:
|
||
project = image_studio.get_project(project_id, path=self.db_path)
|
||
except Exception as exc:
|
||
self._status(f"读取AI工场项目失败：{exc}", "danger")
|
||
return
|
||
if project is None:
|
||
return
|
||
self.current_project = project
|
||
self._set_account_combo(project.account_alias)
|
||
self.item_id_edit.setText(project.item_id)
|
||
self.prompt_edit.blockSignals(True)
|
||
try:
|
||
self.prompt_edit.setPlainText(project.draft_prompt or "")
|
||
finally:
|
||
self.prompt_edit.blockSignals(False)
|
||
self.selected_source_asset_id = None
|
||
self.refresh_project_assets()
|
||
self._sync_project_selection(project.id)
|
||
if not quiet:
|
||
self._status(f"当前AI工场项目：{project.account_alias} / {project.item_id}", "muted")
|
||
|
||
def _sync_project_selection(self, project_id):
|
||
for row in range(self.project_table.rowCount()):
|
||
item = self.project_table.item(row, 0)
|
||
if item is not None and item.data(Qt.UserRole) == int(project_id):
|
||
if not self.project_table.item(row, 0).isSelected():
|
||
self.project_table.selectRow(row)
|
||
break
|
||
|
||
def refresh_project_assets(self):
|
||
if self.current_project is None:
|
||
self.assets = []
|
||
self.jobs = []
|
||
self.selections = []
|
||
else:
|
||
self.assets = image_studio.list_assets(self.current_project.id, path=self.db_path)
|
||
self.jobs = self._list_project_jobs(self.current_project.id)
|
||
self.selections = image_studio.list_selections(self.current_project.id, path=self.db_path)
|
||
self._fill_original_table()
|
||
self._fill_pool_table()
|
||
self._refresh_selection_labels()
|
||
self._refresh_source_label()
|
||
|
||
def _fill_original_table(self):
|
||
originals = [asset for asset in self.assets if asset.kind == image_studio.ASSET_KIND_ORIGINAL]
|
||
self.original_table.setRowCount(len(originals))
|
||
for row, asset in enumerate(originals):
|
||
values = [
|
||
str(asset.source_order or row + 1),
|
||
_asset_status_text(asset),
|
||
asset.remote_url or "",
|
||
]
|
||
for column, value in enumerate(values):
|
||
item = QTableWidgetItem(str(value or ""))
|
||
item.setData(Qt.UserRole, {"type": "asset", "asset_id": int(asset.id)})
|
||
self.original_table.setItem(row, column, item)
|
||
self.original_table.resizeColumnsToContents()
|
||
|
||
def _fill_pool_table(self):
|
||
rows = []
|
||
for asset in self.assets:
|
||
if asset.status == image_studio.ASSET_STATUS_MISSING:
|
||
continue
|
||
if asset.kind not in {"original", "generated_main", "generated_detail"}:
|
||
continue
|
||
rows.append(("asset", asset))
|
||
for job in self.jobs:
|
||
if job.status in {"pending", "submitted", "running", "failed", "expired", "cancelled"}:
|
||
rows.append(("job", job))
|
||
self.pool_table.setRowCount(len(rows))
|
||
for row, (row_type, obj) in enumerate(rows):
|
||
if row_type == "asset":
|
||
values = [
|
||
_asset_badge(obj.kind),
|
||
obj.aspect_ratio or "未知",
|
||
_asset_status_text(obj),
|
||
_source_text(obj, self.assets),
|
||
obj.local_path or "",
|
||
]
|
||
data = {"type": "asset", "asset_id": int(obj.id)}
|
||
else:
|
||
values = [
|
||
"任务",
|
||
"-",
|
||
self.JOB_STATUS_LABELS.get(obj.status, obj.status),
|
||
f"源图 #{obj.source_asset_id or '-'}",
|
||
obj.error or "",
|
||
]
|
||
data = {"type": "job", "job_id": int(obj.id)}
|
||
for column, value in enumerate(values):
|
||
item = QTableWidgetItem(str(value or ""))
|
||
item.setData(Qt.UserRole, data)
|
||
self.pool_table.setItem(row, column, item)
|
||
self.pool_table.resizeColumnsToContents()
|
||
|
||
def _on_original_clicked(self, row, column):
|
||
asset = self._asset_from_table_row(self.original_table, row)
|
||
if asset is not None:
|
||
self._ensure_original_in_pool(asset, open_after=False)
|
||
|
||
def _on_original_double_clicked(self, row, column):
|
||
asset = self._asset_from_table_row(self.original_table, row)
|
||
if asset is not None:
|
||
self._ensure_original_in_pool(asset, open_after=True)
|
||
|
||
def _on_pool_clicked(self, row, column):
|
||
data = self._row_data(self.pool_table, row)
|
||
if not data or data.get("type") != "asset":
|
||
return
|
||
asset = self._asset_by_id(data.get("asset_id"))
|
||
if asset is not None:
|
||
self._select_source_asset(asset)
|
||
|
||
def _on_pool_double_clicked(self, row, column):
|
||
data = self._row_data(self.pool_table, row)
|
||
if not data or data.get("type") != "asset":
|
||
return
|
||
asset = self._asset_by_id(data.get("asset_id"))
|
||
if asset is not None:
|
||
self._open_preview(asset)
|
||
|
||
def _ensure_original_in_pool(self, asset, open_after=False):
|
||
if str(asset.local_path or "").strip() and os.path.isfile(asset.local_path):
|
||
self._select_source_asset(asset)
|
||
if open_after:
|
||
self._open_preview(asset)
|
||
return
|
||
worker = ImageStudioDownloadOriginalWorker(
|
||
asset.id,
|
||
db_path=self.db_path,
|
||
config=self.config,
|
||
open_after=open_after,
|
||
)
|
||
worker.log.connect(self._append_log)
|
||
worker.finished.connect(self._on_download_finished)
|
||
worker.failed.connect(self._on_worker_failed)
|
||
self._start_worker(worker, "AI工场下载原图")
|
||
|
||
def _on_download_finished(self, summary):
|
||
if summary.get("ok") is False:
|
||
if self._running_worker is not None:
|
||
self._on_worker_failed(-1, summary.get("error") or "下载原图失败")
|
||
return
|
||
asset = summary.get("asset")
|
||
self._finish_worker()
|
||
self.refresh_project_assets()
|
||
if asset is not None:
|
||
refreshed = self._asset_by_id(asset.id) or asset
|
||
self._select_source_asset(refreshed)
|
||
if summary.get("open_after"):
|
||
self._open_preview(refreshed)
|
||
self._status("原图已加入照片池", "success")
|
||
|
||
def _select_source_asset(self, asset):
|
||
self.selected_source_asset_id = int(asset.id)
|
||
self._refresh_source_label()
|
||
self._status(f"已选择源图：{_asset_badge(asset.kind)} #{asset.id}", "success")
|
||
|
||
def _refresh_source_label(self):
|
||
asset = self._asset_by_id(self.selected_source_asset_id)
|
||
if asset is None:
|
||
self.source_label.setText("源图：未选择")
|
||
return
|
||
self.source_label.setText(f"源图：{_asset_badge(asset.kind)} #{asset.id}")
|
||
|
||
def _open_preview(self, asset):
|
||
dialog = ImageStudioPreviewDialog(asset, self)
|
||
dialog.exec()
|
||
|
||
def _show_pool_context_menu(self, position):
|
||
row = self.pool_table.rowAt(position.y())
|
||
data = self._row_data(self.pool_table, row)
|
||
if not data or data.get("type") != "asset":
|
||
return
|
||
asset_id = data.get("asset_id")
|
||
menu = QMenu(self)
|
||
try:
|
||
counts = image_studio.asset_reference_counts(asset_id, path=self.db_path)
|
||
referenced = bool(counts.get("total"))
|
||
except Exception:
|
||
referenced = True
|
||
remove_action = menu.addAction(
|
||
"移除照片" if not referenced else "移除照片（已被任务或终选引用）"
|
||
)
|
||
remove_action.setEnabled(not referenced)
|
||
action = menu.exec(self.pool_table.viewport().mapToGlobal(position))
|
||
if action is remove_action and not referenced:
|
||
self.remove_asset(asset_id)
|
||
|
||
def remove_asset(self, asset_id):
|
||
if not self._confirm("移除照片", "只从AI工场照片池移除记录，不删除本地图片文件。"):
|
||
return
|
||
try:
|
||
image_studio.remove_asset_if_unused(asset_id, path=self.db_path)
|
||
except Exception as exc:
|
||
self._message("不能移除照片", str(exc))
|
||
return
|
||
if self.selected_source_asset_id == int(asset_id):
|
||
self.selected_source_asset_id = None
|
||
self.refresh_project_assets()
|
||
self._status("照片已从池中移除", "success")
|
||
|
||
def refresh_templates(self, selected=None):
|
||
current = selected or self.template_combo.currentData()
|
||
self.template_combo.blockSignals(True)
|
||
try:
|
||
self.template_combo.clear()
|
||
self.template_combo.addItem("选择模板", "")
|
||
for name in prompts.list_image_studio_templates(self.prompts_dir):
|
||
self.template_combo.addItem(name, name)
|
||
if current:
|
||
index = self.template_combo.findData(current)
|
||
if index >= 0:
|
||
self.template_combo.setCurrentIndex(index)
|
||
finally:
|
||
self.template_combo.blockSignals(False)
|
||
|
||
def load_selected_template(self, index=None):
|
||
name = self.template_combo.currentData()
|
||
if not name:
|
||
return
|
||
try:
|
||
self.prompt_edit.setPlainText(prompts.load_image_studio_template(name, self.prompts_dir))
|
||
except Exception as exc:
|
||
self._message("加载模板失败", str(exc))
|
||
|
||
def create_template(self, checked=False):
|
||
name, ok = QInputDialog.getText(self, "新建模板", "模板名称")
|
||
if not ok:
|
||
return
|
||
try:
|
||
prompts.save_image_studio_template(name, self.prompt_edit.toPlainText(), self.prompts_dir)
|
||
except Exception as exc:
|
||
self._message("新建模板失败", str(exc))
|
||
return
|
||
self.refresh_templates(selected=name)
|
||
self._status("AI工场模板已新建", "success")
|
||
|
||
def rename_template(self, checked=False):
|
||
old = self.template_combo.currentData()
|
||
if not old:
|
||
self._message("未选择模板", "请先选择要重命名的模板。")
|
||
return
|
||
new, ok = QInputDialog.getText(self, "重命名模板", "新模板名称", text=old)
|
||
if not ok:
|
||
return
|
||
try:
|
||
prompts.rename_image_studio_template(old, new, self.prompts_dir)
|
||
except Exception as exc:
|
||
self._message("重命名模板失败", str(exc))
|
||
return
|
||
self.refresh_templates(selected=new)
|
||
self._status("AI工场模板已重命名", "success")
|
||
|
||
def save_template(self, checked=False):
|
||
name = self.template_combo.currentData()
|
||
if not name:
|
||
name, ok = QInputDialog.getText(self, "保存模板", "模板名称")
|
||
if not ok:
|
||
return
|
||
try:
|
||
prompts.save_image_studio_template(name, self.prompt_edit.toPlainText(), self.prompts_dir)
|
||
except Exception as exc:
|
||
self._message("保存模板失败", str(exc))
|
||
return
|
||
self.refresh_templates(selected=name)
|
||
self._status("AI工场模板已保存", "success")
|
||
|
||
def delete_template(self, checked=False):
|
||
name = self.template_combo.currentData()
|
||
if not name:
|
||
self._message("未选择模板", "请先选择要删除的模板。")
|
||
return
|
||
if not self._confirm("删除模板", f"确定删除模板「{name}」吗？"):
|
||
return
|
||
try:
|
||
prompts.delete_image_studio_template(name, self.prompts_dir)
|
||
except Exception as exc:
|
||
self._message("删除模板失败", str(exc))
|
||
return
|
||
self.refresh_templates()
|
||
self._status("AI工场模板已删除", "success")
|
||
|
||
def _save_project_prompt(self):
|
||
if self.current_project is None:
|
||
return
|
||
try:
|
||
self.current_project = image_studio.update_project_prompt(
|
||
self.current_project.id,
|
||
self.prompt_edit.toPlainText(),
|
||
path=self.db_path,
|
||
)
|
||
except Exception as exc:
|
||
self._status(f"保存AI工场草稿提示词失败：{exc}", "warning")
|
||
|
||
def start_generation(self, checked=False):
|
||
if self.current_project is None:
|
||
self._message("项目未打开", "请先打开一个AI工场项目。")
|
||
return
|
||
source = self._asset_by_id(self.selected_source_asset_id)
|
||
if source is None:
|
||
self._message("未选择源图", "请先在照片池单击选择一张源图。")
|
||
return
|
||
if not str(source.local_path or "").strip() or not os.path.isfile(source.local_path):
|
||
self._message("源图不可用", "请先单击蝦皮原主图下载到本地后再生成。")
|
||
return
|
||
prompt_text = self.prompt_edit.toPlainText().strip()
|
||
if not prompt_text:
|
||
self._message("提示词不能为空", "请输入完整图片生成提示词。")
|
||
return
|
||
count = self.count_spin.value()
|
||
self.progress_bar.setRange(0, count)
|
||
self.progress_bar.setValue(0)
|
||
self.log_view.clear()
|
||
self._append_log(f"[AI工场] 本轮生图开始：{count} 张，来源 cmhub 托管模型")
|
||
worker = ImageStudioGenerateJobsWorker(
|
||
self.current_project.id,
|
||
source.id,
|
||
prompt_text,
|
||
count,
|
||
job_type=self.job_type_combo.currentData(),
|
||
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 stop_generation(self, checked=False):
|
||
if self._running_worker is not None and hasattr(self._running_worker, "cancel"):
|
||
self._running_worker.cancel()
|
||
self._append_log("[AI工场] 已请求停止，正在等待安全边界")
|
||
self._status("AI工场生成已请求停止", "warning")
|
||
|
||
def _on_generate_progress(self, payload):
|
||
total = max(1, int(payload.get("total") or self.progress_bar.maximum() or 1))
|
||
done = min(total, int(payload.get("done") or 0))
|
||
self.progress_bar.setRange(0, total)
|
||
self.progress_bar.setValue(done)
|
||
if payload.get("points_balance") is not None:
|
||
text = f"cmhub 托管模型：余额 {payload.get('points_balance')}"
|
||
if payload.get("points_cost") is not None:
|
||
text += f"，本张扣点 {payload.get('points_cost')}"
|
||
self.billing_label.setText(text)
|
||
|
||
def _on_generation_finished(self, summary):
|
||
if summary.get("ok") is False:
|
||
if self._running_worker is not None:
|
||
self._on_worker_failed(-1, summary.get("error") or "AI工场生成失败")
|
||
return
|
||
self._finish_worker()
|
||
self.refresh_project_assets()
|
||
total = int(summary.get("total") or 0)
|
||
success = int(summary.get("success") or 0)
|
||
failed = int(summary.get("failed") or 0)
|
||
cancelled = int(summary.get("cancelled") or 0)
|
||
self._append_log(f"[AI工场] 本轮完成：总数{total}，成功{success}，失败{failed}，停止{cancelled}")
|
||
level = "warning" if failed or cancelled else "success"
|
||
self._status(f"AI工场生成完成：成功{success}，失败{failed}，停止{cancelled}", level)
|
||
|
||
def _on_worker_failed(self, row, error):
|
||
self._finish_worker()
|
||
message = diagnostics.redact_log_text(str(error or "未知错误"))
|
||
self._append_log(f"[AI工场] 失败：{message}")
|
||
self._status(f"AI工场任务失败：{message}", "danger")
|
||
self._message("AI工场任务失败", message)
|
||
self.refresh_project_assets()
|
||
|
||
def _start_worker(self, worker, thread_name):
|
||
self._set_running(True)
|
||
self._running_worker = worker
|
||
self._running_thread = run_worker(worker, thread_name)
|
||
|
||
def _finish_worker(self):
|
||
self._running_worker = None
|
||
self._running_thread = None
|
||
self._set_running(False)
|
||
|
||
def _set_running(self, running):
|
||
self.open_project_button.setEnabled(not running)
|
||
self.pull_images_button.setEnabled(not running)
|
||
self.open_folder_button.setEnabled(not running)
|
||
self.project_table.setEnabled(not running)
|
||
self.original_table.setEnabled(not running)
|
||
self.pool_table.setEnabled(not running)
|
||
self.template_combo.setEnabled(not running)
|
||
self.template_new_button.setEnabled(not running)
|
||
self.template_rename_button.setEnabled(not running)
|
||
self.template_save_button.setEnabled(not running)
|
||
self.template_delete_button.setEnabled(not running)
|
||
self.prompt_edit.setEnabled(not running)
|
||
self.job_type_combo.setEnabled(not running)
|
||
self.count_spin.setEnabled(not running)
|
||
self.aspect_combo.setEnabled(not running)
|
||
self.start_button.setEnabled(not running)
|
||
self.stop_button.setEnabled(running)
|
||
|
||
def _set_account_combo(self, alias):
|
||
index = self.account_combo.findData(alias)
|
||
if index >= 0:
|
||
self.account_combo.setCurrentIndex(index)
|
||
|
||
def _list_project_jobs(self, project_id):
|
||
conn = db.connect(self.db_path)
|
||
try:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT * FROM image_studio_jobs
|
||
WHERE project_id = ?
|
||
ORDER BY updated_at DESC, id DESC
|
||
""",
|
||
(int(project_id),),
|
||
).fetchall()
|
||
return [image_studio.ImageStudioJob(**dict(row)) for row in rows]
|
||
finally:
|
||
conn.close()
|
||
|
||
def _refresh_selection_labels(self):
|
||
main_count = sum(1 for item in self.selections if item.selection_type == "main")
|
||
detail_count = sum(1 for item in self.selections if item.selection_type == "detail")
|
||
main_target = getattr(self.current_project, "target_main_count", 9) if self.current_project else 9
|
||
detail_target = getattr(self.current_project, "target_detail_count", 12) if self.current_project else 12
|
||
self.main_selection_label.setText(f"主图终选 {main_count}/{main_target}（拖放排序将在 T-592 接入）")
|
||
self.detail_selection_label.setText(f"详情图终选 {detail_count}/{detail_target}（拖放排序将在 T-592 接入）")
|
||
|
||
def _asset_from_table_row(self, table, row):
|
||
data = self._row_data(table, row)
|
||
if not data or data.get("type") != "asset":
|
||
return None
|
||
return self._asset_by_id(data.get("asset_id"))
|
||
|
||
def _row_data(self, table, row):
|
||
if row < 0 or row >= table.rowCount():
|
||
return None
|
||
item = table.item(row, 0)
|
||
if item is None:
|
||
return None
|
||
return item.data(Qt.UserRole)
|
||
|
||
def _asset_by_id(self, asset_id):
|
||
if asset_id is None:
|
||
return None
|
||
for asset in self.assets:
|
||
if int(asset.id) == int(asset_id):
|
||
return asset
|
||
try:
|
||
return image_studio.get_asset(asset_id, path=self.db_path)
|
||
except Exception:
|
||
return None
|
||
|
||
def _append_log(self, message):
|
||
text = diagnostics.redact_log_text(str(message or ""))
|
||
self.log_view.appendPlainText(text)
|
||
scrollbar = self.log_view.verticalScrollBar()
|
||
scrollbar.setValue(scrollbar.maximum())
|
||
|
||
def _message(self, title, text):
|
||
box = QMessageBox(self)
|
||
box.setWindowTitle(str(title or "提示"))
|
||
box.setText(str(text or ""))
|
||
ok_button = box.addButton("确定", QMessageBox.AcceptRole)
|
||
box.setDefaultButton(ok_button)
|
||
box.exec()
|
||
|
||
def _confirm(self, title, text):
|
||
box = QMessageBox(self)
|
||
box.setWindowTitle(str(title or "确认"))
|
||
box.setText(str(text or ""))
|
||
yes_button = box.addButton("确定", QMessageBox.AcceptRole)
|
||
box.addButton("取消", QMessageBox.RejectRole)
|
||
box.setDefaultButton(yes_button)
|
||
box.exec()
|
||
return box.clickedButton() is yes_button
|
||
|
||
def _status(self, message, level="muted"):
|
||
_emit_status(self.status_callback, message, level=level)
|
||
|
||
|
||
def _asset_badge(kind):
|
||
return {
|
||
"original": "原图",
|
||
"generated_main": "主图",
|
||
"generated_detail": "详情图",
|
||
}.get(str(kind or ""), str(kind or "图片"))
|
||
|
||
|
||
def _asset_status_text(asset):
|
||
status = str(getattr(asset, "status", "") or "")
|
||
local_path = str(getattr(asset, "local_path", "") or "")
|
||
if status == image_studio.ASSET_STATUS_MISSING:
|
||
return "文件缺失"
|
||
if local_path and os.path.isfile(local_path):
|
||
return "可用"
|
||
if getattr(asset, "remote_url", None):
|
||
return "远程待下载"
|
||
return "待生成"
|
||
|
||
|
||
def _source_text(asset, assets):
|
||
parent_id = getattr(asset, "parent_asset_id", None)
|
||
if not parent_id:
|
||
return "原始来源"
|
||
for item in assets:
|
||
if int(item.id) == int(parent_id):
|
||
return f"{_asset_badge(item.kind)} #{item.id}"
|
||
return f"源图 #{parent_id}"
|