Method build_model

lmdeploy/vl/model/xcomposer2.py:137–182

Build the vision part of a VLM model when backend is turbomind, or load the whole VLM model when `self.with_llm==True`

(self, trust_remote_code: bool = False)

            self.preprocess_func = self._preprocess_7b

    def build_model(self, trust_remote_code: bool = False):
        """Build the vision part of a VLM model when backend is turbomind, or
        load the whole VLM model when `self.with_llm==True`"""
        from accelerate import init_empty_weights
        with init_empty_weights(), warnings.catch_warnings(), \
                init_empty_vit(self.model_path):
            warnings.simplefilter('ignore')
            config = self.hf_config
            model = AutoModelForCausalLM.from_config(config, trust_remote_code=trust_remote_code)
            model.vit.load_model()
            model.vit.resize_pos()
            if hasattr(self.hf_config, 'img_size'):
                model.vit.vision_tower.vision_model.embeddings.image_size = \
                    self.hf_config.img_size
            model.vit.vision_tower.vision_model.post_layernorm.to_empty(device='cpu').half()
            self.vl_model = model
            if not self.with_llm:
                del model.model
                del model.output

        from accelerate.utils import get_balanced_memory, infer_auto_device_map
        max_memory = get_balanced_memory(model,
                                         max_memory=self.max_memory,
                                         dtype=torch.half,
                                         no_split_module_classes=['CLIPEncoderLayer'])
        device_map = infer_auto_device_map(model,
                                           no_split_module_classes=['CLIPEncoderLayer'],
                                           max_memory=max_memory,
                                           dtype=torch.half)
        # make all tensor on same device for postprocess
        if 'plora_glb_GN' in device_map:
            device_map['plora_sub_GN'] = device_map['plora_glb_GN']

        from accelerate import load_checkpoint_and_dispatch
        with disable_logging():
            load_checkpoint_and_dispatch(model=model,
                                         checkpoint=self.model_path,
                                         device_map=device_map if not self.with_llm else {'': 'cpu'},
                                         no_split_module_classes=['CLIPEncoderLayer'],
                                         dtype=torch.half)

        if 'plora_glb_GN' in device_map:
            add_device_hook(model.vit.vision_tower.vision_model.encoder.layers[-1], device_map['plora_glb_GN'],
                            lambda x: (x[0].to(device=device_map['plora_glb_GN']), ))

        self.model = model.eval()

    def _preprocess_2d5(self, image: Image, params: dict) -> dict:
        """Image preprocessing for internlm-xcomposer2d5-7b."""