[feat] add optional embedding_key argument to LISAForCausalLM.evaluate() method

lisa_on_cuda/LISA.py

@@ -7,13 +7,15 @@ import torch.nn.functional as F
 from .llava.model.language_model.llava_llama import (LlavaLlamaForCausalLM, LlavaLlamaModel)
 from .segment_anything import build_sam_vit_h
 
+embedding_dict = {}
+
 
 def dice_loss(
-    inputs: torch.Tensor,
-    targets: torch.Tensor,
-    num_masks: float,
-    scale=1000,  # 100000.0,
-    eps=1e-6,
+        inputs: torch.Tensor,
+        targets: torch.Tensor,
+        num_masks: float,
+        scale=1000,  # 100000.0,
+        eps=1e-6,
 ):
     """
     Compute the DICE loss, similar to generalized IOU for masks
@@ -35,9 +37,9 @@ def dice_loss(
 
 
 def sigmoid_ce_loss(
-    inputs: torch.Tensor,
-    targets: torch.Tensor,
-    num_masks: float,
+        inputs: torch.Tensor,
+        targets: torch.Tensor,
+        num_masks: float,
 ):
     """
     Args:
@@ -56,9 +58,9 @@ def sigmoid_ce_loss(
 
 class LisaMetaModel:
     def __init__(
-        self,
-        config,
-        **kwargs,
+            self,
+            config,
+            **kwargs,
     ):
         super(LisaMetaModel, self).__init__(config)
 
@@ -98,9 +100,9 @@ class LisaMetaModel:
 
 class LisaModel(LisaMetaModel, LlavaLlamaModel):
     def __init__(
-        self,
-        config,
-        **kwargs,
+            self,
+            config,
+            **kwargs,
     ):
         super(LisaModel, self).__init__(config, **kwargs)
 
@@ -117,9 +119,9 @@ class LisaModel(LisaMetaModel, LlavaLlamaModel):
 
 class LISAForCausalLM(LlavaLlamaForCausalLM):
     def __init__(
-        self,
-        config,
-        **kwargs,
+            self,
+            config,
+            **kwargs,
     ):
         if not hasattr(config, "train_mask_decoder"):
             config.mm_use_im_start_end = kwargs.pop("use_mm_start_end", True)
@@ -131,7 +133,7 @@ class LISAForCausalLM(LlavaLlamaForCausalLM):
             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)
@@ -162,18 +164,18 @@ class LISAForCausalLM(LlavaLlamaForCausalLM):
         return self.model_forward(**kwargs)
 
     def model_forward(
-        self,
-        images: torch.FloatTensor,
-        images_clip: torch.FloatTensor,
-        input_ids: torch.LongTensor,
-        labels: torch.LongTensor,
-        attention_masks: torch.LongTensor,
-        offset: torch.LongTensor,
-        masks_list: List[torch.FloatTensor],
-        label_list: List[torch.Tensor],
-        resize_list: List[tuple],
-        inference: bool = False,
-        **kwargs,
+            self,
+            images: torch.FloatTensor,
+            images_clip: torch.FloatTensor,
+            input_ids: torch.LongTensor,
+            labels: torch.LongTensor,
+            attention_masks: torch.LongTensor,
+            offset: torch.LongTensor,
+            masks_list: List[torch.FloatTensor],
+            label_list: List[torch.Tensor],
+            resize_list: List[tuple],
+            inference: bool = False,
+            **kwargs,
     ):
         image_embeddings = self.get_visual_embs(images)
         batch_size = image_embeddings.shape[0]
