390 lines
16 KiB
Python
390 lines
16 KiB
Python
|
"""
|
|||
|
|
Правая панель: превью, параметры генерации, промты, редакторы
|
|||
|
|
workflow/prompt JSON, рейтинг, избранное и теги.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import random
|
|||
|
|
from typing import Optional
|
|||
|
|
|
|||
|
|
from PyQt6.QtWidgets import (
|
|||
|
|
QWidget, QVBoxLayout, QHBoxLayout, QLabel, QTabWidget,
|
|||
|
|
QTextEdit, QLineEdit, QPushButton, QFormLayout, QComboBox,
|
|||
|
|
QMessageBox, QApplication, QCheckBox
|
|||
|
|
)
|
|||
|
|
from PyQt6.QtGui import QPixmap
|
|||
|
|
from PyQt6.QtCore import Qt, pyqtSignal
|
|||
|
|
|
|||
|
|
from comfy_studio.core.i18n import tr
|
|||
|
|
from comfy_studio.parsers.metadata_parser import MetadataParser
|
|||
|
|
from comfy_studio.ui.widgets.translatable_text_edit import (
|
|||
|
|
TranslatableTextEdit, TranslatableLineEdit,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TagChip(QWidget):
|
|||
|
|
"""Тег с кнопкой удаления."""
|
|||
|
|
removed = pyqtSignal(str)
|
|||
|
|
|
|||
|
|
def __init__(self, tag_name: str, parent=None):
|
|||
|
|
super().__init__(parent)
|
|||
|
|
self.tag_name = tag_name
|
|||
|
|
layout = QHBoxLayout(self)
|
|||
|
|
layout.setContentsMargins(6, 2, 2, 2)
|
|||
|
|
layout.setSpacing(4)
|
|||
|
|
|
|||
|
|
label = QLabel(tag_name)
|
|||
|
|
label.setStyleSheet("background: transparent; color: #cde5ff;")
|
|||
|
|
layout.addWidget(label)
|
|||
|
|
|
|||
|
|
btn = QPushButton("×")
|
|||
|
|
btn.setFixedSize(16, 16)
|
|||
|
|
btn.setStyleSheet(
|
|||
|
|
"QPushButton { background: transparent; border: none; color: #90a4ae;"
|
|||
|
|
" font-weight: bold; padding: 0; }"
|
|||
|
|
"QPushButton:hover { color: #ef5350; }"
|
|||
|
|
)
|
|||
|
|
btn.clicked.connect(lambda: self.removed.emit(self.tag_name))
|
|||
|
|
layout.addWidget(btn)
|
|||
|
|
|
|||
|
|
self.setStyleSheet(
|
|||
|
|
"TagChip { background-color: #1a3a5c; border: 1px solid #2a5a8c;"
|
|||
|
|
" border-radius: 9px; }"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class FlowTagContainer(QWidget):
|
|||
|
|
"""Контейнер тегов в строку с переносом."""
|
|||
|
|
tag_removed = pyqtSignal(str)
|
|||
|
|
|
|||
|
|
def __init__(self, parent=None):
|
|||
|
|
super().__init__(parent)
|
|||
|
|
self._layout = QHBoxLayout(self)
|
|||
|
|
self._layout.setContentsMargins(0, 0, 0, 0)
|
|||
|
|
self._layout.setSpacing(4)
|
|||
|
|
self._layout.addStretch()
|
|||
|
|
|
|||
|
|
def set_tags(self, tags: list):
|
|||
|
|
# Удаляем старые чипы
|
|||
|
|
while self._layout.count() > 1:
|
|||
|
|
item = self._layout.takeAt(0)
|
|||
|
|
if item.widget():
|
|||
|
|
item.widget().deleteLater()
|
|||
|
|
for t in tags:
|
|||
|
|
chip = TagChip(t)
|
|||
|
|
chip.removed.connect(self.tag_removed.emit)
|
|||
|
|
self._layout.insertWidget(self._layout.count() - 1, chip)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class RightPanel(QWidget):
|
|||
|
|
save_clicked = pyqtSignal(int, dict)
|
|||
|
|
send_to_comfy_clicked = pyqtSignal(str)
|
|||
|
|
favorite_toggled = pyqtSignal(int, bool)
|
|||
|
|
tag_added = pyqtSignal(int, str)
|
|||
|
|
tag_removed = pyqtSignal(int, str)
|
|||
|
|
json_saved = pyqtSignal(int, str, str) # (file_id, field: "workflow_json"|"prompt_json", json_text)
|
|||
|
|
|
|||
|
|
def __init__(self, parent=None):
|
|||
|
|
super().__init__(parent)
|
|||
|
|
self.current_file_id = None
|
|||
|
|
self.current_filepath = None
|
|||
|
|
self.current_raw_prompt = None
|
|||
|
|
self.init_ui()
|
|||
|
|
|
|||
|
|
def init_ui(self):
|
|||
|
|
layout = QVBoxLayout(self)
|
|||
|
|
layout.setContentsMargins(5, 5, 5, 5)
|
|||
|
|
|
|||
|
|
self.preview_label = QLabel(tr("no_image_selected"))
|
|||
|
|
self.preview_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
|
|||
|
|
self.preview_label.setMinimumHeight(250)
|
|||
|
|
self.preview_label.setStyleSheet(
|
|||
|
|
"QLabel { background-color: #0f0f0f; border: 1px solid #2d2d2d; border-radius: 4px; }"
|
|||
|
|
)
|
|||
|
|
layout.addWidget(self.preview_label)
|
|||
|
|
|
|||
|
|
# Избранное
|
|||
|
|
self.favorite_btn = QPushButton(tr("favorite"))
|
|||
|
|
self.favorite_btn.setCheckable(True)
|
|||
|
|
self.favorite_btn.clicked.connect(self._on_favorite_clicked)
|
|||
|
|
layout.addWidget(self.favorite_btn)
|
|||
|
|
|
|||
|
|
self.tabs = QTabWidget()
|
|||
|
|
|
|||
|
|
# ── Вкладка Generation ──
|
|||
|
|
self.gen_tab = QWidget()
|
|||
|
|
gen_layout = QVBoxLayout(self.gen_tab)
|
|||
|
|
gen_layout.setContentsMargins(5, 5, 5, 5)
|
|||
|
|
|
|||
|
|
form_layout = QFormLayout()
|
|||
|
|
self.model_input = QLineEdit()
|
|||
|
|
self.model_input.setReadOnly(True)
|
|||
|
|
form_layout.addRow(tr("model"), self.model_input)
|
|||
|
|
|
|||
|
|
self.sampler_input = QLineEdit()
|
|||
|
|
self.sampler_input.setReadOnly(True)
|
|||
|
|
form_layout.addRow(tr("sampler"), self.sampler_input)
|
|||
|
|
|
|||
|
|
params_row = QHBoxLayout()
|
|||
|
|
self.seed_input = QLineEdit()
|
|||
|
|
self.seed_input.setReadOnly(True)
|
|||
|
|
self.steps_input = QLineEdit()
|
|||
|
|
self.steps_input.setReadOnly(True)
|
|||
|
|
self.cfg_input = QLineEdit()
|
|||
|
|
self.cfg_input.setReadOnly(True)
|
|||
|
|
|
|||
|
|
params_row.addWidget(QLabel("Seed:"))
|
|||
|
|
params_row.addWidget(self.seed_input)
|
|||
|
|
params_row.addWidget(QLabel("Steps:"))
|
|||
|
|
params_row.addWidget(self.steps_input)
|
|||
|
|
params_row.addWidget(QLabel("CFG:"))
|
|||
|
|
params_row.addWidget(self.cfg_input)
|
|||
|
|
form_layout.addRow(params_row)
|
|||
|
|
|
|||
|
|
self.rating_combo = QComboBox()
|
|||
|
|
self.rating_combo.addItems([tr("no_rating"), "★", "★★", "★★★", "★★★★", "★★★★★"])
|
|||
|
|
form_layout.addRow(tr("rating"), self.rating_combo)
|
|||
|
|
gen_layout.addLayout(form_layout)
|
|||
|
|
|
|||
|
|
# Теги
|
|||
|
|
gen_layout.addWidget(QLabel(tr("tags")))
|
|||
|
|
self.tag_container = FlowTagContainer()
|
|||
|
|
self.tag_container.tag_removed.connect(self._on_tag_removed)
|
|||
|
|
gen_layout.addWidget(self.tag_container)
|
|||
|
|
|
|||
|
|
self.tag_input = TranslatableLineEdit()
|
|||
|
|
self.tag_input.setPlaceholderText(tr("add_tag_placeholder"))
|
|||
|
|
self.tag_input.returnPressed.connect(self._on_tag_entered)
|
|||
|
|
gen_layout.addWidget(self.tag_input)
|
|||
|
|
|
|||
|
|
gen_layout.addWidget(QLabel(tr("positive_prompt")))
|
|||
|
|
self.positive_text = TranslatableTextEdit()
|
|||
|
|
gen_layout.addWidget(self.positive_text)
|
|||
|
|
|
|||
|
|
gen_layout.addWidget(QLabel(tr("negative_prompt")))
|
|||
|
|
self.negative_text = TranslatableTextEdit()
|
|||
|
|
gen_layout.addWidget(self.negative_text)
|
|||
|
|
|
|||
|
|
self.save_btn = QPushButton(tr("save_changes"))
|
|||
|
|
self.save_btn.clicked.connect(self._on_save_clicked)
|
|||
|
|
gen_layout.addWidget(self.save_btn)
|
|||
|
|
|
|||
|
|
# ── Вкладка Workflow (редактируемая) ──
|
|||
|
|
self.workflow_tab = QWidget()
|
|||
|
|
wf_layout = QVBoxLayout(self.workflow_tab)
|
|||
|
|
wf_layout.setContentsMargins(5, 5, 5, 5)
|
|||
|
|
self.workflow_text = TranslatableTextEdit()
|
|||
|
|
wf_layout.addWidget(self.workflow_text)
|
|||
|
|
|
|||
|
|
wf_btn_layout = QHBoxLayout()
|
|||
|
|
self.format_wf_btn = QPushButton(tr("format_json"))
|
|||
|
|
self.format_wf_btn.clicked.connect(lambda: self._format_json(self.workflow_text))
|
|||
|
|
self.save_wf_btn = QPushButton(tr("save_workflow"))
|
|||
|
|
self.save_wf_btn.clicked.connect(
|
|||
|
|
lambda: self._save_json_field("workflow_json", self.workflow_text))
|
|||
|
|
wf_btn_layout.addWidget(self.format_wf_btn)
|
|||
|
|
wf_btn_layout.addWidget(self.save_wf_btn)
|
|||
|
|
wf_layout.addLayout(wf_btn_layout)
|
|||
|
|
|
|||
|
|
# ── Вкладка Prompt JSON (редактируемая) ──
|
|||
|
|
self.prompt_tab = QWidget()
|
|||
|
|
pr_layout = QVBoxLayout(self.prompt_tab)
|
|||
|
|
pr_layout.setContentsMargins(5, 5, 5, 5)
|
|||
|
|
self.prompt_json_text = TranslatableTextEdit()
|
|||
|
|
pr_layout.addWidget(self.prompt_json_text)
|
|||
|
|
|
|||
|
|
pr_btn_layout = QHBoxLayout()
|
|||
|
|
self.format_pr_btn = QPushButton(tr("format_json"))
|
|||
|
|
self.format_pr_btn.clicked.connect(lambda: self._format_json(self.prompt_json_text))
|
|||
|
|
self.save_pr_btn = QPushButton(tr("save_prompt_json"))
|
|||
|
|
self.save_pr_btn.clicked.connect(
|
|||
|
|
lambda: self._save_json_field("prompt_json", self.prompt_json_text))
|
|||
|
|
pr_btn_layout.addWidget(self.format_pr_btn)
|
|||
|
|
pr_btn_layout.addWidget(self.save_pr_btn)
|
|||
|
|
pr_layout.addLayout(pr_btn_layout)
|
|||
|
|
|
|||
|
|
self.tabs.addTab(self.gen_tab, tr("tab_generation"))
|
|||
|
|
self.tabs.addTab(self.workflow_tab, tr("tab_workflow"))
|
|||
|
|
self.tabs.addTab(self.prompt_tab, tr("tab_prompt_json"))
|
|||
|
|
layout.addWidget(self.tabs)
|
|||
|
|
|
|||
|
|
# Кнопки копирования
|
|||
|
|
export_layout = QHBoxLayout()
|
|||
|
|
self.copy_pos_btn = QPushButton(tr("copy_positive"))
|
|||
|
|
self.copy_pos_btn.clicked.connect(
|
|||
|
|
lambda: QApplication.clipboard().setText(self.positive_text.toPlainText()))
|
|||
|
|
self.copy_neg_btn = QPushButton(tr("copy_negative"))
|
|||
|
|
self.copy_neg_btn.clicked.connect(
|
|||
|
|
lambda: QApplication.clipboard().setText(self.negative_text.toPlainText()))
|
|||
|
|
self.copy_wf_btn = QPushButton(tr("copy_workflow"))
|
|||
|
|
self.copy_wf_btn.clicked.connect(
|
|||
|
|
lambda: QApplication.clipboard().setText(self.workflow_text.toPlainText()))
|
|||
|
|
export_layout.addWidget(self.copy_pos_btn)
|
|||
|
|
export_layout.addWidget(self.copy_neg_btn)
|
|||
|
|
export_layout.addWidget(self.copy_wf_btn)
|
|||
|
|
layout.addLayout(export_layout)
|
|||
|
|
|
|||
|
|
send_layout = QHBoxLayout()
|
|||
|
|
self.send_btn = QPushButton(tr("send_to_comfy"))
|
|||
|
|
self.send_btn.clicked.connect(self._on_send_clicked)
|
|||
|
|
send_layout.addWidget(self.send_btn, 1)
|
|||
|
|
|
|||
|
|
self.randomize_seed_check = QCheckBox(tr("randomize_seed"))
|
|||
|
|
self.randomize_seed_check.setChecked(True)
|
|||
|
|
send_layout.addWidget(self.randomize_seed_check)
|
|||
|
|
layout.addLayout(send_layout)
|
|||
|
|
|
|||
|
|
# ── Отображение данных ──
|
|||
|
|
|
|||
|
|
def display_metadata(self, file_details: Optional[dict]):
|
|||
|
|
if not file_details:
|
|||
|
|
self.clear_fields()
|
|||
|
|
return
|
|||
|
|
self.current_file_id = file_details["id"]
|
|||
|
|
self.current_filepath = file_details["filepath"]
|
|||
|
|
self.current_raw_prompt = file_details.get("prompt_json")
|
|||
|
|
|
|||
|
|
pixmap = QPixmap(self.current_filepath)
|
|||
|
|
if not pixmap.isNull():
|
|||
|
|
scaled = pixmap.scaled(
|
|||
|
|
self.preview_label.width(), self.preview_label.height(),
|
|||
|
|
Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.SmoothTransformation
|
|||
|
|
)
|
|||
|
|
self.preview_label.setPixmap(scaled)
|
|||
|
|
else:
|
|||
|
|
self.preview_label.setText(tr("preview_error"))
|
|||
|
|
|
|||
|
|
self.favorite_btn.setChecked(bool(file_details.get("favorite")))
|
|||
|
|
|
|||
|
|
self.model_input.setText(file_details.get("model_name") or tr("not_specified"))
|
|||
|
|
self.sampler_input.setText(file_details.get("sampler") or tr("not_specified"))
|
|||
|
|
self.seed_input.setText(str(file_details.get("seed") or ""))
|
|||
|
|
self.steps_input.setText(str(file_details.get("steps") or ""))
|
|||
|
|
self.cfg_input.setText(str(file_details.get("cfg") or ""))
|
|||
|
|
|
|||
|
|
rating = file_details.get("rating", 0)
|
|||
|
|
self.rating_combo.setCurrentIndex(rating if 0 <= rating <= 5 else 0)
|
|||
|
|
|
|||
|
|
self.tag_container.set_tags(file_details.get("tags", []))
|
|||
|
|
|
|||
|
|
self.positive_text.setPlainText(file_details.get("positive_prompt") or "")
|
|||
|
|
self.negative_text.setPlainText(file_details.get("negative_prompt") or "")
|
|||
|
|
|
|||
|
|
self._display_json(self.workflow_text, file_details.get("workflow_json"),
|
|||
|
|
tr("workflow_not_found"))
|
|||
|
|
self._display_json(self.prompt_json_text, file_details.get("prompt_json"),
|
|||
|
|
tr("prompt_json_not_found"))
|
|||
|
|
has_wf = bool(file_details.get("workflow_json"))
|
|||
|
|
has_pr = bool(file_details.get("prompt_json"))
|
|||
|
|
self.save_wf_btn.setEnabled(True)
|
|||
|
|
self.save_pr_btn.setEnabled(True)
|
|||
|
|
self.format_wf_btn.setEnabled(has_wf)
|
|||
|
|
self.format_pr_btn.setEnabled(has_pr)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _display_json(widget: QTextEdit, raw: Optional[str], placeholder: str):
|
|||
|
|
if raw:
|
|||
|
|
try:
|
|||
|
|
widget.setPlainText(json.dumps(json.loads(raw), indent=2, ensure_ascii=False))
|
|||
|
|
except Exception:
|
|||
|
|
widget.setPlainText(raw)
|
|||
|
|
else:
|
|||
|
|
widget.clear()
|
|||
|
|
widget.setPlaceholderText(placeholder)
|
|||
|
|
|
|||
|
|
def clear_fields(self):
|
|||
|
|
self.current_file_id = None
|
|||
|
|
self.current_filepath = None
|
|||
|
|
self.current_raw_prompt = None
|
|||
|
|
self.preview_label.setText(tr("no_image_selected"))
|
|||
|
|
self.preview_label.setPixmap(QPixmap())
|
|||
|
|
self.favorite_btn.setChecked(False)
|
|||
|
|
self.model_input.clear()
|
|||
|
|
self.sampler_input.clear()
|
|||
|
|
self.seed_input.clear()
|
|||
|
|
self.steps_input.clear()
|
|||
|
|
self.cfg_input.clear()
|
|||
|
|
self.rating_combo.setCurrentIndex(0)
|
|||
|
|
self.tag_container.set_tags([])
|
|||
|
|
self.tag_input.clear()
|
|||
|
|
self.positive_text.clear()
|
|||
|
|
self.negative_text.clear()
|
|||
|
|
self.workflow_text.clear()
|
|||
|
|
self.prompt_json_text.clear()
|
|||
|
|
self.save_wf_btn.setEnabled(False)
|
|||
|
|
self.save_pr_btn.setEnabled(False)
|
|||
|
|
self.format_wf_btn.setEnabled(False)
|
|||
|
|
self.format_pr_btn.setEnabled(False)
|
|||
|
|
|
|||
|
|
# ── Редактор JSON ──
|
|||
|
|
|
|||
|
|
def _format_json(self, widget: QTextEdit):
|
|||
|
|
"""Валидирует и красиво форматирует JSON в редакторе."""
|
|||
|
|
text = widget.toPlainText().strip()
|
|||
|
|
if not text:
|
|||
|
|
return
|
|||
|
|
try:
|
|||
|
|
widget.setPlainText(json.dumps(json.loads(text), indent=2, ensure_ascii=False))
|
|||
|
|
except json.JSONDecodeError as e:
|
|||
|
|
QMessageBox.warning(self, tr("error"), tr("invalid_json", err=str(e)))
|
|||
|
|
|
|||
|
|
def _save_json_field(self, field: str, widget: QTextEdit):
|
|||
|
|
"""Проверяет JSON и передаёт его на сохранение в БД."""
|
|||
|
|
if self.current_file_id is None:
|
|||
|
|
return
|
|||
|
|
text = widget.toPlainText().strip()
|
|||
|
|
if text:
|
|||
|
|
try:
|
|||
|
|
# Сохраняем компактную форму, как пишет ComfyUI
|
|||
|
|
text = json.dumps(json.loads(text), ensure_ascii=False)
|
|||
|
|
except json.JSONDecodeError as e:
|
|||
|
|
QMessageBox.warning(self, tr("error"), tr("invalid_json", err=str(e)))
|
|||
|
|
return
|
|||
|
|
if field == "prompt_json":
|
|||
|
|
self.current_raw_prompt = text or None
|
|||
|
|
self.json_saved.emit(self.current_file_id, field, text)
|
|||
|
|
|
|||
|
|
# ── Обработчики ──
|
|||
|
|
|
|||
|
|
def _on_save_clicked(self):
|
|||
|
|
if self.current_file_id is None:
|
|||
|
|
return
|
|||
|
|
payload = {
|
|||
|
|
"positive_prompt": self.positive_text.toPlainText(),
|
|||
|
|
"negative_prompt": self.negative_text.toPlainText(),
|
|||
|
|
"rating": self.rating_combo.currentIndex()
|
|||
|
|
}
|
|||
|
|
self.save_clicked.emit(self.current_file_id, payload)
|
|||
|
|
|
|||
|
|
def _on_send_clicked(self):
|
|||
|
|
if not self.current_raw_prompt:
|
|||
|
|
QMessageBox.warning(self, tr("error"), tr("no_prompt_graph"))
|
|||
|
|
return
|
|||
|
|
# Внедряем текущие (возможно отредактированные) промты в граф.
|
|||
|
|
# Без нового seed ComfyUI отдаёт идентичный граф из кэша
|
|||
|
|
# и не генерирует ("Prompt executed in 0.00 seconds").
|
|||
|
|
seed = random.randint(0, 2**48) if self.randomize_seed_check.isChecked() else None
|
|||
|
|
updated_prompt = MetadataParser.inject_parameters(
|
|||
|
|
self.current_raw_prompt,
|
|||
|
|
positive=self.positive_text.toPlainText(),
|
|||
|
|
negative=self.negative_text.toPlainText(),
|
|||
|
|
seed=seed,
|
|||
|
|
)
|
|||
|
|
self.send_to_comfy_clicked.emit(updated_prompt)
|
|||
|
|
|
|||
|
|
def _on_favorite_clicked(self, checked: bool):
|
|||
|
|
if self.current_file_id is not None:
|
|||
|
|
self.favorite_toggled.emit(self.current_file_id, checked)
|
|||
|
|
|
|||
|
|
def _on_tag_entered(self):
|
|||
|
|
if self.current_file_id is None:
|
|||
|
|
return
|
|||
|
|
tag = self.tag_input.text().strip()
|
|||
|
|
if tag:
|
|||
|
|
self.tag_added.emit(self.current_file_id, tag)
|
|||
|
|
self.tag_input.clear()
|
|||
|
|
|
|||
|
|
def _on_tag_removed(self, tag: str):
|
|||
|
|
if self.current_file_id is not None:
|
|||
|
|
self.tag_removed.emit(self.current_file_id, tag)
|