Added no grad and storing embeddings

app.py

@@ -13,6 +13,8 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
 processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 
+embedding = None
+
 def mask_2_dots(mask: np.ndarray) -> List[List[int]]:
     gray = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
     _, thresh = cv2.threshold(gray, 127, 255, 0)
@@ -27,11 +29,16 @@ def mask_2_dots(mask: np.ndarray) -> List[List[int]]:
         points.append([cx, cy])
     return [points]
 
+@torch.no_grad()
 def foward_pass(image_input: np.ndarray, points: List[List[int]]) -> np.ndarray:
+    global embedding
     image_input = Image.fromarray(image_input)
-
     inputs = processor(image_input, input_points=points, return_tensors="pt").to(device)
-    outputs = model(**inputs)
+    if not isinstance(embedding, torch.Tensor):
+        embedding = model.get_image_embeddings(inputs["pixel_values"])
+    del inputs["pixel_values"]
+
+    outputs = model.forward(image_embeddings=embedding, **inputs)
     masks = processor.image_processor.post_process_masks(
         outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
     )
@@ -56,6 +63,9 @@ def main_func(inputs) -> List[Image.Image]:
 
     return pred_masks
 
+def reset_embedding():
+    global embedding
+    embedding = None
 
 with gr.Blocks() as demo:
     gr.Markdown("# How to use")
@@ -71,5 +81,6 @@ with gr.Blocks() as demo:
         image_button = gr.Button("Segment Image")
 
     image_button.click(main_func, inputs=image_input, outputs=image_output)
+    image_input.upload(reset_embedding)
 
 demo.launch()

requirements.txt

@@ -1,3 +1,4 @@
 torch
 git+https://github.com/huggingface/transformers
-opencv-python
+opencv-python
+gradio --upgrade