Convert PyTorch Model to Hugging Face model

Question

I have looked at a lot resources but I still have issues trying to convert a PyTorch model to a hugging face model format. I ultimately want to be able to use inference API with my custom model.

I have a "model.pt" file which I got from fine-tuning the Facebook Musicgen medium model (The Git repo I used to train / Fine tune the model is here). I want to upload this to the hugging face hub so i can use this with inference API. How can I convert the .pt model to files/model that can be used on hugging face hub? I tried looking at other posts but there is no clear answer, or it is poorly explained.

Any help / guidance would be greatly appreciated

This is the code I have right now that is not working:

import torch
from transformers import MusicgenConfig, MusicgenModel
from audiocraft.models import musicgen
import os

os.mkdir('models')

state_dict = musicgen.MusicGen.get_pretrained('facebook/musicgen-medium', device='cuda').lm.load_state_dict(torch.load('NEW_MODEL.pt'))

config = MusicgenConfig.from_pretrained('facebook/musicgen-medium')
model = MusicgenModel(config)
model.load_state_dict(state_dict)

model.save_pretrained('/models')

loaded_model = MusicgenModel.from_pretrained('/models')

Answer 1

In case your model is a (custom) PyTorch model, you can leverage the PyTorchModelHubMixin class available in the huggingface_hub Python library. It is a minimal class which adds from_pretrained and push_to_hub capabilities to any nn.Module, along with download metrics.

import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin


class MyModel(nn.Module, PyTorchModelHubMixin):
    def __init__(self, config: dict):
        super().__init__()
        self.param = nn.Parameter(torch.rand(config["num_channels"], config["hidden_size"]))
        self.linear = nn.Linear(config["hidden_size"], config["num_classes"])

    def forward(self, x):
        return self.linear(x + self.param)

# create model
config = {"num_channels": 3, "hidden_size": 32, "num_classes": 10}
model = MyModel(config=config)

# save locally
model.save_pretrained("my-awesome-model", config=config)

# push to the hub
model.push_to_hub("my-awesome-model", config=config)

# reload
model = MyModel.from_pretrained("username/my-awesome-model")

Here is a link to the huggingface docs, explaining how to push pytorch model to the huggingface hub.

https://huggingface.co/docs/hub/en/models-uploading#upload-a-pytorch-model-using-huggingfacehub