UI improvements and easier photo download

app.py

@@ -4,6 +4,9 @@ import util
 from run_cmd import run_cmd
 from random import randint
 from PIL import Image
+import tempfile
+
+temp_path = tempfile.gettempdir()
 
 is_colab = util.is_google_colab()
 
@@ -11,10 +14,10 @@ run_cmd("pip install pngquant")
 
 def inference(img, size, type):
     _id = randint(1, 10000)
-    INPUT_DIR = "/tmp/input_image" + str(_id) + "/"
-    OUTPUT_DIR = "/tmp/output_image" + str(_id) + "/"
+    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, "1_out.png")
+    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}")
@@ -36,37 +39,58 @@ def inference(img, size, type):
 
     # Remove input and output image
     run_cmd(f"rm -f {img_in_path}")
-    run_cmd(f"rm -f {img_out_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)
 
-    return img_out
+css = '''
+    .file-preview {
+        overflow: hidden !important;
+        margin: 5px 0 !important;
+        padding: 0 10px !important;
+    }
 
+    .file-preview div div:nth-child(2) {
+        flex-grow: 1 !important;
+    }
 
-upscale_type = gr.Radio(["Manga", "Anime", "General"], label="Select the type of picture you want to upscale:", value="Manga")
-upscale_size = gr.Radio(["x4", "x2"], label="Upscale by:", value="x4")
-output_image = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")
+    .file-preview div div:nth-child(3) {
+        text-align: right !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() as demo:
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="pil", label="Input")
-            translate_btn = gr.Button(value="Upscale", variant="primary")
-        with gr.Column():
-            output_image = gr.Image(type="file")
-            save_btn = gr.Button(value="Save image")
-
-
-""" demo = gr.Interface(
-    inference,
-    inputs=[input_image, upscale_size, upscale_type],
-    outputs=[output_image],
-    title=title,
-    description=description,
-    article=article
-) """
+with gr.Blocks(title=title, css=css) as demo:
+    gr.Markdown(
+    f"""
+    # {title}
+    {description}
+    """)
+
+    with gr.Box():
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(type="pil", label="Input")
+                upscale_size = gr.Radio(["x4", "x2"], label="Upscale by:", value="x4")
+                upscale_type = gr.Radio(["Manga", "Anime", "General"], label="Select the type of picture you want to upscale:", value="Manga")
+
+                with gr.Row():
+                    upscale_btn = gr.Button(value="Upscale", variant="primary")
+
+            with gr.Column():
+                output_image = gr.Image(type="filepath", interactive=False, label="Upscaled image", )
+                
+                with gr.Row():
+                    out_file = gr.File(interactive=False, show_label=False, visible=False)
+
+    gr.HTML(value=article)
+        
+    upscale_btn.click(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.py

@@ -14,9 +14,9 @@ def is_cuda():
 model_type = sys.argv[3]
 
 if model_type == "Anime":
-    model_path = "4x-AnimeSharp.pth"
+    model_path = "models/4x-AnimeSharp.pth"
 else:
-    model_path = "4x-UniScaleV2_Sharp.pth"
+    model_path = "models/4x-UniScaleV2_Sharp.pth"
 
 img_path = sys.argv[1]
 output_dir = sys.argv[2]
@@ -39,8 +39,7 @@ model = model.to(device)
 
 base = os.path.splitext(os.path.basename(img_path))[0]
 
-# read image
-print(img_path);
+# Read image
 img = cv2.imread(img_path, cv2.IMREAD_COLOR)
 img = img * 1.0 / 255
 img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()

inference_manga_v2.py

@@ -11,7 +11,7 @@ def is_cuda():
     else:
         return False
 
-model_path = '4x_eula_digimanga_bw_v2_nc1_307k.pth'
+model_path = 'models/4x_eula_digimanga_bw_v2_nc1_307k.pth'
 img_path = sys.argv[1]
 output_dir = sys.argv[2]
 device = torch.device('cuda' if is_cuda() else 'cpu')
@@ -33,8 +33,7 @@ model = model.to(device)
 
 base = os.path.splitext(os.path.basename(img_path))[0]
 
-# read image
-print(img_path);
+# Read image
 img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
 img = img * 1.0 / 255
 img = torch.from_numpy(img[np.newaxis, :, :]).float()