11# SPDX-License-Identifier: Apache-2.0
22# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
33
4- import base64
54from io import BytesIO
65from pathlib import Path
76
7+ import pybase64
88import torch
99from PIL import Image
1010
@@ -55,7 +55,7 @@ def load_bytes(self, data: bytes) -> Image.Image:
5555 return convert_image_mode (image , self .image_mode )
5656
5757 def load_base64 (self , media_type : str , data : str ) -> Image .Image :
58- return self .load_bytes (base64 .b64decode (data ))
58+ return self .load_bytes (pybase64 .b64decode (data , validate = True ))
5959
6060 def load_file (self , filepath : Path ) -> Image .Image :
6161 image = Image .open (filepath )
@@ -75,7 +75,7 @@ def encode_base64(
7575 image .save (buffer , image_format )
7676 data = buffer .getvalue ()
7777
78- return base64 .b64encode (data ).decode ('utf-8' )
78+ return pybase64 .b64encode (data ).decode ('utf-8' )
7979
8080
8181class ImageEmbeddingMediaIO (MediaIO [torch .Tensor ]):
@@ -88,10 +88,10 @@ def load_bytes(self, data: bytes) -> torch.Tensor:
8888 return torch .load (buffer , weights_only = True )
8989
9090 def load_base64 (self , media_type : str , data : str ) -> torch .Tensor :
91- return self .load_bytes (base64 .b64decode (data ))
91+ return self .load_bytes (pybase64 .b64decode (data , validate = True ))
9292
9393 def load_file (self , filepath : Path ) -> torch .Tensor :
9494 return torch .load (filepath , weights_only = True )
9595
9696 def encode_base64 (self , media : torch .Tensor ) -> str :
97- return base64 .b64encode (media .numpy ()).decode ('utf-8' )
97+ return pybase64 .b64encode (media .numpy ()).decode ('utf-8' )
0 commit comments