@@ -309,17 +311,17 @@ class LISAForCausalLM(LlavaLlamaForCausalLM):
             pred_mask = pred_masks[batch_idx]
 
             assert (
-                gt_mask.shape[0] == pred_mask.shape[0]
+                    gt_mask.shape[0] == pred_mask.shape[0]
             ), "gt_mask.shape: {}, pred_mask.shape: {}".format(
                 gt_mask.shape, pred_mask.shape
             )
             mask_bce_loss += (
-                sigmoid_ce_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])
-                * gt_mask.shape[0]
+                    sigmoid_ce_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])
+                    * gt_mask.shape[0]
             )
             mask_dice_loss += (
-                dice_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])
-                * gt_mask.shape[0]
+                    dice_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])
+                    * gt_mask.shape[0]
             )
             num_masks += gt_mask.shape[0]
 
@@ -338,16 +340,22 @@ class LISAForCausalLM(LlavaLlamaForCausalLM):
         }
 
     def evaluate(
-        self,
-        images_clip,
-        images,
-        input_ids,
-        resize_list,
-        original_size_list,
-        max_new_tokens=32,
-        tokenizer=None,
+            self,
+            images_clip,
+            images,
+            input_ids,
+            resize_list,
+            original_size_list,
+            max_new_tokens=32,
+            tokenizer=None,
+            model_logger=None,
+            embedding_key=None
     ):
         with torch.no_grad():
+            if model_logger is None:
+                import logging
+                model_logger = logging
+            model_logger.debug("start output generation...")
             outputs = self.generate(
                 images=images_clip,
                 input_ids=input_ids,
@@ -356,11 +364,13 @@ class LISAForCausalLM(LlavaLlamaForCausalLM):
                 output_hidden_states=True,
                 return_dict_in_generate=True,
             )
+            model_logger.debug("done output generation...")
             output_hidden_states = outputs.hidden_states[-1]
             output_ids = outputs.sequences
 
             seg_token_mask = output_ids[:, 1:] == self.seg_token_idx
             # hack for IMAGE_TOKEN_INDEX (we suppose that there is only one image, and it is in the front)
+            model_logger.debug(f"start torch.cat to seg_token_mask...")
             seg_token_mask = torch.cat(
                 [
                     torch.zeros((seg_token_mask.shape[0], 255)).bool().cuda(),
@@ -368,20 +378,25 @@ class LISAForCausalLM(LlavaLlamaForCausalLM):
                 ],
                 dim=1,
             )
+            model_logger.debug("done torch.cat to seg_token_mask...")
 
             hidden_states = []
 
             assert len(self.model.text_hidden_fcs) == 1
             hidden_states.append(self.model.text_hidden_fcs[0](output_hidden_states))
 
+            model_logger.debug("start torch.stack to last_hidden_state...")
             last_hidden_state = torch.stack(hidden_states, dim=-1).sum(dim=-1)
+            model_logger.debug("done torch.stack to last_hidden_state...")
             pred_embeddings = last_hidden_state[seg_token_mask]
 
             seg_token_counts = seg_token_mask.int().sum(-1)  # [bs, ]
             seg_token_offset = seg_token_counts.cumsum(-1)
+            model_logger.debug(f"start torch.cat to seg_token_offset...")
             seg_token_offset = torch.cat(
                 [torch.zeros(1).long().cuda(), seg_token_offset], dim=0
             )
+            model_logger.debug("done torch.cat to seg_token_offset...")
 
             pred_embeddings_ = []
             for i in range(len(seg_token_offset) - 1):
@@ -389,11 +404,25 @@ class LISAForCausalLM(LlavaLlamaForCausalLM):
                 pred_embeddings_.append(pred_embeddings[start_i:end_i])
             pred_embeddings = pred_embeddings_
 
-            image_embeddings = self.get_visual_embs(images)
+            model_logger.debug(f"start get_visual_embs to image_embeddings with embedding_key {embedding_key}.")
+
+            if embedding_key is None:
+                image_embeddings = self.get_visual_embs(images)
+            else:
+                try:
+                    image_embeddings = embedding_dict[embedding_key]
+                except KeyError:
+                    model_logger.debug(f"embedding_key {embedding_key} not in embedding_dict, creating embedding now!")
+                    image_embeddings = self.get_visual_embs(images)
+                    embedding_dict[embedding_key] = image_embeddings
+                    model_logger.debug(f"image embedding added in embedding_dict with embedding_key {embedding_key}!")
+
+            model_logger.debug("done get_visual_embs to image_embeddings...")
 
             multimask_output = False
             pred_masks = []
             for i in range(len(pred_embeddings)):
