Fix ViT-H builder

app.py

@@ -31,7 +31,7 @@ from tokenize_anything.utils.image import im_vstack
 def parse_args():
     """Parse arguments."""
     parser = argparse.ArgumentParser(description="Launch gradio application")
-    parser.add_argument("--model-type", type=str, default="tap_vit_l")
+    parser.add_argument("--model-type", type=str, default="tap_vit_h")
     parser.add_argument("--checkpoint", type=str, default="models/tap_vit_h_v1_1.pkl")
     parser.add_argument("--concept", type=str, default="concepts/merged_2560.pkl")
     parser.add_argument("--device", nargs="+", type=int, default=[0], help="Index of devices")

tokenize_anything/models/easy_build.py

@@ -106,8 +106,10 @@ def image_tokenizer(image_encoder, checkpoint=None, device=0, dtype="float16", *
 
 vit_b_encoder = partial(vit_encoder, depth=12, embed_dim=768, num_heads=12)
 vit_l_encoder = partial(vit_encoder, depth=24, embed_dim=1024, num_heads=16)
+vit_h_encoder = partial(vit_encoder, depth=32, embed_dim=1280, num_heads=16)
 
 model_registry = {
     "tap_vit_b": partial(image_tokenizer, image_encoder=vit_b_encoder),
     "tap_vit_l": partial(image_tokenizer, image_encoder=vit_l_encoder),
+    "tap_vit_h": partial(image_tokenizer, image_encoder=vit_h_encoder),
 }