Update app.py

app.py

@@ -6,13 +6,15 @@ import torch
 from torchvision import transforms
 import os
 import zipfile
+import numpy as np
+from PIL import Image
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
-birefnet.to("cpu")
+birefnet.to("cuda")
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -25,7 +27,7 @@ def fn(image):
     im = load_img(image, output_type="pil")
     im = im.convert("RGB")
     image_size = im.size
-    input_images = transform_image(im).unsqueeze(0).to("cpu")
+    input_images = transform_image(im).unsqueeze(0).to("cuda")
     
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()
@@ -37,13 +39,16 @@ def fn(image):
     output_file_path = os.path.join("output_images", "output_image_single.png")
     im.save(output_file_path)
 
-    return [mask, im]
+    output_path = os.path.join("output_images", "output_image_processed.png")
+    im.save(output_path, "PNG")
+    
+    return [im, mask], output_path
 
 def fn_url(url):
     im = load_img(url, output_type="pil")
     im = im.convert("RGB")
     image_size = im.size
-    input_images = transform_image(im).unsqueeze(0).to("cpu")
+    input_images = transform_image(im).unsqueeze(0).to("cuda")
     
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()
@@ -55,7 +60,10 @@ def fn_url(url):
     output_file_path = os.path.join("output_images", "output_image_url.png")
     im.save(output_file_path)
 
-    return [mask, im]
+    output_path = os.path.join("output_images", "output_image_url_processed.png")
+    im.save(output_path, "PNG")
+    
+    return [im, mask], output_path
 
 def batch_fn(images):
     output_paths = []
@@ -63,7 +71,7 @@ def batch_fn(images):
         im = load_img(image_path, output_type="pil")
         im = im.convert("RGB")
         image_size = im.size
-        input_images = transform_image(im).unsqueeze(0).to("cpu")
+        input_images = transform_image(im).unsqueeze(0).to("cuda")
         
         with torch.no_grad():
             preds = birefnet(input_images)[-1].sigmoid().cpu()
@@ -71,7 +79,7 @@ def batch_fn(images):
         pred_pil = transforms.ToPILImage()(pred)
         mask = pred_pil.resize(image_size)
         
-        im.putalpha(mask)
+        im.putalpha(mask) 
 
         output_file_path = os.path.join("output_images", f"output_image_batch_{idx + 1}.png")
         im.save(output_file_path)
@@ -95,10 +103,10 @@ chameleon = load_img("chameleon.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, gr.File(label="PNG Output")], examples=[chameleon], api_name="image"
 )
 
-tab2 = gr.Interface(fn_url, inputs=text, outputs=slider2, examples=[url], api_name="text")
+tab2 = gr.Interface(fn_url, inputs=text, outputs=[slider2, gr.File(label="PNG Output")], examples=[url], api_name="text")
 
 tab3 = gr.Interface(
     batch_fn,