Spaces:
Paused
Paused
from PyQt5.QtWidgets import ( | |
QWidget, | |
QPushButton, | |
QHBoxLayout, | |
QVBoxLayout, | |
QLabel, | |
QLineEdit, | |
QMainWindow, | |
QSlider, | |
QTabWidget, | |
QSpacerItem, | |
QSizePolicy, | |
QComboBox, | |
QCheckBox, | |
QTextEdit, | |
QToolButton, | |
QFileDialog, | |
QApplication, | |
) | |
from PyQt5 import QtWidgets, QtCore | |
from PyQt5.QtGui import ( | |
QPixmap, | |
QDesktopServices, | |
QDragEnterEvent, | |
QDropEvent, | |
QMouseEvent, | |
) | |
from PyQt5.QtCore import QSize, QThreadPool, Qt, QUrl, QBuffer, QObject, QEvent | |
import io | |
from PIL import Image | |
from constants import DEVICE | |
from PIL.ImageQt import ImageQt | |
from app_settings import AppSettings | |
from urllib.parse import urlparse, unquote | |
from frontend.gui.base_widget import BaseWidget, ImageLabel | |
from backend.models.lcmdiffusion_setting import DiffusionTask | |
from frontend.gui.image_generator_worker import ImageGeneratorWorker | |
class Img2ImgWidget(BaseWidget): | |
def __init__(self, config: AppSettings, parent): | |
super().__init__(config, parent) | |
# Create init image selection widgets | |
self.img_label = QLabel("Init image:") | |
self.img_path = QLineEdit() | |
self.img_path.setReadOnly(True) | |
self.img_path.setAcceptDrops(True) | |
self.img_path.installEventFilter(self) | |
self.img_browse = QToolButton() | |
self.img_browse.setText("...") | |
self.img_browse.setToolTip("Browse for an init image") | |
self.img_browse.clicked.connect(self.browse_click) | |
# Create the init image selection layout | |
hlayout = QHBoxLayout() | |
hlayout.addWidget(self.img_label) | |
hlayout.addWidget(self.img_path) | |
hlayout.addWidget(self.img_browse) | |
self.strength_label = QLabel("Denoising strength: 0.3") | |
self.strength = QSlider(orientation=Qt.Orientation.Horizontal) | |
self.strength.setMaximum(10) | |
self.strength.setMinimum(1) | |
self.strength.setValue(3) | |
self.strength.valueChanged.connect(self.update_strength_label) | |
# self.layout().insertWidget(1, self.strength_label) | |
# self.layout().insertWidget(2, self.strength) | |
self.layout().addLayout(hlayout) | |
self.layout().addWidget(self.strength_label) | |
self.layout().addWidget(self.strength) | |
def browse_click(self): | |
filename = self.show_file_selection_dialog() | |
if filename[0] != "": | |
self.img_path.setText(filename[0]) | |
def show_file_selection_dialog(self) -> str: | |
filename = QFileDialog.getOpenFileName( | |
self, "Open Image", "results", "Image Files (*.png *.jpg *.bmp)" | |
) | |
return filename | |
def eventFilter(self, source, event: QEvent): | |
"""This is the Drag and Drop event filter for the init image QLineEdit""" | |
if event.type() == QEvent.DragEnter: | |
if event.mimeData().hasFormat("text/plain"): | |
event.acceptProposedAction() | |
return True | |
elif event.type() == QEvent.Drop: | |
event.acceptProposedAction() | |
path = unquote(urlparse(event.mimeData().text()).path) | |
self.img_path.setText(path) | |
return True | |
return False | |
def before_generation(self): | |
super().before_generation() | |
self.img_browse.setEnabled(False) | |
self.img_path.setEnabled(False) | |
def after_generation(self): | |
super().after_generation() | |
self.img_browse.setEnabled(True) | |
self.img_path.setEnabled(True) | |
def generate_image(self): | |
self.parent.prepare_generation_settings(self.config) | |
self.config.settings.lcm_diffusion_setting.diffusion_task = ( | |
DiffusionTask.image_to_image.value | |
) | |
self.config.settings.lcm_diffusion_setting.prompt = self.prompt.toPlainText() | |
self.config.settings.lcm_diffusion_setting.negative_prompt = ( | |
self.neg_prompt.toPlainText() | |
) | |
self.config.settings.lcm_diffusion_setting.init_image = Image.open( | |
self.img_path.text() | |
) | |
self.config.settings.lcm_diffusion_setting.strength = self.strength.value() / 10 | |
images = self.parent.context.generate_text_to_image( | |
self.config.settings, | |
self.config.reshape_required, | |
DEVICE, | |
) | |
self.prepare_images(images) | |
self.after_generation() | |
def update_strength_label(self, value): | |
val = round(int(value) / 10, 1) | |
self.strength_label.setText(f"Denoising strength: {val}") | |
self.config.settings.lcm_diffusion_setting.strength = val | |
# Test the widget | |
if __name__ == "__main__": | |
import sys | |
app = QApplication(sys.argv) | |
widget = Img2ImgWidget(None, None) | |
widget.show() | |
app.exec() | |