Custom models in .ckpt format

[irgolic]

Primarly, enable loading a model by passing a URL to a .ckpt in the model parameter. As diffusers does, these models should be cached locally. Perhaps it would make sense to also cache them with BlobRepo for reliable initialization of new runners.

Also allow preloading models by saving them in a subfolder. AFAIK, diffusers requires the models to be saved in their own format (as output by from_pretrained). There should be an easy way to convert a .ckpt to that format. Or maybe I'm missing something obvious about how to load a .ckpt model.

[MicahZoltu]

IIUC, the .ckpt files are the "big" (GBs) part of all of this? If so, I wonder if it makes the most sense to have them baked into the Docker image (thus in a subfolder as you suggested). If they need to be in a certain format, it seems like that should be done prior to baking them into the image.

You can utilize multistage builds and COPY --link to make it so once we get these files into an image you should never need to fetch or build them again unless they are updated (even as you iterate on other pieces of the system).

[irgolic]

IIUC, the .ckpt files are the "big" (GBs) part of all of this? If they need to be in a certain format, it seems like that should be done prior to baking them into the image.

Correct. We don't use .ckpt's right now.

At the moment, models are downloaded and cached by huggingface. Every time a pipeline is evoked, it's checked whether the model exists in the local cache directory, and if not, it's implicitly downloaded prior to execution. On Mac the directory defaults to ~/.cache/huggingface/diffusers, and looks a bit like this:

If we bake a model into the image, it'd make the most sense to cache the mainstream stable-diffusion-v1-5. However, people may want to run another model altogether; having 4GB taken up in the image is erroneous to them. I'd rather keep the image light, and allow runners to cache the models locally, in a folder accessible by any ephemeral process running huggingface. Maybe a baked-in image could be offered alongside a light one.

Honestly, the docker images need some love:

• all processes (api, runner) use the same docker image. The API wouldn't need the baked-in model
• I've not properly tested GPU runners through docker, it's likely they need separate NVIDIA and AMD Dockerfiles that install the appropriate version of pytorch
• CI rebuilds the docker image from scratch every time

All that said, in the context of this issue, my priority is to allow referencing a model by .ckpt URL, and having it implicitly download and convert to huggingface's format.

[MicahZoltu]

Hmm, good point on not wanting to download every ckpt if all you are going to use is SD 1.5 (or whatever). I think we could have a base image that has no ckpt files and then use docker caching magic to create a combinatorial explosion of images with different combinations of models (e.g., just SD 1.5, all of the SD versions, SD + , just , everything known, etc.).

This can wait though, and being able to download and store in a mounted VOLUME is valuable in the short term.