dev(narugo): more models added

README.md

@@ -9,7 +9,11 @@ app_file: app.py
 pinned: true
 license: mit
 preload_from_hub:
+- SmilingWolf/wd-convnext-tagger-v3 model.safetensors,config.json,sw_jax_cv_config.json,selected_tags.csv
+- SmilingWolf/wd-swinv2-tagger-v3 model.safetensors,config.json,sw_jax_cv_config.json,selected_tags.csv
 - SmilingWolf/wd-vit-tagger-v3 model.safetensors,config.json,sw_jax_cv_config.json,selected_tags.csv
+- SmilingWolf/wd-vit-large-tagger-v3 model.safetensors,config.json,sw_jax_cv_config.json,selected_tags.csv
+- SmilingWolf/wd-eva02-large-tagger-v3 model.safetensors,config.json,sw_jax_cv_config.json,selected_tags.csv
 models:
 - SmilingWolf/wd-convnext-tagger-v3
 - SmilingWolf/wd-swinv2-tagger-v3

tagger/model.py

@@ -128,7 +128,7 @@ def render_heatmap(
     # hmap_dim = int(math.sqrt(pos_embed_dim))
 
     image_hmaps = gradients.mean(2, keepdim=True).mul(image_feats.unsqueeze(0)).squeeze()
-    hmap_dim = int(math.sqrt(image_hmaps.mean(-1).shape[0] / len(image_labels)))
+    hmap_dim = int(math.sqrt(image_hmaps.mean(-1).numel() / len(image_labels)))
     image_hmaps = image_hmaps.mean(-1).reshape(len(image_labels), hmap_dim, hmap_dim)
     image_hmaps = image_hmaps.max(torch.zeros_like(image_hmaps))