dev(narugo): add aibooru

app.py

@@ -14,6 +14,8 @@ from hfutils.operate import get_hf_fs, get_hf_client
 from hfutils.utils import TemporaryDirectory
 from imgutils.tagging import wd14
 
+from pools import AIBooruWebpDataPool
+
 _REPO_ID = 'deepghs/index_experiments'
 
 hf_fs = get_hf_fs()
@@ -33,6 +35,7 @@ _SITE_CLS = {
     'konachan': KonachanWebpDataPool,
     'anime_pictures': AnimePicturesWebpDataPool,
     'rule34': Rule34WebpDataPool,
+    'aibooru': AIBooruWebpDataPool,
 }
 
 

pools.py

@@ -0,0 +1,15 @@
+from cheesechaser.datapool import IncrementIDDataPool
+
+_AIBOORU_REPO_ID = 'deepghs/aibooru-webp-4Mpixel'
+
+
+class AIBooruWebpDataPool(IncrementIDDataPool):
+    def __init__(self, revision: str = 'main'):
+        IncrementIDDataPool.__init__(
+            self,
+            data_repo_id=_AIBOORU_REPO_ID,
+            data_revision=revision,
+            idx_repo_id=_AIBOORU_REPO_ID,
+            idx_revision=revision,
+            base_level=3,
+        )