feat: support new ZeroGPU version addig device_map/device/device2 app_helpers.py functions

lisa_on_cuda/utils/app_helpers.py

@@ -22,6 +22,13 @@ from . import constants, utils
 placeholders = utils.create_placeholder_variables()
 
 
+def get_device_map_kwargs(device_map="auto", device="cuda"):
+    kwargs = {"device_map": device_map}
+    if device != "cuda":
+        kwargs['device_map'] = {"": device}
+    return kwargs
+
+
 def parse_args(args_to_parse, internal_logger=None):
     if internal_logger is None:
         internal_logger = app_logger
@@ -118,12 +125,13 @@ 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, device_map="auto"
+        internal_logger: logging = None, device_map="auto", device="cuda"
 ):
     if internal_logger is None:
         internal_logger = app_logger
     internal_logger.debug(f"prepare kwargs, 4bit:{load_in_4bit}, 8bit:{load_in_8bit}.")
-    kwargs = {"torch_dtype": torch_dtype}
+    kwargs_device_map = get_device_map_kwargs(device_map=device_map, device=device)
+    kwargs = {"torch_dtype": torch_dtype, **kwargs_device_map}
     if load_in_4bit:
         kwargs.update(
             {
@@ -154,15 +162,27 @@ 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,
+        # try to avoid CUDA init RuntimeError on ZeroGPU huggingface hardware (injected into kwargs)
         **kwargs
     )
     internal_logger.debug("model loaded!")
     return _model
 
 
-def get_model(args_to_parse, internal_logger: logging = None, inference_decorator: Callable = None, device_map="auto"):
+def get_model(args_to_parse, internal_logger: logging = None, inference_decorator: Callable = None, device_map="auto", device="cpu", device2="cuda"):
+    """Load model and inference function with arguments. Compatible with ZeroGPU (spaces 0.30.2)
+
+    Args:
+        args_to_parse: default input arguments
+        internal_logger: logger
+        inference_decorator: inference decorator (now it's supported and tested ZeroGPU spaces.GPU decorator)
+        device_map: device type needed for ZeroGPU cuda hw
+        device: device type needed for ZeroGPU cuda hw
+        device2: device type needed for ZeroGPU cuda hw, default to cpu to avoid bug on loading model
+
+    Returns:
+        inference function with LISA model
+    """
     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}.")
@@ -201,7 +221,8 @@ 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  # try to avoid CUDA init RuntimeError on ZeroGPU huggingface hardware
+            device_map=device_map,  # try to avoid CUDA init RuntimeError on ZeroGPU huggingface hardware
+            device=device
         )) if inference_decorator else load_model_for_causal_llm_pretrained(
         args_to_parse.version,
         torch_dtype=torch_dtype,
@@ -226,10 +247,8 @@ def get_model(args_to_parse, internal_logger: logging = None, inference_decorato
     )
     internal_logger.debug(f"_model type:{type(_model)} => {_model}.")
     internal_logger.debug(f"vision_tower type:{type(vision_tower)} => {vision_tower}.")
-    # set device to device_map try to avoid CUDA init RuntimeError on ZeroGPU huggingface hardware
-    device = device_map if device_map else args_to_parse.local_rank
-    internal_logger.debug(f"device to use with vision tower:{device}, device_map:{device_map}, local_rank:{args_to_parse.local_rank}.")
-    vision_tower.to(device=device)
+    # set device to "cuda" try to avoid CUDA init RuntimeError on ZeroGPU huggingface hardware
+    vision_tower.to(device=device2)
     internal_logger.debug("vision tower loaded, prepare clip image processor...")
     _clip_image_processor = CLIPImageProcessor.from_pretrained(_model.config.vision_tower)
     internal_logger.debug("clip image processor done.")
@@ -274,12 +293,24 @@ def prepare_model_vision_tower(_model, args_to_parse, torch_dtype, internal_logg
 
 
 def get_inference_model_by_args(
-        args_to_parse, internal_logger0: logging = None, inference_decorator: Callable = None, device_map="auto"
+        args_to_parse, internal_logger0: logging = None, inference_decorator: Callable = None, device_map="auto", device="cuda"
 ):
+    """Load model and inference function with arguments. Compatible with ZeroGPU (spaces 0.30.2)
+
+    Args:
+        args_to_parse: default input arguments
+        internal_logger0: logger
+        inference_decorator: inference decorator (now it's supported and tested ZeroGPU spaces.GPU decorator)
+        device_map: device type needed for ZeroGPU cuda hw
+        device: device type needed for ZeroGPU cuda hw
+
+    Returns:
+        inference function with LISA model
+    """
     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, device_map=device_map)
+    model, clip_image_processor, tokenizer, transform = get_model(args_to_parse, device_map=device_map, device=device)
     internal_logger0.info("created model, preparing inference function")
     no_seg_out = placeholders["no_seg_out"]