|
@@ -83,7 +83,10 @@ def build_sam2_hf(model_id, **kwargs):
|
|
|
model_id_to_filenames = {
|
|
model_id_to_filenames = {
|
|
|
"facebook/sam2-hiera-tiny": ("sam2_hiera_t.yaml", "sam2_hiera_tiny.pt"),
|
|
"facebook/sam2-hiera-tiny": ("sam2_hiera_t.yaml", "sam2_hiera_tiny.pt"),
|
|
|
"facebook/sam2-hiera-small": ("sam2_hiera_s.yaml", "sam2_hiera_small.pt"),
|
|
"facebook/sam2-hiera-small": ("sam2_hiera_s.yaml", "sam2_hiera_small.pt"),
|
|
|
- "facebook/sam2-hiera-base-plus": ("sam2_hiera_b+.yaml", "sam2_hiera_base_plus.pt"),
|
|
|
|
|
|
|
+ "facebook/sam2-hiera-base-plus": (
|
|
|
|
|
+ "sam2_hiera_b+.yaml",
|
|
|
|
|
+ "sam2_hiera_base_plus.pt",
|
|
|
|
|
+ ),
|
|
|
"facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"),
|
|
"facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"),
|
|
|
}
|
|
}
|
|
|
config_name, checkpoint_name = model_id_to_filenames[model_id]
|
|
config_name, checkpoint_name = model_id_to_filenames[model_id]
|
|
@@ -98,12 +101,17 @@ def build_sam2_video_predictor_hf(model_id, **kwargs):
|
|
|
model_id_to_filenames = {
|
|
model_id_to_filenames = {
|
|
|
"facebook/sam2-hiera-tiny": ("sam2_hiera_t.yaml", "sam2_hiera_tiny.pt"),
|
|
"facebook/sam2-hiera-tiny": ("sam2_hiera_t.yaml", "sam2_hiera_tiny.pt"),
|
|
|
"facebook/sam2-hiera-small": ("sam2_hiera_s.yaml", "sam2_hiera_small.pt"),
|
|
"facebook/sam2-hiera-small": ("sam2_hiera_s.yaml", "sam2_hiera_small.pt"),
|
|
|
- "facebook/sam2-hiera-base-plus": ("sam2_hiera_b+.yaml", "sam2_hiera_base_plus.pt"),
|
|
|
|
|
|
|
+ "facebook/sam2-hiera-base-plus": (
|
|
|
|
|
+ "sam2_hiera_b+.yaml",
|
|
|
|
|
+ "sam2_hiera_base_plus.pt",
|
|
|
|
|
+ ),
|
|
|
"facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"),
|
|
"facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"),
|
|
|
}
|
|
}
|
|
|
config_name, checkpoint_name = model_id_to_filenames[model_id]
|
|
config_name, checkpoint_name = model_id_to_filenames[model_id]
|
|
|
ckpt_path = hf_hub_download(repo_id=model_id, filename=checkpoint_name)
|
|
ckpt_path = hf_hub_download(repo_id=model_id, filename=checkpoint_name)
|
|
|
- return build_sam2_video_predictor(config_file=config_name, ckpt_path=ckpt_path, **kwargs)
|
|
|
|
|
|
|
+ return build_sam2_video_predictor(
|
|
|
|
|
+ config_file=config_name, ckpt_path=ckpt_path, **kwargs
|
|
|
|
|
+ )
|
|
|
|
|
|
|
|
|
|
|
|
|
def _load_checkpoint(model, ckpt_path):
|
|
def _load_checkpoint(model, ckpt_path):
|
|
@@ -116,4 +124,4 @@ def _load_checkpoint(model, ckpt_path):
|
|
|
if unexpected_keys:
|
|
if unexpected_keys:
|
|
|
logging.error(unexpected_keys)
|
|
logging.error(unexpected_keys)
|
|
|
raise RuntimeError()
|
|
raise RuntimeError()
|
|
|
- logging.info("Loaded checkpoint sucessfully")
|
|
|
|
|
|
|
+ logging.info("Loaded checkpoint sucessfully")
|