1import os
2import hashlib
3
4import torch
5
6
7urls = dict(
8 nano='https://hf.co/nasaharvest/galileo/resolve/main/models/nano/encoder.pt',
9 tiny='https://hf.co/nasaharvest/galileo/resolve/main/models/tiny/encoder.pt',
10 base='https://hf.co/nasaharvest/galileo/resolve/main/models/base/encoder.pt',
11)
12for name, url in urls.items():
13 state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu", file_name=f'model_{name}.pth')
14 state_dict = {k.replace('.backbone', ''): v for k, v in state_dict.items()}
15 filename = f'model_{name}.pth'
16 torch.save(state_dict, filename)
17 md5 = hashlib.md5(open(filename, "rb").read()).hexdigest()[:8]
18 os.rename(filename, filename.replace(".pth", f"-{md5}.pth"))