🤗 Add support for downloading Huggingface weights.

[lcxrocks]

Now supports download weight from huggingface directly by calling model.from_pretrained(model_name) as demonstrated in hf_demo_2x.py.

model = Model(-1)
model.from_pretrained(args.model)

[NielsRogge]

Note, it would still be useful to store the checkpoints in separate model repos (one for VFIMamba and one for VFIMamba_S), this to ensure downloads work: https://huggingface.co/docs/hub/models-download-stats (if you're interested in seeing how many times people use your model)

[NielsRogge]

or alternatively, you could open a PR on huggingface.js to specify a file extension to count downloads for.

[lcxrocks]

Thank you for your advice! Pushed a new version.

[Wauplin]

Suggested change

	def from_pretrained(self, model_name):
	try:
	from huggingface_hub import hf_hub_download
	assert model_name in ["VFIMamba", "VFIMamba_S"], "Please select a valid model name from ['VFIMamba', 'VFIMamba_S']"
	ckpt_path = hf_hub_download(
	repo_id=f"MCG-NJU/{model_name}", filename=model_name + ".pkl"
	)
	checkpoint = torch.load(ckpt_path)
	except:
	# In case the model is not hosted on huggingface
	# or the user cannot import huggingface_hub correctly, model_name option: VFIMamba, VFIMamba_S
	_VFIMAMBA_URL = f"https://huggingface.co/MCG-NJU/{model_name}/resolve/main/{model_name}.pkl"
	checkpoint = torch.hub.load_state_dict_from_url(_VFIMAMBA_URL)
	self.net.load_state_dict(convert(checkpoint), strict=True)
	@classmethod
	def from_pretrained(cls, model_id: str, local_rank: int) -> "Model":
	try:
	from huggingface_hub import hf_hub_download
	except ImportError:
	raise ImportError(
	"Model is hosted on the Hugging Face Hub. "
	"Please install huggingface_hub by running `pip install huggingface_hub` to load the weights correctly."
	)
	if "/" not in model_id:
	model_id = "MCG-NJU/" + model_id
	ckpt_path = hf_hub_download(repo_id=model_id, filename="model.pkl")
	checkpoint = torch.load(ckpt_path)
	model = cls(local_rank)
	model.net.load_state_dict(convert(checkpoint), strict=True)
	return model

Hi there, maintainer of huggingface_hub here 👋 May I suggest a more opinionated implementation for this method? In particular:

• I would make huggingface_hub required and raise an explicit error if it's not the case. This way you don't have to maintain different ways to load the model. In particular, end users will benefit from the huggingface_hub cache when running the script several times.
• Now that models are in two separate repos, I would rename the weights files to model.pkl in both cases. This way, the repo id and filename don't have to be correlated. Also, it would simplify the way downloads are counted.
• I would also accept any model_id as input, not only the 2 official one from the repo. This way, your library can be reused by anyone wanting the train, retrain, fine-tuned, etc new models based on your work. This should greatly improve adoption of the library. model_id can be either a full repo id (e.g. MCG-NJU/VFIMamba_S) or simply a model name in which case the MCG-NJU/ org is prefixed.
• I would make from_pretrained a class method to initialize the object directly with model = Model.from_pretrained("VFIMamba_S", local_rank=-1)

Those changes are not an obligation to make things work with the Hub but in my opinion it would greatly improve the integration. Let me know what you think!

(cc @hanouticelina for viz')

[lcxrocks]

Thank you very much! Your suggestion is very helpful. I've made some changes based on your advice.

[Wauplin]

Could these config values be stored in a config.json file alongside the model weights on the Hub? This way they could be downloaded when instantiating the model and therefore reduce the boilerplate code to use Model.

(just a suggestion, I'm not entirely sure how cfg and Model are interacting currently)

[lcxrocks]

Absolutely! Thanks to your advice, we can now directly initialize a model when loading weights from Hugging Face using the following code (as also demonstrated in hf_demo_2x.py):

model = Model.from_pretrained(args.model)

Cheers!

[Wauplin]

Awesome!

[NielsRogge]

Hi @GuozhenZhang1999,

Would be great to merge the following 2 PRs:

• https://huggingface.co/MCG-NJU/VFIMamba_S/discussions/1
• https://huggingface.co/MCG-NJU/VFIMamba/discussions/1

This will enable huggingface/huggingface.js#885 to be merged