From 10c3753d7f9ed7cf2415091160f08ee45e30a681 Mon Sep 17 00:00:00 2001
From: user1 <gregghelt@gmail.com>
Date: Sun, 25 Jun 2023 11:16:39 -0700
Subject: [PATCH] Added SAM preprocessor

---
 .../app/invocations/controlnet_image_processors.py  | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/invokeai/app/invocations/controlnet_image_processors.py b/invokeai/app/invocations/controlnet_image_processors.py
index baf558ac24..0fdfc12905 100644
--- a/invokeai/app/invocations/controlnet_image_processors.py
+++ b/invokeai/app/invocations/controlnet_image_processors.py
@@ -29,6 +29,7 @@ from controlnet_aux import (
     ContentShuffleDetector,
     ZoeDetector,
     MediapipeFaceDetector,
+    SamDetector,
 )
 
 from .image import ImageOutput, PILInvocationConfig
@@ -455,3 +456,15 @@ class MediapipeFaceProcessorInvocation(ImageProcessorInvocation, PILInvocationCo
         mediapipe_face_processor = MediapipeFaceDetector()
         processed_image = mediapipe_face_processor(image, max_faces=self.max_faces, min_confidence=self.min_confidence)
         return processed_image
+
+
+class SegmentAnythingProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
+    """Applies segment anything processing to image"""
+    # fmt: off
+    type: Literal["segment_anything_processor"] = "segment_anything_processor"
+    # fmt: on
+
+    def run_processor(self, image):
+        segment_anything_processor = SamDetector.from_pretrained("ybelkada/segment-anything", subfolder="checkpoints")
+        processed_image = segment_anything_processor(image)
+        return processed_image