aboutsummaryrefslogtreecommitdiffstats
path: root/modules
diff options
context:
space:
mode:
Diffstat (limited to 'modules')
-rw-r--r--modules/realesrgan_model.py80
-rw-r--r--modules/shared.py15
-rw-r--r--modules/ui.py45
3 files changed, 90 insertions, 50 deletions
diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py
index e2cef0c8..c32d6c4c 100644
--- a/modules/realesrgan_model.py
+++ b/modules/realesrgan_model.py
@@ -1,17 +1,66 @@
import sys
import traceback
from collections import namedtuple
+
import numpy as np
from PIL import Image
+from realesrgan import RealESRGANer
import modules.images
from modules.shared import cmd_opts, opts
RealesrganModelInfo = namedtuple("RealesrganModelInfo", ["name", "location", "model", "netscale"])
-
realesrgan_models = []
have_realesrgan = False
-RealESRGANer_constructor = None
+
+
+def get_realesrgan_models():
+ try:
+ from basicsr.archs.rrdbnet_arch import RRDBNet
+ from realesrgan import RealESRGANer
+ from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+ models = [
+ RealesrganModelInfo(
+ name="Real-ESRGAN General x4x3",
+ location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
+ netscale=4,
+ model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+ ),
+ RealesrganModelInfo(
+ name="Real-ESRGAN General WDN x4x3",
+ location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",
+ netscale=4,
+ model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+ ),
+ RealesrganModelInfo(
+ name="Real-ESRGAN AnimeVideo",
+ location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth",
+ netscale=4,
+ model=lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
+ ),
+ RealesrganModelInfo(
+ name="Real-ESRGAN 4x plus",
+ location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
+ netscale=4,
+ model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+ ),
+ RealesrganModelInfo(
+ name="Real-ESRGAN 4x plus anime 6B",
+ location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
+ netscale=4,
+ model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+ ),
+ RealesrganModelInfo(
+ name="Real-ESRGAN 2x plus",
+ location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
+ netscale=2,
+ model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+ ),
+ ]
+ return models
+ except Exception as e:
+ print("Error makeing Real-ESRGAN midels list:", file=sys.stderr)
+ print(traceback.format_exc(), file=sys.stderr)
class UpscalerRealESRGAN(modules.images.Upscaler):
@@ -27,35 +76,18 @@ class UpscalerRealESRGAN(modules.images.Upscaler):
def setup_realesrgan():
global realesrgan_models
global have_realesrgan
- global RealESRGANer_constructor
try:
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from realesrgan.archs.srvgg_arch import SRVGGNetCompact
- realesrgan_models = [
- RealesrganModelInfo(
- name="Real-ESRGAN 4x plus",
- location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
- netscale=4, model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
- ),
- RealesrganModelInfo(
- name="Real-ESRGAN 4x plus anime 6B",
- location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
- netscale=4, model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
- ),
- RealesrganModelInfo(
- name="Real-ESRGAN 2x plus",
- location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
- netscale=2, model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
- ),
- ]
+ realesrgan_models = get_realesrgan_models()
have_realesrgan = True
- RealESRGANer_constructor = RealESRGANer
for i, model in enumerate(realesrgan_models):
- modules.shared.sd_upscalers.append(UpscalerRealESRGAN(model.netscale, i))
+ if model.name in opts.realesrgan_enabled_models:
+ modules.shared.sd_upscalers.append(UpscalerRealESRGAN(model.netscale, i))
except Exception:
print("Error importing Real-ESRGAN:", file=sys.stderr)
@@ -66,13 +98,13 @@ def setup_realesrgan():
def upscale_with_realesrgan(image, RealESRGAN_upscaling, RealESRGAN_model_index):
- if not have_realesrgan or RealESRGANer_constructor is None:
+ if not have_realesrgan:
return image
info = realesrgan_models[RealESRGAN_model_index]
model = info.model()
- upsampler = RealESRGANer_constructor(
+ upsampler = RealESRGANer(
scale=info.netscale,
model_path=info.location,
model=model,
diff --git a/modules/shared.py b/modules/shared.py
index 280d6c4c..0a852838 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -95,6 +95,11 @@ face_restorers = []
modules.sd_models.list_models()
+def realesrgan_models_names():
+ import modules.realesrgan_model
+ return [x.name for x in modules.realesrgan_model.get_realesrgan_models()]
+
+
class Options:
class OptionInfo:
def __init__(self, default=None, label="", component=None, component_args=None, onchange=None):
@@ -142,14 +147,12 @@ class Options:
"save_txt": OptionInfo(False, "Create a text file next to every image with generation parameters."),
"ESRGAN_tile": OptionInfo(192, "Tile size for ESRGAN upscalers. 0 = no tiling.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}),
"ESRGAN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for ESRGAN upscalers. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),
+ "realesrgan_enabled_models": OptionInfo(["Real-ESRGAN 4x plus", "Real-ESRGAN 4x plus anime 6B"],"Select which RealESRGAN models to show in the web UI. (Requires restart)", gr.CheckboxGroup, lambda: {"choices": realesrgan_models_names()}),
"SWIN_tile": OptionInfo(192, "Tile size for all SwinIR.", gr.Slider, {"minimum": 16, "maximum": 512, "step": 16}),
"SWIN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for SwinIR. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),
- "ldsr_steps": OptionInfo(100, "LDSR processing steps. Lower = faster",
- gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}),
- "ldsr_pre_down":OptionInfo(1, "LDSR Pre-process downssample scale. 1 = no down-sampling, 4 = 1/4 scale.",
- gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),
- "ldsr_post_down":OptionInfo(1, "LDSR Post-process down-sample scale. 1 = no down-sampling, 4 = 1/4 scale.",
- gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),
+ "ldsr_steps": OptionInfo(100, "LDSR processing steps. Lower = faster", gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}),
+ "ldsr_pre_down":OptionInfo(1, "LDSR Pre-process downssample scale. 1 = no down-sampling, 4 = 1/4 scale.", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),
+ "ldsr_post_down":OptionInfo(1, "LDSR Post-process down-sample scale. 1 = no down-sampling, 4 = 1/4 scale.", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),
"random_artist_categories": OptionInfo([], "Allowed categories for random artists selection when using the Roll button", gr.CheckboxGroup, {"choices": artist_db.categories()}),
"upscale_at_full_resolution_padding": OptionInfo(16, "Inpainting at full resolution: padding, in pixels, for the masked region.", gr.Slider, {"minimum": 0, "maximum": 128, "step": 4}),
"upscaler_for_hires_fix": OptionInfo(None, "Upscaler for highres. fix", gr.Radio, lambda: {"choices": [x.name for x in sd_upscalers]}),
diff --git a/modules/ui.py b/modules/ui.py
index e3cd7f3a..572dcfb9 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -350,8 +350,15 @@ def create_toprow(is_img2img):
with gr.Column(scale=1):
with gr.Row():
+ interrupt = gr.Button('Interrupt', elem_id="interrupt")
submit = gr.Button('Generate', elem_id="generate", variant='primary')
+ interrupt.click(
+ fn=lambda: shared.state.interrupt(),
+ inputs=[],
+ outputs=[],
+ )
+
with gr.Row():
if is_img2img:
interrogate = gr.Button('Interrogate', elem_id="interrogate")
@@ -386,6 +393,15 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
txt2img_prompt, roll, txt2img_prompt_style, txt2img_negative_prompt, txt2img_prompt_style2, submit, _, txt2img_prompt_style_apply, txt2img_save_style = create_toprow(is_img2img=False)
dummy_component = gr.Label(visible=False)
+ with gr.Row(elem_id='progressRow'):
+ with gr.Column(scale=1):
+ columnEmpty = "Empty"
+
+ with gr.Column(scale=1):
+ progressbar = gr.HTML(elem_id="progressbar")
+ txt2img_preview = gr.Image(elem_id='txt2img_preview', visible=False)
+ setup_progressbar(progressbar, txt2img_preview)
+
with gr.Row().style(equal_height=False):
with gr.Column(variant='panel'):
steps = gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20)
@@ -416,21 +432,17 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
custom_inputs = modules.scripts.scripts_txt2img.setup_ui(is_img2img=False)
with gr.Column(variant='panel'):
- progressbar = gr.HTML(elem_id="progressbar")
with gr.Group():
txt2img_preview = gr.Image(elem_id='txt2img_preview', visible=False)
txt2img_gallery = gr.Gallery(label='Output', elem_id='txt2img_gallery').style(grid=4)
- setup_progressbar(progressbar, txt2img_preview)
-
with gr.Group():
with gr.Row():
save = gr.Button('Save')
send_to_img2img = gr.Button('Send to img2img')
send_to_inpaint = gr.Button('Send to inpaint')
send_to_extras = gr.Button('Send to extras')
- interrupt = gr.Button('Interrupt')
with gr.Group():
html_info = gr.HTML()
@@ -479,12 +491,6 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
outputs=[hr_options],
)
- interrupt.click(
- fn=lambda: shared.state.interrupt(),
- inputs=[],
- outputs=[],
- )
-
save.click(
fn=wrap_gradio_call(save_files),
_js="(x, y, z) => [x, y, selected_gallery_index()]",
@@ -513,6 +519,15 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
with gr.Blocks(analytics_enabled=False) as img2img_interface:
img2img_prompt, roll, img2img_prompt_style, img2img_negative_prompt, img2img_prompt_style2, submit, img2img_interrogate, img2img_prompt_style_apply, img2img_save_style = create_toprow(is_img2img=True)
+ with gr.Row(elem_id='progressRow'):
+ with gr.Column(scale=1):
+ columnEmpty = "Empty"
+
+ with gr.Column(scale=1):
+ progressbar = gr.HTML(elem_id="progressbar")
+ img2img_preview = gr.Image(elem_id='img2img_preview', visible=False)
+ setup_progressbar(progressbar, img2img_preview)
+
with gr.Row().style(equal_height=False):
with gr.Column(variant='panel'):
with gr.Group():
@@ -561,21 +576,17 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
custom_inputs = modules.scripts.scripts_img2img.setup_ui(is_img2img=True)
with gr.Column(variant='panel'):
- progressbar = gr.HTML(elem_id="progressbar")
with gr.Group():
img2img_preview = gr.Image(elem_id='img2img_preview', visible=False)
img2img_gallery = gr.Gallery(label='Output', elem_id='img2img_gallery').style(grid=4)
- setup_progressbar(progressbar, img2img_preview)
-
with gr.Group():
with gr.Row():
save = gr.Button('Save')
img2img_send_to_img2img = gr.Button('Send to img2img')
img2img_send_to_inpaint = gr.Button('Send to inpaint')
img2img_send_to_extras = gr.Button('Send to extras')
- interrupt = gr.Button('Interrupt')
img2img_save_style = gr.Button('Save prompt as style')
@@ -689,12 +700,6 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
outputs=[img2img_prompt],
)
- interrupt.click(
- fn=lambda: shared.state.interrupt(),
- inputs=[],
- outputs=[],
- )
-
save.click(
fn=wrap_gradio_call(save_files),
_js="(x, y, z) => [x, y, selected_gallery_index()]",