diff options
Diffstat (limited to 'g4f/image.py')
| -rw-r--r-- | g4f/image.py | 20 |
1 files changed, 14 insertions, 6 deletions
diff --git a/g4f/image.py b/g4f/image.py index 1a4692b3..3f26f75f 100644 --- a/g4f/image.py +++ b/g4f/image.py @@ -210,20 +210,28 @@ def format_images_markdown(images, alt: str, preview: str = None) -> str: end_flag = "<!-- generated images end -->\n" return f"\n{start_flag}{images}\n{end_flag}\n" -def to_bytes(image: Image) -> bytes: +def to_bytes(image: ImageType) -> bytes: """ Converts the given image to bytes. Args: - image (Image.Image): The image to convert. + image (ImageType): The image to convert. Returns: bytes: The image as bytes. """ - bytes_io = BytesIO() - image.save(bytes_io, image.format) - image.seek(0) - return bytes_io.getvalue() + if isinstance(image, bytes): + return image + elif isinstance(image, str): + is_data_uri_an_image(image) + return extract_data_uri(image) + elif isinstance(image, Image): + bytes_io = BytesIO() + image.save(bytes_io, image.format) + image.seek(0) + return bytes_io.getvalue() + else: + return image.read() class ImageResponse: def __init__( |