Update app.py

app.py

@@ -11,11 +11,15 @@ def load_sam_model():
     checkpoint_path = hf_hub_download(repo_id="facebook/sam-vit-huge", filename="pytorch_model.bin")
     
     # Load checkpoint với map_location=torch.device('cpu')
+    checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
+    
+    # Khởi tạo mô hình SAM
     model_type = "vit_h"
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    # Sử dụng map_location=torch.device('cpu') để tải mô hình trên CPU
-    sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
+    # Truyền checkpoint vào mô hình
+    sam = sam_model_registry[model_type]()
+    sam.load_state_dict(checkpoint)
     sam.to(device=device)
     predictor = SamPredictor(sam)
     return predictor