Skip to content

Add PyTorchModelHubMixin to enable downloads on HF #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

NielsRogge
Copy link

@NielsRogge NielsRogge commented Jan 7, 2025

Hi @gastruc,

Thanks for this nice work and integrating with the hub already!

Fixes #1

I wrote a quick PoC to showcase that makes HF integration even easier by leveraging the PyTorchModelHubMixin class which allows to:

  • track download numbers for your models (similar to models in the Transformers library)
  • leverage safetensors for weights serialization
  • add from_pretrained and push_to_hub automatically.

Usage is as follows:

from hubconf import AnySat

# load model
model = AnySat(...)

# equip with weights
model.load_state_dict_from_url("https://huggingface.co/g-astruc/AnySat/blob/main/models/AnySat.pth")

# push to the hub
model.push_to_hub("nielsr/my-awesome-anysat-model")

# great! We can now use it as follows (+ this increments download numbers)
model = AnySat.from_pretrained("nielsr/anysat-demo")

Would you be interested in this integration?

Kind regards,

Niels

@gastruc
Copy link
Owner

gastruc commented Jan 9, 2025

Hi,
Thank you for your contribution!
It seems that your PR also remove torch hub integration. I'd like to keep that :)

@NielsRogge
Copy link
Author

NielsRogge commented Jan 9, 2025

Hi, thanks for your reply.

The model is still compatible with Torch hub, it can be loaded as follows:

import torch

AnySat = torch.hub.load('gastruc/anysat', 'anysat', pretrained=True, flash_attn=False)

Since this PR does not touch the hubconf.py file. You could theoretically add both .pt and safetensors weights in the gastruc/anysat repository.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants