Open coder168 opened 11 months ago
ok, will check. I never touch that button. Thanks for reporting.
Removing the line will remove the attribute error, and I'll push that as a quick fix, but then the model wont unload. this exact line was one of the open questions I had for @CCRcmcpe, who provided the original pull request and might have a better solution for unloading? The collect() call seems part of the unloading that is commented as dysfunctional, see this comment and properly unloading tensorflow models is not possible but maybe the MLDanbooruInterrogator implementation does not have this flaw?
Thanks for your review and quick fix.
So with this quick fix 5cbf4f5, the button works for unloading all loaded models except the ML-Danbooru models? Is my understanding correct?
Edit, actually it seems related to tensorflow only, so it's the models DeepDanbooruInterrogator and the experimental large_batch_interrogate.
Any of the others should unload properly. Note that with my implementation you re-read from the db.json, even after a shutdown and reload of stablediffusion webui, including the ui interface. and then the query will read from database. Or is that not enough on windows? It will allow you to retrieve/combine the former interrogation output without the model loading.
Let's add a notice upon unloading all button, or unloading a danbooru, exactly for this purpose: that the unloading of this model requires a reload of webui and reload-ui and the interrogation is retrievable from db without load
On SO there is a solution mentioned:
from numba import cuda
cuda.select_device(0)
cuda.close()
But then the numba documentation is unclear how to reopen the same device again. And reading this SO someone mentions a .close() is unrecoverable, suggests a .reset(). more numba docs. where a reset() is listed.
I was thinking about
if use_cpu:
import gc
else:
from numba import cuda
...
class DeepDanbooruInterrogator(Interrogator):
...
def unload(self) -> bool:
unloaded = super().unload()
if unloaded:
if use_cpu:
import tensorflow as tf
tf.keras.backend.clear_session()
gc.collect()
else:
device = cuda.get_current_device()
device.reset()
return unloaded
But I run on CPU.
It is little weird. I run a quick test with the latest version, but VRAM usage does not seem change, only chunk of main memory get released after unloaded. I am using nvidia card.
Hi, you mentioned this https://github.com/toriato/stable-diffusion-webui-wd14-tagger/issues/33 refer to https://github.com/keras-team/keras/issues/2102
In my opinion, it's not able to release GPU memory by tf.keras.backend.clear_session()
I don‘t think numba
is a good idea. Because it seems that numba
requires Nvidia cuda's dependency
If you really want to release tensorflow
memory, I will suggest using multiprocessing.Process
to load model and release they by os
Thanks, this is why I placed it behind an experimental option in the settings. The Nvidia dependency for numba did occur to me, but at least for nvidia numba could be an option? AMD (ROCm) and cpu are the others, or do even drivers like nouveau or nv function? I'm also not exactly sure what configs do not release propperly, and trying to get an impression, and whether it is vram or ram.
thanks for the links, I'll do some more digging.
One thing to note is that some users do not install Nvidia cuda tookit, but use Torch cuda. But Nvidia cuda tooki is required by numba. That’s why I think numba is not a good idea.
In fact, any model related to tensorflow will encounter the problem of not being able to release GPU memory.
The reason why other models, such as WD14, can be released normally is because they are onnx models (check this ), and onnxruntime can release them properly.
gc.collection()
is just an explicit garbage collection of python, which cannot release GPU memory.
toriato seems to have tried to release tensorflow memory as well, and I have tried it too, but this seems to be unsolvable. https://github.com/toriato/stable-diffusion-webui-wd14-tagger/issues/31
My solution
I usually prefer the second option, because onnx model is faster to start and run. The downside is that you can’t use the latest models, unless you convert new models in time after the model author releases a new version.
You are welcome. I have a Colab notebook about tf2onnx that might be helpful, if you need it. But it is code commented in Chinese, and it needs some adjustments
https://github.com/WSH032/wd-v1-4-tagger-feature-extractor-tutorials/blob/main/wd14_tf2onnx.ipynb
One option I am considering is implementing a Settings -> Tagger -> checkbox: convert tensorflow models to onnx. There might be checks necessary for updates, maybe keep the tensorflow next to the onnx model; rerun if the tensorflow's mtime, or sha256 changes.
When click the "Unload all interrogate models" button, the following errors occurred:
After a quick check, this error was caused by a code change (in line 136, \tagger\interrogator.py) within an recent pull request merge - "Manually merged: Support ML-Danbooru https://github.com/picobyte/stable-diffusion-webui-wd14-tagger/pull/6, changes amended from CCRcmcpe's".
Please review the relevant code and fix it, currently tagger can not clean its loaded models up. Thanks. BTW, my webui env versions: version: v1.4.1 • python: 3.10.8 • torch: 2.0.0+cu118 • xformers: 0.0.20 • gradio: 3.32.0