| 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) |
| |
