PyTorchModelHubMixin Help

I need some help with uploading models to the Hugging Face repo to later access them with inference endpoints. I’m planning on training a custom DreamBooth model, but I’m confused about how to upload it. I guess I should be using PyTorchModelHubMixin, but I don’t really know how to use it in my notebook for this kind of model. Can someone please help me out?

Hi,

Sure happy to help out. What does your model look like?

It’s basically this: Google Colab

I managed to upload something just for testing, but I’m not really sure what to do with the above Colab notebook

Hi @nielsr,

I would also like to know how to upload a ControlNet model to the hub and then send an image through the inference API as an input to the model

Control Net colab : Google Colab

Hi,

As explained here, if your model is available in a native library (like Diffusers or Transformers), you can use the following:

from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(model_path, safety_checker=None, torch_dtype=torch.float16).to("cuda")

# push to hub
pipe.push_to_hub("...")

This is explained here: Push files to the Hub.

However if you have a custom nn.Module, you can leverage the PyTorchModelHubMixin. I don’t see anything custom in your notebook which means that you can just use the native methods.