Update app.py
Browse files
app.py
CHANGED
@@ -8,12 +8,19 @@ from rich.traceback import install as traceback_install
|
|
8 |
from tagger.common import Heatmap, ImageLabels, LabelData, load_labels_hf, preprocess_image
|
9 |
from tagger.model import load_model_and_transform, process_heatmap
|
10 |
|
11 |
-
TITLE = "WD Tagger Heatmap"
|
12 |
DESCRIPTION = """WD Tagger v3 Heatmap Generator."""
|
13 |
# get HF token
|
14 |
HF_TOKEN = getenv("HF_TOKEN", None)
|
15 |
|
16 |
# model repo and cache
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
17 |
MODEL_REPO = "SmilingWolf/wd-vit-tagger-v3"
|
18 |
# get the repo root (or the current working directory if running in ipython)
|
19 |
WORK_DIR = Path(__file__).parent.resolve() if "__file__" in globals() else Path().resolve()
|
|
|
8 |
from tagger.common import Heatmap, ImageLabels, LabelData, load_labels_hf, preprocess_image
|
9 |
from tagger.model import load_model_and_transform, process_heatmap
|
10 |
|
11 |
+
TITLE = "WD Tagger Heatmap For More Models"
|
12 |
DESCRIPTION = """WD Tagger v3 Heatmap Generator."""
|
13 |
# get HF token
|
14 |
HF_TOKEN = getenv("HF_TOKEN", None)
|
15 |
|
16 |
# model repo and cache
|
17 |
+
AVAILABLE_MODEL_REPOS = [
|
18 |
+
'SmilingWolf/wd-convnext-tagger-v3',
|
19 |
+
'SmilingWolf/wd-swinv2-tagger-v3',
|
20 |
+
'SmilingWolf/wd-vit-tagger-v3',
|
21 |
+
'SmilingWolf/wd-vit-large-tagger-v3',
|
22 |
+
"SmilingWolf/wd-eva02-large-tagger-v3",
|
23 |
+
]
|
24 |
MODEL_REPO = "SmilingWolf/wd-vit-tagger-v3"
|
25 |
# get the repo root (or the current working directory if running in ipython)
|
26 |
WORK_DIR = Path(__file__).parent.resolve() if "__file__" in globals() else Path().resolve()
|