diff options
Diffstat (limited to 'g4f/Provider/bing/create_images.py')
| -rw-r--r-- | g4f/Provider/bing/create_images.py | 20 |
1 files changed, 13 insertions, 7 deletions
diff --git a/g4f/Provider/bing/create_images.py b/g4f/Provider/bing/create_images.py index c868a6f0..7c9a8815 100644 --- a/g4f/Provider/bing/create_images.py +++ b/g4f/Provider/bing/create_images.py @@ -1,5 +1,3 @@ - - import asyncio import time, json, os from aiohttp import ClientSession @@ -9,10 +7,12 @@ from typing import Generator from ...webdriver import WebDriver, get_driver_cookies, get_browser from ...Provider.helper import get_event_loop +from ...base_provider import ProviderType +from ...Provider.create_images import CreateImagesProvider BING_URL = "https://www.bing.com" -def wait_for_login(driver: WebDriver, timeout: int = 1200): +def wait_for_login(driver: WebDriver, timeout: int = 1200) -> Generator: driver.get(f"{BING_URL}/") value = driver.get_cookie("_U") if value: @@ -29,7 +29,7 @@ def wait_for_login(driver: WebDriver, timeout: int = 1200): return time.sleep(0.1) -def create_session(cookies: dict): +def create_session(cookies: dict) -> ClientSession: headers = { "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7", "accept-encoding": "gzip, deflate, br", @@ -51,7 +51,7 @@ def create_session(cookies: dict): headers["cookie"] = "; ".join(f"{k}={v}" for k, v in cookies.items()) return ClientSession(headers=headers) -async def create_images(session: ClientSession, prompt: str, proxy: str = None, timeout: int = 200): +async def create_images(session: ClientSession, prompt: str, proxy: str = None, timeout: int = 200) -> list: url_encoded_prompt = quote(prompt) payload = f"q={url_encoded_prompt}&rt=4&FORM=GENCRE" url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=4&FORM=GENCRE" @@ -111,7 +111,10 @@ async def create_images(session: ClientSession, prompt: str, proxy: str = None, def format_images_markdown(images: list, prompt: str) -> str: images = [f"[]({image})" for idx, image in enumerate(images)] - return f"\n\n<img data-prompt=\"{prompt}\">\n<!-- generated images start -->\n" + ("\n".join(images)) + "\n<!-- generated images end -->\n\n" + images = "\n".join(images) + start_flag = "<!-- generated images start -->\n" + end_flag = "<!-- generated images end -->\n" + return f"\n\n<img data-prompt=\"{prompt}\">\n{start_flag}{images}\n{end_flag}\n" def get_images(text: str) -> list: html_soup = BeautifulSoup(text, "html.parser") @@ -143,4 +146,7 @@ def create_completion(prompt: str, proxy: str = None) -> Generator: images = loop.run_until_complete(run_session()) yield format_images_markdown(images, prompt) finally: - driver.quit()
\ No newline at end of file + driver.quit() + +def patch_provider(provider: ProviderType) -> CreateImagesProvider: + return CreateImagesProvider(provider, create_completion)
\ No newline at end of file |