Update app.py

app.py

@@ -522,8 +522,12 @@ def start_controller():
     return subprocess.Popen(controller_command)
 
 
-def start_worker(model_path: str):
+def start_worker(model_path: str, bits=16):
     logger.info(f"Starting the model worker for the model {model_path}")
+    model_name = model_path.strip('/').split('/')[-1]
+    assert bits in [4, 8, 16], "It can be only loaded with 16-bit, 8-bit, and 4-bit."
+    if bits != 16:
+        model_name += f'-{bits}bit'
     worker_command = [
         "python",
         "-m",
@@ -535,9 +539,10 @@ def start_worker(model_path: str):
         "--model-path",
         model_path,
         "--model-name",
-        "llava-v1.5-13b-4bit",
-        "--load-4bit",
+        model_name,
     ]
+    if bits != 16:
+        worker_command += [f'--load-{bits}bit']
     return subprocess.Popen(worker_command)
 
 
@@ -586,8 +591,10 @@ if __name__ == "__main__":
     logger.info(f"args: {args}")
 
     model_path = "liuhaotian/llava-v1.5-13b"
+    bits = 4
+    # set bits=4 for T4, bits=8 for A10G (24G), and bits=16 for A100 (40G)
 
-    preload_models(model_path)
+    preload_models(model_path, bits=bits)
 
     controller_proc = start_controller()
     worker_proc = start_worker(model_path)