sahl-express / preload_models.py
arij155's picture
Update preload_models.py
5bf3e85 verified
raw
history blame contribute delete
621 Bytes
import os
import torch
import torch.hub
os.environ['TORCH_HOME'] = '/tmp/torch_cache'
os.environ['YOLO_CONFIG_DIR'] = '/tmp/ultralytics_config'
# Permanently disable the trust check for ALL hub.load() calls,
# including the nested one inside MiDaS's hubconf.py
def _trust_all(*args, **kwargs):
pass
torch.hub._check_repo_is_trusted = _trust_all
print("Pre-downloading MiDaS and dependencies...")
midas = torch.hub.load(
"intel-isl/MiDaS",
"MiDaS_small",
trust_repo=True
)
transforms = torch.hub.load(
"intel-isl/MiDaS",
"transforms",
trust_repo=True
)
print("MiDaS cached successfully.")