add auto download for tools

Author: ControlNet

hydra_vl4ai/tool/blip.py

@@ -1,3 +1,4 @@
+import os
 import torch
 import re
 
@@ -24,6 +25,9 @@ def __init__(self, gpu_number=0, blip_v2_model_type="blip2-flan-t5-xxl"):
 
         max_memory = {gpu_number: torch.cuda.mem_get_info(self.dev)[0]}
 
+        if not os.path.exists(get_root_folder() / "pretrained_models" / "blip2" / blip_v2_model_type):
+            self.prepare()
+
         self.processor = Blip2Processor.from_pretrained(get_root_folder() / "pretrained_models" / "blip2" / blip_v2_model_type)
         # Device_map must be sequential for manual GPU selection
         try:

hydra_vl4ai/tool/efficient_sam.py

@@ -1,3 +1,4 @@
+import os
 import cv2
 import numpy as np
 from ..util.misc import get_root_folder
@@ -12,6 +13,8 @@ class EfficientSam(BaseModel):
 
     def __init__(self, gpu_number):
         super().__init__(gpu_number)
+        if not os.path.exists(get_root_folder() / "pretrained_models" / "efficient_sam" / "efficientsam_s_gpu.jit"):
+            self.prepare()
         self.model = torch.jit.load(str(get_root_folder() / "pretrained_models" / "efficient_sam" / "efficientsam_s_gpu.jit"), map_location=self.dev)
         self.to_tensor = ToTensor()
 

hydra_vl4ai/tool/glip.py

@@ -1,3 +1,4 @@
+import os
 from typing import Union
 
 from ..util.config import Config
@@ -27,6 +28,9 @@ def __init__(self, model_size='large', gpu_number=0, *args):
             config_file = working_dir / "configs/glip_Swin_L.yaml"
             weight_file = str(working_dir / "checkpoints/glip_large_model.pth")
 
+        if not os.path.exists(weight_file):
+            self.prepare(model_size)
+
         class OurGLIPDemo(GLIPDemo):
 
             def __init__(self, dev, *args_demo):

hydra_vl4ai/tool/grounding_dino.py

@@ -1,3 +1,4 @@
+import os
 import cv2
 import numpy as np
 import torch
@@ -24,6 +25,8 @@ def __init__(self, gpu_number=0):
 
         super().__init__(gpu_number)
         path_checkpoint = str(get_root_folder() / "pretrained_models" / "grounding_dino" / "groundingdino_swint_ogc.pth")
+        if not os.path.exists(path_checkpoint):
+            self.prepare()
         config_file = str(get_root_folder() / "module_repos" / "Grounded-Segment-Anything" / "GroundingDINO" / "groundingdino" / "config" / "GroundingDINO_SwinT_OGC.py")
         args = SLConfig.fromfile(config_file) 
         args.device = self.dev

hydra_vl4ai/tool/llava.py

@@ -1,10 +1,12 @@
+import os
 import torch
 import re
 from torchvision import transforms
 from llava.conversation import conv_templates
 from llava.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path
 from llava.model.builder import load_pretrained_model
 from llava.utils import disable_torch_init
+from huggingface_hub import snapshot_download
 
 from ._base import BaseModel, module_registry
 from ..util.misc import get_root_folder
@@ -17,6 +19,8 @@ def __init__(self, gpu_number=0, model_name: str = "liuhaotian/llava-v1.5-7b"):
         super().__init__(gpu_number)
         self.model_path = get_root_folder() / "pretrained_models" / "llava" / model_name.split("/")[-1]
         self.model_name = get_model_name_from_path(str(self.model_path))
+        if not os.path.exists(self.model_path):
+            self.prepare(model_name)
         disable_torch_init()
         self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(self.model_path,
             model_name=self.model_name, model_base=None, load_8bit=False, load_4bit=False)

hydra_vl4ai/tool/xvlm.py

@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 
 from transformers import BertTokenizer
@@ -39,6 +40,8 @@ def __init__(self, gpu_number=0,
             'num_heads': [4, 8, 16, 32]
         }
         model = XVLMBase(config_xvlm, use_contrastive_loss=True, vision_config=vision_config)
+        if not os.path.exists(path_checkpoint):
+            self.prepare()
         checkpoint = torch.load(path_checkpoint, map_location='cpu')
         state_dict = checkpoint['model'] if 'model' in checkpoint.keys() else checkpoint
         msg = model.load_state_dict(state_dict, strict=False)