| | from PIL.ImageQt import ImageQt
|
| | from PyQt5 import QtCore
|
| | from PyQt5.QtCore import QSize, Qt, QUrl
|
| | from PyQt5.QtGui import (
|
| | QDesktopServices,
|
| | QPixmap,
|
| | )
|
| | from PyQt5.QtWidgets import (
|
| | QApplication,
|
| | QHBoxLayout,
|
| | QLabel,
|
| | QPushButton,
|
| | QSizePolicy,
|
| | QTextEdit,
|
| | QToolButton,
|
| | QVBoxLayout,
|
| | QWidget,
|
| | )
|
| |
|
| | from app_settings import AppSettings
|
| | from constants import DEVICE
|
| | from frontend.gui.image_generator_worker import ImageGeneratorWorker
|
| |
|
| |
|
| | class ImageLabel(QLabel):
|
| | """Defines a simple QLabel widget"""
|
| |
|
| | changed = QtCore.pyqtSignal()
|
| |
|
| | def __init__(self, text: str):
|
| | super().__init__(text)
|
| | self.setAlignment(Qt.AlignCenter)
|
| | self.resize(512, 512)
|
| | self.setSizePolicy(QSizePolicy.MinimumExpanding, QSizePolicy.MinimumExpanding)
|
| | self.sizeHint = QSize(512, 512)
|
| | self.setAcceptDrops(False)
|
| |
|
| | def show_image(self, pixmap: QPixmap = None):
|
| | """Updates the widget pixamp"""
|
| | if pixmap == None or pixmap.isNull():
|
| | return
|
| | self.current_pixmap = pixmap
|
| | self.changed.emit()
|
| |
|
| |
|
| | image_width = self.current_pixmap.width()
|
| | image_height = self.current_pixmap.height()
|
| | if image_width > 512 or image_height > 512:
|
| | new_width = 512 if image_width > 512 else image_width
|
| | new_height = 512 if image_height > 512 else image_height
|
| | self.setPixmap(
|
| | self.current_pixmap.scaled(
|
| | new_width,
|
| | new_height,
|
| | Qt.KeepAspectRatio,
|
| | )
|
| | )
|
| | else:
|
| | self.setPixmap(self.current_pixmap)
|
| |
|
| |
|
| | class BaseWidget(QWidget):
|
| | def __init__(self, config: AppSettings, parent):
|
| | super().__init__()
|
| | self.config = config
|
| | self.gen_images = []
|
| | self.image_index = 0
|
| | self.config = config
|
| | self.parent = parent
|
| |
|
| |
|
| | self.prev_btn = QToolButton()
|
| | self.prev_btn.setText("<")
|
| | self.prev_btn.clicked.connect(self.on_show_previous_image)
|
| | self.img = ImageLabel("<<Image>>")
|
| | self.next_btn = QToolButton()
|
| | self.next_btn.setText(">")
|
| | self.next_btn.clicked.connect(self.on_show_next_image)
|
| | self.prompt = QTextEdit()
|
| | self.prompt.setPlaceholderText("A fantasy landscape")
|
| | self.prompt.setAcceptRichText(False)
|
| | self.prompt.setFixedHeight(40)
|
| | self.neg_prompt = QTextEdit()
|
| | self.neg_prompt.setPlaceholderText("")
|
| | self.neg_prompt.setAcceptRichText(False)
|
| | self.neg_prompt_label = QLabel("Negative prompt (Set guidance scale > 1.0):")
|
| | self.neg_prompt.setFixedHeight(35)
|
| | self.neg_prompt.setEnabled(False)
|
| | self.generate = QPushButton("Generate")
|
| | self.generate.clicked.connect(self.generate_click)
|
| | self.browse_results = QPushButton("...")
|
| | self.browse_results.setFixedWidth(30)
|
| | self.browse_results.clicked.connect(self.on_open_results_folder)
|
| | self.browse_results.setToolTip("Open output folder")
|
| |
|
| |
|
| | ilayout = QHBoxLayout()
|
| | ilayout.addWidget(self.prev_btn)
|
| | ilayout.addWidget(self.img)
|
| | ilayout.addWidget(self.next_btn)
|
| |
|
| |
|
| | hlayout = QHBoxLayout()
|
| | hlayout.addWidget(self.neg_prompt)
|
| | hlayout.addWidget(self.generate)
|
| | hlayout.addWidget(self.browse_results)
|
| |
|
| |
|
| | vlayout = QVBoxLayout()
|
| | vlayout.addLayout(ilayout)
|
| |
|
| | vlayout.addWidget(self.prompt)
|
| | vlayout.addWidget(self.neg_prompt_label)
|
| | vlayout.addLayout(hlayout)
|
| | self.setLayout(vlayout)
|
| |
|
| | self.parent.settings_changed.connect(self.on_settings_changed)
|
| |
|
| | def generate_image(self):
|
| | self.parent.prepare_generation_settings(self.config)
|
| | self.config.settings.lcm_diffusion_setting.prompt = self.prompt.toPlainText()
|
| | self.config.settings.lcm_diffusion_setting.negative_prompt = (
|
| | self.neg_prompt.toPlainText()
|
| | )
|
| | images = self.parent.context.generate_text_to_image(
|
| | self.config.settings,
|
| | self.config.reshape_required,
|
| | DEVICE,
|
| | )
|
| | self.parent.context.save_images(
|
| | images,
|
| | self.config.settings,
|
| | )
|
| | self.prepare_images(images)
|
| | self.after_generation()
|
| |
|
| | def prepare_images(self, images):
|
| | """Prepares the generated images to be displayed in the Qt widget"""
|
| | self.image_index = 0
|
| | self.gen_images = []
|
| | for img in images:
|
| | im = ImageQt(img).copy()
|
| | pixmap = QPixmap.fromImage(im)
|
| | self.gen_images.append(pixmap)
|
| |
|
| | if len(self.gen_images) > 1:
|
| | self.next_btn.setEnabled(True)
|
| | self.prev_btn.setEnabled(False)
|
| | else:
|
| | self.next_btn.setEnabled(False)
|
| | self.prev_btn.setEnabled(False)
|
| |
|
| | self.img.show_image(pixmap=self.gen_images[0])
|
| |
|
| | def on_show_next_image(self):
|
| | if self.image_index != len(self.gen_images) - 1 and len(self.gen_images) > 0:
|
| | self.prev_btn.setEnabled(True)
|
| | self.image_index += 1
|
| | self.img.show_image(pixmap=self.gen_images[self.image_index])
|
| | if self.image_index == len(self.gen_images) - 1:
|
| | self.next_btn.setEnabled(False)
|
| |
|
| | def on_show_previous_image(self):
|
| | if self.image_index != 0:
|
| | self.next_btn.setEnabled(True)
|
| | self.image_index -= 1
|
| | self.img.show_image(pixmap=self.gen_images[self.image_index])
|
| | if self.image_index == 0:
|
| | self.prev_btn.setEnabled(False)
|
| |
|
| | def on_open_results_folder(self):
|
| | QDesktopServices.openUrl(
|
| | QUrl.fromLocalFile(self.config.settings.generated_images.path)
|
| | )
|
| |
|
| | def generate_click(self):
|
| | self.img.setText("Please wait...")
|
| | self.before_generation()
|
| | worker = ImageGeneratorWorker(self.generate_image)
|
| | self.parent.threadpool.start(worker)
|
| |
|
| | def before_generation(self):
|
| | """Call this function before running a generation task"""
|
| | self.img.setEnabled(False)
|
| | self.generate.setEnabled(False)
|
| | self.browse_results.setEnabled(False)
|
| |
|
| | def after_generation(self):
|
| | """Call this function after running a generation task"""
|
| | self.img.setEnabled(True)
|
| | self.generate.setEnabled(True)
|
| | self.browse_results.setEnabled(True)
|
| | self.parent.store_dimension_settings()
|
| |
|
| | def on_settings_changed(self):
|
| | self.neg_prompt.setEnabled(
|
| | self.config.settings.lcm_diffusion_setting.use_openvino
|
| | or self.config.settings.lcm_diffusion_setting.use_lcm_lora
|
| | )
|
| |
|
