[CI/Build] vLLM cache directory for images (#6444)

This commit is contained in:
Cyrus Leung
2024-07-16 14:12:25 +08:00
committed by GitHub
parent 37d776606f
commit d97011512e
13 changed files with 123 additions and 140 deletions

View File

@ -1,12 +1,5 @@
import os
import subprocess
from PIL import Image
from vllm import LLM
# The assets are located at `s3://air-example-data-2/vllm_opensource_llava/`.
# You can use `.buildkite/download-images.sh` to download them
from vllm.assets.image import ImageAsset
def run_llava():
@ -14,7 +7,7 @@ def run_llava():
prompt = "USER: <image>\nWhat is the content of this image?\nASSISTANT:"
image = Image.open("images/stop_sign.jpg")
image = ImageAsset("stop_sign").pil_image
outputs = llm.generate({
"prompt": prompt,
@ -28,25 +21,5 @@ def run_llava():
print(generated_text)
def main():
run_llava()
if __name__ == "__main__":
# Download from s3
s3_bucket_path = "s3://air-example-data-2/vllm_opensource_llava/"
local_directory = "images"
# Make sure the local directory exists or create it
os.makedirs(local_directory, exist_ok=True)
# Use AWS CLI to sync the directory, assume anonymous access
subprocess.check_call([
"aws",
"s3",
"sync",
s3_bucket_path,
local_directory,
"--no-sign-request",
])
main()
run_llava()