File size: 1,215 Bytes
44039c6 |
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 |
import torch
from huggingface_hub import hf_hub_download
# Download the model file from Hugging Face Hub
repo_name = "roughness_model"
downloaded_file = hf_hub_download(
repo_id=f"Pra-tham/{repo_name}", # Replace with your Hugging Face username
filename="roughness_model.pth"
)
print(f"Model downloaded from Hugging Face Hub: {downloaded_file}")
# Initialize the model and load the state_dict
model.load_state_dict(torch.load(downloaded_file))
model.eval() # Set to evaluation mode
print("Model loaded successfully from Hugging Face Hub!")
# Set the computation device
device0 = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Load the pretrained model checkpoint
load_path = "/kaggle/working/metric_depth_vit_large_800k.pth"
checkpoint = torch.load(load_path, map_location="cpu")
# Load the model configuration
cfg_large = Config.fromfile('/kaggle/working/Texture_training/training/mono/configs/RAFTDecoder/vit.raft5.large.py')
# Initialize the DepthModel
model = DepthModel(cfg_large, None)
# Load the model's state dictionary
ckpt_state_dict = checkpoint['model_state_dict']
model.load_state_dict(ckpt_state_dict, strict=False)
# Print the model architecture
#print(model)
|