dev(narugo): add support for pixai taggers

.github/workflows/doc.yml

@@ -67,6 +67,7 @@ jobs:
           ENV_PROD: 'true'
           PLANTUML_HOST: http://localhost:18080
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          CI: 'true'
         run: |
           plantumlcli -c
           make docs
@@ -145,6 +146,7 @@ jobs:
           ENV_PROD: 'true'
           PLANTUML_HOST: http://localhost:18080
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          CI: 'true'
         run: |
           git fetch --all --tags
           git branch -av

Makefile

@@ -41,7 +41,7 @@ unittest:
 		--cov="${RANGE_SRC_DIR}" \
 		$(if ${MIN_COVERAGE},--cov-fail-under=${MIN_COVERAGE},) \
 		$(if ${WORKERS},-n ${WORKERS},) \
-		--reruns 8 --reruns-delay 2 --only-rerun '(OSError|Timeout|HTTPError.*429|HTTPError.*502|HTTPError.*504|check your connection|429 error)'
+		--reruns 8 --reruns-delay 2 --only-rerun '(OSError|Timeout|HTTPError.*429|HTTPError.*502|HTTPError.*504|check your connection|429 error|CAS service error)'
 
 docs:
 	$(MAKE) -C "${DOC_DIR}" build

docs/source/_libs/benchmark.py

@@ -34,6 +34,9 @@ def load(self):
     def unload(self):
         raise NotImplementedError
 
+    def after_unload(self):
+        pass
+
     def run(self):
         raise NotImplementedError
 
@@ -60,6 +63,7 @@ def _record(name):
 
         self.unload()
         _record('<unload>')
+        self.after_unload()
 
         mems = np.array([mem for _, mem, _ in logs])
         mems -= mems[0]

docs/source/api_doc/tagging/index.rst

@@ -12,6 +12,7 @@ imgutils.tagging
     mldanbooru
     wd14
     camie
+    pixai
     deepdanbooru
     deepgelbooru
     format

docs/source/api_doc/tagging/pixai.rst

@@ -0,0 +1,15 @@
+imgutils.tagging.pixai
+====================================
+
+.. currentmodule:: imgutils.tagging.pixai
+
+.. automodule:: imgutils.tagging.pixai
+
+
+get_pixai_tags
+----------------------
+
+.. autofunction:: get_pixai_tags
+
+
+

docs/source/api_doc/tagging/tagging_benchmark.plot.py

@@ -1,25 +1,36 @@
+import os
 import random
 
+from hfutils.cache import delete_cache
+
 from benchmark import BaseBenchmark, create_plot_cli
 from imgutils.tagging import get_deepdanbooru_tags, get_wd14_tags, get_mldanbooru_tags, get_deepgelbooru_tags, \
-    get_camie_tags
+    get_camie_tags, get_pixai_tags
+
 
+class CleanModelStorageBenchmark(BaseBenchmark):
+    def after_unload(self):
+        if os.environ.get('CI'):
+            delete_cache()
 
-class DeepdanbooruBenchmark(BaseBenchmark):
+
+class DeepdanbooruBenchmark(CleanModelStorageBenchmark):
     def load(self):
-        from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model
+        from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model, _get_deepdanbooru_labels
         _ = _get_deepdanbooru_model()
+        _ = _get_deepdanbooru_labels
 
     def unload(self):
-        from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model
+        from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model, _get_deepdanbooru_labels
         _get_deepdanbooru_model.cache_clear()
+        _get_deepdanbooru_labels.cache_clear()
 
     def run(self):
         image_file = random.choice(self.all_images)
         _ = get_deepdanbooru_tags(image_file)
 
 
-class DeepgelbooruBenchmark(BaseBenchmark):
+class DeepgelbooruBenchmark(CleanModelStorageBenchmark):
     def load(self):
         from imgutils.tagging.deepgelbooru import _open_tags, _open_model, _open_preprocessor
         _ = _open_tags()
@@ -37,39 +48,43 @@ def run(self):
         _ = get_deepgelbooru_tags(image_file)
 
 
-class Wd14Benchmark(BaseBenchmark):
+class Wd14Benchmark(CleanModelStorageBenchmark):
     def __init__(self, model):
         BaseBenchmark.__init__(self)
         self.model = model
 
     def load(self):
-        from imgutils.tagging.wd14 import _get_wd14_model
+        from imgutils.tagging.wd14 import _get_wd14_model, _get_wd14_labels
         _ = _get_wd14_model(self.model)
+        _ = _get_wd14_labels(self.model)
 
     def unload(self):
-        from imgutils.tagging.wd14 import _get_wd14_model
+        from imgutils.tagging.wd14 import _get_wd14_model, _get_wd14_labels
         _get_wd14_model.cache_clear()
+        _get_wd14_labels.cache_clear()
 
     def run(self):
         image_file = random.choice(self.all_images)
         _ = get_wd14_tags(image_file, model_name=self.model)
 
 
-class MLDanbooruBenchmark(BaseBenchmark):
+class MLDanbooruBenchmark(CleanModelStorageBenchmark):
     def load(self):
-        from imgutils.tagging.mldanbooru import _open_mldanbooru_model
+        from imgutils.tagging.mldanbooru import _open_mldanbooru_model, _get_mldanbooru_labels
         _ = _open_mldanbooru_model()
+        _ = _get_mldanbooru_labels()
 
     def unload(self):
-        from imgutils.tagging.mldanbooru import _open_mldanbooru_model
+        from imgutils.tagging.mldanbooru import _open_mldanbooru_model, _get_mldanbooru_labels
         _open_mldanbooru_model.cache_clear()
+        _get_mldanbooru_labels.cache_clear()
 
     def run(self):
         image_file = random.choice(self.all_images)
         _ = get_mldanbooru_tags(image_file)
 
 
-class CamieBenchmark(BaseBenchmark):
+class CamieBenchmark(CleanModelStorageBenchmark):
     def __init__(self, model_name):
         BaseBenchmark.__init__(self)
         self.model_name = model_name
@@ -95,6 +110,32 @@ def run(self):
         _ = get_camie_tags(image_file, model_name=self.model_name)
 
 
+class PixAITaggerBenchmark(CleanModelStorageBenchmark):
+    def __init__(self, model_name: str):
+        BaseBenchmark.__init__(self)
+        self.model_name = model_name
+
+    def load(self):
+        from imgutils.tagging.pixai import _open_tags, _open_preprocess, _open_onnx_model, \
+            _open_default_category_thresholds
+        _ = _open_tags(self.model_name)
+        _ = _open_preprocess(self.model_name)
+        _ = _open_onnx_model(self.model_name)
+        _ = _open_default_category_thresholds(self.model_name)
+
+    def unload(self):
+        from imgutils.tagging.pixai import _open_tags, _open_preprocess, _open_onnx_model, \
+            _open_default_category_thresholds
+        _open_tags.cache_clear()
+        _open_preprocess.cache_clear()
+        _open_onnx_model.cache_clear()
+        _open_default_category_thresholds.cache_clear()
+
+    def run(self):
+        image_file = random.choice(self.all_images)
+        _ = get_pixai_tags(image_file, model_name=self.model_name)
+
+
 if __name__ == '__main__':
     create_plot_cli(
         [
@@ -113,6 +154,7 @@ def run(self):
             ('mldanbooru', MLDanbooruBenchmark()),
             ('camie-initial', CamieBenchmark('initial')),
             ('camie-refined', CamieBenchmark('refined')),
+            ('pixai-tagger-v0.9', PixAITaggerBenchmark('v0.9')),
         ],
         title='Benchmark for Tagging Models',
         run_times=10,