1
1
# SPDX-License-Identifier: Apache-2.0
2
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
3
4
- import base64
5
4
from io import BytesIO
6
5
from pathlib import Path
7
6
7
+ import pybase64
8
8
import torch
9
9
from PIL import Image
10
10
@@ -55,7 +55,7 @@ def load_bytes(self, data: bytes) -> Image.Image:
55
55
return convert_image_mode (image , self .image_mode )
56
56
57
57
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 ))
59
59
60
60
def load_file (self , filepath : Path ) -> Image .Image :
61
61
image = Image .open (filepath )
@@ -75,7 +75,7 @@ def encode_base64(
75
75
image .save (buffer , image_format )
76
76
data = buffer .getvalue ()
77
77
78
- return base64 .b64encode (data ).decode ('utf-8' )
78
+ return pybase64 .b64encode (data ).decode ('utf-8' )
79
79
80
80
81
81
class ImageEmbeddingMediaIO (MediaIO [torch .Tensor ]):
@@ -88,10 +88,10 @@ def load_bytes(self, data: bytes) -> torch.Tensor:
88
88
return torch .load (buffer , weights_only = True )
89
89
90
90
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 ))
92
92
93
93
def load_file (self , filepath : Path ) -> torch .Tensor :
94
94
return torch .load (filepath , weights_only = True )
95
95
96
96
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