[Bugfix] Fix feature size calculation for LLaVA-NeXT (#6982)

2024-07-31 23:46:17 +08:00
parent 2f4e108f75
commit daed30c4a9
5 changed files with 98 additions and 50 deletions
--- a/vllm/model_executor/models/fuyu.py
+++ b/vllm/model_executor/models/fuyu.py
@@ -169,7 +169,7 @@ def input_processor_for_fuyu(ctx: InputContext, llm_inputs: LLMInputs):
        raise TypeError(f"Invalid image type: {type(image_data)}")

    # process prompts
-    prompt = llm_inputs["prompt"]
+    prompt = llm_inputs.get("prompt")
    prompt_token_ids = llm_inputs["prompt_token_ids"]
    tokenizer = cached_get_tokenizer(model_config.model)
    # dim0 is batch_size, dim1 is subseq_size which will always be 1