|
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 |
|
) |
|
|
