Merge pull request #163 from jhc13/llava-next-prompts

Fix LLaVA-NeXT Vicuna prompt postprocessing
jhc13 · Jun 4, 2024 · 204a58e · 204a58e
2 parents 5edb6f1 + e6d5d55
commit 204a58e
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/taggui/auto_captioning/prompts.py b/taggui/auto_captioning/prompts.py
@@ -61,9 +61,10 @@ def postprocess_prompt_and_generated_text(model_type: CaptionModelType,
             generated_text)
         prompt = prompt.replace('<grounding>', '')
     elif model_type in (CaptionModelType.LLAVA_1_5,
-                        CaptionModelType.LLAVA_NEXT_MISTRAL,
-                        CaptionModelType.LLAVA_NEXT_VICUNA):
+                        CaptionModelType.LLAVA_NEXT_MISTRAL):
         prompt = prompt.replace('<image>', ' ')
+    elif model_type == CaptionModelType.LLAVA_NEXT_VICUNA:
+        prompt = prompt.replace('<image>', '')
     elif model_type == CaptionModelType.LLAVA_LLAMA_3:
         prompt = prompt.replace('<|start_header_id|>', '')
         prompt = prompt.replace('<|end_header_id|>', '')