MegaLoc/hubconf.py at main · gmberton/MegaLoc · GitHub

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
"""PyTorch Hub configuration for MegaLoc.

Usage:
    import torch
    model = torch.hub.load("gmberton/MegaLoc", "get_trained_model")
"""

dependencies = ["torch", "torchvision", "huggingface_hub"]

import torch  # noqa: E402

from megaloc_model import MegaLoc  # noqa: E402


def get_trained_model() -> torch.nn.Module:
    """Load the pretrained MegaLoc model.

    Returns:
        MegaLoc model with pretrained weights.
    """
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file

    model = MegaLoc()

    weights_path = hf_hub_download(
        repo_id="gberton/MegaLoc",
        filename="model.safetensors",
    )

    state_dict = load_file(weights_path)
    model.load_state_dict(state_dict)

    return model