Merge pull request #2 from TRI-ML/master

interactive_demo.py

@@ -47,20 +47,12 @@ def heart_beat_worker(controller):
 
 
 class ModelWorker:
-    def __init__(self, controller_addr, worker_addr, worker_id, no_register, vlm, model_base, model_name):
+    def __init__(self, controller_addr, worker_addr, worker_id, no_register, vlm, model_name):
         self.controller_addr = controller_addr
         self.worker_addr = worker_addr
         self.worker_id = worker_id
         self.model_name = model_name
-
-        # logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
         self.vlm = vlm
-        self.tokenizer, self.model, self.image_processor, self.context_len = (
-            vlm.tokenizer,
-            vlm.model,
-            vlm.image_processor,
-            vlm.max_length,
-        )
 
         if not no_register:
             self.register_to_controller()
@@ -68,18 +60,12 @@ class ModelWorker:
             self.heart_beat_thread.start()
 
     def register_to_controller(self):
-        # logger.info("Register to controller")
-
         url = self.controller_addr + "/register_worker"
         data = {"worker_name": self.worker_addr, "check_heart_beat": True, "worker_status": self.get_status()}
         r = requests.post(url, json=data)
         assert r.status_code == 200
 
     def send_heart_beat(self):
-        # logger.info(f"Send heart beat. Models: {[self.model_name]}. "
-        #             f"Semaphore: {pretty_print_semaphore(model_semaphore)}. "
-        #             f"global_counter: {global_counter}")
-
         url = self.controller_addr + "/receive_heart_beat"
 
         while True:
@@ -91,7 +77,6 @@ class ModelWorker:
                 break
             except requests.exceptions.RequestException:
                 pass
-                # logger.error(f"heart beat error: {e}")
             time.sleep(5)
 
         if not exist:
@@ -145,12 +130,12 @@ class ModelWorker:
             else:
                 question_prompt = [prompt_fn()]
 
-        if isinstance(self.image_processor, Compose) or hasattr(self.image_processor, "is_prismatic"):
+        if isinstance(self.vlm.image_processor, Compose) or hasattr(self.vlm.image_processor, "is_prismatic"):
             # This is a standard `torchvision.transforms` object or custom PrismaticVLM wrapper
-            pixel_values = self.image_processor(images[0].convert("RGB"))
+            pixel_values = self.vlm.image_processor(images[0].convert("RGB"))
         else:
             # Assume `image_transform` is a HF ImageProcessor...
-            pixel_values = self.image_processor(images[0].convert("RGB"), return_tensors="pt")["pixel_values"][0]
+            pixel_values = self.vlm.image_processor(images[0].convert("RGB"), return_tensors="pt")["pixel_values"][0]
 
         if type(pixel_values) is dict:
             for k in pixel_values.keys():
@@ -227,31 +212,29 @@ overwatch = initialize_overwatch(__name__)
 class DemoConfig:
     # fmt: off
 
-    # === Model Parameters =>> Quartz ===
-    model_family: str = "quartz"                    # Model family to load from in < `quartz` | `llava-v15` | ... >
-    model_id: str = "llava-v1.5-7b"     # Model ID to load and run (instance of `model_family`)
-    model_dir: Path = (                             # Path to model checkpoint to load --> should be self-contained
-        "resize-naive-siglip-vit-l-16-384px-no-align-2-epochs+13b+stage-finetune+x7"
-    )
+    # === Model Parameters =>> Prismatic ===
+    model_family: str = "prismatic"           # Model family to load from in < `prismatic` | `llava-v15` | ... >
+    model_id: str = "prism-dinosiglip+7b"     # Model ID to load and run (instance of `model_family`)
+    model_dir: str = None                     # Can optionally supply model_dir instead of model_id
 
     # === Model Parameters =>> Official LLaVa ===
     # model_family: str = "llava-v15"
     # model_id: str = "llava-v1.5-13b"
     # model_dir: Path = "liuhaotian/llava-v1.5-13b"
 
+    # === Model Parameters =>> Official InstructBLIP ===
+    # model_family: str = "instruct-blip"
+    # model_id: str = "instructblip-vicuna-7b"
+    # model_dir: Path = "Salesforce/instructblip-vicuna-7b"
+
     # Model Worker Parameters
     host: str = "0.0.0.0"
     port: int = 40000
     controller_address: str = "http://localhost:10000"
-    model_base: str = "llava-v15"
     limit_model_concurrency: int = 5
     stream_interval: int = 1
     no_register: bool = False
 
-    # Inference Parameters
-    device_batch_size: int = 1                      # Device Batch Size set to 1 until LLaVa/HF LLaMa fixes bugs!
-    num_workers: int = 2                            # Number of Dataloader Workers (on each process)
-
     # HF Hub Credentials (for LLaMa-2)
     hf_token: Union[str, Path] = Path(".hf_token")  # Environment variable or Path to HF Token
 
@@ -259,14 +242,8 @@ class DemoConfig:
     seed: int = 21                                  # Random Seed (for reproducibility)
 
     def __post_init__(self) -> None:
