Renamed GFPGAN to extras

Added Real-ESRGAN to extras tab

README.md

@@ -62,14 +62,25 @@ Open the URL in browser, and you are good to go.
 The script creates a web UI for Stable Diffusion's txt2img and img2img scripts. Following are features added
 that are not in original script.
 
-### GFPGAN
+### Extras tab
+Additional neural network image improvement methods unrelated to stable diffusion.
+
+#### GFPGAN
 Lets you improve faces in pictures using the GFPGAN model. There is a checkbox in every tab to use GFPGAN at 100%, and
 also a separate tab that just allows you to use GFPGAN on any picture, with a slider that controls how strongthe effect is.
 
 ![](images/GFPGAN.png)
 
+#### Real-ESRGAN
+Image upscaler. You can choose from multiple models by original author, and specify by how much the image should be upscaled.
+Requires `realesrgan` librarty:
+
+```commandline
+pip install realesrgan
+```
+
 ### Sampling method selection
-Pick out of three sampling methods for txt2img: DDIM, PLMS, k-diffusion:
+Pick out of multiple sampling methods for txt2img:
 
 ![](images/sampling.png)
 

webui.py

@@ -76,6 +76,38 @@ samplers = [
     SamplerData('PLMS', lambda model: PLMSSampler(model)),
 ]
 
+RealesrganModelInfo = namedtuple("RealesrganModelInfo", ["name", "location", "model", "netscale"])
+
+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 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)
+        ),
+        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)
+        ),
+    ]
+    have_realesrgan = True
+except:
+    print("Error loading Real-ESRGAN:", file=sys.stderr)
+    print(traceback.format_exc(), file=sys.stderr)
+
+    realesrgan_models = [RealesrganModelInfo('None', '', 0, None)]
+    have_realesrgan = False
+
 
 class Options:
     data = None
@@ -196,10 +228,6 @@ def torch_gc():
     torch.cuda.ipc_collect()
 
 
-def sanitize_filename_part(text):
-    return text.replace(' ', '_').translate({ord(x): '' for x in invalid_filename_chars})[:128]
-
-
 def save_image(image, path, basename, seed, prompt, extension, info=None, short_filename=False):
     prompt = sanitize_filename_part(prompt)
 
@@ -208,7 +236,7 @@ def save_image(image, path, basename, seed, prompt, extension, info=None, short_
     else:
         filename = f"{basename}-{seed}-{prompt[:128]}.{extension}"
 
-    if extension == 'png' and opts.enable_pnginfo:
+    if extension == 'png' and opts.enable_pnginfo and info is not None:
         pnginfo = PngImagePlugin.PngInfo()
         pnginfo.add_text("parameters", info)
     else:
@@ -217,6 +245,10 @@ def save_image(image, path, basename, seed, prompt, extension, info=None, short_
     image.save(os.path.join(path, filename), quality=opts.jpeg_quality, pnginfo=pnginfo)
 
 
+def sanitize_filename_part(text):
+    return text.replace(' ', '_').translate({ord(x): '' for x in invalid_filename_chars})[:128]
+
+
 def plaintext_to_html(text):
     text = "".join([f"<p>{html.escape(x)}</p>\n" for x in text.split('\n')])
     return text
@@ -835,7 +867,7 @@ def img2img(prompt: str, init_img, ddim_steps: int, use_GFPGAN: bool, prompt_mat
                 prompt_matrix=prompt_matrix,
                 use_GFPGAN=use_GFPGAN,
                 do_not_save_grid=True,
-                extra_generation_params = {"Denoising Strength": denoising_strength},
+                extra_generation_params={"Denoising Strength": denoising_strength},
             )
 
             if initial_seed is None:
@@ -870,7 +902,7 @@ def img2img(prompt: str, init_img, ddim_steps: int, use_GFPGAN: bool, prompt_mat
             height=height,
             prompt_matrix=prompt_matrix,
             use_GFPGAN=use_GFPGAN,
-            extra_generation_params = {"Denoising Strength": denoising_strength},
+            extra_generation_params={"Denoising Strength": denoising_strength},
         )
 
     del sampler
@@ -908,30 +940,56 @@ img2img_interface = gr.Interface(
 )
 
 
-def run_GFPGAN(image, strength):
+def run_extras(image, GFPGAN_strength, RealESRGAN_upscaling, RealESRGAN_model_index):
     image = image.convert("RGB")
 
-    cropped_faces, restored_faces, restored_img = GFPGAN.enhance(np.array(image, dtype=np.uint8), has_aligned=False, only_center_face=False, paste_back=True)
-    res = Image.fromarray(restored_img)
+    outpath = opts.outdir or "outputs/extras-samples"
+
+    if GFPGAN is not None and GFPGAN_strength > 0:
+        cropped_faces, restored_faces, restored_img = GFPGAN.enhance(np.array(image, dtype=np.uint8), has_aligned=False, only_center_face=False, paste_back=True)
+        res = Image.fromarray(restored_img)
+
+        if GFPGAN_strength < 1.0:
+            res = Image.blend(image, res, GFPGAN_strength)
+
+        image = res
+
+    if have_realesrgan and RealESRGAN_upscaling != 1.0:
+        info = realesrgan_models[RealESRGAN_model_index]
+
+        model = info.model()
+        upsampler = RealESRGANer(
+            scale=info.netscale,
+            model_path=info.location,
+            model=model,
+            half=True
+        )
+
+        upsampled = upsampler.enhance(np.array(image), outscale=RealESRGAN_upscaling)[0]
+
+        image = Image.fromarray(upsampled)
 
-    if strength < 1.0:
-        res = Image.blend(image, res, strength)
+    os.makedirs(outpath, exist_ok=True)
+    base_count = len(os.listdir(outpath))
+
+    save_image(image, outpath, f"{base_count:05}", None, '', opts.samples_format, short_filename=True)
 
-    return res, 0, ''
+    return image, 0, ''
 
 
-gfpgan_interface = gr.Interface(
-    run_GFPGAN,
+extras_interface = gr.Interface(
+    wrap_gradio_call(run_extras),
     inputs=[
         gr.Image(label="Source", source="upload", interactive=True, type="pil"),
-        gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Effect strength", value=100),
+        gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="GFPGAN strength", value=1, interactive=GFPGAN is not None),
+        gr.Slider(minimum=1.0, maximum=4.0, step=0.05, label="Real-ESRGAN upscaling", value=2, interactive=have_realesrgan),
+        gr.Radio(label='Real-ESRGAN model', choices=[x.name for x in realesrgan_models], value=realesrgan_models[0].name, type="index", interactive=have_realesrgan),
     ],
     outputs=[
         gr.Image(label="Result"),
         gr.Number(label='Seed', visible=False),
         gr.HTML(),
     ],
-    description="Fix faces on images",
     allow_flagging="never",
 )
 
@@ -989,7 +1047,7 @@ settings_interface = gr.Interface(
 interfaces = [
     (txt2img_interface, "txt2img"),
     (img2img_interface, "img2img"),
-    (gfpgan_interface, "GFPGAN"),
+    (extras_interface, "Extras"),
     (settings_interface, "Settings"),
 ]
 
@@ -1003,9 +1061,6 @@ text_inversion_embeddings = TextInversionEmbeddings()
 if os.path.exists(cmd_opts.embeddings_dir):
     text_inversion_embeddings.hijack(model)
 
-if GFPGAN is None:
-    interfaces = [x for x in interfaces if x[0] != gfpgan_interface]
-
 demo = gr.TabbedInterface(
     interface_list=[x[0] for x in interfaces],
     tab_names=[x[1] for x in interfaces],