File size: 4,731 Bytes
99b955f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
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()