Merge pull request haotian-liu#548 from eggry/main

fix: some args are not used in cli.py
Chomolungma · Oct 12, 2023 · 04d102f · 04d102f
2 parents b4f4e01 + 2484462
commit 04d102f
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/llava/serve/cli.py b/llava/serve/cli.py
@@ -29,7 +29,7 @@ def main(args):
     disable_torch_init()
 
     model_name = get_model_name_from_path(args.model_path)
-    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
 
     if 'llama-2' in model_name.lower():
         conv_mode = "llava_llama_2"
@@ -90,8 +90,8 @@ def main(args):
                 input_ids,
                 images=image_tensor,
                 do_sample=True,
-                temperature=0.2,
-                max_new_tokens=1024,
+                temperature=args.temperature,
+                max_new_tokens=args.max_new_tokens,
                 streamer=streamer,
                 use_cache=True,
                 stopping_criteria=[stopping_criteria])
@@ -108,7 +108,7 @@ def main(args):
     parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
     parser.add_argument("--model-base", type=str, default=None)
     parser.add_argument("--image-file", type=str, required=True)
-    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--device", type=str, default="cuda")
     parser.add_argument("--conv-mode", type=str, default=None)
     parser.add_argument("--temperature", type=float, default=0.2)
     parser.add_argument("--max-new-tokens", type=int, default=512)