Format as input

app.py

@@ -14,6 +14,7 @@ import math
 import time
 import random
 import spaces
+import re
 from huggingface_hub import hf_hub_download
 
 hf_hub_download(repo_id="laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", filename="open_clip_pytorch_model.bin", local_dir="laion_CLIP-ViT-bigG-14-laion2B-39B-b160k")
@@ -86,7 +87,7 @@ def stage1_process(
         gr.Warning('Set this space to GPU config to make it work.')
         return None, None
     torch.cuda.set_device(SUPIR_device)
-    LQ = HWC3(np.array(Image.open(input_image)))
+    LQ = HWC3(Image.open(input_image))
     LQ = fix_resize(LQ, 512)
     # stage1
     LQ = np.array(LQ) / 255 * 2 - 1
@@ -170,7 +171,7 @@ def stage2_process(
         if noisy_image is None:
             output_format = "png"
         else:
-            output_format = noisy_image.format
+            output_format = re.sub(r"^.*\.([^\.]+)$", "\1", noisy_image)
 
     if prompt is None:
         prompt = ""
@@ -328,7 +329,7 @@ def restore(
         elif model_select == 'v0-F':
             model.load_state_dict(ckpt_F, strict=False)
         model.current_model = model_select
-    input_image = HWC3(np.array(input_image))
+    input_image = HWC3(input_image)
     input_image = upscale_image(input_image, upscale, unit_resolution=32,
                                 min_size=min_size)
 
@@ -379,7 +380,7 @@ def restore(
     print(information)
 
     # Only one image can be shown in the slider
-    return [input_image] + [results[0]], gr.update(format = output_format, value = results), gr.update(value = information, visible = True)
+    return [input_image] + [results[0]], gr.update(label="Downloadable results in *." + output_format + " format", format = output_format, value = results), gr.update(value = information, visible = True)
 
 def load_and_reset(param_setting):
     print('load_and_reset ==>>')
@@ -474,7 +475,7 @@ with gr.Blocks() as interface:
         prompt_hint = gr.HTML("You can use a <a href='"'https://huggingface.co/spaces/MaziyarPanahi/llava-llama-3-8b'"'>LlaVa space</a> to auto-generate the description of your image.")
         upscale = gr.Radio([["x1", 1], ["x2", 2], ["x3", 3], ["x4", 4], ["x5", 5], ["x6", 6], ["x7", 7], ["x8", 8]], label="Upscale factor", info="Resolution x1 to x8", value=2, interactive=True)
         allocation = gr.Radio([["1 min", 1], ["2 min", 2], ["3 min", 3], ["4 min", 4], ["5 min", 5], ["6 min", 6], ["7 min", 7], ["8 min", 8]], label="GPU allocation time", info="lower=May abort run, higher=Quota penalty for next runs", value=6, interactive=True)
-        output_format = gr.Radio([["*.png", "png"], ["*.webp", "webp"], ["*.jpeg", "jpeg"], ["*.gif", "gif"], ["*.bmp", "bmp"]], label="Image format for result", info="File extention", value="png", interactive=True)
+        output_format = gr.Radio([["As input", "input"], ["*.png", "png"], ["*.webp", "webp"], ["*.jpeg", "jpeg"], ["*.gif", "gif"], ["*.bmp", "bmp"]], label="Image format for result", info="File extention", value="png", interactive=True)
 
     with gr.Accordion("Pre-denoising (optional)", open=False):
         gamma_correction = gr.Slider(label="Gamma Correction", info = "lower=lighter, higher=darker", minimum=0.1, maximum=2.0, value=1.0, step=0.1)