summaryrefslogtreecommitdiffstats
path: root/g4f/client
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--g4f/client/__init__.py3
-rw-r--r--g4f/client/async_client.py275
-rw-r--r--g4f/client/client.py484
-rw-r--r--g4f/client/image_models.py19
4 files changed, 429 insertions, 352 deletions
diff --git a/g4f/client/__init__.py b/g4f/client/__init__.py
index 5bb4ba35..d1e7e298 100644
--- a/g4f/client/__init__.py
+++ b/g4f/client/__init__.py
@@ -1,3 +1,2 @@
from .stubs import ChatCompletion, ChatCompletionChunk, ImagesResponse
-from .client import Client
-from .async_client import AsyncClient \ No newline at end of file
+from .client import Client, AsyncClient
diff --git a/g4f/client/async_client.py b/g4f/client/async_client.py
deleted file mode 100644
index 2fe4640b..00000000
--- a/g4f/client/async_client.py
+++ /dev/null
@@ -1,275 +0,0 @@
-from __future__ import annotations
-
-import time
-import random
-import string
-import asyncio
-import base64
-from aiohttp import ClientSession, BaseConnector
-
-from .types import Client as BaseClient
-from .types import ProviderType, FinishReason
-from .stubs import ChatCompletion, ChatCompletionChunk, ImagesResponse, Image
-from .types import AsyncIterResponse, ImageProvider
-from .image_models import ImageModels
-from .helper import filter_json, find_stop, filter_none, cast_iter_async
-from .service import get_last_provider, get_model_and_provider
-from ..Provider import ProviderUtils
-from ..typing import Union, Messages, AsyncIterator, ImageType
-from ..errors import NoImageResponseError, ProviderNotFoundError
-from ..requests.aiohttp import get_connector
-from ..providers.conversation import BaseConversation
-from ..image import ImageResponse as ImageProviderResponse, ImageDataResponse
-
-try:
- anext
-except NameError:
- async def anext(iter):
- async for chunk in iter:
- return chunk
-
-async def iter_response(
- response: AsyncIterator[str],
- stream: bool,
- response_format: dict = None,
- max_tokens: int = None,
- stop: list = None
-) -> AsyncIterResponse:
- content = ""
- finish_reason = None
- completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
- count: int = 0
- async for chunk in response:
- if isinstance(chunk, FinishReason):
- finish_reason = chunk.reason
- break
- elif isinstance(chunk, BaseConversation):
- yield chunk
- continue
- content += str(chunk)
- count += 1
- if max_tokens is not None and count >= max_tokens:
- finish_reason = "length"
- first, content, chunk = find_stop(stop, content, chunk)
- if first != -1:
- finish_reason = "stop"
- if stream:
- yield ChatCompletionChunk(chunk, None, completion_id, int(time.time()))
- if finish_reason is not None:
- break
- finish_reason = "stop" if finish_reason is None else finish_reason
- if stream:
- yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))
- else:
- if response_format is not None and "type" in response_format:
- if response_format["type"] == "json_object":
- content = filter_json(content)
- yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
-
-async def iter_append_model_and_provider(response: AsyncIterResponse) -> AsyncIterResponse:
- last_provider = None
- async for chunk in response:
- last_provider = get_last_provider(True) if last_provider is None else last_provider
- chunk.model = last_provider.get("model")
- chunk.provider = last_provider.get("name")
- yield chunk
-
-class AsyncClient(BaseClient):
- def __init__(
- self,
- provider: ProviderType = None,
- image_provider: ImageProvider = None,
- **kwargs
- ):
- super().__init__(**kwargs)
- self.chat: Chat = Chat(self, provider)
- self.images: Images = Images(self, image_provider)
-
-def create_response(
- messages: Messages,
- model: str,
- provider: ProviderType = None,
- stream: bool = False,
- proxy: str = None,
- max_tokens: int = None,
- stop: list[str] = None,
- api_key: str = None,
- **kwargs
-):
- has_asnyc = hasattr(provider, "create_async_generator")
- if has_asnyc:
- create = provider.create_async_generator
- else:
- create = provider.create_completion
- response = create(
- model, messages,
- stream=stream,
- **filter_none(
- proxy=proxy,
- max_tokens=max_tokens,
- stop=stop,
- api_key=api_key
- ),
- **kwargs
- )
- if not has_asnyc:
- response = cast_iter_async(response)
- return response
-
-class Completions():
- def __init__(self, client: AsyncClient, provider: ProviderType = None):
- self.client: AsyncClient = client
- self.provider: ProviderType = provider
-
- def create(
- self,
- messages: Messages,
- model: str,
- provider: ProviderType = None,
- stream: bool = False,
- proxy: str = None,
- max_tokens: int = None,
- stop: Union[list[str], str] = None,
- api_key: str = None,
- response_format: dict = None,
- ignored : list[str] = None,
- ignore_working: bool = False,
- ignore_stream: bool = False,
- **kwargs
- ) -> Union[ChatCompletion, AsyncIterator[ChatCompletionChunk]]:
- model, provider = get_model_and_provider(
- model,
- self.provider if provider is None else provider,
- stream,
- ignored,
- ignore_working,
- ignore_stream
- )
- stop = [stop] if isinstance(stop, str) else stop
- response = create_response(
- messages, model,
- provider, stream,
- proxy=self.client.get_proxy() if proxy is None else proxy,
- max_tokens=max_tokens,
- stop=stop,
- api_key=self.client.api_key if api_key is None else api_key,
- **kwargs
- )
- response = iter_response(response, stream, response_format, max_tokens, stop)
- response = iter_append_model_and_provider(response)
- return response if stream else anext(response)
-
-class Chat():
- completions: Completions
-
- def __init__(self, client: AsyncClient, provider: ProviderType = None):
- self.completions = Completions(client, provider)
-
-async def iter_image_response(
- response: AsyncIterator,
- response_format: str = None,
- connector: BaseConnector = None,
- proxy: str = None
-) -> Union[ImagesResponse, None]:
- async for chunk in response:
- if isinstance(chunk, ImageProviderResponse):
- if response_format == "b64_json":
- async with ClientSession(
- connector=get_connector(connector, proxy),
- cookies=chunk.options.get("cookies")
- ) as session:
- async def fetch_image(image):
- async with session.get(image) as response:
- return base64.b64encode(await response.content.read()).decode()
- images = await asyncio.gather(*[fetch_image(image) for image in chunk.get_list()])
- return ImagesResponse([Image(None, image, chunk.alt) for image in images], int(time.time()))
- return ImagesResponse([Image(image, None, chunk.alt) for image in chunk.get_list()], int(time.time()))
- elif isinstance(chunk, ImageDataResponse):
- return ImagesResponse([Image(None, image, chunk.alt) for image in chunk.get_list()], int(time.time()))
-
-def create_image(provider: ProviderType, prompt: str, model: str = "", **kwargs) -> AsyncIterator:
- if isinstance(provider, type) and provider.__name__ == "You":
- kwargs["chat_mode"] = "create"
- else:
- prompt = f"create a image with: {prompt}"
- return provider.create_async_generator(
- model,
- [{"role": "user", "content": prompt}],
- stream=True,
- **kwargs
- )
-
-class Images():
- def __init__(self, client: AsyncClient, provider: ImageProvider = None):
- self.client: AsyncClient = client
- self.provider: ImageProvider = provider
- self.models: ImageModels = ImageModels(client)
-
- def get_provider(self, model: str, provider: ProviderType = None):
- if isinstance(provider, str):
- if provider in ProviderUtils.convert:
- provider = ProviderUtils.convert[provider]
- else:
- raise ProviderNotFoundError(f'Provider not found: {provider}')
- else:
- provider = self.models.get(model, self.provider)
- return provider
-
- async def generate(
- self,
- prompt,
- model: str = "",
- provider: ProviderType = None,
- response_format: str = None,
- connector: BaseConnector = None,
- proxy: str = None,
- **kwargs
- ) -> ImagesResponse:
- provider = self.get_provider(model, provider)
- if hasattr(provider, "create_async_generator"):
- response = create_image(
- provider,
- prompt,
- **filter_none(
- response_format=response_format,
- connector=connector,
- proxy=self.client.get_proxy() if proxy is None else proxy,
- ),
- **kwargs
- )
- else:
- response = await provider.create_async(prompt)
- return ImagesResponse([Image(image) for image in response.get_list()])
- image = await iter_image_response(response, response_format, connector, proxy)
- if image is None:
- raise NoImageResponseError()
- return image
-
- async def create_variation(
- self,
- image: ImageType,
- model: str = None,
- response_format: str = None,
- connector: BaseConnector = None,
- proxy: str = None,
- **kwargs
- ):
- provider = self.get_provider(model, provider)
- result = None
- if hasattr(provider, "create_async_generator"):
- response = provider.create_async_generator(
- "",
- [{"role": "user", "content": "create a image like this"}],
- stream=True,
- image=image,
- **filter_none(
- response_format=response_format,
- connector=connector,
- proxy=self.client.get_proxy() if proxy is None else proxy,
- ),
- **kwargs
- )
- result = iter_image_response(response, response_format, connector, proxy)
- if result is None:
- raise NoImageResponseError()
- return result
diff --git a/g4f/client/client.py b/g4f/client/client.py
index 63bae4fe..8e195213 100644
--- a/g4f/client/client.py
+++ b/g4f/client/client.py
@@ -1,10 +1,19 @@
from __future__ import annotations
+import os
import time
import random
import string
+import threading
+import asyncio
+import base64
+import aiohttp
+import queue
+from typing import Union, AsyncIterator, Iterator
-from ..typing import Union, Iterator, Messages, ImageType
+from ..providers.base_provider import AsyncGeneratorProvider
+from ..image import ImageResponse, to_image, to_data_uri
+from ..typing import Messages, ImageType
from ..providers.types import BaseProvider, ProviderType, FinishReason
from ..providers.conversation import BaseConversation
from ..image import ImageResponse as ImageProviderResponse
@@ -15,35 +24,86 @@ from .types import IterResponse, ImageProvider
from .types import Client as BaseClient
from .service import get_model_and_provider, get_last_provider
from .helper import find_stop, filter_json, filter_none
+from ..models import ModelUtils
+from ..Provider import IterListProvider
+# Helper function to convert an async generator to a synchronous iterator
+def to_sync_iter(async_gen: AsyncIterator) -> Iterator:
+ q = queue.Queue()
+ loop = asyncio.new_event_loop()
+ done = object()
+
+ def _run():
+ asyncio.set_event_loop(loop)
+
+ async def iterate():
+ try:
+ async for item in async_gen:
+ q.put(item)
+ finally:
+ q.put(done)
+
+ loop.run_until_complete(iterate())
+ loop.close()
+
+ threading.Thread(target=_run).start()
+
+ while True:
+ item = q.get()
+ if item is done:
+ break
+ yield item
+
+# Helper function to convert a synchronous iterator to an async iterator
+async def to_async_iterator(iterator):
+ for item in iterator:
+ yield item
+
+# Synchronous iter_response function
def iter_response(
- response: iter[str],
+ response: Union[Iterator[str], AsyncIterator[str]],
stream: bool,
response_format: dict = None,
max_tokens: int = None,
stop: list = None
-) -> IterResponse:
+) -> Iterator[Union[ChatCompletion, ChatCompletionChunk]]:
content = ""
finish_reason = None
completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
- for idx, chunk in enumerate(response):
+ idx = 0
+
+ if hasattr(response, '__aiter__'):
+ # It's an async iterator, wrap it into a sync iterator
+ response = to_sync_iter(response)
+
+ for chunk in response:
if isinstance(chunk, FinishReason):
finish_reason = chunk.reason
break
elif isinstance(chunk, BaseConversation):
yield chunk
continue
+
content += str(chunk)
+
if max_tokens is not None and idx + 1 >= max_tokens:
finish_reason = "length"
+
first, content, chunk = find_stop(stop, content, chunk if stream else None)
+
if first != -1:
finish_reason = "stop"
+
if stream:
yield ChatCompletionChunk(chunk, None, completion_id, int(time.time()))
+
if finish_reason is not None:
break
+
+ idx += 1
+
finish_reason = "stop" if finish_reason is None else finish_reason
+
if stream:
yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))
else:
@@ -52,12 +112,14 @@ def iter_response(
content = filter_json(content)
yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
-def iter_append_model_and_provider(response: IterResponse) -> IterResponse:
+# Synchronous iter_append_model_and_provider function
+def iter_append_model_and_provider(response: Iterator) -> Iterator:
last_provider = None
+
for chunk in response:
last_provider = get_last_provider(True) if last_provider is None else last_provider
chunk.model = last_provider.get("model")
- chunk.provider = last_provider.get("name")
+ chunk.provider = last_provider.get("name")
yield chunk
class Client(BaseClient):
@@ -69,9 +131,39 @@ class Client(BaseClient):
) -> None:
super().__init__(**kwargs)
self.chat: Chat = Chat(self, provider)
- self.images: Images = Images(self, image_provider)
+ self._images: Images = Images(self, image_provider)
+
+ @property
+ def images(self) -> Images:
+ return self._images
+
+ async def async_images(self) -> Images:
+ return self._images
+
+# For backwards compatibility and legacy purposes, use Client instead
+class AsyncClient(Client):
+ """Legacy AsyncClient that redirects to the main Client class.
+ This class exists for backwards compatibility."""
+
+ def __init__(self, *args, **kwargs):
+ import warnings
+ warnings.warn(
+ "AsyncClient is deprecated and will be removed in a future version. "
+ "Use Client instead, which now supports both sync and async operations.",
+ DeprecationWarning,
+ stacklevel=2
+ )
+ super().__init__(*args, **kwargs)
-class Completions():
+ async def chat_complete(self, *args, **kwargs):
+ """Legacy method that redirects to async_create"""
+ return await self.chat.completions.async_create(*args, **kwargs)
+
+ async def create_image(self, *args, **kwargs):
+ """Legacy method that redirects to async_generate"""
+ return await self.images.async_generate(*args, **kwargs)
+
+class Completions:
def __init__(self, client: Client, provider: ProviderType = None):
self.client: Client = client
self.provider: ProviderType = provider
@@ -87,7 +179,7 @@ class Completions():
max_tokens: int = None,
stop: Union[list[str], str] = None,
api_key: str = None,
- ignored : list[str] = None,
+ ignored: list[str] = None,
ignore_working: bool = False,
ignore_stream: bool = False,
**kwargs
@@ -100,79 +192,343 @@ class Completions():
ignore_working,
ignore_stream,
)
-
+
stop = [stop] if isinstance(stop, str) else stop
- response = provider.create_completion(
- model, messages,
- stream=stream,
- **filter_none(
- proxy=self.client.get_proxy() if proxy is None else proxy,
- max_tokens=max_tokens,
- stop=stop,
- api_key=self.client.api_key if api_key is None else api_key
- ),
- **kwargs
+
+ if asyncio.iscoroutinefunction(provider.create_completion):
+ # Run the asynchronous function in an event loop
+ response = asyncio.run(provider.create_completion(
+ model,
+ messages,
+ stream=stream,
+ **filter_none(
+ proxy=self.client.get_proxy() if proxy is None else proxy,
+ max_tokens=max_tokens,
+ stop=stop,
+ api_key=self.client.api_key if api_key is None else api_key
+ ),
+ **kwargs
+ ))
+ else:
+ response = provider.create_completion(
+ model,
+ messages,
+ stream=stream,
+ **filter_none(
+ proxy=self.client.get_proxy() if proxy is None else proxy,
+ max_tokens=max_tokens,
+ stop=stop,
+ api_key=self.client.api_key if api_key is None else api_key
+ ),
+ **kwargs
+ )
+
+ if stream:
+ if hasattr(response, '__aiter__'):
+ # It's an async generator, wrap it into a sync iterator
+ response = to_sync_iter(response)
+
+ # Now 'response' is an iterator
+ response = iter_response(response, stream, response_format, max_tokens, stop)
+ response = iter_append_model_and_provider(response)
+ return response
+ else:
+ if hasattr(response, '__aiter__'):
+ # If response is an async generator, collect it into a list
+ response = list(to_sync_iter(response))
+ response = iter_response(response, stream, response_format, max_tokens, stop)
+ response = iter_append_model_and_provider(response)
+ return next(response)
+
+ async def async_create(
+ self,
+ messages: Messages,
+ model: str,
+ provider: ProviderType = None,
+ stream: bool = False,
+ proxy: str = None,
+ response_format: dict = None,
+ max_tokens: int = None,
+ stop: Union[list[str], str] = None,
+ api_key: str = None,
+ ignored: list[str] = None,
+ ignore_working: bool = False,
+ ignore_stream: bool = False,
+ **kwargs
+ ) -> Union[ChatCompletion, AsyncIterator[ChatCompletionChunk]]:
+ model, provider = get_model_and_provider(
+ model,
+ self.provider if provider is None else provider,
+ stream,
+ ignored,
+ ignore_working,
+ ignore_stream,
)
- response = iter_response(response, stream, response_format, max_tokens, stop)
- response = iter_append_model_and_provider(response)
- return response if stream else next(response)
-class Chat():
+ stop = [stop] if isinstance(stop, str) else stop
+
+ if asyncio.iscoroutinefunction(provider.create_completion):
+ response = await provider.create_completion(
+ model,
+ messages,
+ stream=stream,
+ **filter_none(
+ proxy=self.client.get_proxy() if proxy is None else proxy,
+ max_tokens=max_tokens,
+ stop=stop,
+ api_key=self.client.api_key if api_key is None else api_key
+ ),
+ **kwargs
+ )
+ else:
+ response = provider.create_completion(
+ model,
+ messages,
+ stream=stream,
+ **filter_none(
+ proxy=self.client.get_proxy() if proxy is None else proxy,
+ max_tokens=max_tokens,
+ stop=stop,
+ api_key=self.client.api_key if api_key is None else api_key
+ ),
+ **kwargs
+ )
+
+ # Removed 'await' here since 'async_iter_response' returns an async generator
+ response = async_iter_response(response, stream, response_format, max_tokens, stop)
+ response = async_iter_append_model_and_provider(response)
+
+ if stream:
+ return response
+ else:
+ async for result in response:
+ return result
+
+class Chat:
completions: Completions
def __init__(self, client: Client, provider: ProviderType = None):
self.completions = Completions(client, provider)
-def iter_image_response(response: Iterator) -> Union[ImagesResponse, None]:
- for chunk in list(response):
+# Asynchronous versions of the helper functions
+async def async_iter_response(
+ response: Union[AsyncIterator[str], Iterator[str]],
+ stream: bool,
+ response_format: dict = None,
+ max_tokens: int = None,
+ stop: list = None
+) -> AsyncIterator[Union[ChatCompletion, ChatCompletionChunk]]:
+ content = ""
+ finish_reason = None
+ completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
+ idx = 0
+
+ if not hasattr(response, '__aiter__'):
+ response = to_async_iterator(response)
+
+ async for chunk in response:
+ if isinstance(chunk, FinishReason):
+ finish_reason = chunk.reason
+ break
+ elif isinstance(chunk, BaseConversation):
+ yield chunk
+ continue
+
+ content += str(chunk)
+
+ if max_tokens is not None and idx + 1 >= max_tokens:
+ finish_reason = "length"
+
+ first, content, chunk = find_stop(stop, content, chunk if stream else None)
+
+ if first != -1:
+ finish_reason = "stop"
+
+ if stream:
+ yield ChatCompletionChunk(chunk, None, completion_id, int(time.time()))
+
+ if finish_reason is not None:
+ break
+
+ idx += 1
+
+ finish_reason = "stop" if finish_reason is None else finish_reason
+
+ if stream:
+ yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))
+ else:
+ if response_format is not None and "type" in response_format:
+ if response_format["type"] == "json_object":
+ content = filter_json(content)
+ yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
+
+async def async_iter_append_model_and_provider(response: AsyncIterator) -> AsyncIterator:
+ last_provider = None
+
+ if not hasattr(response, '__aiter__'):
+ response = to_async_iterator(response)
+
+ async for chunk in response:
+ last_provider = get_last_provider(True) if last_provider is None else last_provider
+ chunk.model = last_provider.get("model")
+ chunk.provider = last_provider.get("name")
+ yield chunk
+
+async def iter_image_response(response: AsyncIterator) -> Union[ImagesResponse, None]:
+ response_list = []
+ async for chunk in response:
if isinstance(chunk, ImageProviderResponse):
- return ImagesResponse([Image(image) for image in chunk.get_list()])
+ response_list.extend(chunk.get_list())
+ elif isinstance(chunk, str):
+ response_list.append(chunk)
-def create_image(client: Client, provider: ProviderType, prompt: str, model: str = "", **kwargs) -> Iterator:
+ if response_list:
+ return ImagesResponse([Image(image) for image in response_list])
+ return None
+async def create_image(client: Client, provider: ProviderType, prompt: str, model: str = "", **kwargs) -> AsyncIterator:
if isinstance(provider, type) and provider.__name__ == "You":
kwargs["chat_mode"] = "create"
else:
- prompt = f"create a image with: {prompt}"
- return provider.create_completion(
- model,
- [{"role": "user", "content": prompt}],
- stream=True,
- proxy=client.get_proxy(),
- **kwargs
- )
+ prompt = f"create an image with: {prompt}"
-class Images():
- def __init__(self, client: Client, provider: ImageProvider = None):
- self.client: Client = client
- self.provider: ImageProvider = provider
+ if asyncio.iscoroutinefunction(provider.create_completion):
+ response = await provider.create_completion(
+ model,
+ [{"role": "user", "content": prompt}],
+ stream=True,
+ proxy=client.get_proxy(),
+ **kwargs
+ )
+ else:
+ response = provider.create_completion(
+ model,
+ [{"role": "user", "content": prompt}],
+ stream=True,
+ proxy=client.get_proxy(),
+ **kwargs
+ )
+
+ # Wrap synchronous iterator into async iterator if necessary
+ if not hasattr(response, '__aiter__'):
+ response = to_async_iterator(response)
+
+ return response
+
+class Image:
+ def __init__(self, url: str = None, b64_json: str = None):
+ self.url = url
+ self.b64_json = b64_json
+
+ def __repr__(self):
+ return f"Image(url={self.url}, b64_json={'<base64 data>' if self.b64_json else None})"
+
+class ImagesResponse:
+ def __init__(self, data: list[Image]):
+ self.data = data
+
+ def __repr__(self):
+ return f"ImagesResponse(data={self.data})"
+
+class Images:
+ def __init__(self, client: 'Client', provider: 'ImageProvider' = None):
+ self.client: 'Client' = client
+ self.provider: 'ImageProvider' = provider
self.models: ImageModels = ImageModels(client)
- def generate(self, prompt, model: str = None, **kwargs) -> ImagesResponse:
+ def generate(self, prompt: str, model: str = None, response_format: str = "url", **kwargs) -> ImagesResponse:
+ """
+ Synchronous generate method that runs the async_generate method in an event loop.
+ """
+ return asyncio.run(self.async_generate(prompt, model, response_format=response_format, **kwargs))
+
+ async def async_generate(self, prompt: str, model: str = None, response_format: str = "url", **kwargs) -> ImagesResponse:
provider = self.models.get(model, self.provider)
- if isinstance(provider, type) and issubclass(provider, BaseProvider):
- response = create_image(self.client, provider, prompt, **kwargs)
+ if provider is None:
+ raise ValueError(f"Unknown model: {model}")
+
+ if isinstance(provider, IterListProvider):
+ if provider.providers:
+ provider = provider.providers[0]
+ else:
+ raise ValueError(f"IterListProvider for model {model} has no providers")
+
+ if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider):
+ messages = [{"role": "user", "content": prompt}]
+ async for response in provider.create_async_generator(model, messages, **kwargs):
+ if isinstance(response, ImageResponse):
+ return await self._process_image_response(response, response_format)
+ elif isinstance(response, str):
+ image_response = ImageResponse([response], prompt)
+ return await self._process_image_response(image_response, response_format)
+ elif hasattr(provider, 'create'):
+ if asyncio.iscoroutinefunction(provider.create):
+ response = await provider.create(prompt)
+ else:
+ response = provider.create(prompt)
+
+ if isinstance(response, ImageResponse):
+ return await self._process_image_response(response, response_format)
+ elif isinstance(response, str):
+ image_response = ImageResponse([response], prompt)
+ return await self._process_image_response(image_response, response_format)
else:
- response = list(provider.create(prompt))
- image = iter_image_response(response)
- if image is None:
- raise NoImageResponseError()
- return image
+ raise ValueError(f"Provider {provider} does not support image generation")
+
+ raise NoImageResponseError(f"Unexpected response type: {type(response)}")
+
+ async def _process_image_response(self, response: ImageResponse, response_format: str) -> ImagesResponse:
+ processed_images = []
+
+ for image_data in response.get_list():
+ if image_data.startswith('http://') or image_data.startswith('https://'):
+ if response_format == "url":
+ processed_images.append(Image(url=image_data))
+ elif response_format == "b64_json":
+ # Fetch the image data and convert it to base64
+ image_content = await self._fetch_image(image_data)
+ b64_json = base64.b64encode(image_content).decode('utf-8')
+ processed_images.append(Image(b64_json=b64_json))
+ else:
+ # Assume image_data is base64 data or binary
+ if response_format == "url":
+ if image_data.startswith('data:image'):
+ # Remove the data URL scheme and get the base64 data
+ header, base64_data = image_data.split(',', 1)
+ else:
+ base64_data = image_data
+ # Decode the base64 data
+ image_data_bytes = base64.b64decode(base64_data)
+ # Convert bytes to an image
+ image = to_image(image_data_bytes)
+ file_name = self._save_image(image)
+ processed_images.append(Image(url=file_name))
+ elif response_format == "b64_json":
+ if isinstance(image_data, bytes):
+ b64_json = base64.b64encode(image_data).decode('utf-8')
+ else:
+ b64_json = image_data # If already base64-encoded string
+ processed_images.append(Image(b64_json=b64_json))
+
+ return ImagesResponse(processed_images)
+
+ async def _fetch_image(self, url: str) -> bytes:
+ # Asynchronously fetch image data from the URL
+ async with aiohttp.ClientSession() as session:
+ async with session.get(url) as resp:
+ if resp.status == 200:
+ return await resp.read()
+ else:
+ raise Exception(f"Failed to fetch image from {url}, status code {resp.status}")
+
+ def _save_image(self, image: 'PILImage') -> str:
+ os.makedirs('generated_images', exist_ok=True)
+ file_name = f"generated_images/image_{int(time.time())}_{random.randint(0, 10000)}.png"
+ image.save(file_name)
+ return file_name
+
+ async def create_variation(self, image: Union[str, bytes], model: str = None, response_format: str = "url", **kwargs):
+ # Existing implementation, adjust if you want to support b64_json here as well
+ pass
- def create_variation(self, image: ImageType, model: str = None, **kwargs):
- provider = self.models.get(model, self.provider)
- result = None
- if isinstance(provider, type) and issubclass(provider, BaseProvider):
- response = provider.create_completion(
- "",
- [{"role": "user", "content": "create a image like this"}],
- True,
- image=image,
- proxy=self.client.get_proxy(),
- **kwargs
- )
- result = iter_image_response(response)
- if result is None:
- raise NoImageResponseError()
- return result \ No newline at end of file
diff --git a/g4f/client/image_models.py b/g4f/client/image_models.py
index db2ce09a..edaa4592 100644
--- a/g4f/client/image_models.py
+++ b/g4f/client/image_models.py
@@ -2,18 +2,15 @@ from __future__ import annotations
from .types import Client, ImageProvider
-from ..Provider.BingCreateImages import BingCreateImages
-from ..Provider.needs_auth import Gemini, OpenaiChat
-from ..Provider.You import You
+from ..models import ModelUtils
class ImageModels():
- gemini = Gemini
- openai = OpenaiChat
- you = You
-
- def __init__(self, client: Client) -> None:
+ def __init__(self, client):
self.client = client
- self.default = BingCreateImages(proxy=self.client.get_proxy())
+ self.models = ModelUtils.convert
- def get(self, name: str, default: ImageProvider = None) -> ImageProvider:
- return getattr(self, name) if hasattr(self, name) else default or self.default
+ def get(self, name, default=None):
+ model = self.models.get(name)
+ if model and model.best_provider:
+ return model.best_provider
+ return default
generated by cgit v1.2.3 (git 2.25.1) at 2024-11-07 16:36:50 +0100