Update llava_llama3/serve/cli.py

llava_llama3/serve/cli.py

@@ -26,58 +26,67 @@ def load_image(image_file):
     return image
 
 
-def chat_llava(args, image_file, text, tokenizer, model, streamer, image_processor, context_len):
+def chat_llava(args, image_file, text, tokenizer, model, image_processor, context_len, streamer=None):
     # Model
     disable_torch_init()
 
     conv = conv_templates[args.conv_mode].copy()
     roles = conv.roles
     
-    print(f"\033[91m{image_file}, {type(image_file)}\033[0m")
-    image = load_image(image_file)
-    print(f"\033[91m{image}, {type(image)}\033[0m")
-    
-    image_size = image.size
-    image_tensor = process_images([image], image_processor, model.config)
-    if type(image_tensor) is list:
-        image_tensor = [image.to(model.device, dtype=torch.float16) for image in image_tensor]
-    else:
-        image_tensor = image_tensor.to(model.device, dtype=torch.float16)
-        
     inp = text
 
-    # print(f"{roles[1]}: ", end="")
-
-    if image is not None:
-        # first message
+    if image_file is not None:
+        print(image_file, type(image_file))
+        image = load_image(image_file)
+        print(image, type(image))
+        image_size = image.size
+        image_tensor = process_images([image], image_processor, model.config)
+        if type(image_tensor) is list:
+            image_tensor = [image.to(model.device, dtype=torch.float16) for image in image_tensor]
+        else:
+            image_tensor = image_tensor.to(model.device, dtype=torch.float16)
+        
         if model.config.mm_use_im_start_end:
             inp = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + inp
         else:
             inp = DEFAULT_IMAGE_TOKEN + '\n' + inp
-        image = None
-    
-    conv.append_message(conv.roles[0], inp)
-    conv.append_message(conv.roles[1], None)
-    prompt = conv.get_prompt()
 
-    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
-    stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
-    keywords = [stop_str]
-    # streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        conv.append_message(conv.roles[0], inp)
+        conv.append_message(conv.roles[1], None)
+        prompt = conv.get_prompt()
+
+        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
+        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+        keywords = [stop_str]
 
-    with torch.inference_mode():
-        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,
-            streamer=streamer,
-            use_cache=True)
+        with torch.inference_mode():
+            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,
+                streamer=streamer,
+                use_cache=True)
+    else:
+        conv.append_message(conv.roles[0], inp)
+        conv.append_message(conv.roles[1], None)
+        prompt = conv.get_prompt()
+
+        input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(model.device)
+        
+        with torch.inference_mode():
+            output_ids = model.generate(
+                input_ids,
+                do_sample=True if args.temperature > 0 else False,
+                temperature=args.temperature,
+                max_new_tokens=args.max_new_tokens,
+                use_cache=True)
 
     outputs = tokenizer.decode(output_ids[0]).strip()
     conv.messages[-1][-1] = outputs
 
     # Return the model's output as a string
+    # return outputs
     return outputs.replace('<|end_of_text|>', '\n').lstrip()