clean

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")