Duplicate from SmilingWolf/wd-swinv2-tagger-v3

Co-authored-by: Smiling Wolf <SmilingWolf@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,49 @@
+---
+license: apache-2.0
+library_name: timm
+---
+# WD SwinV2 Tagger v3
+
+Supports ratings, characters and general tags.
+
+Trained using https://github.com/SmilingWolf/JAX-CV.  
+TPUs used for training kindly provided by the [TRC program](https://sites.research.google/trc/about/).
+
+## Dataset
+Last image id: 7220105  
+Trained on Danbooru images with IDs modulo 0000-0899.  
+Validated on images with IDs modulo 0950-0999.  
+Images with less than 10 general tags were filtered out.  
+Tags with less than 600 images were filtered out.
+
+## Validation results
+`v2.0: P=R: threshold = 0.2653, F1 = 0.4541`  
+`v1.0: P=R: threshold = 0.2521, F1 = 0.4411`
+
+## What's new
+Model v2.0/Dataset v3:  
+Trained for a few more epochs.  
+Used tag frequency-based loss scaling to combat class imbalance.
+
+Model v1.1/Dataset v3:  
+Amended the JAX model config file: add image size.  
+No change to the trained weights.
+
+Model v1.0/Dataset v3:  
+More training images, more and up-to-date tags (up to 2024-02-28).  
+Now `timm` compatible! Load it up and give it a spin using the canonical one-liner!  
+ONNX model is compatible with code developed for the v2 series of models.  
+The batch dimension of the ONNX model is not fixed to 1 anymore. Now you can go crazy with batch inference.  
+Switched to Macro-F1 to measure model performance since it gives me a better gauge of overall training progress.
+
+# Runtime deps
+ONNX model requires `onnxruntime >= 1.17.0`
+
+# Inference code examples
+For timm: https://github.com/neggles/wdv3-timm  
+For ONNX: https://huggingface.co/spaces/SmilingWolf/wd-tagger  
+For JAX: https://github.com/SmilingWolf/wdv3-jax
+
+## Final words
+Subject to change and updates.
+Downstream users are encouraged to use tagged releases rather than relying on the head of the repo.

config.json

@@ -0,0 +1,37 @@
+{
+  "architecture": "swinv2_base_window8_256",
+  "num_classes": 10861,
+  "num_features": 1024,
+  "global_pool": "avg",
+  "model_args": {
+    "act_layer": "gelu_tanh",
+    "img_size": 448,
+    "window_size": 14
+  },
+  "pretrained_cfg": {
+    "custom_load": false,
+    "input_size": [
+      3,
+      448,
+      448
+    ],
+    "fixed_input_size": false,
+    "interpolation": "bicubic",
+    "crop_pct": 1.0,
+    "crop_mode": "center",
+    "mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "num_classes": 10861,
+    "pool_size": null,
+    "first_conv": null,
+    "classifier": null
+  }
+}

model.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a26ae4edbe45fa155ede34257153642b3bf476044a13e0ae6a0c4c131198f444
+size 413777297

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6774bff34d43bd49f75a47db4ef217dce701c9847b546523eb85ff6dbba1db1
+size 467460978

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39fe04000d5907a0a7616df54af744d43fbd70b5fba7abc3135d3675734c03ef
+size 392149220

sw_jax_cv_config.json

@@ -0,0 +1,19 @@
+{
+    "image_size": 448,
+    "model_name": "swinv2_base",
+    "model_args": {
+        "image_size": 448,
+        "patch_size": 4,
+        "in_chans": 3,
+        "num_classes": 10861,
+        "embed_dim": 128,
+        "window_size": 14,
+        "mlp_ratio": 4.0,
+        "qkv_bias": true,
+        "drop_rate": 0.0,
+        "attn_drop_rate": 0.0,
+        "drop_path_rate": 0.1,
+        "patch_norm": true,
+        "layer_norm_eps": 1e-05
+    }
+}