+                model_logger.debug(f"start ({i}nth time) visual_model.prompt_encoder to sparse/dense")
                 (
                     sparse_embeddings,
                     dense_embeddings,
@@ -403,8 +432,9 @@ class LISAForCausalLM(LlavaLlamaForCausalLM):
                     masks=None,
                     text_embeds=pred_embeddings[i].unsqueeze(1),
                 )
-
+                model_logger.debug(f"done ({i}nth) visual_model.prompt_encoder to sparse/dense, start sparse2sparse")
                 sparse_embeddings = sparse_embeddings.to(pred_embeddings[i].dtype)
+                model_logger.debug(f"done ({i}nth) sparse2sparse, start visual_model.mask_decoder")
                 low_res_masks, iou_predictions = self.model.visual_model.mask_decoder(
                     image_embeddings=image_embeddings[i].unsqueeze(0),
                     image_pe=self.model.visual_model.prompt_encoder.get_dense_pe(),
@@ -412,11 +442,14 @@ class LISAForCausalLM(LlavaLlamaForCausalLM):
                     dense_prompt_embeddings=dense_embeddings,
                     multimask_output=multimask_output,
                 )
+                model_logger.debug(f"done ({i}nth) visual_model.mask_decoder, start postprocess_masks")
                 pred_mask = self.model.visual_model.postprocess_masks(
                     low_res_masks,
                     input_size=resize_list[i],
                     original_size=original_size_list[i],
                 )
+                model_logger.debug(f"done ({i}nth) postprocess_masks")
                 pred_masks.append(pred_mask[:, 0])
 
+        model_logger.debug(f"env evaluate! ")
         return output_ids, pred_masks

lisa_on_cuda/utils/app_helpers.py

@@ -211,8 +211,12 @@ def get_inference_model_by_args(args_to_parse):
     no_seg_out = placeholders["no_seg_out"]
 
     @session_logger.set_uuid_logging
-    def inference(input_str: str, input_image: str | np.ndarray, internal_logger: logging = None):
-
+    def inference(
+            input_str: str,
+            input_image: str | np.ndarray,
+            internal_logger: logging = None,
+            embedding_key: str = None
+    ):
         if internal_logger is None:
             internal_logger = app_logger
 
@@ -255,7 +259,7 @@ def get_inference_model_by_args(args_to_parse):
             image_np = cv2.imread(input_image)
             image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
         original_size_list = [image_np.shape[:2]]
