Update app.py

app.py

@@ -65,37 +65,27 @@ predictor = Predictor()
 
 with gr.Blocks(title="Face Swap & Enhance") as demo:
     gr.Markdown("## Face Swap & Enhancement Tool")
-
     with gr.Row():
         input_img = gr.Image(type="file", label="Foto tjeter")
         swap_img = gr.Image(type="file", label="Foto e juaj")
 
     with gr.Row():
-        output_img = gr.Image(type="filepath", label="Rezultati")
-    
-    with gr.Row():
-        submit_btn = gr.Button("Përpunoni")
-        download_btn = gr.Button("Shkarko", visible=False)
-        hidden_file = gr.File(visible=False)
+        output_img = gr.Image(type="filepath", label="Rezultati", visible=False)
+        download_btn = gr.File(label="Shkarko", visible=False)
+
+    submit_btn = gr.Button("Përpunoni")
 
     def predict_and_update_ui(img1, img2):
         result = predictor.predict(img1, img2)
         if result:
-            return result, gr.update(visible=True), gr.update(value=result)
+            return result, result, gr.update(visible=True), gr.update(visible=True)
         else:
-            return None, gr.update(visible=False), gr.update(value=None)
+            return None, None, gr.update(visible=False), gr.update(visible=False)
 
     submit_btn.click(
         fn=predict_and_update_ui,
         inputs=[input_img, swap_img],
-        outputs=[output_img, download_btn, hidden_file]
-    )
-
-    # Make download button trigger download of the hidden file
-    download_btn.click(
-        fn=lambda f: f,
-        inputs=[hidden_file],
-        outputs=[hidden_file]
+        outputs=[output_img, download_btn, output_img, download_btn]
     )
 
-demo.launch()
+demo.launch()