| import os |
| import sys |
|
|
| _SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__)) |
| _SUBMIT_ROOT = os.path.dirname(os.path.dirname(_SCRIPTS_DIR)) |
|
|
| if _SCRIPTS_DIR not in sys.path: |
| sys.path.insert(0, _SCRIPTS_DIR) |
| if _SUBMIT_ROOT not in sys.path: |
| sys.path.insert(0, _SUBMIT_ROOT) |
|
|
| _LM_PATH = os.path.join(_SUBMIT_ROOT, "model", "LM") |
| _VL_B_PATH = os.path.join(_SUBMIT_ROOT, "model", "VL-B") |
| _VL_F_PATH = os.path.join(_SUBMIT_ROOT, "model", "VL-F") |
| _DIRECTION_MAP_PATH = os.path.join(_SCRIPTS_DIR, "object_direction_map.yaml") |
|
|
|
|
| def load_robospatialBrain_models(low_memory=False): |
| import yaml |
| import torch |
| from lm_classifier import load_lm_classifier |
| from robobrain_runner import load_robobrain |
|
|
| lm_clf = load_lm_classifier(_LM_PATH) |
| torch.cuda.empty_cache() |
|
|
| with open(_DIRECTION_MAP_PATH) as f: |
| direction_map = yaml.safe_load(f) |
|
|
| if low_memory: |
| return { |
| "lm_clf": lm_clf, |
| "direction_map": direction_map, |
| "low_memory": True, |
| "_vl_b_path": _VL_B_PATH, |
| "_vl_f_path": _VL_F_PATH, |
| "_current_vl": None, |
| "_current_vl_type": None, |
| } |
|
|
| vl_b = load_robobrain(_VL_B_PATH) |
| torch.cuda.empty_cache() |
|
|
| vl_f = load_robobrain(_VL_F_PATH) |
| torch.cuda.empty_cache() |
|
|
| return { |
| "lm_clf": lm_clf, |
| "vl_b": vl_b, |
| "vl_f": vl_f, |
| "direction_map": direction_map, |
| "low_memory": False, |
| } |
|
|
|
|
| def _swap_vl(kwargs, needed_type, load_robobrain): |
| import torch |
| if kwargs["_current_vl_type"] != needed_type: |
| if kwargs["_current_vl"] is not None: |
| kwargs["_current_vl"].clear() |
| kwargs["_current_vl"] = None |
| torch.cuda.empty_cache() |
| path = kwargs["_vl_f_path"] if needed_type == "F" else kwargs["_vl_b_path"] |
| kwargs["_current_vl"] = load_robobrain(path) |
| kwargs["_current_vl_type"] = needed_type |
|
|
|
|
| def _run_compatibility(question, image_path, kwargs, low_memory, run_robobrain, load_robobrain): |
| if low_memory: |
| _swap_vl(kwargs, "F", load_robobrain) |
| result_f = run_robobrain(question, image_path, None, kwargs["_current_vl"], LM_classify="compatibility") |
| if result_f != "yes": |
| return result_f |
| _swap_vl(kwargs, "B", load_robobrain) |
| return run_robobrain(question, image_path, None, kwargs["_current_vl"], LM_classify="compatibility") |
| else: |
| result_f = run_robobrain(question, image_path, None, kwargs["vl_f"], LM_classify="compatibility") |
| if result_f != "yes": |
| return result_f |
| return run_robobrain(question, image_path, None, kwargs["vl_b"], LM_classify="compatibility") |
|
|
|
|
| def inference_single(question, image_path, kwargs): |
| import torch |
| from lm_classifier import classify_single |
| from object_direction_remap import should_remap, run_context_with_remap |
| from robobrain_runner import run_robobrain, load_robobrain |
|
|
| lm_clf = kwargs["lm_clf"] |
| direction_map = kwargs["direction_map"] |
| low_memory = kwargs.get("low_memory", False) |
|
|
| category = classify_single(question, lm_clf) |
| torch.cuda.empty_cache() |
|
|
| if category == "compatibility": |
| return _run_compatibility(question, image_path, kwargs, low_memory, run_robobrain, load_robobrain) |
|
|
| if low_memory: |
| _swap_vl(kwargs, "B", load_robobrain) |
| model_kw = kwargs["_current_vl"] |
| else: |
| model_kw = kwargs["vl_b"] |
|
|
| if category == "context": |
| do_remap, _, _ = should_remap(question) |
| if do_remap: |
| return run_context_with_remap( |
| question, image_path, None, model_kw, lm_clf, direction_map |
| ) |
|
|
| return run_robobrain(question, image_path, None, model_kw, LM_classify=category) |
|
|
