Spaces:
Paused
Paused
File size: 4,385 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 |
from PyQt5.QtWidgets import (
QWidget,
QPushButton,
QHBoxLayout,
QVBoxLayout,
QLabel,
QLineEdit,
QMainWindow,
QSlider,
QTabWidget,
QSpacerItem,
QSizePolicy,
QComboBox,
QCheckBox,
QTextEdit,
QToolButton,
QFileDialog,
QApplication,
QRadioButton,
QFrame,
)
from PyQt5 import QtWidgets, QtCore
from PyQt5.QtGui import QPixmap, QDesktopServices, QDragEnterEvent, QDropEvent
from PyQt5.QtCore import QSize, QThreadPool, Qt, QUrl, QBuffer
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 backend.models.upscale import UpscaleMode
from backend.upscale.upscaler import upscale_image
from frontend.gui.img2img_widget import Img2ImgWidget
from paths import FastStableDiffusionPaths, join_paths
from backend.models.lcmdiffusion_setting import DiffusionTask
from frontend.gui.image_generator_worker import ImageGeneratorWorker
from frontend.webui.image_variations_ui import generate_image_variations
class UpscalerWidget(Img2ImgWidget):
scale_factor = 2.0
def __init__(self, config: AppSettings, parent):
super().__init__(config, parent)
# Hide prompt and negative prompt widgets
self.prompt.hide()
self.neg_prompt.hide()
# self.neg_prompt.deleteLater()
self.strength_label.hide()
self.strength.hide()
self.generate.setText("Upscale")
# Create upscaler widgets
self.edsr_button = QRadioButton("EDSR", self)
self.edsr_button.toggled.connect(self.on_mode_change)
self.edsr_button.toggle()
self.sd_button = QRadioButton("SD", self)
self.sd_button.toggled.connect(self.on_mode_change)
self.aura_button = QRadioButton("AURA-SR", self)
self.aura_button.toggled.connect(self.on_mode_change)
self.neg_prompt_label.setText("Upscale mode (2x) | AURA-SR (4x):")
# Create upscaler buttons layout
hlayout = QHBoxLayout()
hlayout.addWidget(self.edsr_button)
hlayout.addWidget(self.sd_button)
hlayout.addWidget(self.aura_button)
# Can't use a layout with replaceWidget(), so the layout is assigned
# to a dummy widget used to replace the negative prompt button and
# obtain the desired GUI design
self.container = QWidget()
self.container.setLayout(hlayout)
self.layout().replaceWidget(self.neg_prompt, self.container)
def generate_image(self):
self.parent.prepare_generation_settings(self.config)
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
upscaled_filepath = FastStableDiffusionPaths.get_upscale_filepath(
None,
self.scale_factor,
self.config.settings.generated_images.format,
)
images = upscale_image(
context=self.parent.context,
src_image_path=self.img_path.text(),
dst_image_path=upscaled_filepath,
upscale_mode=self.upscale_mode,
strength=0.1,
)
self.prepare_images(images)
self.after_generation()
def before_generation(self):
super().before_generation()
self.container.setEnabled(False)
def after_generation(self):
super().after_generation()
self.container.setEnabled(True)
# TODO For some reason, if init_image is not set to None, saving settings
# when closing the main window will fail, even though the save() method
# explicitly sets init_image to None?
self.config.settings.lcm_diffusion_setting.init_image = None
def on_mode_change(self):
self.scale_factor = 2.0
if self.edsr_button.isChecked():
self.upscale_mode = UpscaleMode.normal.value
elif self.sd_button.isChecked():
self.upscale_mode = UpscaleMode.sd_upscale.value
else:
self.upscale_mode = UpscaleMode.aura_sr.value
self.scale_factor = 4.0
# Test the widget
if __name__ == "__main__":
import sys
app = QApplication(sys.argv)
widget = ImageVariationsWidget(None, None)
widget.show()
app.exec()
|