narugo commited on
Commit
aeac786
·
verified ·
1 Parent(s): dbb6240

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -1
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()