Changed checking of cached data to take into account people using the app at the same time

app.py

@@ -8,12 +8,13 @@ import cv2
 from typing import List
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+device = 'cpu'
 
 # Load model and processor
 model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
 processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 
-embedding = None
+cache_data = None
 
 def mask_2_dots(mask: np.ndarray) -> List[List[int]]:
     gray = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
@@ -31,14 +32,16 @@ def mask_2_dots(mask: np.ndarray) -> List[List[int]]:
 
 @torch.no_grad()
 def foward_pass(image_input: np.ndarray, points: List[List[int]]) -> np.ndarray:
-    global embedding
+    global cache_data
     image_input = Image.fromarray(image_input)
     inputs = processor(image_input, input_points=points, return_tensors="pt").to(device)
-    if not isinstance(embedding, torch.Tensor):
+    if not cache_data or not torch.equal(inputs['pixel_values'],cache_data[0]):
         embedding = model.get_image_embeddings(inputs["pixel_values"])
+        pixels = inputs["pixel_values"]
+        cache_data = [pixels, embedding]
     del inputs["pixel_values"]
 
-    outputs = model.forward(image_embeddings=embedding, **inputs)
+    outputs = model.forward(image_embeddings=cache_data[1], **inputs)
     masks = processor.image_processor.post_process_masks(
         outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
     )
@@ -63,9 +66,9 @@ def main_func(inputs) -> List[Image.Image]:
 
     return pred_masks
 
-def reset_embedding():
-    global embedding
-    embedding = None
+def reset_data():
+    global cache_data
+    cache_data = None
 
 with gr.Blocks() as demo:
     gr.Markdown("# How to use")
@@ -81,6 +84,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)
+    image_input.upload(reset_data)
 
 demo.launch()