feat: add device_map argument to try avoid CUDA init RuntimeError, bump lisa-on-cuda to version 1.4.5

lisa_on_cuda/utils/app_helpers.py

@@ -118,7 +118,7 @@ def preprocess(
 
 def load_model_for_causal_llm_pretrained(
         version, torch_dtype, load_in_8bit, load_in_4bit, seg_token_idx, vision_tower,
-        internal_logger: logging = None
+        internal_logger: logging = None, device_map="auto"
 ):
     if internal_logger is None:
         internal_logger = app_logger
@@ -128,8 +128,7 @@ def load_model_for_causal_llm_pretrained(
         kwargs.update(
             {
                 "torch_dtype": torch.half,
-                # commentare?
-                "load_in_4bit": True,
+                # "load_in_4bit": True,
                 "quantization_config": BitsAndBytesConfig(
                     load_in_4bit=True,
                     bnb_4bit_compute_dtype=torch.float16,
@@ -155,13 +154,15 @@ def load_model_for_causal_llm_pretrained(
         low_cpu_mem_usage=True,
         vision_tower=vision_tower,
         seg_token_idx=seg_token_idx,
+        # try to avoid CUDA init RuntimeError on ZeroGPU huggingface hardware
+        device_map=device_map,
         **kwargs
     )
-    internal_logger.debug(f"model loaded!")
+    internal_logger.debug("model loaded!")
     return _model
 
 
-def get_model(args_to_parse, internal_logger: logging = None, inference_decorator: Callable = None):
+def get_model(args_to_parse, internal_logger: logging = None, inference_decorator: Callable = None, device_map="auto"):
     if internal_logger is None:
         internal_logger = app_logger
     internal_logger.info(f"starting model preparation, folder creation for path: {args_to_parse.vis_save_path}.")
@@ -183,7 +184,7 @@ def get_model(args_to_parse, internal_logger: logging = None, inference_decorato
         use_fast=False,
     )
     _tokenizer.pad_token = _tokenizer.unk_token
-    internal_logger.info(f"tokenizer ok")
+    internal_logger.info("tokenizer ok")
     args_to_parse.seg_token_idx = _tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
     torch_dtype = torch.float32
     if args_to_parse.precision == "bf16":
@@ -199,7 +200,8 @@ def get_model(args_to_parse, internal_logger: logging = None, inference_decorato
             load_in_8bit=args_to_parse.load_in_8bit,
             load_in_4bit=args_to_parse.load_in_4bit,
             seg_token_idx=args_to_parse.seg_token_idx,
-            vision_tower=args_to_parse.vision_tower
+            vision_tower=args_to_parse.vision_tower,
+            device_map=device_map  # try to avoid CUDA init RuntimeError on ZeroGPU huggingface hardware
         )) if inference_decorator else load_model_for_causal_llm_pretrained(
         args_to_parse.version,
         torch_dtype=torch_dtype,
@@ -207,8 +209,9 @@ def get_model(args_to_parse, internal_logger: logging = None, inference_decorato
         load_in_4bit=args_to_parse.load_in_4bit,
         seg_token_idx=args_to_parse.seg_token_idx,
         vision_tower=args_to_parse.vision_tower,
+        device_map=device_map
     )
-    internal_logger.debug(f"causal llm loaded!")
+    internal_logger.debug("causal llm loaded!")
 
     _model.config.eos_token_id = _tokenizer.eos_token_id
     _model.config.bos_token_id = _tokenizer.bos_token_id
@@ -222,11 +225,11 @@ def get_model(args_to_parse, internal_logger: logging = None, inference_decorato
         _model, args_to_parse, torch_dtype
     )
     vision_tower.to(device=args_to_parse.local_rank)
-    internal_logger.debug(f"vision tower loaded, prepare clip image processor...")
+    internal_logger.debug("vision tower loaded, prepare clip image processor...")
     _clip_image_processor = CLIPImageProcessor.from_pretrained(_model.config.vision_tower)
-    internal_logger.debug(f"clip image processor done.")
+    internal_logger.debug("clip image processor done.")
     _transform = ResizeLongestSide(args_to_parse.image_size)
-    internal_logger.debug(f"start model evaluation...")
+    internal_logger.debug("start model evaluation...")
     inference_decorator(_model.eval()) if inference_decorator else _model.eval()
     internal_logger.info("model preparation ok!")
     return _model, _clip_image_processor, _tokenizer, _transform
@@ -261,15 +264,17 @@ def prepare_model_vision_tower(_model, args_to_parse, torch_dtype, internal_logg
         internal_logger.debug(f"vision tower precision fp32? {args_to_parse.precision}, 3.")
         _model = _model.float().cuda()
     vision_tower = _model.get_model().get_vision_tower()
-    internal_logger.debug(f"vision tower ok!")
+    internal_logger.debug("vision tower ok!")
     return _model, vision_tower
 
 
-def get_inference_model_by_args(args_to_parse, internal_logger0: logging = None, inference_decorator: Callable = None):
+def get_inference_model_by_args(
+        args_to_parse, internal_logger0: logging = None, inference_decorator: Callable = None, device_map="auto"
+):
     if internal_logger0 is None:
         internal_logger0 = app_logger
     internal_logger0.info(f"args_to_parse:{args_to_parse}, creating model...")
-    model, clip_image_processor, tokenizer, transform = get_model(args_to_parse)
+    model, clip_image_processor, tokenizer, transform = get_model(args_to_parse, device_map=device_map)
     internal_logger0.info("created model, preparing inference function")
     no_seg_out = placeholders["no_seg_out"]
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "lisa-on-cuda"
-version = "1.4.4"
+version = "1.4.5"
 description = "LISA (Reasoning Segmentation via Large Language Model) on cuda, now with huggingface ZeroGPU support!"
 authors = ["alessandro trinca tornidor <[email protected]>"]
 license = "Apache 2.0"
@@ -8,7 +8,7 @@ readme = "README.md"
 
 [metadata]
 name = "lisa-on-cuda"
-version = "1.4.4"
+version = "1.4.5"
 
 [tool.poetry.urls]
 Source = "https://huggingface.co/spaces/aletrn/lisa-on-cuda/"