udpate

models/mllava/processing_llava.py

@@ -122,6 +122,10 @@ class MLlavaProcessor(ProcessorMixin):
                         # prepend empty image tokens to text
                         if "USER:" in t:
                             t = t.replace("USER:", "USER:" + "<image>" * (num_images - num_image_tokens), 1)
+                        elif "Human:" in t:
+                            t = t.replace("Human:", "Human:" + "<image>" * (num_images - num_image_tokens), 1)
+                        elif "HUMAN:" in t:
+                            t = t.replace("HUMAN:", "HUMAN:" + "<image>" * (num_images - num_image_tokens), 1)
                         else:
                             t = "<image>" * (num_images - num_image_tokens) + t
                         # logger.warning("Image Tokens <image> are not provided in the text. Automatically prepending them before the text. This might cause model to behave unexpectedly.")
@@ -133,6 +137,7 @@ class MLlavaProcessor(ProcessorMixin):
                         t = "".join(t)
                         logger.warning("Number of <image> tokens exceeds number of images. Automatically removing extra tokens at the end of the text.")
                         # raise ValueError("Invalid input text. Number of <image> tokens exceeds number of images.")
+                    text[i] = t
                 texts = text
             else:
                 raise ValueError("Invalid input text. text must be a string or a list of strings.")