Update app.py

app.py

@@ -205,9 +205,7 @@ def restore_in_Xmin(
 
     if input_format.lower() not in ['png', 'webp', 'jpg', 'jpeg', 'gif', 'bmp', 'heic']:
         gr.Warning('Invalid image format. Please use a supported image format.')
-        return None, None, None, None
-
-    output_format = "png"
+        return None, None, None
 
     if prompt is None:
         prompt = ""
@@ -234,7 +232,7 @@ def restore_in_Xmin(
 
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')
-        return [input_image_path, denoise_image], gr.update(label="Downloadable results", format=output_format, value=[denoise_image]), None, gr.update(visible=True)
+        return [input_image_path, denoise_image], gr.update(value="GPU not available."), gr.update(visible=True)
 
     if model_select != model.current_model:
         print('Loading model: ' + model_select)
@@ -339,14 +337,11 @@ def restore_on_gpu(
     start = time.time()
     print('Starting GPU restoration...')
 
-    output_format = "png"
-
     torch.cuda.set_device(SUPIR_device)
 
     with torch.no_grad():
         # Convert input image to NumPy array and ensure it has 3 channels
         input_image = HWC3(np.array(Image.open(input_image_path)))
-
         input_image = upscale_image(input_image, upscale, unit_resolution=32, min_size=min_size)
         LQ = input_image / 255.0
         LQ = np.power(LQ, gamma_correction)
@@ -382,8 +377,13 @@ def restore_on_gpu(
     information = ("Start the process again if you want a different result. " if seed is not None else "") + \
         "The image has been enhanced successfully."
 
-    # Only one image can be shown in the slider
-    return [input_image_path] + [results[0]], gr.update(label="Downloadable results", format=output_format, value=results), gr.update(value=information, visible=True), gr.update(visible=True)
+    # Save the result image to a temporary file for downloading
+    result_image = Image.fromarray(results[0])
+    result_image_path = f"result_{uuid.uuid4().hex}.png"
+    result_image.save(result_image_path)
+
+    # Update the result slider with the before and after images
+    return [input_image_path, result_image_path], gr.update(value=information, visible=True), gr.update(visible=True)
 
 def load_and_reset(param_setting):
     print('Resetting parameters...')
@@ -420,18 +420,10 @@ def load_and_reset(param_setting):
     return edm_steps, s_cfg, s_stage2, s_stage1, s_churn, s_noise, a_prompt, n_prompt, color_fix_type, linear_CFG, \
         linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select
 
-def log_information(result_gallery):
+def log_information(result_slider):
     print('Logging information...')
-    if result_gallery is not None:
-        for i, result in enumerate(result_gallery):
-            print(result[0])
-
-def on_select_result(result_slider, result_gallery, evt: gr.SelectData):
-    print('Result selected.')
-    if result_gallery is not None:
-        for i, result in enumerate(result_gallery):
-            print(result[0])
-    return [result_slider[0], result_gallery[evt.index][0]]
+    if result_slider is not None:
+        print(result_slider)
 
 title_html = """
     <h1><center>Maree's Magical Photo Tool</center></h1>
@@ -647,17 +639,11 @@ with gr.Blocks() as interface:
             visible=False
         )
         result_slider = ImageSlider(
-            label='Comparator',
+            label='Result',
             show_label=False,
             interactive=False,
             elem_id="slider1",
-            show_download_button=False
-        )
-        result_gallery = gr.Gallery(
-            label='Downloadable results',
-            show_label=True,
-            interactive=False,
-            elem_id="gallery1"
+            show_download_button=True  # Enable the download button
         )
 
     input_image.upload(
@@ -718,21 +704,17 @@ with gr.Blocks() as interface:
         ],
         outputs=[
             result_slider,
-            result_gallery,
             restore_information,
             reset_btn
         ]
     ).success(
         fn=log_information,
-        inputs=[result_gallery],
+        inputs=[result_slider],
         outputs=[],
         queue=False,
         show_progress=False
     )
 
-    result_gallery.change(on_select_result, [result_slider, result_gallery], result_slider)
-    result_gallery.select(on_select_result, [result_slider, result_gallery], result_slider)
-
     restart_button.click(
         fn=load_and_reset,
         inputs=[param_setting],