Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
-from transformers import AutoProcessor, AutoModelForMaskGeneration
+# from transformers import AutoProcessor, AutoModelForMaskGeneration
+from transformers import SamModel, SamProcessor
 from transformers import pipeline
 from PIL import Image, ImageOps
 # from PIL import Image
@@ -11,11 +12,13 @@ from io import BytesIO
 
 def main():
     st.title("Image Segmentation w/ Object Detection")
-    st.subheader("models : facebook/sam-vit-huge, facebook/detr-resnet-50")
+    st.subheader("models : Zigeng/SlimSAM-uniform-77, facebook/detr-resnet-50")
     
     # Load SAM by Facebook
-    processor = AutoProcessor.from_pretrained("facebook/sam-vit-huge")
-    model = AutoModelForMaskGeneration.from_pretrained("facebook/sam-vit-huge")
+    # processor = AutoProcessor.from_pretrained("facebook/sam-vit-huge")
+    # model = AutoModelForMaskGeneration.from_pretrained("facebook/sam-vit-huge")
+    model = SamModel.from_pretrained("Zigeng/SlimSAM-uniform-77")
+    processor = SamProcessor.from_pretrained("Zigeng/SlimSAM-uniform-77")
     # Load Object Detection
     od_pipe = pipeline("object-detection", "facebook/detr-resnet-50")