sunshangquan / Histoformer

[ECCV 2024] Histoformer: Restoring Images in Adverse Weather Conditions via Histogram Transformer
143 stars 10 forks source link

Release artifacts (models, datasets) on Hugging Face #12

Closed NielsRogge closed 1 month ago

NielsRogge commented 1 month ago

Hi @sunshangquan,

Niels here from the open-source team at Hugging Face. I discovered your work through ECCV, I indexed your paper: https://huggingface.co/papers/2407.10172. Congrats on getting it accepted! I work together with AK on improving the visibility of researchers' work on the hub.

It'd be great to make the checkpoints and data available on the πŸ€— hub, rather than Google Drive, to improve their discoverability/visibility. We can add tags so that people find them when filtering https://huggingface.co/models.

Uploading models

See here for a guide: https://huggingface.co/docs/hub/models-uploading.

In this case, we could leverage the PyTorchModelHubMixin class which adds from_pretrained and push_to_hub to any custom nn.Module. Alternatively, one can leverages the hf_hub_download one-liner to download a checkpoint from the hub.

We encourage researchers to push each model checkpoint to a separate model repository, so that things like download stats also work. We can then also link the checkpoints to the paper page.

Uploading dataset

Would be awesome to make the training dataset available on πŸ€— , so that people can do:

from datasets import load_dataset

dataset = load_dataset("your-hf-org/your-dataset")

See here for a guide: https://huggingface.co/docs/datasets/image_dataset

Besides that, there's the dataset viewer which allows people to quickly explore the first few rows of the data in the browser.

Let me know if you're interested/need any help regarding this!

Cheers,

Niels ML Engineer @ HF πŸ€—

sunshangquan commented 1 month ago

Hi @NielsRogge, thank you for your interest in our work! It’s a great pleasure to make our model available on HuggingFace. Please allow me a few days to learn how to proceed, as well as to attend to some personal matters.

sunshangquan commented 1 month ago

Thank you for the invitation again. I am uploading our work on HuggingFace at sunsean/Histoformer.

NielsRogge commented 1 month ago

Great, thanks a lot!

I see you uploaded the pickle files (.pth) to the hub, just a note that even better would be to inherit from PyTorchModelHubMixin. Inheriting from this class ensures downloads are tracked for your model, and safetensors is leveraged (a safer format to store weights - also adopted by recent models like llama).

It seems the main class for the model is ImageCleanModel, which seems like it should work by inheriting. Basically you can do:

from basicsr.model.image_restoration_model import ImageCleanModel
from huggingface_hub import hf_hub_download

# define model
model = ImageCleanModel(...)

# equip with weights
filepath = hf_hub_download(repo_id="sunsean/Histoformer", filename="net_g_best.pth")
model.load_state_dict(filepath, map_location="cpu")

# push to hub
model.push_to_hub("sunsean/Histoformer-best")

# reload
model = ImageCleanModel.from_pretrained("sunsean/Histoformer-best")

This way, model checkpoints (for best and real in this case) get pushed to separate model repositories, and downloads are tracked for both.

Would be great to try this :)

NielsRogge commented 1 month ago

Btw if you're interested, would be great to create a demo for your model on πŸ€— Spaces!

We have the image slider feature in Gradio, which seems useful here, see this recent example: https://www.linkedin.com/posts/gradio_introducing-lotus-a-diffusion-based-activity-7249013930369425408-3mo_?utm_source=share&utm_medium=member_desktop.

sunshangquan commented 1 month ago

Really appreciate your helpful advice. I will try them as soon as possible.

sunshangquan commented 1 month ago

Glad that Histoformer has been added to Hugging Face checkpoint and space, able to inference online! But now it can only run on CPU and may be very slow.

NielsRogge commented 1 month ago

You may be interested in using ZeroGPU: https://huggingface.co/zero-gpu-explorers. It allows to dynamically allocate A100 GPUs to your Space.