Improve HF integration

[NielsRogge]

Hi @feizc,

Thanks for this nice work! I see the weights are already on the hub (nice!) but things like download counts don't work.

This PR aims to improve the HF integration:

• automatically load the model using from_pretrained (and push it using push_to_hub)
• track download numbers for your models (similar to models in the Transformers library)
• have nice model cards on a per-model basis along with appropriate tags (so that people find them when filtering https://huggingface.co/models?pipeline_tag=text-to-audio)
• perhaps most importantly, leverage safetensors for the weights in favor of pickle.

It leverages the PyTorchModelHubMixin class which allows to inherits these methods.

Usage is as follows:

from model import Flux

# load from hub
model = Flux.from_pretrained("feizhengcong/FluxMusic-small")

# optionally, save locally or push to the hub
model.save_pretrained("...")
model.push_to_hub("your-hf-org-or-username/flux-small")

This means people don't need to manually download a checkpoint first in their local environment, it just loads automatically from the hub.

Would you be interested in this integration?

Kind regards,

Niels

Note

Please don't merge this PR before pushing the model to the hub :)