Update app.py

app.py

@@ -1,12 +1,11 @@
 import gradio as gr
 from gradio_imageslider import ImageSlider
 from loadimg import load_img
-import spaces
 from transformers import AutoModelForImageSegmentation
 import torch
 from torchvision import transforms
-import zipfile
 import os
+import zipfile
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
@@ -22,12 +21,10 @@ transform_image = transforms.Compose(
     ]
 )
 
-@spaces.GPU
 def fn(image):
     im = load_img(image, output_type="pil")
     im = im.convert("RGB")
     image_size = im.size
-    origin = im.copy()
     input_images = transform_image(im).unsqueeze(0).to("cpu")
     
     with torch.no_grad():
@@ -35,34 +32,32 @@ def fn(image):
     pred = preds[0].squeeze()
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
+    
     im.putalpha(mask)
-
     output_file_path = os.path.join("output_images", "output_image_single.png")
     im.save(output_file_path)
 
-    return (im, origin)
+    return [mask, im]
 
-@spaces.GPU
 def fn_url(url):
     im = load_img(url, output_type="pil")
     im = im.convert("RGB")
-    origin = im.copy() 
     image_size = im.size
     input_images = transform_image(im).unsqueeze(0).to("cpu")
     
+    # Prediction
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()
     pred = preds[0].squeeze()
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
+    
     im.putalpha(mask)
-
     output_file_path = os.path.join("output_images", "output_image_url.png")
     im.save(output_file_path)
 
-    return [im, origin] 
+    return [mask, im]
 
-@spaces.GPU
 def batch_fn(images):
     output_paths = []
     for idx, image_path in enumerate(images):
@@ -76,6 +71,7 @@ def batch_fn(images):
         pred = preds[0].squeeze()
         pred_pil = transforms.ToPILImage()(pred)
         mask = pred_pil.resize(image_size)
+        
         im.putalpha(mask)
 
         output_file_path = os.path.join("output_images", f"output_image_batch_{idx + 1}.png")
@@ -89,7 +85,7 @@ def batch_fn(images):
 
     return zip_file_path
 
-batch_image = gr.File(label="Upload multiple images", type="filepath", file_count="multiple")  # 複数画像のアップロードを許可
+batch_image = gr.File(label="Upload multiple images", type="filepath", file_count="multiple")
 
 slider1 = ImageSlider(label="Processed Image", type="pil")
 slider2 = ImageSlider(label="Processed Image from URL", type="pil")
@@ -109,12 +105,7 @@ tab3 = gr.Interface(
     batch_fn, 
     inputs=batch_image, 
     outputs=gr.File(label="Download Processed Files"), 
-    api_name="batch",
-    css="""
-    #component-37 {
-        display: none;
-    }
-    """
+    api_name="batch"
 )
 
 demo = gr.TabbedInterface(