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,