572 lines
31 KiB
Python
572 lines
31 KiB
Python
import sys
|
||
import os
|
||
import json
|
||
import time
|
||
import requests
|
||
import copy
|
||
import uuid
|
||
import configparser
|
||
from pathlib import Path
|
||
from urllib.parse import quote
|
||
|
||
from PySide6.QtWidgets import (
|
||
QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
|
||
QLineEdit, QPushButton, QComboBox, QPlainTextEdit, QGroupBox,
|
||
QFormLayout, QFileDialog, QMessageBox, QStatusBar, QProgressBar,
|
||
QDialog, QLabel, QDialogButtonBox, QStyle, QListWidget, QListWidgetItem
|
||
)
|
||
from PySide6.QtCore import Qt, QThread, Signal, QObject, QSettings
|
||
from PySide6.QtGui import QIcon, QPixmap
|
||
|
||
|
||
# --- 配置管理类 ---
|
||
class ConfigManager:
|
||
def __init__(self, filename='config.ini'):
|
||
self.filename, self.config = filename, configparser.ConfigParser()
|
||
self.config.read(self.filename)
|
||
|
||
def get(self, section, key, fallback=None): return self.config.get(section, key, fallback=fallback)
|
||
|
||
def set(self, section, key, value):
|
||
if not self.config.has_section(section): self.config.add_section(section)
|
||
self.config.set(section, key, str(value))
|
||
|
||
def save(self):
|
||
with open(self.filename, 'w') as configfile: self.config.write(configfile)
|
||
|
||
|
||
# --- 辅助函数：将工作流图谱格式转换为API提示格式 ---
|
||
def convert_graph_to_prompt(workflow_graph):
|
||
if 'nodes' not in workflow_graph or 'links' not in workflow_graph: return None
|
||
prompt, nodes_by_id = {}, {str(node['id']): node for node in workflow_graph['nodes']}
|
||
for node_id, node_data in nodes_by_id.items():
|
||
inputs, widget_values, widget_index = {}, node_data.get('widgets_values', []), 0
|
||
if 'inputs' in node_data:
|
||
for an_input in node_data['inputs']:
|
||
if an_input.get('widget') and widget_index < len(widget_values):
|
||
inputs[an_input['name']] = widget_values[widget_index]
|
||
widget_index += 1
|
||
for link_info in workflow_graph['links']:
|
||
target_node_id_str, source_node_id_str = str(link_info[3]), str(link_info[1])
|
||
if target_node_id_str == node_id:
|
||
source_slot_idx, target_slot_idx = link_info[2], link_info[4]
|
||
if 'inputs' in node_data and target_slot_idx < len(node_data['inputs']):
|
||
inputs[node_data['inputs'][target_slot_idx]['name']] = [source_node_id_str, source_slot_idx]
|
||
prompt[node_id] = {"class_type": node_data['type'], "inputs": inputs}
|
||
return prompt
|
||
|
||
|
||
# --- API通信类 ---
|
||
class ComfyAPI:
|
||
def __init__(self, host, port):
|
||
self.base_url, self.client_id = f"http://{host}:{port}", str(uuid.uuid4())
|
||
|
||
def _request(self, method, endpoint, **kwargs):
|
||
try:
|
||
response = requests.request(method, f"{self.base_url}{endpoint}", timeout=10, **kwargs)
|
||
response.raise_for_status()
|
||
return response.json(), None
|
||
except requests.exceptions.RequestException as e:
|
||
return None, str(e)
|
||
|
||
def get_workflows(self):
|
||
return self._request("get", "/api/userdata?dir=workflows&recurse=true&split=false&full_info=true")
|
||
|
||
def get_workflow_details(self, workflow_name):
|
||
return self._request("get", f"/api/userdata/workflows%2F{quote(workflow_name)}")
|
||
|
||
def get_queue(self):
|
||
return self._request("get", "/api/queue")
|
||
|
||
def queue_prompt(self, prompt, workflow_graph):
|
||
payload = {"prompt": prompt, "extra_data": {"extra_pnginfo": {"workflow": workflow_graph}},
|
||
"client_id": self.client_id}
|
||
return self._request("post", "/api/prompt", json=payload)
|
||
|
||
|
||
# --- 后台执行线程 ---
|
||
class ExecutionWorker(QObject):
|
||
log, progress, finished = Signal(str), Signal(int, int), Signal(bool, str)
|
||
status_update = Signal(str) # FIX 2: New signal for status bar
|
||
|
||
def __init__(self, api, workflow_prompt, workflow_graph, model_image_path, model_node_id,
|
||
clothing_dir, clothing_node_id, output_dir, output_node_id):
|
||
super().__init__()
|
||
self.api, self.workflow_prompt, self.workflow_graph = api, workflow_prompt, workflow_graph
|
||
self.model_image_path, self.model_node_id = Path(model_image_path), model_node_id
|
||
self.clothing_dir, self.clothing_node_id = Path(clothing_dir), clothing_node_id
|
||
self.output_dir, self.output_node_id = Path(output_dir), output_node_id
|
||
self._is_running = True
|
||
|
||
def stop(self):
|
||
self._is_running = False; self.log.emit("执行停止请求已发送...")
|
||
|
||
def run(self):
|
||
image_extensions = ['.png', '.jpg', '.jpeg', '.webp', '.bmp']
|
||
try:
|
||
clothing_files = [f for f in self.clothing_dir.iterdir() if
|
||
f.is_file() and f.suffix.lower() in image_extensions]
|
||
except FileNotFoundError:
|
||
return self.finished.emit(False, f"错误：找不到衣物文件夹 {self.clothing_dir}")
|
||
if not clothing_files: return self.finished.emit(False, f"在 {self.clothing_dir} 中没有找到有效的图片文件。")
|
||
total_files = len(clothing_files);
|
||
self.log.emit(f"找到 {total_files} 件衣物图片待处理。");
|
||
self.progress.emit(0, total_files)
|
||
for i, clothing_path in enumerate(clothing_files):
|
||
if not self._is_running: break
|
||
self.log.emit(
|
||
f"--- 开始处理: {self.model_image_path.name} + {clothing_path.name} ({i + 1}/{total_files}) ---")
|
||
prompt = copy.deepcopy(self.workflow_prompt)
|
||
if self.model_node_id in prompt:
|
||
prompt[self.model_node_id]["inputs"]["image"] = str(self.model_image_path)
|
||
else:
|
||
return self.finished.emit(False, f"错误：模特节点ID '{self.model_node_id}' 在工作流中不存在！")
|
||
if self.clothing_node_id in prompt:
|
||
prompt[self.clothing_node_id]["inputs"]["image"] = str(clothing_path)
|
||
else:
|
||
return self.finished.emit(False, f"错误：衣物节点ID '{self.clothing_node_id}' 在工作流中不存在！")
|
||
if self.output_node_id in prompt:
|
||
prompt[self.output_node_id]["inputs"]["output_dir"] = str(self.output_dir)
|
||
prompt[self.output_node_id]["inputs"][
|
||
"filename_prefix"] = f"{self.model_image_path.stem}_{clothing_path.stem}"
|
||
else:
|
||
return self.finished.emit(False, f"错误：输出节点ID '{self.output_node_id}' 在工作流中不存在！")
|
||
response, error = self.api.queue_prompt(prompt, self.workflow_graph)
|
||
if error: self.log.emit(f"错误: 任务排队失败 - {error}"); continue
|
||
prompt_id = response['prompt_id'];
|
||
self.log.emit(f"任务已入队，ID: {prompt_id}")
|
||
|
||
# FIX 2: New wait logic
|
||
logged_wait_for_task = False
|
||
while self._is_running:
|
||
queue_info, error = self.api.get_queue()
|
||
if error: self.status_update.emit(f"无法获取队列状态: {error}"); time.sleep(2); continue
|
||
|
||
queue_running = queue_info.get('queue_running', [])
|
||
queue_pending = queue_info.get('queue_pending', [])
|
||
|
||
is_running = any(item[1] == prompt_id for item in queue_running)
|
||
is_pending = any(item[1] == prompt_id for item in queue_pending)
|
||
|
||
if not is_running and not is_pending:
|
||
self.status_update.emit(f"任务 {prompt_id} 已完成")
|
||
self.log.emit(f"任务 {prompt_id} 处理完成。")
|
||
break
|
||
|
||
if not logged_wait_for_task:
|
||
self.log.emit(f"等待任务 {prompt_id} 完成...")
|
||
logged_wait_for_task = True
|
||
|
||
if is_running:
|
||
self.status_update.emit(f"正在执行任务: {prompt_id}")
|
||
elif is_pending:
|
||
try:
|
||
# Find position in the pending queue
|
||
pending_ids = [item[1] for item in queue_pending]
|
||
position = pending_ids.index(prompt_id) + 1
|
||
self.status_update.emit(f"队列等候中 (位置 {position}/{len(pending_ids)})")
|
||
except ValueError:
|
||
self.status_update.emit("正在更新队列状态...")
|
||
time.sleep(1)
|
||
|
||
self.progress.emit(i + 1, total_files)
|
||
if self._is_running:
|
||
self.finished.emit(True, "全部任务执行完毕。")
|
||
else:
|
||
self.finished.emit(False, "执行被用户中止。")
|
||
|
||
|
||
# --- UI 组件类 ---
|
||
class ImagePreviewDialog(QDialog):
|
||
def __init__(self, image_path, parent=None):
|
||
super().__init__(parent);
|
||
self.setWindowTitle("模特预览")
|
||
self.image_label = QLabel(self);
|
||
self.image_label.setAlignment(Qt.AlignCenter)
|
||
pixmap = QPixmap(image_path);
|
||
scaled_pixmap = pixmap.scaled(800, 800, Qt.KeepAspectRatio, Qt.SmoothTransformation)
|
||
self.image_label.setPixmap(scaled_pixmap);
|
||
layout = QVBoxLayout(self);
|
||
layout.addWidget(self.image_label)
|
||
button_box = QDialogButtonBox(QDialogButtonBox.Ok);
|
||
button_box.accepted.connect(self.accept)
|
||
layout.addWidget(button_box);
|
||
self.setLayout(layout);
|
||
self.adjustSize()
|
||
|
||
|
||
class SettingsDialog(QDialog):
|
||
def __init__(self, workflow_prompt, current_settings, parent=None):
|
||
super().__init__(parent);
|
||
self.setWindowTitle("工作流输入输出设置");
|
||
self.setMinimumWidth(450)
|
||
self.workflow_prompt = workflow_prompt;
|
||
layout = QFormLayout(self)
|
||
self.model_node_combo = QComboBox();
|
||
self.clothing_node_combo = QComboBox();
|
||
self.output_node_combo = QComboBox()
|
||
for node_id, node_info in self.workflow_prompt.items():
|
||
text = f"ID: {node_id} (类型: {node_info['class_type']})"
|
||
if "LoadImage" in node_info['class_type']: self.model_node_combo.addItem(text,
|
||
userData=node_id); self.clothing_node_combo.addItem(
|
||
text, userData=node_id)
|
||
if "Save" in node_info['class_type']: self.output_node_combo.addItem(text, userData=node_id)
|
||
if current_settings.get('model_node_id'): self.model_node_combo.setCurrentIndex(
|
||
self.model_node_combo.findData(current_settings['model_node_id']))
|
||
if current_settings.get('clothing_node_id'): self.clothing_node_combo.setCurrentIndex(
|
||
self.clothing_node_combo.findData(current_settings['clothing_node_id']))
|
||
if current_settings.get('output_node_id'): self.output_node_combo.setCurrentIndex(
|
||
self.output_node_combo.findData(current_settings['output_node_id']))
|
||
layout.addRow("1. 模特加载节点 (LoadImage):", self.model_node_combo);
|
||
layout.addRow("2. 衣物加载节点 (LoadImage):", self.clothing_node_combo);
|
||
layout.addRow("3. 图像保存节点 (SaveImageAnywhere):", self.output_node_combo)
|
||
button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel);
|
||
button_box.accepted.connect(self.accept);
|
||
button_box.rejected.connect(self.reject);
|
||
layout.addRow(button_box)
|
||
|
||
def get_settings(self):
|
||
return {"model_node_id": self.model_node_combo.currentData(),
|
||
"clothing_node_id": self.clothing_node_combo.currentData(),
|
||
"output_node_id": self.output_node_combo.currentData()}
|
||
|
||
|
||
# --- 主窗口 ---
|
||
class ComfyUIController(QMainWindow):
|
||
def __init__(self):
|
||
super().__init__();
|
||
self.setWindowTitle("ComfyUI 智能试衣控制器");
|
||
self.setWindowIcon(QIcon(os.path.join(os.path.dirname(__file__), "comfyui_hz.ico")));
|
||
self.setGeometry(100, 100, 1000, 700)
|
||
self.config = ConfigManager(os.path.join(os.path.dirname(os.path.dirname(__file__)), "config_hz.ini"));
|
||
self.api, self.workflow_settings = None, {}
|
||
self.current_workflow_prompt, self.current_workflow_graph = None, None
|
||
self.execution_thread, self.execution_worker = None, None
|
||
self._create_widgets();
|
||
self._create_layouts();
|
||
self._connect_signals();
|
||
self._load_settings();
|
||
self._load_models();
|
||
self._update_ui_state()
|
||
|
||
def _create_widgets(self):
|
||
self.model_group = QGroupBox("模特库");
|
||
self.model_list_widget = QListWidget();
|
||
self.add_model_button = QPushButton("✚ 添加模特");
|
||
self.delete_model_button = QPushButton("✖ 删除选中")
|
||
self.server_group = QGroupBox("① 服务器连接");
|
||
self.host_input = QLineEdit();
|
||
self.port_input = QLineEdit();
|
||
self.connect_button = QPushButton("连接")
|
||
self.workflow_group = QGroupBox("② 工作流选择");
|
||
self.workflow_combo = QComboBox();
|
||
self.refresh_button = QPushButton("🔄");
|
||
self.settings_button = QPushButton("⚙️");
|
||
self.workflow_tip_label = QLabel("提示：请选择一个包含两个LoadImage节点和1个SaveImageAnywhere节点的换装工作流。");
|
||
self.workflow_tip_label.setStyleSheet("color: #e67e22;")
|
||
self.exec_group = QGroupBox("③ 执行配置");
|
||
self.clothing_dir_input = QLineEdit();
|
||
self.clothing_dir_button = QPushButton("浏览...")
|
||
self.output_group = QGroupBox("④ 输出设置");
|
||
self.output_dir_input = QLineEdit();
|
||
self.output_dir_button = QPushButton("浏览...")
|
||
self.action_group = QGroupBox("⑤ 开始执行");
|
||
self.execute_button = QPushButton("▶️ 开始批量换装");
|
||
self.stop_button = QPushButton("⏹️ 停止执行")
|
||
self.log_group = QGroupBox("执行日志");
|
||
self.log_box = QPlainTextEdit();
|
||
self.log_box.setReadOnly(True)
|
||
self.status_bar = QStatusBar();
|
||
self.setStatusBar(self.status_bar);
|
||
self.progress_bar = QProgressBar();
|
||
self.status_bar.addPermanentWidget(self.progress_bar)
|
||
|
||
def _create_layouts(self):
|
||
model_layout, model_buttons_layout = QVBoxLayout(), QHBoxLayout();
|
||
model_buttons_layout.addWidget(self.add_model_button);
|
||
model_buttons_layout.addWidget(self.delete_model_button);
|
||
model_layout.addWidget(self.model_list_widget);
|
||
model_layout.addLayout(model_buttons_layout);
|
||
self.model_group.setLayout(model_layout)
|
||
right_v_layout = QVBoxLayout()
|
||
server_h_layout = QHBoxLayout();
|
||
server_h_layout.addWidget(QLabel("主机:"));
|
||
server_h_layout.addWidget(self.host_input, 1);
|
||
server_h_layout.addWidget(QLabel("端口:"));
|
||
server_h_layout.addWidget(self.port_input);
|
||
server_h_layout.addWidget(self.connect_button);
|
||
self.server_group.setLayout(server_h_layout)
|
||
workflow_layout, workflow_controls_layout = QVBoxLayout(), QHBoxLayout();
|
||
workflow_controls_layout.addWidget(self.workflow_combo, 1);
|
||
workflow_controls_layout.addWidget(self.refresh_button);
|
||
workflow_controls_layout.addWidget(self.settings_button);
|
||
workflow_layout.addLayout(workflow_controls_layout);
|
||
workflow_layout.addWidget(self.workflow_tip_label);
|
||
self.workflow_group.setLayout(workflow_layout)
|
||
exec_form_layout = QFormLayout();
|
||
clothing_layout = QHBoxLayout();
|
||
clothing_layout.addWidget(self.clothing_dir_input, 1);
|
||
clothing_layout.addWidget(self.clothing_dir_button);
|
||
exec_form_layout.addRow("衣物文件夹:", clothing_layout);
|
||
self.exec_group.setLayout(exec_form_layout)
|
||
output_form_layout = QFormLayout();
|
||
output_layout = QHBoxLayout();
|
||
output_layout.addWidget(self.output_dir_input, 1);
|
||
output_layout.addWidget(self.output_dir_button);
|
||
output_form_layout.addRow("输出文件夹:", output_layout);
|
||
self.output_group.setLayout(output_form_layout)
|
||
action_layout = QHBoxLayout();
|
||
action_layout.addStretch();
|
||
action_layout.addWidget(self.execute_button);
|
||
action_layout.addWidget(self.stop_button);
|
||
self.action_group.setLayout(action_layout)
|
||
log_layout = QVBoxLayout();
|
||
log_layout.addWidget(self.log_box);
|
||
self.log_group.setLayout(log_layout)
|
||
right_v_layout.addWidget(self.server_group);
|
||
right_v_layout.addWidget(self.workflow_group);
|
||
right_v_layout.addWidget(self.exec_group);
|
||
right_v_layout.addWidget(self.output_group);
|
||
right_v_layout.addWidget(self.action_group);
|
||
right_v_layout.addWidget(self.log_group, 1)
|
||
main_h_layout = QHBoxLayout();
|
||
main_h_layout.setSpacing(15);
|
||
main_h_layout.addWidget(self.model_group, 1);
|
||
main_h_layout.addLayout(right_v_layout, 3)
|
||
central_widget = QWidget();
|
||
central_widget.setLayout(main_h_layout);
|
||
self.setCentralWidget(central_widget)
|
||
|
||
def _connect_signals(self):
|
||
self.add_model_button.clicked.connect(self._add_model);
|
||
self.delete_model_button.clicked.connect(self._delete_model);
|
||
self.model_list_widget.itemDoubleClicked.connect(self.show_model_preview)
|
||
self.connect_button.clicked.connect(self.connect_and_fetch_workflows);
|
||
self.refresh_button.clicked.connect(self.fetch_workflows);
|
||
self.workflow_combo.currentTextChanged.connect(self.fetch_workflow_details)
|
||
self.settings_button.clicked.connect(self.open_settings_dialog)
|
||
self.clothing_dir_button.clicked.connect(
|
||
lambda: self._select_directory(self.clothing_dir_input, "选择衣物文件夹"))
|
||
self.output_dir_button.clicked.connect(lambda: self._select_directory(self.output_dir_input, "选择输出文件夹"))
|
||
self.execute_button.clicked.connect(self.start_execution);
|
||
self.stop_button.clicked.connect(self.stop_execution)
|
||
|
||
def _update_ui_state(self, is_running=False):
|
||
# FIX 1: Finer control over UI elements
|
||
self.server_group.setEnabled(not is_running)
|
||
self.workflow_group.setEnabled(not is_running)
|
||
self.exec_group.setEnabled(not is_running)
|
||
self.output_group.setEnabled(not is_running)
|
||
self.add_model_button.setEnabled(not is_running)
|
||
self.delete_model_button.setEnabled(not is_running)
|
||
self.model_list_widget.setEnabled(True) # Always enabled for preview
|
||
|
||
self.execute_button.setEnabled(not is_running and self.api is not None)
|
||
self.stop_button.setEnabled(is_running)
|
||
|
||
def _load_settings(self):
|
||
self.host_input.setText(self.config.get("server", "host", "127.0.0.1"));
|
||
self.port_input.setText(self.config.get("server", "port", "8188"));
|
||
self.clothing_dir_input.setText(self.config.get("paths", "clothing", ""));
|
||
self.output_dir_input.setText(self.config.get("paths", "output", ""));
|
||
self.log_message("欢迎使用ComfyUI智能试衣控制器！配置已从 config.ini 加载。")
|
||
|
||
def _save_settings(self):
|
||
self.config.set("server", "host", self.host_input.text());
|
||
self.config.set("server", "port", self.port_input.text());
|
||
self.config.set("paths", "clothing", self.clothing_dir_input.text());
|
||
self.config.set("paths", "output", self.output_dir_input.text());
|
||
self._save_models()
|
||
self.config.save()
|
||
|
||
def closeEvent(self, event):
|
||
reply = QMessageBox.question(self, '确认退出', '您确定要退出程序吗？\n所有配置将自动保存。',
|
||
QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
|
||
if reply == QMessageBox.Yes:
|
||
self._save_settings(); self.log_message("配置已保存，程序退出。"); event.accept()
|
||
else:
|
||
event.ignore()
|
||
|
||
def log_message(self, message):
|
||
self.log_box.appendPlainText(f"[{time.strftime('%H:%M:%S')}] {message}"); self.status_bar.showMessage(
|
||
message.split('\n')[-1], 5000)
|
||
|
||
# FIX 2: Slot for status bar updates
|
||
def _update_status_bar(self, message):
|
||
self.status_bar.showMessage(message, 0) # 0 means persistent message
|
||
|
||
def show_model_preview(self, item):
|
||
image_path = item.data(Qt.UserRole)
|
||
if image_path and os.path.exists(image_path):
|
||
ImagePreviewDialog(image_path, self).exec()
|
||
else:
|
||
self.show_error("文件错误", f"找不到图片文件：\n{image_path}")
|
||
|
||
def _add_model(self):
|
||
file_path, _ = QFileDialog.getOpenFileName(self, "选择模特图片", self.config.get("paths", "last_used", ""),
|
||
"图片文件 (*.png *.jpg *.jpeg *.webp)");
|
||
if file_path: item = QListWidgetItem(QIcon(file_path), os.path.basename(file_path)); item.setData(Qt.UserRole,
|
||
file_path); self.model_list_widget.addItem(
|
||
item); self.log_message(f"已添加模特: {os.path.basename(file_path)}")
|
||
|
||
def _delete_model(self):
|
||
selected_items = self.model_list_widget.selectedItems()
|
||
if not selected_items: return self.show_error("操作无效", "请先在左侧列表中选择一个要删除的模特。")
|
||
for item in selected_items: self.model_list_widget.takeItem(self.model_list_widget.row(item)); self.log_message(
|
||
f"已删除模特: {item.text()}")
|
||
|
||
def _load_models(self):
|
||
models_str = self.config.get("models", "list", "");
|
||
if models_str:
|
||
models = models_str.splitlines()
|
||
for path in models:
|
||
if path and os.path.exists(path): item = QListWidgetItem(QIcon(path),
|
||
os.path.basename(path)); item.setData(
|
||
Qt.UserRole, path); self.model_list_widget.addItem(item)
|
||
self.log_message(f"已加载 {len(models)} 个模特。")
|
||
|
||
def _save_models(self):
|
||
models = [self.model_list_widget.item(i).data(Qt.UserRole) for i in range(self.model_list_widget.count())]
|
||
self.config.set("models", "list", "\n".join(models))
|
||
|
||
def connect_and_fetch_workflows(self):
|
||
host, port = self.host_input.text(), self.port_input.text();
|
||
self.log_message(f"正在连接到 http://{host}:{port}...");
|
||
self.api = ComfyAPI(host, port);
|
||
self.fetch_workflows()
|
||
|
||
def fetch_workflows(self):
|
||
if not self.api: return
|
||
self.log_message("正在获取工作流列表...")
|
||
workflows_data, error = self.api.get_workflows();
|
||
self.workflow_combo.clear()
|
||
if error:
|
||
self.show_error("连接错误", f"无法获取工作流列表。\n原因: {error}"); self.api = None
|
||
elif workflows_data:
|
||
workflows = [item['path'] for item in workflows_data if item.get('path', '').endswith('.json')]
|
||
if workflows:
|
||
self.workflow_combo.addItems(workflows); self.log_message(f"成功找到 {len(workflows)} 个工作流。")
|
||
else:
|
||
self.log_message("连接成功，但在workflows文件夹中未找到.json文件。")
|
||
self._update_ui_state()
|
||
|
||
def fetch_workflow_details(self, workflow_name):
|
||
self.current_workflow_prompt, self.current_workflow_graph = None, None
|
||
if not workflow_name or not self.api: self._update_ui_state(); return
|
||
self.log_message(f"正在加载工作流: {workflow_name}")
|
||
workflow_graph, error = self.api.get_workflow_details(workflow_name)
|
||
if error: self._update_ui_state(); return self.show_error("加载失败",
|
||
f"无法获取工作流 '{workflow_name}' 的详细信息。\n\n原因: {error}")
|
||
if workflow_graph:
|
||
self.current_workflow_graph = copy.deepcopy(workflow_graph)
|
||
original_node_count = len(workflow_graph['nodes'])
|
||
workflow_graph['nodes'] = [node for node in workflow_graph['nodes'] if node.get('type') != 'Note']
|
||
removed_count = original_node_count - len(workflow_graph['nodes'])
|
||
if removed_count > 0: self.log_message(f"已自动移除 {removed_count} 个 'Note' 节点。")
|
||
self.log_message("正在转换工作流为可执行格式...")
|
||
self.current_workflow_prompt = convert_graph_to_prompt(workflow_graph)
|
||
if self.current_workflow_prompt:
|
||
self.log_message("转换成功，工作流已就绪。");
|
||
settings_key = f"workflow/{workflow_name.replace('/', '_').replace('.', '_')}"
|
||
self.workflow_settings = json.loads(self.config.get("workflows", settings_key, "{}"))
|
||
else:
|
||
self.show_error("转换失败", "无法将工作流转换为API格式。")
|
||
self._update_ui_state()
|
||
|
||
def _select_directory(self, line_edit, title="选择文件夹"):
|
||
start_dir = line_edit.text() or self.config.get("paths", "last_used", "")
|
||
directory = QFileDialog.getExistingDirectory(self, title, start_dir);
|
||
if directory: line_edit.setText(directory); self.config.set("paths", "last_used", directory)
|
||
|
||
def open_settings_dialog(self):
|
||
if not self.current_workflow_prompt: return self.show_error("未加载工作流", "请先选择一个工作流。")
|
||
dialog = SettingsDialog(self.current_workflow_prompt, self.workflow_settings, self)
|
||
if dialog.exec():
|
||
self.workflow_settings = dialog.get_settings();
|
||
workflow_name = self.workflow_combo.currentText()
|
||
settings_key = f"workflow/{workflow_name.replace('/', '_').replace('.', '_')}";
|
||
self.config.set("workflows", settings_key, json.dumps(self.workflow_settings))
|
||
self.log_message(f"已为'{workflow_name}'保存节点设置。")
|
||
|
||
def start_execution(self):
|
||
if self.model_list_widget.currentItem() is None: return self.show_error("缺少输入",
|
||
"请在左侧模特库中选择一位模特。")
|
||
if not self.clothing_dir_input.text(): return self.show_error("缺少输入", "请选择包含衣物图片的文件夹。")
|
||
if not self.output_dir_input.text(): return self.show_error("缺少输入", "请选择输出文件夹。")
|
||
if not self.current_workflow_prompt: return self.show_error("缺少工作流", "请选择一个有效的工作流。")
|
||
ws = self.workflow_settings
|
||
if not all(
|
||
[ws.get('model_node_id'), ws.get('clothing_node_id'), ws.get('output_node_id')]): return self.show_error(
|
||
"配置不完整", "请点击设置按钮(⚙️)来指定模特、衣物和输出节点。")
|
||
|
||
self._update_ui_state(is_running=True)
|
||
self.execution_thread = QThread()
|
||
self.execution_worker = ExecutionWorker(api=self.api, workflow_prompt=self.current_workflow_prompt,
|
||
workflow_graph=self.current_workflow_graph,
|
||
model_image_path=self.model_list_widget.currentItem().data(Qt.UserRole),
|
||
model_node_id=ws['model_node_id'],
|
||
clothing_dir=self.clothing_dir_input.text(),
|
||
clothing_node_id=ws['clothing_node_id'],
|
||
output_dir=self.output_dir_input.text(),
|
||
output_node_id=ws['output_node_id'])
|
||
self.execution_worker.moveToThread(self.execution_thread);
|
||
self.execution_worker.log.connect(self.log_message)
|
||
self.execution_worker.progress.connect(self.update_progress)
|
||
self.execution_worker.finished.connect(self.execution_finished)
|
||
self.execution_worker.status_update.connect(self._update_status_bar) # FIX 2
|
||
self.execution_thread.started.connect(self.execution_worker.run)
|
||
self.execution_thread.start();
|
||
self.log_message("=== 开始批量换装任务 ===")
|
||
|
||
def stop_execution(self):
|
||
if self.execution_worker:
|
||
self.execution_worker.stop()
|
||
else:
|
||
self.log_message("没有正在执行的任务可停止。")
|
||
|
||
def update_progress(self, current, total):
|
||
self.progress_bar.setMaximum(total); self.progress_bar.setValue(current)
|
||
|
||
def execution_finished(self, success, message):
|
||
self.log_message(f"=== {message} ===");
|
||
if not success and "中止" not in message: self.show_error("执行失败", message)
|
||
self.progress_bar.setValue(0);
|
||
self.status_bar.clearMessage() # Clear status bar on finish
|
||
if self.execution_thread:
|
||
if self.execution_thread.isRunning(): self.execution_thread.quit()
|
||
self.execution_thread.wait(1000)
|
||
self.execution_thread, self.execution_worker = None, None
|
||
self._update_ui_state(is_running=False)
|
||
|
||
def show_error(self, title, message):
|
||
self.log_message(f"错误: {title} - {message}");
|
||
QMessageBox.critical(self, title, message)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
app = QApplication(sys.argv)
|
||
LIGHT_STYLE_SHEET = """
|
||
QMainWindow, QDialog { background-color: #f4f6f8; color: #34495e; font-family: "Microsoft YaHei", "Segoe UI", "Roboto", sans-serif; }
|
||
QGroupBox { background-color: #ffffff; border: 1px solid #dcdfe6; border-radius: 8px; margin-top: 10px; padding: 15px; }
|
||
QGroupBox::title { subcontrol-origin: margin; subcontrol-position: top left; padding: 5px 15px; background-color: #5dade2; color: white; border-top-left-radius: 8px; border-bottom-right-radius: 8px; font-weight: bold; }
|
||
QLabel { font-size: 13px; font-weight: bold; color: #34495e; padding: 5px 0; }
|
||
QLineEdit, QPlainTextEdit, QComboBox { background-color: #ffffff; border: 1px solid #dcdfe6; border-radius: 4px; padding: 6px; color: #34495e; font-size: 13px; }
|
||
QLineEdit:focus, QPlainTextEdit:focus, QComboBox:focus { border-color: #5dade2; }
|
||
QPushButton { background-color: #5dade2; color: white; border: none; border-radius: 4px; padding: 8px 16px; font-size: 13px; font-weight: bold; }
|
||
QPushButton:hover { background-color: #85c1e9; }
|
||
QPushButton:pressed { background-color: #2e86c1; }
|
||
QPushButton:disabled { background-color: #bdc3c7; color: #7f8c8d; }
|
||
QComboBox::drop-down { border: none; }
|
||
QComboBox QAbstractItemView { background-color: #ffffff; border: 1px solid #dcdfe6; selection-background-color: #5dade2; selection-color: white; }
|
||
QListWidget { background-color: #ffffff; border: 1px solid #dcdfe6; border-radius: 4px; }
|
||
QListWidget::item { padding: 8px; border-bottom: 1px solid #f4f6f8; }
|
||
QListWidget::item:hover { background-color: #ecf5ff; }
|
||
QListWidget::item:selected { background-color: #5dade2; color: white; }
|
||
QProgressBar { border: 1px solid #dcdfe6; border-radius: 4px; text-align: center; background-color: #e9ecef; color: #495057; }
|
||
QProgressBar::chunk { background-color: #28a745; border-radius: 3px; }
|
||
QStatusBar { background-color: #ffffff; border-top: 1px solid #dcdfe6; }
|
||
"""
|
||
app.setStyleSheet(LIGHT_STYLE_SHEET)
|
||
window = ComfyUIController()
|
||
window.show()
|
||
sys.exit(app.exec()) |