-        if self.model_family == "quartz":
-            self.model_name = MODEL_ID_TO_NAME[str(self.model_dir)]
-            self.run_dir = Path("/mnt/fsx/x-onyx-vlms/runs") / self.model_dir
-        elif self.model_family in {"instruct-blip", "llava", "llava-v15"}:
-            self.model_name = MODEL_ID_TO_NAME[self.model_id]
-            self.run_dir = self.model_dir
-        else:
-            raise ValueError(f"Run Directory for `{self.model_family = }` does not exist!")
+        self.run_dir = self.model_dir
+        self.model_name = MODEL_ID_TO_NAME[str(self.model_id)]
         self.worker_address = f"http://localhost:{self.port}"
 
     # fmt: on
@@ -286,7 +263,7 @@ def interactive_demo(cfg: DemoConfig):
     global limit_model_concurrency
     limit_model_concurrency = cfg.limit_model_concurrency
     worker = ModelWorker(
-        cfg.controller_address, cfg.worker_address, worker_id, cfg.no_register, vlm, cfg.model_base, cfg.model_name
+        cfg.controller_address, cfg.worker_address, worker_id, cfg.no_register, vlm, cfg.model_name
     )
     uvicorn.run(app, host=cfg.host, port=cfg.port, log_level="info")
 

serve/__init__.py

@@ -5,31 +5,24 @@ from collections import OrderedDict
 MODEL_ID_TO_NAME = OrderedDict(
     [
         (
-            "llava-lvis4v-lrv+redux-lvis4v-lrv-resize-naive-dinosiglip-vit-so-14-384px-no-align+13b+stage-finetune+x7",
-            "PrismaticVLM 13B - Chat",
-        ),
-        (
-            "llava-lvis4v-lrv+redux-lvis4v-lrv-resize-naive-dinosiglip-vit-so-14-384px-no-align+7b+stage-finetune+x7",
-            "PrismaticVLM 7B - Chat",
-        ),
-        (
-            "llava-lvis4v-lrv+redux-lvis4v-lrv-resize-naive-dinosiglip-vit-so-14-384px-no-align-llama2pure+13b+stage-finetune+x7",
+            "prism-dinosiglip+13b",
             "PrismaticVLM 13B",
         ),
         (
-            "llava-lvis4v-lrv+redux-lvis4v-lrv-resize-naive-dinosiglip-vit-so-14-384px-no-align-llama2pure+7b+stage-finetune+x7",
+            "prism-dinosiglip+7b",
             "PrismaticVLM 7B",
         ),
         (
-            "redux-resize-naive-dinosiglip-vit-so-14-384px-no-align-llama2pure+13b+stage-finetune+x7",
+            "prism-dinosiglip-controlled+13b",
             "PrismaticVLM 13B (Controlled)",
         ),
         (
-            "redux-resize-naive-dinosiglip-vit-so-14-384px-no-align-llama2pure+7b+stage-finetune+x7",
+            "prism-dinosiglip-controlled+7b",
             "PrismaticVLM 7B (Controlled)",
         ),
-        ("llava-v1.5-13b", "LLaVA 1.5: 13B"),
-        ("llava-v1.5-7b", "LLaVA 1.5: 7B"),
+        ("llava-v1.5-13b", "LLaVA 1.5 13B"),
+        ("llava-v1.5-7b", "LLaVA 1.5 7B"),
+        ("instructblip-vicuna-7b", "InstructBLIP 7B"),
     ]
 )
 

serve/gradio_web_server.py

@@ -93,24 +93,6 @@ def vote_last_response(state, vote_type, model_selector, request: gr.Request):
         fout.write(json.dumps(data) + "\n")
 
 
-# def upvote_last_response(state, model_selector, request: gr.Request):
-#     logger.info(f"upvote. ip: {request.client.host}")
-#     vote_last_response(state, "upvote", model_selector, request)
-#     return ("",) + (disable_btn,) * 3
-
-
-# def downvote_last_response(state, model_selector, request: gr.Request):
-#     logger.info(f"downvote. ip: {request.client.host}")
-#     vote_last_response(state, "downvote", model_selector, request)
-#     return ("",) + (disable_btn,) * 3
-
-
-# def flag_last_response(state, model_selector, request: gr.Request):
-#     logger.info(f"flag. ip: {request.client.host}")
-#     vote_last_response(state, "flag", model_selector, request)
-#     return ("",) + (disable_btn,) * 3
-
-
 def regenerate(state, image_process_mode, request: gr.Request):
     logger.info(f"regenerate. ip: {request.client.host}")
     state.messages[-1][-1] = None
@@ -388,15 +370,6 @@ def build_demo(embed_mode):
 
         # Register listeners
         btn_list = [regenerate_btn, clear_btn]
-        # upvote_btn.click(
-        #     upvote_last_response, [state, model_selector], [textbox, upvote_btn, downvote_btn, flag_btn], queue=False
-        # )
-        # downvote_btn.click(
-        #     downvote_last_response, [state, model_selector], [textbox, upvote_btn, downvote_btn, flag_btn], queue=False
-        # )
-        # flag_btn.click(
-        #     flag_last_response, [state, model_selector], [textbox, upvote_btn, downvote_btn, flag_btn], queue=False
-        # )
 
         regenerate_btn.click(
             regenerate, [state, image_process_mode], [state, chatbot, textbox, imagebox, *btn_list], queue=False