[refactor] start reducing complexity of chat.py

app/chat.py

@@ -1,70 +1,21 @@
-import argparse
+import logging
 import os
 import sys
 
 import cv2
 import numpy as np
 import torch
-import torch.nn.functional as F
 from transformers import AutoTokenizer, BitsAndBytesConfig, CLIPImageProcessor
 
 from model.LISA import LISAForCausalLM
 from model.llava import conversation as conversation_lib
 from model.llava.mm_utils import tokenizer_image_token
 from model.segment_anything.utils.transforms import ResizeLongestSide
-from utils.utils import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,
-                         DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX)
-
-
-def parse_args(args):
-    parser = argparse.ArgumentParser(description="LISA chat")
-    parser.add_argument("--version", default="xinlai/LISA-13B-llama2-v1")
-    parser.add_argument("--vis_save_path", default="./vis_output", type=str)
-    parser.add_argument(
-        "--precision",
-        default="bf16",
-        type=str,
-        choices=["fp32", "bf16", "fp16"],
-        help="precision for inference",
-    )
-    parser.add_argument("--image_size", default=1024, type=int, help="image size")
-    parser.add_argument("--model_max_length", default=512, type=int)
-    parser.add_argument("--lora_r", default=8, type=int)
-    parser.add_argument(
-        "--vision-tower", default="openai/clip-vit-large-patch14", type=str
-    )
-    parser.add_argument("--local-rank", default=0, type=int, help="node rank")
-    parser.add_argument("--load_in_8bit", action="store_true", default=False)
-    parser.add_argument("--load_in_4bit", action="store_true", default=False)
-    parser.add_argument("--use_mm_start_end", action="store_true", default=True)
-    parser.add_argument(
-        "--conv_type",
-        default="llava_v1",
-        type=str,
-        choices=["llava_v1", "llava_llama_2"],
-    )
-    return parser.parse_args(args)
-
-
-def preprocess(
-    x,
-    pixel_mean=torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1),
-    pixel_std=torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1),
-    img_size=1024,
-) -> torch.Tensor:
-    """Normalize pixel values and pad to a square input."""
-    # Normalize colors
-    x = (x - pixel_mean) / pixel_std
-    # Pad
-    h, w = x.shape[-2:]
-    padh = img_size - h
-    padw = img_size - w
-    x = F.pad(x, (0, padw, 0, padh))
-    return x
+from utils import app_helpers, utils
 
 
 def main(args):
-    args = parse_args(args)
+    args = app_helpers.parse_args(args)
     os.makedirs(args.vis_save_path, exist_ok=True)
 
     # Create model
@@ -78,12 +29,7 @@ def main(args):
     tokenizer.pad_token = tokenizer.unk_token
     args.seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
 
-
-    torch_dtype = torch.float32
-    if args.precision == "bf16":
-        torch_dtype = torch.bfloat16
-    elif args.precision == "fp16":
-        torch_dtype = torch.half
+    torch_dtype = change_torch_dtype_by_precision(args.precision)
 
     kwargs = {"torch_dtype": torch_dtype}
     if args.load_in_4bit:
@@ -156,12 +102,12 @@ def main(args):
         conv.messages = []
 
         prompt = input("Please input your prompt: ")
-        prompt = DEFAULT_IMAGE_TOKEN + "\n" + prompt
+        prompt = utils.DEFAULT_IMAGE_TOKEN + "\n" + prompt
         if args.use_mm_start_end:
             replace_token = (
-                DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN
+                utils.DEFAULT_IM_START_TOKEN + utils.DEFAULT_IMAGE_TOKEN + utils.DEFAULT_IM_END_TOKEN
             )
-            prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)
+            prompt = prompt.replace(utils.DEFAULT_IMAGE_TOKEN, replace_token)
 
         conv.append_message(conv.roles[0], prompt)
         conv.append_message(conv.roles[1], "")
@@ -183,27 +129,19 @@ def main(args):
             .unsqueeze(0)
             .cuda()
         )
-        if args.precision == "bf16":
-            image_clip = image_clip.bfloat16()
-        elif args.precision == "fp16":
-            image_clip = image_clip.half()
-        else:
-            image_clip = image_clip.float()
+        logging.info(f"image_clip type: {type(image_clip)}.")
+        image_clip = app_helpers.set_image_precision_by_args(image_clip, args.precision)
 
         image = transform.apply_image(image_np)
         resize_list = [image.shape[:2]]
 
         image = (
-            preprocess(torch.from_numpy(image).permute(2, 0, 1).contiguous())
+            app_helpers.preprocess(torch.from_numpy(image).permute(2, 0, 1).contiguous())
             .unsqueeze(0)
             .cuda()
         )
-        if args.precision == "bf16":
-            image = image.bfloat16()
-        elif args.precision == "fp16":
-            image = image.half()
-        else:
-            image = image.float()
+        logging.info(f"image_clip type: {type(image_clip)}.")
+        image = app_helpers.set_image_precision_by_args(image, args.precision)
 
         input_ids = tokenizer_image_token(prompt, tokenizer, return_tensors="pt")
         input_ids = input_ids.unsqueeze(0).cuda()
@@ -217,11 +155,11 @@ def main(args):
             max_new_tokens=512,
             tokenizer=tokenizer,
         )
-        output_ids = output_ids[0][output_ids[0] != IMAGE_TOKEN_INDEX]
+        output_ids = output_ids[0][output_ids[0] != utils.IMAGE_TOKEN_INDEX]
 
         text_output = tokenizer.decode(output_ids, skip_special_tokens=False)
         text_output = text_output.replace("\n", "").replace("  ", " ")
-        print("text_output: ", text_output)
+        logging.info(f"text_output: {text_output}.")
 
         for i, pred_mask in enumerate(pred_masks):
             if pred_mask.shape[0] == 0:
@@ -249,5 +187,14 @@ def main(args):
             print("{} has been saved.".format(save_path))
 
 
+def change_torch_dtype_by_precision(precision):
+    torch_dtype = torch.float32
+    if precision == "bf16":
+        torch_dtype = torch.bfloat16
+    elif precision == "fp16":
+        torch_dtype = torch.half
+    return torch_dtype
+
+
 if __name__ == "__main__":
     main(sys.argv[1:])