Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,578 Bytes
5f9d349 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
from typing import Union
import torch
from .roma_models import roma_model, tiny_roma_v1_model
weight_urls = {
"romatch": {
"outdoor": "https://github.com/Parskatt/storage/releases/download/roma/roma_outdoor.pth",
"indoor": "https://github.com/Parskatt/storage/releases/download/roma/roma_indoor.pth",
},
"tiny_roma_v1": {
"outdoor": "https://github.com/Parskatt/storage/releases/download/roma/tiny_roma_v1_outdoor.pth",
},
"dinov2": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_pretrain.pth", #hopefully this doesnt change :D
}
def tiny_roma_v1_outdoor(device, weights = None, xfeat = None):
if weights is None:
weights = torch.hub.load_state_dict_from_url(
weight_urls["tiny_roma_v1"]["outdoor"],
map_location=device)
if xfeat is None:
xfeat = torch.hub.load(
'verlab/accelerated_features',
'XFeat',
pretrained = True,
top_k = 4096).net
return tiny_roma_v1_model(weights = weights, xfeat = xfeat).to(device)
def roma_outdoor(device, weights=None, dinov2_weights=None, coarse_res: Union[int,tuple[int,int]] = 560, upsample_res: Union[int,tuple[int,int]] = 864, amp_dtype: torch.dtype = torch.float16):
if isinstance(coarse_res, int):
coarse_res = (coarse_res, coarse_res)
if isinstance(upsample_res, int):
upsample_res = (upsample_res, upsample_res)
if str(device) == 'cpu':
amp_dtype = torch.float32
assert coarse_res[0] % 14 == 0, "Needs to be multiple of 14 for backbone"
assert coarse_res[1] % 14 == 0, "Needs to be multiple of 14 for backbone"
if weights is None:
weights = torch.hub.load_state_dict_from_url(weight_urls["romatch"]["outdoor"],
map_location=device)
if dinov2_weights is None:
dinov2_weights = torch.hub.load_state_dict_from_url(weight_urls["dinov2"],
map_location=device)
model = roma_model(resolution=coarse_res, upsample_preds=True,
weights=weights,dinov2_weights = dinov2_weights,device=device, amp_dtype=amp_dtype)
model.upsample_res = upsample_res
print(f"Using coarse resolution {coarse_res}, and upsample res {model.upsample_res}")
return model
def roma_indoor(device, weights=None, dinov2_weights=None, coarse_res: Union[int,tuple[int,int]] = 560, upsample_res: Union[int,tuple[int,int]] = 864, amp_dtype: torch.dtype = torch.float16):
if isinstance(coarse_res, int):
coarse_res = (coarse_res, coarse_res)
if isinstance(upsample_res, int):
upsample_res = (upsample_res, upsample_res)
assert coarse_res[0] % 14 == 0, "Needs to be multiple of 14 for backbone"
assert coarse_res[1] % 14 == 0, "Needs to be multiple of 14 for backbone"
if weights is None:
weights = torch.hub.load_state_dict_from_url(weight_urls["romatch"]["indoor"],
map_location=device)
if dinov2_weights is None:
dinov2_weights = torch.hub.load_state_dict_from_url(weight_urls["dinov2"],
map_location=device)
model = roma_model(resolution=coarse_res, upsample_preds=True,
weights=weights,dinov2_weights = dinov2_weights,device=device, amp_dtype=amp_dtype)
model.upsample_res = upsample_res
print(f"Using coarse resolution {coarse_res}, and upsample res {model.upsample_res}")
return model
|