Update app.py

app.py

@@ -24,7 +24,7 @@ transform_image = transforms.Compose(
 @spaces.GPU
 def fn(image):
     if image is None or len(image) == 0:
-        return image
+        return image, None  # 원본 이미지도 반환
     im = load_img(image, output_type="pil")
     im = im.convert("RGB")
     image_size = im.size
@@ -38,24 +38,38 @@ def fn(image):
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
     image.putalpha(mask)
-    return (image, origin)
+    return image, origin  # 변환된 이미지와 원본 이미지 반환
+
+
+def save_image(image):
+    if image is not None:
+        image.save("output.png")
+        return "output.png"
+    return None
 
 
 slider1 = ImageSlider(label="birefnet", type="pil")
 slider2 = ImageSlider(label="birefnet", type="pil")
 image = gr.Image(label="Upload an image")
 text = gr.Textbox(label="Paste an image URL")
-
+download_button = gr.Button("Download Image")  # 다운로드 버튼 추가
+output_file = gr.File()  # 파일 출력 추가
 
 chameleon = load_img("butterfly.jpg", output_type="pil")
 
 url = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
 tab1 = gr.Interface(
-    fn, inputs=image, outputs=slider1, examples=[chameleon], api_name="image"
+    fn, inputs=image, outputs=[slider1, output_file], examples=[chameleon], api_name="image"
 )
 
-tab2 = gr.Interface(fn, inputs=text, outputs=slider2, examples=[url], api_name="text")
+tab2 = gr.Interface(fn, inputs=text, outputs=[slider2, output_file], examples=[url], api_name="text")
+
+
+def process_download(image):
+    return save_image(image[0])
+
 
+download_button.click(process_download, inputs=slider1, outputs=output_file)
 
 demo = gr.TabbedInterface(
     [tab1, tab2], ["image", "text"], title="birefnet for background removal"