vlm: enforce pybase64 for image and str encode/decode (#10700)

This commit is contained in:
b8zhong
2025-10-21 04:05:32 -07:00
committed by GitHub
parent 05d3667ab9
commit d0a64c7e2c
8 changed files with 20 additions and 36 deletions

View File

@@ -872,9 +872,9 @@ def get_image_bytes(image_file: Union[str, bytes]):
return f.read()
elif image_file.startswith("data:"):
image_file = image_file.split(",")[1]
return pybase64.b64decode(image_file)
return pybase64.b64decode(image_file, validate=True)
elif isinstance(image_file, str):
return pybase64.b64decode(image_file)
return pybase64.b64decode(image_file, validate=True)
else:
raise NotImplementedError(f"Invalid image: {image_file}")
@@ -911,7 +911,7 @@ def load_video(video_file: Union[str, bytes], use_gpu: bool = True):
vr = VideoReader(tmp_file.name, ctx=ctx)
elif video_file.startswith("data:"):
_, encoded = video_file.split(",", 1)
video_bytes = pybase64.b64decode(encoded)
video_bytes = pybase64.b64decode(encoded, validate=True)
tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
tmp_file.write(video_bytes)
tmp_file.close()
@@ -919,7 +919,7 @@ def load_video(video_file: Union[str, bytes], use_gpu: bool = True):
elif os.path.isfile(video_file):
vr = VideoReader(video_file, ctx=ctx)
else:
video_bytes = pybase64.b64decode(video_file)
video_bytes = pybase64.b64decode(video_file, validate=True)
tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
tmp_file.write(video_bytes)
tmp_file.close()
@@ -2083,7 +2083,7 @@ class MultiprocessingSerializer:
if output_str:
# Convert bytes to base64-encoded string
output = pybase64.b64encode(output).decode("utf-8")
pybase64.b64encode(output).decode("utf-8")
return output