"""Tab 6: AI image studio UI.""" from __future__ import annotations import os from PySide6.QtCore import QMimeData from PySide6.QtWidgets import QListWidget, QListWidgetItem from ... import accounts, appconfig, cmhub_models, db, image_studio, image_studio_export, prompts from .. import file_manager from ..widgets import * from ..workers import ( ImageStudioDownloadOriginalWorker as _RealImageStudioDownloadOriginalWorker, ) 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" 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, ) def ImageStudioExportWorker(*args, **kwargs): return _call_package_attr( "ImageStudioExportWorker", _RealImageStudioExportWorker, *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() return event.pos() def _mime_asset_id(mime_data): if mime_data is None or not mime_data.hasFormat(ASSET_MIME_TYPE): return None raw = bytes(mime_data.data(ASSET_MIME_TYPE)).decode("utf-8", errors="ignore") try: return int(str(raw).strip()) except (TypeError, ValueError): return None class ImageStudioPoolTable(QTableWidget): """Photo pool table that can drag available asset IDs into final trays.""" def __init__(self, parent=None): super().__init__(parent) self.setDragEnabled(True) self.setDragDropMode(QAbstractItemView.DragOnly) self.setDefaultDropAction(Qt.CopyAction) def mimeData(self, items): mime = QMimeData() if not items: return mime row = items[0].row() item = self.item(row, 0) if item is None: return mime data = item.data(Qt.UserRole) or {} if data.get("type") != "asset" or not data.get("draggable"): return mime mime.setData(ASSET_MIME_TYPE, str(data.get("asset_id")).encode("utf-8")) return mime class ImageStudioSelectionList(QListWidget): """Ordered final tray for main/detail selected assets.""" def __init__(self, selection_type, tab, parent=None): super().__init__(parent) self.selection_type = selection_type self.tab = tab self.setAcceptDrops(True) self.setDragEnabled(True) self.setDefaultDropAction(Qt.MoveAction) self.setDragDropMode(QAbstractItemView.DragDrop) self.setSelectionMode(QAbstractItemView.SingleSelection) self.setContextMenuPolicy(Qt.CustomContextMenu) self.customContextMenuRequested.connect(self._show_context_menu) def mimeData(self, items): mime = QMimeData() if not items: return mime asset_id = items[0].data(Qt.UserRole) if asset_id is not None: mime.setData(ASSET_MIME_TYPE, str(asset_id).encode("utf-8")) return mime def dragEnterEvent(self, event): if _mime_asset_id(event.mimeData()) is not None: event.acceptProposedAction() return super().dragEnterEvent(event) def dragMoveEvent(self, event): if _mime_asset_id(event.mimeData()) is not None: event.acceptProposedAction() return super().dragMoveEvent(event) def dropEvent(self, event): asset_id = _mime_asset_id(event.mimeData()) if asset_id is None: super().dropEvent(event) return index = self.indexAt(_drop_event_position(event)) insert_index = index.row() if index.isValid() else self.count() if event.source() is self: current_index = self._row_for_asset(asset_id) if current_index is None: event.ignore() return self.tab.move_selection_asset(self.selection_type, current_index, insert_index) else: self.tab.add_asset_to_selection(self.selection_type, asset_id, insert_index=insert_index) event.acceptProposedAction() def keyPressEvent(self, event): if event.key() == Qt.Key_Delete: self.remove_current_asset() event.accept() return super().keyPressEvent(event) def remove_current_asset(self): item = self.currentItem() if item is None: return asset_id = item.data(Qt.UserRole) self.tab.remove_asset_from_selection(self.selection_type, asset_id) def _show_context_menu(self, position): item = self.itemAt(position) if item is None: return menu = QMenu(self) remove_action = menu.addAction("移出终选") action = menu.exec(self.viewport().mapToGlobal(position)) if action is remove_action: self.setCurrentItem(item) self.remove_current_asset() def _row_for_asset(self, asset_id): for row in range(self.count()): item = self.item(row) if item is not None and int(item.data(Qt.UserRole)) == int(asset_id): return row return None 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_model_hint() 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 = ImageStudioPoolTable() self.pool_table.setColumnCount(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.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) 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) layout.setSpacing(10) main_panel = QWidget() main_layout = QVBoxLayout(main_panel) main_layout.setContentsMargins(0, 0, 0, 0) self.main_selection_label = QLabel("主图终选 0/9") self.main_selection_label.setObjectName("imageStudioMainSelectionLabel") self.main_selection_list = ImageStudioSelectionList("main", self) self.main_selection_list.setObjectName("imageStudioMainSelectionList") self.main_selection_list.setMinimumHeight(96) main_layout.addWidget(self.main_selection_label) main_layout.addWidget(self.main_selection_list) detail_panel = QWidget() detail_layout = QVBoxLayout(detail_panel) detail_layout.setContentsMargins(0, 0, 0, 0) self.detail_selection_label = QLabel("详情图终选 0/12") self.detail_selection_label.setObjectName("imageStudioDetailSelectionLabel") self.detail_selection_list = ImageStudioSelectionList("detail", self) self.detail_selection_list.setObjectName("imageStudioDetailSelectionList") self.detail_selection_list.setMinimumHeight(96) detail_layout.addWidget(self.detail_selection_label) detail_layout.addWidget(self.detail_selection_list) action_panel = QWidget() action_layout = QVBoxLayout(action_panel) action_layout.setContentsMargins(0, 0, 0, 0) self.export_button = QPushButton("导出终选") self.export_button.setObjectName("imageStudioExportButton") self.export_hint_label = QLabel("可部分导出，不要求主图/详情图满额") self.export_hint_label.setObjectName("imageStudioExportHintLabel") self.export_hint_label.setWordWrap(True) action_layout.addWidget(self.export_button) action_layout.addWidget(self.export_hint_label) action_layout.addStretch(1) layout.addWidget(main_panel, 1) layout.addWidget(detail_panel, 1) layout.addWidget(action_panel, 0) 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.resume_button.clicked.connect(self.resume_generation_jobs) self.stop_button.clicked.connect(self.stop_generation) self.export_button.clicked.connect(self.export_selections) 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": draggable = _asset_is_usable(obj) 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), "draggable": draggable} else: values = [ "任务", "-", _job_status_text(obj, self.JOB_STATUS_LABELS), 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 _selection_asset_ids(self, selection_type): return [ int(selection.asset_id) for selection in self.selections if selection.selection_type == selection_type ] def add_asset_to_selection(self, selection_type, asset_id, insert_index=None): if self.current_project is None: return False asset = self._asset_by_id(asset_id) if asset is None: self._message("不能加入终选", "照片不存在。") return False if not _asset_is_usable(asset): self._message("不能加入终选", "照片尚未下载或本地文件缺失。") return False ids = self._selection_asset_ids(selection_type) if int(asset.id) in ids: self._message("不能重复加入", "同一张照片在同一类终选中只能出现一次。") return False target = self._selection_target_count(selection_type) if len(ids) >= target: self._message("终选已满", f"{_selection_label(selection_type)}最多 {target} 张。") return False index = _clamp_insert_index(insert_index, len(ids)) ids.insert(index, int(asset.id)) return self._persist_selection(selection_type, ids, success_message="已加入终选") def move_selection_asset(self, selection_type, from_index, to_index): ids = self._selection_asset_ids(selection_type) if not ids: return False from_index = int(from_index) if from_index < 0 or from_index >= len(ids): return False asset_id = ids.pop(from_index) target_index = _clamp_insert_index(to_index, len(ids)) ids.insert(target_index, asset_id) return self._persist_selection(selection_type, ids, success_message="终选顺序已更新") def remove_asset_from_selection(self, selection_type, asset_id): ids = self._selection_asset_ids(selection_type) filtered = [value for value in ids if int(value) != int(asset_id)] if len(filtered) == len(ids): return False return self._persist_selection(selection_type, filtered, success_message="已移出终选") def _persist_selection(self, selection_type, asset_ids, success_message): try: image_studio.replace_selections( self.current_project.id, selection_type, asset_ids, path=self.db_path, ) except Exception as exc: self._message("终选保存失败", str(exc)) self.refresh_project_assets() return False self.refresh_project_assets() self._status(success_message, "success") return True def _selection_target_count(self, selection_type): if self.current_project is None: return 0 if selection_type == "detail": return int(self.current_project.target_detail_count or 12) return int(self.current_project.target_main_count or 9) 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} 张，使用 {self._cmhub_image_model_summary()}") 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 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工场项目。") return main_count = len(self._selection_asset_ids("main")) detail_count = len(self._selection_asset_ids("detail")) if main_count + detail_count <= 0: self._message("没有可导出图片", "主图和详情图终选都为空。") return parent_dir = QFileDialog.getExistingDirectory(self, "选择导出父目录") if not parent_dir: return mode = self._export_existing_mode(parent_dir, main_count, detail_count) if mode is None: return worker = ImageStudioExportWorker( self.current_project.id, parent_dir, existing_mode=mode, db_path=self.db_path, config=self.config, ) worker.log.connect(self._append_log) worker.finished.connect(self._on_export_finished) worker.failed.connect(self._on_worker_failed) self._start_worker(worker, "AI工场导出终选") def _export_existing_mode(self, parent_dir, main_count, detail_count): try: target = image_studio_export.target_dir_for_project(self.current_project, parent_dir) except Exception as exc: self._message("导出目录不可用", str(exc)) return None if not os.path.exists(target): return image_studio_export.EXISTING_FAIL box = QMessageBox(self) box.setWindowTitle("商品目录已存在") box.setText( f"目标目录已存在：{target}\n" f"本次将导出主图 {main_count} 张、详情图 {detail_count} 张。请选择处理方式。" ) overwrite_button = box.addButton("覆盖本软件导出的图片", QMessageBox.AcceptRole) timestamp_button = box.addButton("新建带时间目录", QMessageBox.ActionRole) box.addButton("取消", QMessageBox.RejectRole) box.setDefaultButton(timestamp_button) box.exec() clicked = box.clickedButton() if clicked is overwrite_button: return image_studio_export.EXISTING_OVERWRITE_MANAGED if clicked is timestamp_button: return image_studio_export.EXISTING_TIMESTAMP return None def _on_export_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 self._finish_worker() target_dir = summary.get("target_dir") or "" main_count = int(summary.get("main_count") or 0) detail_count = int(summary.get("detail_count") or 0) file_count = int(summary.get("file_count") or 0) self._append_log(f"[AI工场] 导出完成：主图{main_count}，详情图{detail_count}，文件{file_count}") self._status(f"AI工场终选已导出：{file_count} 个文件", "success") self._show_export_success(target_dir, main_count, detail_count, file_count) def _show_export_success(self, target_dir, main_count, detail_count, file_count): box = QMessageBox(self) box.setWindowTitle("导出完成") box.setText( f"已导出 {file_count} 个 JPEG 文件。\n" f"主图 {main_count} 张，详情图 {detail_count} 张。\n" f"目录：{target_dir}" ) open_button = box.addButton("打开目录", QMessageBox.AcceptRole) box.addButton("确定", QMessageBox.RejectRole) box.setDefaultButton(open_button) box.exec() if box.clickedButton() is open_button: try: file_manager.open_in_file_manager(target_dir) except Exception as exc: self._message("打开目录失败", str(exc)) 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"{self._cmhub_image_model_summary()}；余额 {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.main_selection_list.setEnabled(not running) self.detail_selection_list.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.resume_button.setEnabled(not running) self.stop_button.setEnabled(running) self.export_button.setEnabled(not 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}") self.detail_selection_label.setText(f"详情图终选 {detail_count}/{detail_target}") self._fill_selection_list(self.main_selection_list, "main") self._fill_selection_list(self.detail_selection_list, "detail") def _fill_selection_list(self, widget, selection_type): widget.clear() for index, selection in enumerate( [item for item in self.selections if item.selection_type == selection_type], start=1, ): asset = self._asset_by_id(selection.asset_id) if asset is None: text = f"{index}. 缺失照片 #{selection.asset_id}" item = QListWidgetItem(text) item.setData(Qt.UserRole, int(selection.asset_id)) item.setToolTip("终选引用的照片记录不存在") item.setForeground(_qcolor(COLOR_DANGER)) widget.addItem(item) continue text = f"{index}. {_asset_badge(asset.kind)} #{asset.id} · {asset.aspect_ratio or '未知比例'}" item = QListWidgetItem(text) item.setData(Qt.UserRole, int(asset.id)) tooltip = _selection_tooltip(selection_type, asset) item.setToolTip(tooltip) if _selection_ratio_warning(selection_type, asset): item.setBackground(QColor("#fff8c5")) item.setToolTip(tooltip + "\n比例与主图推荐比例不一致，仅提示，不阻止导出。") if not _asset_is_usable(asset): item.setForeground(_qcolor(COLOR_MUTED)) item.setToolTip(tooltip + "\n本地文件缺失，不能用于导出。") widget.addItem(item) 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 _refresh_model_hint(self): self.billing_label.setText(self._cmhub_image_model_summary()) def _cmhub_image_model_summary(self): alias = appconfig.cmhub_config(self.config).get("image_alias", "") return cmhub_models.configured_alias_summary(alias) 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 _asset_is_usable(asset): if asset is None: return False if str(getattr(asset, "status", "") or "") == image_studio.ASSET_STATUS_MISSING: return False local_path = str(getattr(asset, "local_path", "") or "") return bool(local_path and os.path.isfile(local_path)) 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}" def _selection_label(selection_type): return "详情图终选" if selection_type == "detail" else "主图终选" def _clamp_insert_index(value, length): try: index = int(value) except (TypeError, ValueError): index = length return max(0, min(index, int(length))) def _selection_ratio_warning(selection_type, asset): if selection_type != "main": return False ratio = str(getattr(asset, "aspect_ratio", "") or "").strip() return bool(ratio and ratio != "1:1") def _selection_tooltip(selection_type, asset): return ( 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)