From babad145427a90611794c283b6ea8f8eec603f89 Mon Sep 17 00:00:00 2001 From: corkborg <62742115+corkborg@users.noreply.github.com> Date: Tue, 23 Apr 2024 01:00:28 +0900 Subject: [PATCH 1/6] Support gpu. --- README.md | 21 +++++++++++++++++++-- requirements.txt | 3 +-- run.py | 9 +++++++++ tagger/interrogator.py | 21 ++++++++------------- 4 files changed, 37 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index c74c475..877b335 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,8 @@ pip install -r requirements.txt ## usage ``` -usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}] +usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--gpu] + [--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}] options: -h, --help show this help message and exit @@ -22,7 +23,8 @@ options: --threshold THRESHOLD Prediction threshold (default is 0.35) --ext EXT Extension to add to caption file in case of dir option (default is .txt) - --overwrite Overwrite the file if it exists. + --overwrite Overwrite the file if it exists + --gpu Use GPU (Compatible CUDA and cuDNN required) --model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000} modelname to use for prediction (default is wd14-convnextv2.v1) ``` @@ -52,3 +54,18 @@ python run.py --file image.jpg --model wd-v1-4-moat-tagger.v2 python run.py --file image.jpg --model mld-caformer.dec-5-97527 python run.py --file image.jpg --model mld-tresnetd.6-30000 ``` + +## Use GPU + +Requires CUDA 12.2 and cuDNN8.x. + +``` +pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ +``` + +``` +python run.py --file image.jpg --gpu +``` + +https://onnxruntime.ai/docs/install/ +https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements diff --git a/requirements.txt b/requirements.txt index d9c3aa1..c9d47c0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,6 @@ pandas==2.1.1 huggingface-hub==0.17.3 -torch==2.0.1 Pillow==10.0.1 deepdanbooru==1.0.2 -onnxruntime==1.16.0 +onnxruntime==1.17.3 opencv-python==4.8.0.76 \ No newline at end of file diff --git a/run.py b/run.py index cc055f1..1d4747d 100644 --- a/run.py +++ b/run.py @@ -1,6 +1,8 @@ import sys + +from sympy import use from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator from PIL import Image from pathlib import Path @@ -27,6 +29,10 @@ parser.add_argument( '--overwrite', action='store_true', help='Overwrite the file if it exists') +parser.add_argument( + '--gpu', + action='store_true', + help='Use GPU (Compatible CUDA and cuDNN required)') parser.add_argument( '--model', default='wd14-convnextv2.v1', @@ -37,6 +43,9 @@ args = parser.parse_args() # get interrogator configs interrogator = interrogators[args.model] +if args.gpu: + interrogator.use_gpu() + def image_interrogate(image_path: Path): """ Predictions from a image path diff --git a/tagger/interrogator.py b/tagger/interrogator.py index dc5dbe2..e85112f 100644 --- a/tagger/interrogator.py +++ b/tagger/interrogator.py @@ -17,7 +17,6 @@ tag_escape_pattern = re.compile(r'([\\()])') import tagger.dbimutils as dbimutils -use_cpu = True class Interrogator: @staticmethod @@ -73,6 +72,7 @@ class Interrogator: def __init__(self, name: str) -> None: self.name = name + self.providers = ['CPUExecutionProvider'] def load(self): raise NotImplementedError() @@ -90,6 +90,9 @@ class Interrogator: return unloaded + def use_gpu(self) -> None: + self.providers = ['CUDAExecutionProvider'] + def interrogate( self, image: Image @@ -111,6 +114,7 @@ class WaifuDiffusionInterrogator(Interrogator): self.model_path = model_path self.tags_path = tags_path self.kwargs = kwargs + self.use_cpu = True def download(self) -> Tuple[os.PathLike, os.PathLike]: print(f"Loading {self.name} model file from {self.kwargs['repo_id']}") @@ -125,14 +129,7 @@ class WaifuDiffusionInterrogator(Interrogator): model_path, tags_path = self.download() from onnxruntime import InferenceSession - - # https://onnxruntime.ai/docs/execution-providers/ - # https://github.com/toriato/stable-diffusion-webui-wd14-tagger/commit/e4ec460122cf674bbf984df30cdb10b4370c1224#r92654958 - providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] - if use_cpu: - providers.pop(0) - - self.model = InferenceSession(str(model_path), providers=providers) + self.model = InferenceSession(str(model_path), providers=self.providers) print(f'Loaded {self.name} model from {model_path}') @@ -202,6 +199,7 @@ class MLDanbooruInterrogator(Interrogator): self.repo_id = repo_id self.tags = None self.model = None + self.use_cpu = False def download(self) -> Tuple[str, str]: print(f"Loading {self.name} model file from {self.repo_id}") @@ -220,11 +218,8 @@ class MLDanbooruInterrogator(Interrogator): model_path, tags_path = self.download() from onnxruntime import InferenceSession - providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] - if use_cpu: - providers.pop(0) self.model = InferenceSession(model_path, - providers=providers) + providers=self.providers) print(f'Loaded {self.name} model from {model_path}') with open(tags_path, 'r', encoding='utf-8') as filen: From 00083b1819880c5c93805d3c8a52bc56fd3438df Mon Sep 17 00:00:00 2001 From: corkborg <62742115+corkborg@users.noreply.github.com> Date: Tue, 23 Apr 2024 01:04:26 +0900 Subject: [PATCH 2/6] Fix README.md. --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 877b335..23016db 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,7 @@ python run.py --file image.jpg --model mld-caformer.dec-5-97527 python run.py --file image.jpg --model mld-tresnetd.6-30000 ``` -## Use GPU +## Using GPU Requires CUDA 12.2 and cuDNN8.x. From f452c56ea69d57d69ef99f563242ced8f73d6176 Mon Sep 17 00:00:00 2001 From: corkborg <62742115+corkborg@users.noreply.github.com> Date: Tue, 23 Apr 2024 01:20:42 +0900 Subject: [PATCH 3/6] Remove useless code. --- tagger/interrogator.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tagger/interrogator.py b/tagger/interrogator.py index e85112f..fa88b9d 100644 --- a/tagger/interrogator.py +++ b/tagger/interrogator.py @@ -114,7 +114,6 @@ class WaifuDiffusionInterrogator(Interrogator): self.model_path = model_path self.tags_path = tags_path self.kwargs = kwargs - self.use_cpu = True def download(self) -> Tuple[os.PathLike, os.PathLike]: print(f"Loading {self.name} model file from {self.kwargs['repo_id']}") @@ -199,7 +198,6 @@ class MLDanbooruInterrogator(Interrogator): self.repo_id = repo_id self.tags = None self.model = None - self.use_cpu = False def download(self) -> Tuple[str, str]: print(f"Loading {self.name} model file from {self.repo_id}") From 3d954884073befdbcb68fa8c7a907b2e7cc01b77 Mon Sep 17 00:00:00 2001 From: corkborg <62742115+corkborg@users.noreply.github.com> Date: Tue, 23 Apr 2024 01:27:15 +0900 Subject: [PATCH 4/6] Version up dependencies. --- requirements.txt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/requirements.txt b/requirements.txt index c9d47c0..8e8fd86 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ -pandas==2.1.1 -huggingface-hub==0.17.3 +pandas==2.2.2 +huggingface-hub==0.22.2 Pillow==10.0.1 deepdanbooru==1.0.2 onnxruntime==1.17.3 -opencv-python==4.8.0.76 \ No newline at end of file +opencv-python==4.9.0.80 \ No newline at end of file From 7ecb4ed52604918015f9c294f0158750e1ff4dd3 Mon Sep 17 00:00:00 2001 From: corkborg <62742115+corkborg@users.noreply.github.com> Date: Tue, 23 Apr 2024 01:59:22 +0900 Subject: [PATCH 5/6] Add cpu only option instead of gpu option. --- README.md | 6 +----- run.py | 15 +++++---------- tagger/interrogator.py | 6 +++--- 3 files changed, 9 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 23016db..d4f96ac 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,7 @@ options: Prediction threshold (default is 0.35) --ext EXT Extension to add to caption file in case of dir option (default is .txt) --overwrite Overwrite the file if it exists - --gpu Use GPU (Compatible CUDA and cuDNN required) + --cpu Use CPU only --model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000} modelname to use for prediction (default is wd14-convnextv2.v1) ``` @@ -63,9 +63,5 @@ Requires CUDA 12.2 and cuDNN8.x. pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ ``` -``` -python run.py --file image.jpg --gpu -``` - https://onnxruntime.ai/docs/install/ https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements diff --git a/run.py b/run.py index 1d4747d..99cba59 100644 --- a/run.py +++ b/run.py @@ -1,8 +1,3 @@ - - -import sys - -from sympy import use from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator from PIL import Image from pathlib import Path @@ -30,9 +25,9 @@ parser.add_argument( action='store_true', help='Overwrite the file if it exists') parser.add_argument( - '--gpu', + '--cpu', action='store_true', - help='Use GPU (Compatible CUDA and cuDNN required)') + help='Use CPU only') parser.add_argument( '--model', default='wd14-convnextv2.v1', @@ -43,8 +38,8 @@ args = parser.parse_args() # get interrogator configs interrogator = interrogators[args.model] -if args.gpu: - interrogator.use_gpu() +if args.cpu: + interrogator.use_cpu() def image_interrogate(image_path: Path): """ @@ -72,7 +67,7 @@ if args.dir: tags = image_interrogate(image_path) tags_str = ", ".join(tags.keys()) - with open(caption_path, "w") as fp: + with open(caption_path, 'w') as fp: fp.write(tags_str) if args.file: diff --git a/tagger/interrogator.py b/tagger/interrogator.py index fa88b9d..c48caf0 100644 --- a/tagger/interrogator.py +++ b/tagger/interrogator.py @@ -72,7 +72,7 @@ class Interrogator: def __init__(self, name: str) -> None: self.name = name - self.providers = ['CPUExecutionProvider'] + self.providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] def load(self): raise NotImplementedError() @@ -90,8 +90,8 @@ class Interrogator: return unloaded - def use_gpu(self) -> None: - self.providers = ['CUDAExecutionProvider'] + def use_cpu(self) -> None: + self.providers = ['CPUExecutionProvider'] def interrogate( self, From b6d257f55dfcf5e96e0c750f0a01b624ec1dcb49 Mon Sep 17 00:00:00 2001 From: corkborg <62742115+corkborg@users.noreply.github.com> Date: Tue, 23 Apr 2024 02:02:10 +0900 Subject: [PATCH 6/6] Fix description. --- README.md | 5 ++--- run.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index d4f96ac..5449f9b 100644 --- a/README.md +++ b/README.md @@ -13,8 +13,7 @@ pip install -r requirements.txt ## usage ``` -usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--gpu] - [--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}] +usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--cpu] [--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}] options: -h, --help show this help message and exit @@ -23,7 +22,7 @@ options: --threshold THRESHOLD Prediction threshold (default is 0.35) --ext EXT Extension to add to caption file in case of dir option (default is .txt) - --overwrite Overwrite the file if it exists + --overwrite Overwrite caption file if it exists --cpu Use CPU only --model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000} modelname to use for prediction (default is wd14-convnextv2.v1) diff --git a/run.py b/run.py index 99cba59..a3f45a8 100644 --- a/run.py +++ b/run.py @@ -23,7 +23,7 @@ parser.add_argument( parser.add_argument( '--overwrite', action='store_true', - help='Overwrite the file if it exists') + help='Overwrite caption file if it exists') parser.add_argument( '--cpu', action='store_true',