diff options
Diffstat (limited to 'g4f/image.py')
| -rw-r--r-- | g4f/image.py | 23 |
1 files changed, 13 insertions, 10 deletions
diff --git a/g4f/image.py b/g4f/image.py index 01d6ae50..6370a06f 100644 --- a/g4f/image.py +++ b/g4f/image.py @@ -11,7 +11,7 @@ try: has_requirements = True except ImportError: has_requirements = False - + from .errors import MissingRequirementsError ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp', 'svg'} @@ -28,9 +28,11 @@ def to_image(image: ImageType, is_svg: bool = False) -> Image: """ if not has_requirements: raise MissingRequirementsError('Install "pillow" package for images') + if isinstance(image, str): is_data_uri_an_image(image) image = extract_data_uri(image) + if is_svg: try: import cairosvg @@ -41,6 +43,7 @@ def to_image(image: ImageType, is_svg: bool = False) -> Image: buffer = BytesIO() cairosvg.svg2png(image, write_to=buffer) return open_image(buffer) + if isinstance(image, bytes): is_accepted_format(image) return open_image(BytesIO(image)) @@ -48,6 +51,7 @@ def to_image(image: ImageType, is_svg: bool = False) -> Image: image = open_image(image) image.load() return image + return image def is_allowed_extension(filename: str) -> bool: @@ -200,17 +204,16 @@ def format_images_markdown(images: Union[str, list], alt: str, preview: Union[st str: The formatted markdown string. """ if isinstance(images, str): - images = f"[ if preview else images})]({images})" + result = f"[ if preview else images})]({images})" else: if not isinstance(preview, list): preview = [preview.replace('{image}', image) if preview else image for image in images] - images = [ + result = "\n".join( f"[]({image})" for idx, image in enumerate(images) - ] - images = "\n".join(images) + ) start_flag = "<!-- generated images start -->\n" end_flag = "<!-- generated images end -->\n" - return f"\n{start_flag}{images}\n{end_flag}\n" + return f"\n{start_flag}{result}\n{end_flag}\n" def to_bytes(image: ImageType) -> bytes: """ @@ -245,19 +248,19 @@ class ImageResponse: self.images = images self.alt = alt self.options = options - + def __str__(self) -> str: return format_images_markdown(self.images, self.alt, self.get("preview")) - + def get(self, key: str): return self.options.get(key) - + class ImageRequest: def __init__( self, options: dict = {} ): self.options = options - + def get(self, key: str): return self.options.get(key)
\ No newline at end of file |