UI and memory improvements.

app.py

@@ -1,50 +1,12 @@
-import os
 import gradio as gr
 import util
+import process_image
 from run_cmd import run_cmd
-from random import randint
-from PIL import Image
-import tempfile
 
-temp_path = tempfile.gettempdir()
+run_cmd("pip install split-image")
 
 is_colab = util.is_google_colab()
 
-run_cmd("pip install pngquant")
-
-def inference(img, size, type):
-    _id = randint(1, 10000)
-    INPUT_DIR = os.path.join(temp_path, f"input_image{str(_id)}")
-    OUTPUT_DIR = os.path.join(temp_path, f"output_image{str(_id)}")
-    img_in_path = os.path.join(INPUT_DIR, "1.jpg")
-    img_out_path = os.path.join(OUTPUT_DIR, f"1_{size}.png")
-    run_cmd(f"rm -rf {INPUT_DIR}")
-    run_cmd(f"rm -rf {OUTPUT_DIR}")
-    run_cmd(f"mkdir {INPUT_DIR}")
-    run_cmd(f"mkdir {OUTPUT_DIR}")
-    img.save(img_in_path, "PNG")
-
-    if type == "Manga":
-        run_cmd(f"python inference_manga_v2.py {img_in_path} {img_out_path}")
-    else:
-        run_cmd(f"python inference.py {img_in_path} {img_out_path} {type}")
-
-    img_out = Image.open(img_out_path)
-
-    if size == "x2":
-        img_out = img_out.resize((img_out.width // 2, img_out.height // 2), resample=Image.BICUBIC)
-
-    #img_out.save(img_out_path, optimize=True) # Add more optimizations
-    #img_out = Image.open(img_out_path)
-
-    # Remove input and output image
-    run_cmd(f"rm -f {img_in_path}")
-    #run_cmd(f"rm -f {img_out_path}")
-
-    out_file.update(value=img_out_path, visible=True)
-
-    return img_out, gr.File.update(value=img_out_path, visible=True)
-
 css = '''
     .file-preview {
         overflow: hidden !important;
@@ -59,17 +21,23 @@ css = '''
     .file-preview div div:nth-child(3) {
         text-align: right !important;
     }
+
+    #preview_img {
+        user-select: none !important;
+        touch-action: none !important;
+        pointer-events: none !important;
+    }
 '''
 
 title = "ESRGAN Upscaling With Custom Models"
-description = "This space uses old ESRGAN architecture to upscale images, using models made by the community."
-article = "<p><a href='https://upscale.wiki/wiki/Model_Database'>Model Database</a></p>"
 
 with gr.Blocks(title=title, css=css) as demo:
     gr.Markdown(
     f"""
     # {title}
-    {description}
+    This space uses old ESRGAN architecture to upscale images, using models made by the community.
+
+    Once upscaled, click or tap the download button under the image to download it.
     """)
 
     with gr.Box():
@@ -83,14 +51,14 @@ with gr.Blocks(title=title, css=css) as demo:
                     upscale_btn = gr.Button(value="Upscale", variant="primary")
 
             with gr.Column():
-                output_image = gr.Image(type="filepath", interactive=False, label="Upscaled image", )
+                output_image = gr.Image(type="filepath", interactive=False, label="Upscaled image", elem_id="preview_img")
                 
                 with gr.Row():
                     out_file = gr.File(interactive=False, show_label=False, visible=False)
 
-    gr.HTML(value=article)
+    gr.HTML(value="<p><a href='https://upscale.wiki/wiki/Model_Database'>Model Database</a></p>")
         
-    upscale_btn.click(inference, inputs=[input_image, upscale_size, upscale_type], outputs=[output_image, out_file])
+    upscale_btn.click(process_image.inference, inputs=[input_image, upscale_size, upscale_type], outputs=[output_image, out_file])
 
 demo.queue()
 demo.launch(debug=is_colab, share=is_colab, inline=is_colab)

inference_manga_v2.py

@@ -4,6 +4,10 @@ import cv2
 import numpy as np
 import torch
 import architecture as arch
+from split_image import split_image
+from run_cmd import run_cmd
+import re
+from pathlib import Path
 
 def is_cuda():
     if torch.cuda.is_available():
@@ -31,20 +35,39 @@ for k, v in model.named_parameters():
     v.requires_grad = False
 model = model.to(device)
 
-base = os.path.splitext(os.path.basename(img_path))[0]
-
-# Read image
-img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
-img = img * 1.0 / 255
-img = torch.from_numpy(img[np.newaxis, :, :]).float()
-img_LR = img.unsqueeze(0)
-img_LR = img_LR.to(device)
-
-print('Start upscaling...')
-with torch.no_grad():
-    output = model(img_LR).squeeze(dim=0).float().cpu().clamp_(0, 1).numpy()
-output = np.transpose(output, (1, 2, 0))
-output = (output * 255.0).round()
-print('Finished upscaling, saving image.')
-print(output_dir)
-cv2.imwrite(output_dir, output, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])
+base = os.path.dirname(img_path)
+
+# Split image
+run_cmd(f"split-image {img_path} 7 7 --output-dir {base} --quiet")
+
+for root, dirs, files in os.walk(base, topdown=True):
+    for x, name in enumerate(files):
+        file_path = os.path.join(root, name)
+
+        if file_path == img_path:
+            print(file_path)
+            continue
+
+        # Read image
+        img = cv2.imread(file_path, cv2.IMREAD_GRAYSCALE)
+        img = img * 1.0 / 255
+        img = torch.from_numpy(img[np.newaxis, :, :]).float()
+        img_LR = img.unsqueeze(0)
+        img_LR = img_LR.to(device)
+
+        #print(f"Start upscaling tile {x}...")
+        with torch.no_grad():
+            output = model(img_LR).squeeze(dim=0).float().cpu().clamp_(0, 1).numpy()
+        output = np.transpose(output, (1, 2, 0))
+        output = (output * 255.0).round()
+        #print(f"Finished upscaling tile {x}, saving tile.")
+        #print(output_dir)
+        cv2.imwrite(file_path, output)
+
+# Join all tiles
+run_cmd(f"cd {base} && split-image input.jpg 7 7 -r --quiet")
+
+# Open image and save as png with the ouput name
+img_out = cv2.imread(img_path);
+cv2.imwrite(output_dir, img_out, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])
+

process_image.py

@@ -0,0 +1,43 @@
+import os
+import gradio as gr
+from run_cmd import run_cmd
+from random import randint
+from PIL import Image
+import tempfile
+
+temp_path = tempfile.gettempdir()
+
+def inference(img, size, type):
+    if not img:
+        raise Exception("No image!")
+
+    _id = randint(1, 10000)
+    INPUT_DIR = os.path.join(temp_path, f"input_image{str(_id)}")
+    OUTPUT_DIR = os.path.join(temp_path, f"output_image{str(_id)}")
+    img_in_path = os.path.join(INPUT_DIR, "input.jpg")
+    img_out_path = os.path.join(OUTPUT_DIR, f"output_{size}.png")
+    run_cmd(f"rm -rf {INPUT_DIR}")
+    run_cmd(f"rm -rf {OUTPUT_DIR}")
+    run_cmd(f"mkdir {INPUT_DIR}")
+    run_cmd(f"mkdir {OUTPUT_DIR}")
+    img.save(img_in_path, "PNG")
+
+    if type == "Manga":
+        run_cmd(f"python inference_manga_v2.py {img_in_path} {img_out_path}")
+    else:
+        run_cmd(f"python inference.py {img_in_path} {img_out_path} {type}")
+
+    img_out = Image.open(img_out_path)
+
+    if size == "x2":
+        img_out = img_out.resize((img_out.width // 2, img_out.height // 2), resample=Image.BICUBIC)
+
+    #img_out.save(img_out_path, optimize=True) # Add more optimizations
+    #img_out = Image.open(img_out_path)
+
+    # Remove input and output image
+    #run_cmd(f"rm -f {img_in_path}")
+
+    img_out.thumbnail((600, 600), Image.ANTIALIAS)
+
+    return img_out, gr.File.update(value=img_out_path, visible=True)