diff --git a/README.md b/README.md index c74c475..5449f9b 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ 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] [--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 @@ -22,7 +22,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 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) ``` @@ -52,3 +53,14 @@ 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 ``` + +## Using 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/ +``` + +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..8e8fd86 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,6 @@ -pandas==2.1.1 -huggingface-hub==0.17.3 -torch==2.0.1 +pandas==2.2.2 +huggingface-hub==0.22.2 Pillow==10.0.1 deepdanbooru==1.0.2 -onnxruntime==1.16.0 -opencv-python==4.8.0.76 \ No newline at end of file +onnxruntime==1.17.3 +opencv-python==4.9.0.80 \ No newline at end of file diff --git a/run.py b/run.py index cc055f1..a3f45a8 100644 --- a/run.py +++ b/run.py @@ -1,6 +1,3 @@ - - -import sys from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator from PIL import Image from pathlib import Path @@ -26,7 +23,11 @@ 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', + help='Use CPU only') parser.add_argument( '--model', default='wd14-convnextv2.v1', @@ -37,6 +38,9 @@ args = parser.parse_args() # get interrogator configs interrogator = interrogators[args.model] +if args.cpu: + interrogator.use_cpu() + def image_interrogate(image_path: Path): """ Predictions from a image path @@ -63,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 dc5dbe2..c48caf0 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 = ['CUDAExecutionProvider', 'CPUExecutionProvider'] def load(self): raise NotImplementedError() @@ -90,6 +90,9 @@ class Interrogator: return unloaded + def use_cpu(self) -> None: + self.providers = ['CPUExecutionProvider'] + def interrogate( self, image: Image @@ -125,14 +128,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}') @@ -220,11 +216,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: