fix bug in inference

Former-commit-id: 4b1776203d3410cd71b2d4720fc7d9cc61d1db3c

app.py

@@ -92,7 +92,6 @@ if args.load_in_4bit:
     kwargs.update(
         {
             "torch_dtype": torch.half,
-            "device_map": "auto",
             "load_in_4bit": True,
             "quantization_config": BitsAndBytesConfig(
                 load_in_4bit=True,
@@ -107,7 +106,6 @@ elif args.load_in_8bit:
     kwargs.update(
         {
             "torch_dtype": torch.half,
-            "device_map": "auto",
             "quantization_config": BitsAndBytesConfig(
                 llm_int8_skip_modules=["visual_model"],
                 load_in_8bit=True,
@@ -116,7 +114,7 @@ elif args.load_in_8bit:
     )
 
 model = LISAForCausalLM.from_pretrained(
-    args.version, low_cpu_mem_usage=True, seg_token_idx=args.seg_token_idx, **kwargs
+    args.version, low_cpu_mem_usage=True, vision_tower=args.vision_tower, seg_token_idx=args.seg_token_idx, **kwargs
 )
 
 model.config.eos_token_id = tokenizer.eos_token_id

chat.py

@@ -90,7 +90,6 @@ def main(args):
         kwargs.update(
             {
                 "torch_dtype": torch.half,
-                "device_map": "auto",
                 "load_in_4bit": True,
                 "quantization_config": BitsAndBytesConfig(
                     load_in_4bit=True,
@@ -105,7 +104,6 @@ def main(args):
         kwargs.update(
             {
                 "torch_dtype": torch.half,
-                "device_map": "auto",
                 "quantization_config": BitsAndBytesConfig(
                     llm_int8_skip_modules=["visual_model"],
                     load_in_8bit=True,
@@ -114,7 +112,7 @@ def main(args):
         )
 
     model = LISAForCausalLM.from_pretrained(
-        args.version, low_cpu_mem_usage=True, seg_token_idx=args.seg_token_idx, **kwargs
+        args.version, low_cpu_mem_usage=True, vision_tower=args.vision_tower, seg_token_idx=args.seg_token_idx, **kwargs
     )
 
     model.config.eos_token_id = tokenizer.eos_token_id
@@ -223,6 +221,7 @@ def main(args):
 
         text_output = tokenizer.decode(output_ids, skip_special_tokens=False)
         text_output = text_output.replace("\n", "").replace("  ", " ")
+        print("text_output: ", text_output)
 
         for i, pred_mask in enumerate(pred_masks):
             if pred_mask.shape[0] == 0:

model/LISA.py

@@ -134,7 +134,9 @@ class LISAForCausalLM(LlavaLlamaForCausalLM):
             self.ce_loss_weight = kwargs.pop("ce_loss_weight", None)
             self.dice_loss_weight = kwargs.pop("dice_loss_weight", None)
             self.bce_loss_weight = kwargs.pop("bce_loss_weight", None)
-
+        else:
+            config.mm_vision_tower = config.vision_tower
+            
         self.seg_token_idx = kwargs.pop("seg_token_idx")
 
         super().__init__(config)

train_ds.py

@@ -90,7 +90,6 @@ def parse_args(args):
     parser.add_argument("--eval_only", action="store_true", default=False)
     parser.add_argument("--vision_pretrained", default="PATH_TO_SAM_ViT-H", type=str)
     parser.add_argument("--out_dim", default=256, type=int)
-    parser.add_argument("--weight", default="", type=str)
     parser.add_argument("--resume", default="", type=str)
     parser.add_argument("--print_freq", default=1, type=int)
     parser.add_argument("--start_epoch", default=0, type=int)