Update libra/model/builder.py

libra/model/builder.py

@@ -23,14 +23,13 @@ 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=False, device_map="auto", device="cpu"):
-    kwargs = {"device_map": device_map}
+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, **kwargs}
 
     
     if device != "cuda":
         kwargs['device_map'] = {"": device}
         
-        
     if load_8bit:
         kwargs['load_in_8bit'] = True
     elif load_4bit:
@@ -125,7 +124,7 @@ 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()
+            vision_tower.load_model(device_map=device_map)
         vision_tower.to(device=device, dtype=torch.float16)
         image_processor = vision_tower.image_processor