[CI/Build] vLLM cache directory for images (#6444)
This commit is contained in:
@ -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()
|
||||
|
||||
Reference in New Issue
Block a user