Update app.py

app.py

@@ -1,19 +1,20 @@
-from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
 import gradio as gr
 from PIL import Image
 import torch
 import matplotlib.pyplot as plt
-import torch
 import numpy as np
+import spaces
+from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
 
 processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
-model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
-
+model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined").cuda()
 
+@spaces.GPU
 def process_image(image, prompt):
     inputs = processor(
         text=prompt, images=image, padding="max_length", return_tensors="pt"
     )
+    inputs = {k: v.cuda() for k, v in inputs.items()}
 
     # predict
     with torch.no_grad():
@@ -33,22 +34,22 @@ def process_image(image, prompt):
     mask = (mask - mask_min) / (mask_max - mask_min)
     return mask
 
-
-def get_masks(prompts, img, threhsold):
+@spaces.GPU
+def get_masks(prompts, img, threshold):
     prompts = prompts.split(",")
     masks = []
     for prompt in prompts:
         mask = process_image(img, prompt)
-        mask = mask > threhsold
+        mask = mask > threshold
         masks.append(mask)
     return masks
 
-
-def extract_image(pos_prompts, neg_prompts, img, threhsold):
+@spaces.GPU
+def extract_image(pos_prompts, neg_prompts, img, threshold):
     positive_masks = get_masks(pos_prompts, img, 0.5)
     negative_masks = get_masks(neg_prompts, img, 0.5)
 
-    # combine masks into one masks, logic OR
+    # combine masks into one mask, logic OR
     pos_mask = np.any(np.stack(positive_masks), axis=0)
     neg_mask = np.any(np.stack(negative_masks), axis=0)
     final_mask = pos_mask & ~neg_mask
@@ -59,12 +60,10 @@ def extract_image(pos_prompts, neg_prompts, img, threhsold):
     output_image.paste(img, mask=final_mask)
     return output_image, final_mask
 
-
 title = "Interactive demo: zero-shot image segmentation with CLIPSeg"
 description = "Demo for using CLIPSeg, a CLIP-based model for zero- and one-shot image segmentation. To use it, simply upload an image and add a text to mask (identify in the image), or use one of the examples below and click 'submit'. Results will show up in a few seconds."
 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.10003'>CLIPSeg: Image Segmentation Using Text and Image Prompts</a> | <a href='https://huggingface.co/docs/transformers/main/en/model_doc/clipseg'>HuggingFace docs</a></p>"
 
-
 with gr.Blocks() as demo:
     gr.Markdown("# CLIPSeg: Image Segmentation Using Text and Image Prompts")
     gr.Markdown(article)
@@ -100,5 +99,4 @@ with gr.Blocks() as demo:
         outputs=[output_image, output_mask],
     )
 
-
 demo.launch()