Improve discoverability on Hugging Face

[NielsRogge]

Hi @mv-lab and team,

Thanks for this nice work!

I wrote a quick PoC to showcase that you can easily have integration with the 🤗 hub so that you can automatically:

• load the various DarkIR models using from_pretrained (and push them using push_to_hub
• track download numbers for your models (similar to models in the Transformers library)
• leverage safetensors for weights serialization
• have nice model cards on a per-model basis. It leverages the PyTorchModelHubMixin class which allows to inherits these methods.

Usage is as follows:

from archs import DarkIR

model = DarkIR.from_pretrained("cidautai/darkir")

One can push a trained model as follows:

from archs import DarkIR

model = DarkIR()
state_dict = ...
model.load_state_dict(state_dict)

#one can push a trained model to the hub
model.push_to_hub("cidautai/darkir-example")

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

[PCMCI-Tonight]

Hi， As you take distributed training, so I am reporting error when I use single threaded training, I took some methods to change to single threaded but still reporting error, would like to ask any solution? Thank you
The image shows my changes