llava_example.py 1.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647
  1. import os
  2. from PIL import Image
  3. from aphrodite import LLM, SamplingParams
  4. # The assets are located at `s3://air-example-data-2/vllm_opensource_llava/`.
  5. # You can use `.buildkite/download-images.sh` to download them
  6. def run_llava():
  7. llm = LLM(
  8. model="llava-hf/llava-1.5-7b-hf",
  9. image_token_id=32000,
  10. image_input_shape="1,3,336,336",
  11. image_feature_size=576,
  12. )
  13. prompt = "USER: <image>\nWhat is the content of this image?\nASSISTANT:"
  14. image_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
  15. "burg.jpg")
  16. image = Image.open(image_path)
  17. sampling_params = SamplingParams(temperature=1.1,
  18. min_p=0.06,
  19. max_tokens=512)
  20. outputs = llm.generate(
  21. {
  22. "prompt": prompt,
  23. "multi_modal_data": {
  24. "image": image
  25. }
  26. },
  27. sampling_params=sampling_params)
  28. for o in outputs:
  29. generated_text = o.outputs[0].text
  30. print(generated_text)
  31. def main():
  32. run_llava()
  33. if __name__ == "__main__":
  34. main()