0.3.2.6

1 files changed, 28 insertions, 13 deletions

diff --git a/g4f/gui/server/api.py b/g4f/gui/server/api.py
index 3da0fe17..c984abec 100644
--- a/g4f/gui/server/api.py
+++ b/g4f/gui/server/api.py
@@ -6,6 +6,7 @@ import os.path
 import uuid
 import asyncio
 import time
+import base64
 from aiohttp import ClientSession
 from typing import Iterator, Optional
 from flask import send_from_directory
@@ -195,18 +196,32 @@ class Api():
                             cookies=cookies
                         ) as session:
                             async def copy_image(image):
-                                async with session.get(image) as response:
-                                    target = os.path.join(images_dir, f"{int(time.time())}_{str(uuid.uuid4())}")
-                                    with open(target, "wb") as f:
-                                        async for chunk in response.content.iter_any():
-                                            f.write(chunk)
-                                    with open(target, "rb") as f:
-                                        extension = is_accepted_format(f.read(12)).split("/")[-1]
-                                        extension = "jpg" if extension == "jpeg" else extension
-                                    new_target = f"{target}.{extension}"
-                                    os.rename(target, new_target)
-                                    return f"/images/{os.path.basename(new_target)}"
-                            return await asyncio.gather(*[copy_image(image) for image in images])                                
+                                if image.startswith("data:"):
+                                    # Processing the data URL
+                                    data_uri_parts = image.split(",")
+                                    if len(data_uri_parts) == 2:
+                                        content_type, base64_data = data_uri_parts
+                                        extension = content_type.split("/")[-1].split(";")[0]
+                                        target = os.path.join(images_dir, f"{int(time.time())}_{str(uuid.uuid4())}.{extension}")
+                                        with open(target, "wb") as f:
+                                            f.write(base64.b64decode(base64_data))
+                                        return f"/images/{os.path.basename(target)}"
+                                    else:
+                                        return None
+                                else:
+                                    # Обробка звичайної URL-адреси
+                                    async with session.get(image) as response:
+                                        target = os.path.join(images_dir, f"{int(time.time())}_{str(uuid.uuid4())}")
+                                        with open(target, "wb") as f:
+                                            async for chunk in response.content.iter_any():
+                                                f.write(chunk)
+                                        with open(target, "rb") as f:
+                                            extension = is_accepted_format(f.read(12)).split("/")[-1]
+                                            extension = "jpg" if extension == "jpeg" else extension
+                                        new_target = f"{target}.{extension}"
+                                        os.rename(target, new_target)
+                                        return f"/images/{os.path.basename(new_target)}"
+                            return await asyncio.gather(*[copy_image(image) for image in images])
                     images = asyncio.run(copy_images(chunk.get_list(), chunk.options.get("cookies")))
                     yield self._format_json("content", str(ImageResponse(images, chunk.alt)))
                 elif not isinstance(chunk, FinishReason):
@@ -245,4 +260,4 @@ def get_error_message(exception: Exception) -> str:
     provider = get_last_provider()
     if provider is None:
         return message
-    return f"{provider.__name__}: {message}"
\ No newline at end of file
+    return f"{provider.__name__}: {message}"