auto-sd-paint-ext/frontends/krita/krita_diff/pages/common.py

124 lines
4.2 KiB
Python

from krita import QHBoxLayout, QPushButton, QVBoxLayout, QWidget
from ..script import script
from ..widgets import QCheckBox, QComboBoxLayout, QLabel, QSpinBoxLayout
# Notes:
# - move tiling mode to config?
# - move upscaler/face restorer to config?
class SDCommonPage(QWidget):
name = "SD Common Options"
def __init__(self, *args, **kwargs):
super(SDCommonPage, self).__init__(*args, **kwargs)
self.title = QLabel("<em>Quick Config</em>")
# Model list
self.sd_model_layout = QComboBoxLayout(
script.cfg, "sd_model_list", "sd_model", label="SD model:"
)
# batch size & count
self.batch_count_layout = QSpinBoxLayout(
script.cfg, "sd_batch_count", label="Batch count:", min=1, max=9999, step=1
)
self.batch_size_layout = QSpinBoxLayout(
script.cfg, "sd_batch_size", label="Batch size:", min=1, max=9999, step=1
)
batch_layout = QHBoxLayout()
batch_layout.addLayout(self.batch_count_layout)
batch_layout.addLayout(self.batch_size_layout)
# base/max size adjustment
self.base_size_layout = QSpinBoxLayout(
script.cfg, "sd_base_size", label="Base size:", min=64, max=8192, step=64
)
self.max_size_layout = QSpinBoxLayout(
script.cfg, "sd_max_size", label="Max size:", min=64, max=8192, step=64
)
size_layout = QHBoxLayout()
size_layout.addLayout(self.base_size_layout)
size_layout.addLayout(self.max_size_layout)
# global upscaler
self.upscaler_layout = QComboBoxLayout(
script.cfg, "upscaler_list", "upscaler_name", label="Upscaler:"
)
# Restore faces
self.face_restorer_layout = QComboBoxLayout(
script.cfg,
"face_restorer_model_list",
"face_restorer_model",
label="Face restorer:",
)
self.codeformer_weight_layout = QSpinBoxLayout(
script.cfg,
"codeformer_weight",
label="CodeFormer weight (max 0, min 1):",
step=0.01,
)
# Tiling mode
self.tiling = QCheckBox(script.cfg, "sd_tiling", "Tiling mode")
# Interrupt button
self.interrupt_btn = QPushButton("Interrupt")
layout = QVBoxLayout()
layout.setContentsMargins(0, 0, 0, 0)
layout.addWidget(self.title)
layout.addLayout(self.upscaler_layout)
layout.addLayout(self.face_restorer_layout)
layout.addLayout(self.codeformer_weight_layout)
layout.addWidget(self.tiling)
layout.addLayout(self.sd_model_layout)
layout.addLayout(batch_layout)
layout.addLayout(size_layout)
layout.addWidget(self.interrupt_btn)
layout.addStretch()
self.setLayout(layout)
def cfg_init(self):
self.sd_model_layout.cfg_init()
self.batch_count_layout.cfg_init()
self.batch_size_layout.cfg_init()
self.base_size_layout.cfg_init()
self.max_size_layout.cfg_init()
self.upscaler_layout.cfg_init()
self.face_restorer_layout.cfg_init()
self.codeformer_weight_layout.cfg_init()
self.tiling.cfg_init()
self.title.setVisible(not script.cfg("minimize_ui", bool))
def cfg_connect(self):
self.sd_model_layout.cfg_connect()
self.batch_count_layout.cfg_connect()
self.batch_size_layout.cfg_connect()
self.base_size_layout.cfg_connect()
self.max_size_layout.cfg_connect()
self.upscaler_layout.cfg_connect()
self.face_restorer_layout.cfg_connect()
self.codeformer_weight_layout.cfg_connect()
self.tiling.cfg_connect()
# Hide codeformer_weight when model isnt codeformer
def toggle_codeformer_weights(visible):
self.codeformer_weight_layout.qspin.setVisible(visible)
self.codeformer_weight_layout.qlabel.setVisible(visible)
self.face_restorer_layout.qcombo.currentTextChanged.connect(
lambda t: toggle_codeformer_weights(t == "CodeFormer")
)
toggle_codeformer_weights(
self.face_restorer_layout.qcombo.currentText() == "CodeFormer"
)
self.interrupt_btn.released.connect(lambda: script.action_interrupt())