aboutsummaryrefslogtreecommitdiff
path: root/modules
diff options
context:
space:
mode:
authorGreendayle <Greendayle>2022-10-07 20:46:38 +0200
committerGreendayle <Greendayle>2022-10-07 20:46:38 +0200
commitfa2ea648db81f5723bb5d722f2fe0ebd7dfc319a (patch)
treed5ae0c2626cc0691ea264fb3f372198b3b2eed51 /modules
parent54fa613c8391e3973cca9d94cdf539061932508b (diff)
even more powerfull fix
Diffstat (limited to 'modules')
-rw-r--r--modules/deepbooru.py9
1 files changed, 7 insertions, 2 deletions
diff --git a/modules/deepbooru.py b/modules/deepbooru.py
index 79dc59bd..60094336 100644
--- a/modules/deepbooru.py
+++ b/modules/deepbooru.py
@@ -60,8 +60,13 @@ def _load_tf_and_return_tags(pil_image, threshold):
return ', '.join(result_tags_out).replace('_', ' ').replace(':', ' ')
+def subprocess_init_no_cuda():
+ import os
+ os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+
def get_deepbooru_tags(pil_image, threshold=0.5):
- with ProcessPoolExecutor() as executor:
- f = executor.submit(_load_tf_and_return_tags, pil_image, threshold)
+ with ProcessPoolExecutor(initializer=subprocess_init_no_cuda) as executor:
+ f = executor.submit(_load_tf_and_return_tags, pil_image, threshold, )
ret = f.result() # will rethrow any exceptions
return ret \ No newline at end of file