Update libra/model/builder.py

libra/model/builder.py

@@ -23,25 +23,14 @@ from libra.model import *
 from libra.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
 
 
-def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=True, device_map="auto", device="cpu"):
-    kwargs = {"device_map": device_map}
+def load_pretrained_model(model_path, model_base, model_name, device="cpu"):
+    device_map = {"": device}
+    kwargs = {
+        "device_map": device_map,
+        "torch_dtype": torch.float32  # 对于 CPU，建议使用 float32 或 bfloat16
+    }
 
-    
-    if device != "cuda":
-        kwargs['device_map'] = {"": device}
         
-    if load_8bit:
-        kwargs['load_in_8bit'] = True
-    elif load_4bit:
-        kwargs['load_in_4bit'] = True
-        kwargs['quantization_config'] = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.float16,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type='nf4'
-        )
-    else:
-        kwargs['torch_dtype'] = torch.float16
 
     if 'libra' in model_name.lower():
         # Load Libra model
@@ -92,7 +81,7 @@ def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, l
             model.load_state_dict(mm_projector_weights, strict=False)
         else:
             tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
-            model = LibraLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+            model = LibraLlamaForCausalLM.from_pretrained(model_path, **kwargs)
     else:
         # Load language model
         if model_base is not None:
@@ -124,8 +113,8 @@ def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, l
 
         vision_tower = model.get_vision_tower()
         if not vision_tower.is_loaded:
-            vision_tower.load_model(device_map=device_map)
-        vision_tower.to(device=device, dtype=torch.float16)
+            vision_tower.load_model()
+        vision_tower.to(device=device, dtype=torch.float32)
         image_processor = vision_tower.image_processor
 
     if hasattr(model.config, "max_sequence_length"):