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] 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: