yoinked commited on
Commit
8a3e614
·
verified ·
1 Parent(s): 0286385

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -2
app.py CHANGED
@@ -42,6 +42,7 @@ class WaifuDiffusionInterrogator:
42
  repo='SmilingWolf/wd-v1-4-vit-tagger',
43
  model_path='model.onnx',
44
  tags_path='selected_tags.csv',
 
45
  mode: str = "auto"
46
  ) -> None:
47
  self.__repo = repo
@@ -56,8 +57,10 @@ class WaifuDiffusionInterrogator:
56
  return
57
 
58
  model_path = hf_hub_download(self.__repo, filename=self.__model_path)
59
- tags_path = hf_hub_download(self.__repo, filename=self.__tags_path)
60
-
 
 
61
  self._model = InferenceSession(str(model_path))
62
  self._tags = pd.read_csv(tags_path)
63
 
@@ -141,6 +144,12 @@ WAIFU_MODELS: Mapping[str, WaifuDiffusionInterrogator] = {
141
  'chen-evangelion': WaifuDiffusionInterrogator(
142
  repo='SmilingWolf/wd-eva02-large-tagger-v3'
143
  ),
 
 
 
 
 
 
144
  'chenkaku-evangelion': WaifuDiffusionInterrogator(
145
  repo='deepghs/idolsankaku-eva02-large-tagger-v1'
146
  ),
 
42
  repo='SmilingWolf/wd-v1-4-vit-tagger',
43
  model_path='model.onnx',
44
  tags_path='selected_tags.csv',
45
+ tags_path_is_local=False,
46
  mode: str = "auto"
47
  ) -> None:
48
  self.__repo = repo
 
57
  return
58
 
59
  model_path = hf_hub_download(self.__repo, filename=self.__model_path)
60
+ if not tags_path_is_local:
61
+ tags_path = hf_hub_download(self.__repo, filename=self.__tags_path)
62
+ else:
63
+ tags_path = self.__tags_path
64
  self._model = InferenceSession(str(model_path))
65
  self._tags = pd.read_csv(tags_path)
66
 
 
144
  'chen-evangelion': WaifuDiffusionInterrogator(
145
  repo='SmilingWolf/wd-eva02-large-tagger-v3'
146
  ),
147
+ 'chen-cltagger-evangelion-optimized': WaifuDiffusionInterrogator(
148
+ repo='cella110n/cl_tagger',
149
+ model_path='cl_tagger_1_00/model_optimized.onnx',
150
+ tags_path='cltags.csv',
151
+ tags_path_is_local=True,
152
+ ),
153
  'chenkaku-evangelion': WaifuDiffusionInterrogator(
154
  repo='deepghs/idolsankaku-eva02-large-tagger-v1'
155
  ),