src/frontend/gui/base_widget.py

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

        # Resize the pixmap to the widget dimensions
        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

        # Initialize GUI widgets
        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")

        # Create the image navigation layout
        ilayout = QHBoxLayout()
        ilayout.addWidget(self.prev_btn)
        ilayout.addWidget(self.img)
        ilayout.addWidget(self.next_btn)

        # Create the generate button layout
        hlayout = QHBoxLayout()
        hlayout.addWidget(self.neg_prompt)
        hlayout.addWidget(self.generate)
        hlayout.addWidget(self.browse_results)

        # Create the actual widget layout
        vlayout = QVBoxLayout()
        vlayout.addLayout(ilayout)
        # vlayout.addItem(self.vspacer)
        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
        )