Skip to content

Commit

Permalink
fix(vllm): images and videos are base64 by default
Browse files Browse the repository at this point in the history
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
  • Loading branch information
mudler committed Oct 17, 2024
1 parent 54c0f15 commit ef963ba
Showing 1 changed file with 15 additions and 5 deletions.
20 changes: 15 additions & 5 deletions backend/python/vllm/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
from vllm.transformers_utils.tokenizer import get_tokenizer
from vllm.multimodal.utils import fetch_image
from vllm.assets.video import VideoAsset
import base64
import io

_ONE_DAY_IN_SECONDS = 60 * 60 * 24

Expand Down Expand Up @@ -262,16 +264,19 @@ async def _predict(self, request, context, streaming=False):

def load_image(self, image_path: str):
"""
Load an image from the given file path.
Load an image from the given file path or base64 encoded data.
Args:
image_path (str): The path to the image file.
image_path (str): The path to the image file or base64 encoded data.
Returns:
Image: The loaded image.
"""
try:
return Image.open(image_path)

image_data = base64.b64decode(image_path)
image = Image.open(io.BytesIO(image_data))
return image
except Exception as e:
print(f"Error loading image {image_path}: {e}", file=sys.stderr)
return self.load_video(image_path)
Expand All @@ -287,10 +292,15 @@ def load_video(self, video_path: str):
Video: The loaded video.
"""
try:
video = VideoAsset(name=video_path).np_ndarrays
timestamp = str(int(time.time() * 1000)) # Generate timestamp
p = f"/tmp/vl-{timestamp}.data" # Use timestamp in filename
with open(p, "wb") as f:
f.write(base64.b64decode(video_path))
video = VideoAsset(name=p).np_ndarrays
os.remove(p)
return video
except Exception as e:
print(f"Error loading video {image_path}: {e}", file=sys.stderr)
print(f"Error loading video {video_path}: {e}", file=sys.stderr)
return None

async def serve(address):
Expand Down

0 comments on commit ef963ba

Please sign in to comment.