Merge pull request haotian-liu#1038 from heltrix/patch-1

[Fix] Add image_sizes for cli.py for LLaVA-v1.6
harithoppil · Feb 1, 2024 · c878cc3 · c878cc3
2 parents ac6b1e7 + 4718611
commit c878cc3
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/llava/serve/cli.py b/llava/serve/cli.py
@@ -52,6 +52,7 @@ def main(args):
         roles = conv.roles
 
     image = load_image(args.image_file)
+    image_size = image.size
     # Similar operation in model_worker.py
     image_tensor = process_images([image], image_processor, model.config)
     if type(image_tensor) is list:
@@ -94,6 +95,7 @@ def main(args):
             output_ids = model.generate(
                 input_ids,
                 images=image_tensor,
+                image_sizes=[image_size],
                 do_sample=True if args.temperature > 0 else False,
                 temperature=args.temperature,
                 max_new_tokens=args.max_new_tokens,