llava_example.py 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. import os
  2. import subprocess
  3. from PIL import Image
  4. from aphrodite import LLM
  5. # The assets are located at `s3://air-example-data-2/vllm_opensource_llava/`.
  6. # You can use `.buildkite/download-images.sh` to download them
  7. def run_llava():
  8. llm = LLM(
  9. model="llava-hf/llava-1.5-7b-hf",
  10. image_token_id=32000,
  11. image_input_shape="1,3,336,336",
  12. image_feature_size=576,
  13. )
  14. prompt = "<image>" * 576 + (
  15. "\nUSER: What is the content of this image?\nASSISTANT:")
  16. image = Image.open("images/stop_sign.jpg")
  17. outputs = llm.generate({
  18. "prompt": prompt,
  19. "multi_modal_data": {
  20. "image": image
  21. },
  22. })
  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. # Download from s3
  30. s3_bucket_path = "s3://air-example-data-2/vllm_opensource_llava/"
  31. local_directory = "images"
  32. # Make sure the local directory exists or create it
  33. os.makedirs(local_directory, exist_ok=True)
  34. # Use AWS CLI to sync the directory, assume anonymous access
  35. subprocess.check_call([
  36. "aws",
  37. "s3",
  38. "sync",
  39. s3_bucket_path,
  40. local_directory,
  41. "--no-sign-request",
  42. ])
  43. main()