vlm: enforce pybase64 for image and str encode/decode (#10700)
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user