dev(narugo): quick webp

app.py

@@ -14,7 +14,7 @@ from hfutils.operate import get_hf_fs, get_hf_client
 from hfutils.utils import TemporaryDirectory
 from imgutils.tagging import wd14
 
-from pools import AIBooruWebpDataPool
+from pools import quick_webp_pool
 
 _REPO_ID = 'deepghs/index_experiments'
 
@@ -35,13 +35,13 @@ _SITE_CLS = {
     'konachan': KonachanWebpDataPool,
     'anime_pictures': AnimePicturesWebpDataPool,
     'rule34': Rule34WebpDataPool,
-    'aibooru': AIBooruWebpDataPool,
 }
 
 
 def _get_from_ids(site_name: str, ids: List[int]) -> Dict[int, Image.Image]:
     with TemporaryDirectory() as td:
-        datapool = _SITE_CLS[site_name]()
+        site_cls = _SITE_CLS.get(site_name) or quick_webp_pool(site_name, 3)
+        datapool = site_cls()
         datapool.batch_download_to_directory(
             resource_ids=ids,
             dst_dir=td,

pools.py

@@ -1,3 +1,5 @@
+from typing import Type
+
 from cheesechaser.datapool import IncrementIDDataPool
 
 _AIBOORU_REPO_ID = 'deepghs/aibooru-webp-4Mpixel'
@@ -13,3 +15,20 @@ class AIBooruWebpDataPool(IncrementIDDataPool):
             idx_revision=revision,
             base_level=3,
         )
+
+
+def quick_webp_pool(site_name: str, level: int = 3) -> Type[IncrementIDDataPool]:
+    repo_id = f'deepghs/{site_name}-webp-4Mpixel'
+
+    class _QuickWebpDataPool(IncrementIDDataPool):
+        def __init__(self, revision: str = 'main'):
+            IncrementIDDataPool.__init__(
+                self,
+                data_repo_id=repo_id,
+                data_revision=revision,
+                idx_repo_id=repo_id,
+                idx_revision=revision,
+                base_level=level,
+            )
+
+    return _QuickWebpDataPool