Update app.py

app.py

@@ -8,37 +8,37 @@ from transformers import AutoProcessor, CLIPSegForImageSegmentation
 processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
 model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
 
+# Ensure that the model uses GPU if available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model.to(device)
+
 def segment_everything(image):
-    # Check if image is a list and extract the actual image data
     if isinstance(image, list):
         image = image[0]
     
-    # Convert numpy array to PIL Image
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
     
-    inputs = processor(text=["object"], images=[image], padding="max_length", return_tensors="pt")
+    inputs = processor(text=["object"], images=[image], padding="max_length", return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = model(**inputs)
-    preds = outputs.logits.squeeze().sigmoid()
+    preds = outputs.logits.squeeze().sigmoid().cpu()
     segmentation = (preds.numpy() * 255).astype(np.uint8)
     return Image.fromarray(segmentation)
 
 def segment_box(image, x1, y1, x2, y2):
-    # Check if image is a list and extract the actual image data
     if isinstance(image, list):
         image = image[0]
     
-    # Convert PIL Image to numpy array if necessary
     if isinstance(image, Image.Image):
         image = np.array(image)
     
     x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
     cropped_image = image[y1:y2, x1:x2]
-    inputs = processor(text=["object"], images=[Image.fromarray(cropped_image)], padding="max_length", return_tensors="pt")
+    inputs = processor(text=["object"], images=[Image.fromarray(cropped_image)], padding="max_length", return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = model(**inputs)
-    preds = outputs.logits.squeeze().sigmoid()
+    preds = outputs.logits.squeeze().sigmoid().cpu()
     segmentation = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
     segmentation[y1:y2, x1:x2] = (preds.numpy() * 255).astype(np.uint8)
     return Image.fromarray(segmentation)
@@ -47,24 +47,19 @@ def update_image(image, segmentation):
     if segmentation is None:
         return image
     
-    # Check if image is a list and extract the actual image data
     if isinstance(image, list):
         image = image[0]
     
-    # Ensure image is in the correct format (PIL Image)
     if isinstance(image, np.ndarray):
         image_pil = Image.fromarray(image)
     else:
         image_pil = image
     
-    # Convert segmentation to RGBA
     seg_pil = Image.fromarray(segmentation).convert('RGBA')
     
-    # Resize segmentation to match input image if necessary
     if image_pil.size != seg_pil.size:
         seg_pil = seg_pil.resize(image_pil.size, Image.NEAREST)
     
-    # Blend images
     blended = Image.blend(image_pil.convert('RGBA'), seg_pil, 0.5)
     
     return np.array(blended)
@@ -101,4 +96,4 @@ with gr.Blocks() as demo:
         outputs=[output_image]
     )
 
-demo.launch()
+demo.launch()