Make sure Hugging Face download stats work, better discoverability

[NielsRogge]

Hi @vsehwag-sony,

Thanks for this nice work! Great to see the models being released on the hub :)

This PR proposes to improve the 🤗 integration by leveraging the PyTorchModelHubMixin class. This allows:

• to automatically load the various DiT models using from_pretrained (and push them using push_to_hub)
• better discoverability using the text-to-image tag
• track download numbers for your models (similar to models in the Transformers library)
• use safetensors instead of pickle for weights serialization
• have nice model cards on a per-model basis.

Usage is as follows:

from micro_diffusion.models.dit import DiT

model = DiT.from_pretrained("sony/MicroDiT-16-channel")

See also the push_and_load_from_hf.py script.

Would you be interested in this integration?

Kind regards,

Niels

[NielsRogge]

Hi @vsehwag-sony is this something that interests you?

[NielsRogge]

HI @vsehwag-sony I guess I can close this PR?