From 2a8c67fe9e3532361e222ca52f470e3ff1ed34bc Mon Sep 17 00:00:00 2001 From: miraodasilva Date: Wed, 1 Nov 2023 15:11:38 +0900 Subject: [PATCH] download imagebind_huge checkpoint to cache instead of local dir --- imagebind/models/imagebind_model.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/imagebind/models/imagebind_model.py b/imagebind/models/imagebind_model.py index c560945f..e564fdea 100644 --- a/imagebind/models/imagebind_model.py +++ b/imagebind/models/imagebind_model.py @@ -490,17 +490,18 @@ def imagebind_huge(pretrained=False): ) if pretrained: - if not os.path.exists(".checkpoints/imagebind_huge.pth"): + cache_dir = os.path.expanduser("~/.cache/imagebind_checkpoints") + if not os.path.exists("%s/imagebind_huge.pth" % (cache_dir)): print( - "Downloading imagebind weights to .checkpoints/imagebind_huge.pth ..." + "Downloading imagebind weights to %s/imagebind_huge.pth ..." % (cache_dir) ) - os.makedirs(".checkpoints", exist_ok=True) + os.makedirs(cache_dir, exist_ok=True) torch.hub.download_url_to_file( "https://dl.fbaipublicfiles.com/imagebind/imagebind_huge.pth", - ".checkpoints/imagebind_huge.pth", + "%s/imagebind_huge.pth" % (cache_dir), progress=True, ) - model.load_state_dict(torch.load(".checkpoints/imagebind_huge.pth")) + model.load_state_dict(torch.load("%s/imagebind_huge.pth" % (cache_dir))) return model