Change the call interface and adjust the program execution logic

mlcd_seg.py

@@ -38,6 +38,7 @@ from PIL import Image
 from transformers import Qwen2Config, Qwen2Model, Qwen2ForCausalLM, AutoConfig, AutoModelForCausalLM, PreTrainedTokenizer
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.generation.utils import GenerateOutput
+from transformers.utils import cached_file
 from safetensors.torch import load_file as safetensors_load
 from .vision_tower import build_vision_tower
 from .vision_resampler import build_vision_resampler
@@ -141,10 +142,8 @@ class MLCDSegMetaModel:
 
     def dispatch_weight(self, config):
         safetensors_set = set()
-        index_folder = Path(getattr(config, "_name_or_path", "./"))
-        index_file = index_folder / "model.safetensors.index.json"
-        if not index_file.exists():
-            os.getenv("")
+        repo = getattr(config, "_name_or_path", "'DeepGlint-AI/MLCD-Seg'")
+        index_file = cached_file(repo, "model.safetensors.index.json")
         with open(index_file, "r") as safetensors_index:
             safetensors_map = json.loads(safetensors_index.read())
             for key, value in safetensors_map["weight_map"].items():
@@ -156,7 +155,7 @@ class MLCDSegMetaModel:
         projector_weight = {}
         text2sam_projection_weight = {}
         for safetensors_file in safetensors_set:
-            temp_load = safetensors_load(index_folder / safetensors_file)
+            temp_load = safetensors_load(cached_file(repo, safetensors_file))
             for key, value in temp_load.items():
                 if key.startswith("model.sam."):
                     sam_weight[key.replace("model.sam.", "")] = value
@@ -174,86 +173,6 @@ class MLCDSegMetaModel:
             vision_tower = vision_tower[0]
         return vision_tower
 
-    # def initialize_vision_modules(self, model_args, fsdp=None):
-    #     vision_tower = model_args.vision_tower
-    #     mm_vision_select_layer = model_args.mm_vision_select_layer
-    #     mm_vision_select_feature = model_args.mm_vision_select_feature
-    #     pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
-    #     mm_patch_merge_type = model_args.mm_patch_merge_type
-
-    #     self.config.mm_vision_tower = vision_tower
-    #     self.config.vision_tower_pretrained = getattr(model_args, "vision_tower_pretrained", "")
-
-    #     if self.get_vision_tower() is None:
-    #         vision_tower = build_vision_tower(model_args)
-    #         vision_resampler = build_vision_resampler(model_args, vision_tower=vision_tower)
-    #         for k, v in vision_resampler.config.items():
-    #             setattr(self.config, k, v)
-
-    #         if fsdp is not None and len(fsdp) > 0:
-    #             self.vision_tower = [vision_tower]
-    #             self.vision_resampler = [vision_resampler]
-    #         else:
-    #             self.vision_tower = vision_tower
-    #             self.vision_resampler = vision_resampler
-    #     else:
-    #         if fsdp is not None and len(fsdp) > 0:
-    #             vision_resampler = self.vision_resampler[0]
-    #             vision_tower = self.vision_tower[0]
-    #         else:
-    #             vision_resampler = self.vision_resampler
-    #             vision_tower = self.vision_tower
-    #         vision_tower.load_model()
-
-    #         # In case it is frozen by LoRA
-    #         for p in self.vision_resampler.parameters():
-    #             p.requires_grad = True
-
-    #     self.config.use_mm_proj = True
-    #     self.config.mm_projector_type = getattr(model_args, "mm_projector_type", "linear")
-    #     self.config.mm_hidden_size = getattr(vision_resampler, "hidden_size", vision_tower.hidden_size)
-    #     self.config.mm_vision_select_layer = mm_vision_select_layer
-    #     self.config.mm_vision_select_feature = mm_vision_select_feature
-    #     self.config.mm_patch_merge_type = mm_patch_merge_type
-        
-    #     for key in vars(model_args):  
-    #         if key.startswith('sam_'):
-    #             setattr(self.config, key, getattr(model_args, key))
-        
-    #     if not hasattr(self.config, 'add_faster_video'):
-    #         if model_args.add_faster_video:
-    #             embed_std = 1 / torch.sqrt(torch.tensor(self.config.hidden_size, dtype=self.dtype))
-    #             self.faster_token = nn.Parameter(
-    #                 torch.randn(self.config.hidden_size, dtype=self.dtype) * embed_std
-    #             )
-
-    #     if getattr(self, "mm_projector", None) is None:
-    #         self.mm_projector = build_vision_projector(self.config, vision_cfg=vision_tower.config)
-
-    #         if "unpad" in mm_patch_merge_type:
-    #             embed_std = 1 / torch.sqrt(torch.tensor(self.config.hidden_size, dtype=self.dtype))
-    #             self.image_newline = nn.Parameter(torch.randn(self.config.hidden_size, dtype=self.dtype) * embed_std)
-                
-    #         if getattr(self.config, 'sam_path', None) is not None:
-    #             self.sam = build_sam_vit_h(self.config.sam_path)
-    #             self.text2sam_projection = text2sam_projection_layer(self.config)
-    #     else:
-    #         if getattr(self.config, 'sam_path', None) is not None and self.config.sam_path !="":
-    #             self.sam = build_sam_vit_h(self.config.sam_path)
-    #             self.text2sam_projection = text2sam_projection_layer(self.config)
-    #         # In case it is frozen by LoRA
-    #         for p in self.mm_projector.parameters():
-    #             p.requires_grad = True
-
-    #     if pretrain_mm_mlp_adapter is not None:
-    #         mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location="cpu")
-
-    #         def get_w(weights, keyword):
-    #             return {k.split(keyword + ".")[1]: v for k, v in weights.items() if keyword in k}
-
-    #         incompatible_keys = self.mm_projector.load_state_dict(get_w(mm_projector_weights, "mm_projector"))
-    #         incompatible_keys = self.vision_resampler.load_state_dict(get_w(mm_projector_weights, "vision_resampler"), strict=False)
-
 
 def unpad_image(tensor, original_size):
     """