[MM] Make controlnet image processors use MM cache system
Summary
This PR changes the loading of most controlnet image processors to use the model manager caching system. This should result in a performance improvement when multiple different control images are processed using the same processor. However, the application of the processed image during image generation will not be affected.
A couple of notes here.
- The
controlnet_auxmodule, which defines the image processors, uses an oddfrom_pretrained()pattern in which the same HuggingFace repo is used for each of the processor classes. Internally, each processor class looks for a different hard-coded model checkpoint file. To use the MM caching system, I figured out which model was being loaded, and created a mapping from the processor class to the checkpoint file incontrolnet_image_processors.py. If at some point one of the processors incontrolnet_auxis updated to use a different file name, a compensatory change will need to be made tocontrolnet_image_processors.py. - I had to make a small change to the model installer's
download_and_cache_model()method in order to accommodate the way that the SegmentAnything class looks for its model in a subfolder. - I was unable to migrate the Leres processor due to the fact that its
from_pretrained()method loads two different models, butdownload_and_cache_model()doesn't accept more than a single model file at a time.
Related Issues / Discussions
QA Instructions
In the workflow editor, execute the following controlnet processors and confirm that they load and work as expected:
- Midas Depth Processor
- MLSD Processor
- Pidi Processor
- Zoe (Depth) Processor
- Leres (Depth) Processor
- Segment Anything Processor
- Lineart Processor
- Lineart Anime Processor
- Hed Processor
The remaining processors (canny, shuffle, etc) do not download models from HF.
Merge Plan
Merge when approved.
Checklist
- [X] The PR has a short but descriptive title, suitable for a changelog
- [X] Tests added / updated (if applicable)
- [X] Documentation added / updated (if applicable)
This is a good step in the right direction. Eventually we should finish the move away from controlnet_aux as it is not actively maintained.
Overall, makes sense to me. A few comments:
- I smoke tested the lineart processor. I trust that you have tested them all?
Yes, all tested successfully.
- The lineart processor (and probably others) are treated as having size 0 by
calc_model_size_by_data(...). Once you merge main into your branch, this will result in error logs (as of Fix some models treated as having size 0 in the model cache #6571).
I've added memory calculations for the controlnet processors as well as the scheduler and tokenizer. Hopefully we won't see any more complaints about size 0.
- The lineart processor (and probably others) don't run on the GPU, because they don't have a
.to(...)method. Is that expected?
Some of the controlnet processors run in GPU and others in CPU. The CPU models only runs once when the encoded controlnet image is first calculated and cached, so it's not a high priority to fix at the moment.
- @brandonrising can you review for potential commercial implications?
Otherwise, this all worked really well and commercial would be improved by having these in the cache
Hmm I did get this error when testing each of the controlnets the first time, but I don't seem to be able to reproduce it again.
Otherwise, this all worked really well and commercial would be improved by having these in the cache
oh ignore that error message, I thought it was downloading a preprocessor but it was probably just from when I kicked off installs in MM
Thanks @lstein, I picked up on this for the Canvas v2 release and got all the nodes using the MM cache in #6831. Well, technically they don't exactly use the cache bc I wasn't sure how to get them to report on their sizes, but we have that on the TODO list.
Hmm I did get this error when testing each of the controlnets the first time, but I don't seem to be able to reproduce it again.