Add PyTorchModelHubMixin for better HF integration

[NielsRogge]

Hi @haofeixu,

Thanks for this nice work, great to see a demo for UniMatch being released on the hub! I indexed the paper here: https://huggingface.co/papers/2211.05783. Right now, it does not show linked models/Spaces.

This PR proposes to integrate the PyTorchModelHubMixin class developed by the 🤗 team to make sure a custom PyTorch model like yours:

• automatically inherits from_pretrained and push_to_hub methods
• has better metadata
• has working download stats.

Usage is as follows:

from unimatch.unimatch import UniMatch
from huggingface_hub import hf_hub_download

# Load a pre-trained model
model = UniMatch(...)
filepath = hf_hub_download(repo_id="jhaofeixu/unimatch", filename="pretrained/gmflow-scale2-regrefine6-mixdata-train320x576-4e7b215d.pth", repo_type="space")
model.load_state_dict(filepath, map_location="cpu")

# One can optionally push this to the hub
model.push_to_hub("jhaofeixu/unimatch-gmflow-scale2-regrefine6-mixdata-train320x576")

# Can now be reloaded as follows (and will increment download count)
model = UniMatch.from_pretrained("jhaofeixu/unimatch-gmflow-scale2-regrefine6-mixdata-train320x576")

So to push each of the UniMatch checkpoints, one would need to run the script as shown above.

Besides that, one can link the models/Spaces demo to the paper page (by including this link: https://huggingface.co/papers/2211.05783 in the respective READMEs).

Let me know whether you need any help!

Kind regards,

Niels from HF