Spaces:
Running
Running
Update app.py
Browse files
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 |
-
|
| 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 |
),
|