Spaces:
Running
Running
File size: 3,469 Bytes
c7f95a6 |
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 |
from PIL import Image
from PyQt5.QtWidgets import (
QApplication,
QHBoxLayout,
QRadioButton,
QWidget,
)
from app_settings import AppSettings
from backend.models.upscale import UpscaleMode
from backend.upscale.upscaler import upscale_image
from frontend.gui.img2img_widget import Img2ImgWidget
from paths import FastStableDiffusionPaths
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
|