Update app.py

app.py

@@ -68,6 +68,11 @@ transform = transforms.Compose([
 # INFERENCE FUNCTION
 # ---------------------------
 
+import gradio as gr
+import numpy as np
+from PIL import Image
+import tempfile
+
 def enhance_image(input_image):
     img = input_image.convert("RGB")
 
@@ -79,20 +84,27 @@ def enhance_image(input_image):
     output_img = output.squeeze().permute(1,2,0).cpu().numpy()
     output_img = (output_img * 255).astype(np.uint8)
 
-    return output_img
+    # 🔥 SAVE TEMP FILE FOR DOWNLOAD
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    Image.fromarray(output_img).save(temp_file.name)
+
+    return output_img, temp_file.name
 
 
 with gr.Blocks() as demo:
     gr.Markdown("# 🔍 Image Super Resolution")
-    
+
     input_img = gr.Image(type="pil", label="Upload Image")
-    output_img = gr.Image(type="numpy", label="Enhanced Image")
+    output_img = gr.Image(label="Enhanced Image")
+    download_file = gr.File(label="Download Image")
 
     btn = gr.Button("Enhance Image")
 
-    btn.click(fn=enhance_image, inputs=input_img, outputs=output_img)
-
-    gr.DownloadButton(label="Download Enhanced Image", data=output_img)
+    btn.click(
+        fn=enhance_image,
+        inputs=input_img,
+        outputs=[output_img, download_file]
+    )
 
 demo.launch()