llava_example.py 1.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142
  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(model="llava-hf/llava-1.5-7b-hf")
  8. prompt = "USER: <image>\nWhat is the content of this image?\nASSISTANT:"
  9. image_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
  10. "burg.jpg")
  11. image = Image.open(image_path)
  12. sampling_params = SamplingParams(temperature=1.1,
  13. min_p=0.06,
  14. max_tokens=512)
  15. outputs = llm.generate(
  16. {
  17. "prompt": prompt,
  18. "multi_modal_data": {
  19. "image": image
  20. }
  21. },
  22. sampling_params=sampling_params)
  23. for o in outputs:
  24. generated_text = o.outputs[0].text
  25. print(generated_text)
  26. def main():
  27. run_llava()
  28. if __name__ == "__main__":
  29. main()