-
+        app_logger.debug("start clip_image_processor.preprocess")
         image_clip = (
             clip_image_processor.preprocess(image_np, return_tensors="pt")[
                 "pixel_values"
@@ -263,24 +267,27 @@ def get_inference_model_by_args(args_to_parse):
             .unsqueeze(0)
             .cuda()
         )
+        app_logger.debug("done clip_image_processor.preprocess")
         internal_logger.info(f"image_clip type: {type(image_clip)}.")
         image_clip = set_image_precision_by_args(image_clip, args_to_parse.precision)
 
         image = transform.apply_image(image_np)
         resize_list = [image.shape[:2]]
 
+        internal_logger.debug(f"starting preprocess image: {type(image_clip)}.")
         image = (
             preprocess(torch.from_numpy(image).permute(2, 0, 1).contiguous())
             .unsqueeze(0)
             .cuda()
         )
-        internal_logger.info(f"image_clip type: {type(image_clip)}.")
+        internal_logger.info(f"done preprocess image:{type(image)}, image_clip type: {type(image_clip)}.")
         image = set_image_precision_by_args(image, args_to_parse.precision)
 
         input_ids = tokenizer_image_token(prompt, tokenizer, return_tensors="pt")
         input_ids = input_ids.unsqueeze(0).cuda()
 
-        internal_logger.info("start model evaluation...")
+        embedding_key = get_hash_array(embedding_key, image, internal_logger)
+        internal_logger.info(f"start model evaluation with embedding_key {embedding_key}.")
         output_ids, pred_masks = model.evaluate(
             image_clip,
             image,
@@ -289,6 +296,8 @@ def get_inference_model_by_args(args_to_parse):
             original_size_list,
             max_new_tokens=512,
             tokenizer=tokenizer,
+            model_logger=internal_logger,
+            embedding_key=embedding_key
         )
         internal_logger.info("model evaluation done, start token decoding...")
         output_ids = output_ids[0][output_ids[0] != utils.IMAGE_TOKEN_INDEX]
@@ -347,6 +356,25 @@ def get_gradio_interface(
     )
 
 
+def get_hash_array(embedding_key: str, arr: np.ndarray | torch.Tensor, model_logger: logging):
+    from base64 import b64encode
+    from hashlib import sha256
+
+    model_logger.debug(f"embedding_key {embedding_key} is None? {embedding_key is None}.")
+    if embedding_key is None:
+        img2hash = arr
+        if isinstance(arr, torch.Tensor):
+            model_logger.debug("images variable is a Tensor, start converting back to numpy")
+            img2hash = arr.numpy(force=True)
+            model_logger.debug("done Tensor converted back to numpy")
+        model_logger.debug("start image hashing")
+        img2hash_fn = sha256(img2hash)
+        embedding_key = b64encode(img2hash_fn.digest())
+        embedding_key = embedding_key.decode("utf-8")
+        model_logger.debug(f"done image hashing, now embedding_key is {embedding_key}.")
+    return embedding_key
+
+
 if __name__ == '__main__':
     parsed_args = parse_args([])
     print("arrrrg:", parsed_args)

scripts/baremetal_entrypoint.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+
+if [ -z "${WORKDIR}" ];
+then
+    WORKDIR=$1
+fi
+
+if [ -z "${XDG_CACHE_HOME}" ];
+then
+    XDG_CACHE_HOME=$HOME/.cache
+fi
+
+echo "WORKDIR: ${WORKDIR} ..."
+echo "XDG_CACHE_HOME: ${XDG_CACHE_HOME} ..."
+
+cd ${WORKDIR}
+
+if [ ! -f "${WORKDIR}/.env_source" ];
+then
+    echo "missing ${WORKDIR}/.env_source file, exit now..."
+    exit 1
+fi
+
+source ${WORKDIR}/.env_source
+echo "FOLDERS_MAP: ${FOLDERS_MAP} ..."
+
+which python
+python --version
+python ${WORKDIR}/scripts/create_folders_and_variables_if_not_exists.py
+
+cd ${WORKDIR}/static
+npm install -g npm pnpm
+pnpm install
+pnpm build
+pnpm tailwindcss -i ${WORKDIR}/static/src/input.css -o ${WORKDIR}/static/dist/output.css
+cd ${WORKDIR}
+
+chmod +x ${WORKDIR}/scripts/entrypoint.sh
+bash ${WORKDIR}/scripts/entrypoint.sh "baremetal"
+
+exit 0

scripts/create_folders_and_variables_if_not_exists.py

@@ -0,0 +1,51 @@
+import json
+import logging
+import os
+from pathlib import Path
+
+
+def stats_pathname(pathname: Path | str):
+    current_pathname = Path(pathname)
+    return current_pathname.is_dir()
+
+
+def create_folder_if_not_exists(pathname: Path | str):
+    current_pathname = Path(pathname)
+    try:
+        print(f"Pathname exists? {current_pathname.exists()}, That's a folder? {current_pathname.is_dir()}...")
+        logging.info(f"Pathname exists? {current_pathname.exists()}, That's a folder? {current_pathname.is_dir()}...")
+        current_pathname.unlink(missing_ok=True)
+    except PermissionError as pe:
+        print(f"permission denied on removing pathname before folder creation:{pe}.")
+        logging.error(f"permission denied on removing pathname before folder creation:{pe}.")
+    except IsADirectoryError as errdir:
+        print(f"that's a directory:{errdir}.")
+        logging.error(f"that's a directory:{errdir}.")
+
+    print(f"Creating pathname: {current_pathname} ...")
+    logging.info(f"Creating pathname: {current_pathname} ...")
+    current_pathname.mkdir(mode=0o770, parents=True, exist_ok=True)
+
+    print(f"assertion: pathname exists and is a folder: {current_pathname} ...")
+    logging.info(f"assertion: pathname exists and is a folder: {current_pathname} ...")
+    assert current_pathname.is_dir()
+
+
+if __name__ == '__main__':
+    folders_string = os.getenv("FOLDERS_MAP")
+    try:
+        folders_dict = json.loads(folders_string)
+        for folder_env_ref, folder_env_path in folders_dict.items():
+            print(f"folder_env_ref:{folder_env_ref}, folder_env_path:{folder_env_path}.")
+            logging.info(f"folder_env_ref:{folder_env_ref}, folder_env_path:{folder_env_path}.")
+            create_folder_if_not_exists(folder_env_path)
+            print("========")
+            assert os.getenv(folder_env_ref) == folder_env_path
+    except (json.JSONDecodeError, TypeError) as jde:
+        print(f"jde:{jde}.")
+        logging.error(f"jde:{jde}.")
+        print("double check your variables, e.g. for mispelling like 'FOLDER_MAP'...")
+        logging.info("double check your variables, e.g. for mispelling like 'FOLDER_MAP' instead than 'FOLDERS_MAP'...")
+        for k_env, v_env in dict(os.environ).items():
+            print(f"{k_env}, v_env:{v_env}.")
+            logging.info(f"{k_env}, v_env:{v_env}.")

scripts/entrypoint.sh

@@ -1,7 +1,11 @@
 #!/usr/bin/env bash
 
-WORKDIR="/var/task"
-XDG_CACHE_HOME="/data"
+if [ -z "$1" ];
+then
+  echo "use no \$1 variable, set WORKDIR and XDG_CACHE_HOME as for docker container mode"
+  WORKDIR="/var/task"
+  XDG_CACHE_HOME="/data"
+fi
 MPLCONFIGDIR=${XDG_CACHE_HOME}/.cache/matplotlib
 TRANSFORMERS_CACHE=${XDG_CACHE_HOME}/.cache/transformers
 FASTAPI_STATIC=${XDG_CACHE_HOME}/static
@@ -45,13 +49,22 @@ echo "WORKDIR - /var/task"
 ls -l ${WORKDIR}
 
 echo "XDG_CACHE_HOME - /data"
-find ${XDG_CACHE_HOME}
+if [ -z "$1" ];
+then
+  echo "use no \$1 variable, show folder ${XDG_CACHE_HOME} content"
+  find ${XDG_CACHE_HOME}
+fi
 
 CUDA_VISIBLE_DEVICES=$(nvidia-smi --query-gpu=memory.free,index --format=csv,nounits,noheader | sort -nr | head -1 | awk '{ print $NF }')
 echo "calculated CUDA_VISIBLE_DEVICES env variable: ${CUDA_VISIBLE_DEVICES}."
 export CUDA_VISIBLE_DEVICES
 
-echo "running command 'uvicorn app.main:app --host 0.0.0.0 --port 7860'..."
-uvicorn app.main:app --host 0.0.0.0 --port 7860
+PYTHONFILE="lisa_on_cuda.app.main"
+if [ -z "$1" ];
+then
+  PYTHONFILE="app.main"
+fi
+echo "running command 'uvicorn ${PYTHONFILE}:app --host 0.0.0.0 --port 7860'..."
+uvicorn ${PYTHONFILE}:app --host 0.0.0.0 --port 7860
 
 exit 0