Update app.py

app.py

@@ -232,17 +232,17 @@ def bot(history):
     our_chatbot.conversation.append_message(our_chatbot.conversation.roles[1], None)
     prompt = our_chatbot.conversation.get_prompt()
 
-    # input_ids = (
-    #     tokenizer_image_token(
-    #         prompt, our_chatbot.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
-    #     )
-    #     .unsqueeze(0)
-    #     .to(our_chatbot.model.device)
-    # )
-    input_ids = tokenizer_image_token(
+    input_ids = (
+        tokenizer_image_token(
             prompt, our_chatbot.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
-        ).unsqueeze(0).to(our_chatbot.model.device)
-    print("### input_id",input_ids)
+        )
+        .unsqueeze(0)
+        .to(our_chatbot.model.device)
+    )
+    # input_ids = tokenizer_image_token(
+    #         prompt, our_chatbot.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
+    #     ).unsqueeze(0).to(our_chatbot.model.device)
+    # print("### input_id",input_ids)
     stop_str = (
         our_chatbot.conversation.sep
         if our_chatbot.conversation.sep_style != SeparatorStyle.TWO
@@ -252,58 +252,58 @@ def bot(history):
     stopping_criteria = KeywordsStoppingCriteria(
         keywords, our_chatbot.tokenizer, input_ids
     )
-    # streamer = TextStreamer(
-    #     our_chatbot.tokenizer, skip_prompt=True, skip_special_tokens=True
-    # )
-    streamer = TextIteratorStreamer(
+    streamer = TextStreamer(
         our_chatbot.tokenizer, skip_prompt=True, skip_special_tokens=True
     )
+    # streamer = TextIteratorStreamer(
+    #     our_chatbot.tokenizer, skip_prompt=True, skip_special_tokens=True
+    # )
     print(our_chatbot.model.device)
     print(input_ids.device)
     print(image_tensor.device)
     # import pdb;pdb.set_trace()
-    # with torch.inference_mode():
-    #     output_ids = our_chatbot.model.generate(
-    #         input_ids,
-    #         images=image_tensor,
-    #         do_sample=True,
-    #         temperature=0.2,
-    #         max_new_tokens=1024,
-    #         streamer=streamer,
-    #         use_cache=False,
-    #         stopping_criteria=[stopping_criteria],
-    #     )
-
-    # outputs = our_chatbot.tokenizer.decode(output_ids[0]).strip()
-    # if outputs.endswith(stop_str):
-    #     outputs = outputs[: -len(stop_str)]
-    # our_chatbot.conversation.messages[-1][-1] = outputs
-
-    # history[-1] = [text, outputs]
-
-    # return history
-    generate_kwargs = dict(
-        inputs=input_ids,
-        streamer=streamer,
-        images=image_tensor,
-        max_new_tokens=1024,
-        do_sample=True,
-        temperature=0.2,
-        num_beams=1,
-        use_cache=False,
-        stopping_criteria=[stopping_criteria],
-    )
+    with torch.inference_mode():
+        output_ids = our_chatbot.model.generate(
+            input_ids,
+            images=image_tensor,
+            do_sample=True,
+            temperature=0.2,
+            max_new_tokens=1024,
+            streamer=streamer,
+            use_cache=False,
+            stopping_criteria=[stopping_criteria],
+        )
+
+    outputs = our_chatbot.tokenizer.decode(output_ids[0]).strip()
+    if outputs.endswith(stop_str):
+        outputs = outputs[: -len(stop_str)]
+    our_chatbot.conversation.messages[-1][-1] = outputs
+
+    history[-1] = [text, outputs]
+
+    return history
+    # generate_kwargs = dict(
+    #     inputs=input_ids,
+    #     streamer=streamer,
+    #     images=image_tensor,
+    #     max_new_tokens=1024,
+    #     do_sample=True,
+    #     temperature=0.2,
+    #     num_beams=1,
+    #     use_cache=False,
+    #     stopping_criteria=[stopping_criteria],
+    # )
     
-    t = Thread(target=our_chatbot.model.generate, kwargs=generate_kwargs)
-    t.start()
+    # t = Thread(target=our_chatbot.model.generate, kwargs=generate_kwargs)
+    # t.start()
     
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
+    # outputs = []
+    # for text in streamer:
+    #     outputs.append(text)
+    #     yield "".join(outputs)
     
-    our_chatbot.conversation.messages[-1][-1] = "".join(outputs)
-    history[-1] = [text, "".join(outputs)]
+    # our_chatbot.conversation.messages[-1][-1] = "".join(outputs)
+    # history[-1] = [text, "".join(outputs)]
     
 
 txt = gr.Textbox(