narugo commited on
Commit
9ecc7a5
·
1 Parent(s): 681a350

dev(narugo): support multiple sites

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -8,7 +8,7 @@ import faiss
8
  import gradio as gr
9
  import numpy as np
10
  from PIL import Image
11
- from cheesechaser.datapool import DanbooruWebpDataPool
12
  from hfutils.operate import get_hf_fs, get_hf_client
13
  from hfutils.utils import TemporaryDirectory
14
  from imgutils.tagging import wd14
@@ -26,6 +26,7 @@ _ALL_MODEL_NAMES = [
26
 
27
  _SITE_CLS = {
28
  'danbooru': DanbooruWebpDataPool,
 
29
  }
30
 
31
 
 
8
  import gradio as gr
9
  import numpy as np
10
  from PIL import Image
11
+ from cheesechaser.datapool import DanbooruWebpDataPool, YandeWebpDataPool
12
  from hfutils.operate import get_hf_fs, get_hf_client
13
  from hfutils.utils import TemporaryDirectory
14
  from imgutils.tagging import wd14
 
26
 
27
  _SITE_CLS = {
28
  'danbooru': DanbooruWebpDataPool,
29
+ 'yandere': YandeWebpDataPool,
30
  }
31
 
32