diff --git a/samgeo/text_sam.py b/samgeo/text_sam.py
index 1c010111..897d7c5c 100644
--- a/samgeo/text_sam.py
+++ b/samgeo/text_sam.py
@@ -229,7 +229,12 @@ def predict_sam(self, image, boxes):
             )
             return masks.cpu()
         elif self._sam_version == 2:
-            self.sam.set_image(self.source)
+
+            if isinstance(self.source, str):
+                self.sam.set_image(self.source)
+            # If no source is set provide PIL image
+            if self.source is None:
+                self.sam.set_image(image)
             self.sam.boxes = boxes.numpy().tolist()
             masks, _, _ = self.sam.predict(
                 boxes=boxes.numpy().tolist(),