From 4b8202f493b3c3f1dc828b212190ebbab934f7a3 Mon Sep 17 00:00:00 2001
From: Luna <git@l4.pm>
Date: Fri, 4 Aug 2023 20:19:20 -0300
Subject: [PATCH] update model ids

---
 README.md        |  2 +-
 main.py          | 16 ++++++++++------
 requirements.txt |  3 +--
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 26af782..5fe5b0a 100644
--- a/README.md
+++ b/README.md
@@ -40,7 +40,7 @@ cp config.example.json config.json
 env/bin/python3 ./main.py fight
 
 # score mode -- crank the final numbers, generates graphs under plots/ folder
-env/bin/python3 ./main.py fight
+env/bin/python3 ./main.py scores
 
 # keep in mind that you can download more images, run fight mode, and then
 # run score mode! the commands are aware of work that's been already done and
diff --git a/main.py b/main.py
index 3a43999..69be1ad 100644
--- a/main.py
+++ b/main.py
@@ -24,13 +24,17 @@ from dataclasses import dataclass, field
 log = logging.getLogger(__name__)
 
 
+# from https://github.com/picobyte/stable-diffusion-webui-wd14-tagger/blob/49cfdcd6be77086e6cc0d26b5ec26605b702be6e/tagger/utils.py#L15C1-L58C1
 DEFAULTS = [
-    "wd14-vit-v2-git",
-    "wd14-vit",
-    "wd14-swinv2-v2-git",
-    "wd14-convnextv2-v2-git",
-    "wd14-convnext-v2-git",
-    "wd14-convnext",
+    "wd14-vit.v1",
+    "wd14-vit.v2",
+    "wd14-convnext.v1",
+    "wd14-convnext.v2",
+    "wd14-convnextv2.v1",
+    "wd14-swinv2-v1",
+    "wd-v1-4-moat-tagger.v2",
+    "mld-caformer.dec-5-97527",
+    # broken model: "mld-tresnetd.6-30000",
 ]
 
 
diff --git a/requirements.txt b/requirements.txt
index 391859f..6bb4e45 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,8 +1,7 @@
 aiosqlite==0.19.0
 aiohttp==3.8.4
 aiolimiter>1.1.0<2.0
-
-
+aiofiles==23.1.0
 plotly>5.15.0<6.0
 pandas==2.0.2
 kaleido==0.2.1
\ No newline at end of file