Repository: Integration-Automation/ReEdgeGPT (language: python)

Issue:
Image generation
It was in original lib. Not supported here?


Discussion / hints:
No original image generation is broken too and i don't want to fix it.
If original image generation only have ip problem i will try it.
I will add image gernation feature next PR.

Fix patch:
diff --git a/.github/workflows/re_edge_gpt_dev_python3_10.yml b/.github/workflows/re_edge_gpt_dev_python3_10.yml
index c110881..2689524 100644
--- a/.github/workflows/re_edge_gpt_dev_python3_10.yml
+++ b/.github/workflows/re_edge_gpt_dev_python3_10.yml
@@ -28,4 +28,8 @@ jobs:
     - name: Test Bot
       env:
         EDGE_COOKIES: ${{ secrets.EDGE_COOKIES }}
-      run: python ./test/unit_test/test_bot.py
\ No newline at end of file
+      run: python ./test/unit_test/test_bot/test_bot.py
+    - name: Test Generation Image
+      env:
+        IMAGE_COOKIE: ${{ secrets.IMAGE_COOKIE }}
+      run: python ./test/unit_test/test_generate_image/test_generation.py
\ No newline at end of file
diff --git a/.github/workflows/re_edge_gpt_dev_python3_11.yml b/.github/workflows/re_edge_gpt_dev_python3_11.yml
index 234ef62..1107220 100644
--- a/.github/workflows/re_edge_gpt_dev_python3_11.yml
+++ b/.github/workflows/re_edge_gpt_dev_python3_11.yml
@@ -28,4 +28,8 @@ jobs:
     - name: Test Bot
       env:
         EDGE_COOKIES: ${{ secrets.EDGE_COOKIES }}
-      run: python ./test/unit_test/test_bot.py
\ No newline at end of file
+      run: python ./test/unit_test/test_bot/test_bot.py
+    - name: Test Generation Image
+      env:
+        IMAGE_COOKIE: ${{ secrets.IMAGE_COOKIE }}
+      run: python ./test/unit_test/test_generate_image/test_generation.py
\ No newline at end of file
diff --git a/.github/workflows/re_edge_gpt_dev_python3_8.yml b/.github/workflows/re_edge_gpt_dev_python3_8.yml
deleted file mode 100644
index 6ecb130..0000000
--- a/.github/workflows/re_edge_gpt_dev_python3_8.yml
+++ /dev/null
@@ -1,31 +0,0 @@
-name: ReEdgeGPT Dev Python3.8
-
-on:
-  push:
-    branches: [ "dev" ]
-  pull_request:
-    branches: [ "dev" ]
-  schedule:
-    - cron: "0 8 * * *"
-
-permissions:
-  contents: read
-
-jobs:
-  build_dev_version:
-    runs-on: windows-2019
-
-    steps:
-    - uses: actions/checkout@v3
-    - name: Set up Python 3.8
-      uses: actions/setup-python@v3
-      with:
-        python-version: "3.8"
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip wheel
-        pip install -r dev_requirements.txt
-    - name: Test Bot
-      env:
-        EDGE_COOKIES: ${{ secrets.EDGE_COOKIES }}
-      run: python ./test/unit_test/test_bot.py
diff --git a/.github/workflows/re_edge_gpt_dev_python3_9.yml b/.github/workflows/re_edge_gpt_dev_python3_9.yml
index 4dafcbe..e3fbb99 100644
--- a/.github/workflows/re_edge_gpt_dev_python3_9.yml
+++ b/.github/workflows/re_edge_gpt_dev_python3_9.yml
@@ -28,4 +28,8 @@ jobs:
     - name: Test Bot
       env:
         EDGE_COOKIES: ${{ secrets.EDGE_COOKIES }}
-      run: python ./test/unit_test/test_bot.py
\ No newline at end of file
+      run: python ./test/unit_test/test_bot/test_bot.py
+    - name: Test Generation Image
+      env:
+        IMAGE_COOKIE: ${{ secrets.IMAGE_COOKIE }}
+      run: python ./test/unit_test/test_generate_image/test_generation.py
\ No newline at end of file
diff --git a/.github/workflows/re_edge_gpt_stable_python3_10.yml b/.github/workflows/re_edge_gpt_stable_python3_10.yml
index 7d674fb..3fada23 100644
--- a/.github/workflows/re_edge_gpt_stable_python3_10.yml
+++ b/.github/workflows/re_edge_gpt_stable_python3_10.yml
@@ -28,4 +28,8 @@ jobs:
     - name: Test Bot
       env:
         EDGE_COOKIES: ${{ secrets.EDGE_COOKIES }}
-      run: python ./test/unit_test/test_bot.py
\ No newline at end of file
+      run: python ./test/unit_test/test_bot/test_bot.py
+    - name: Test Generation Image
+      env:
+        IMAGE_COOKIE: ${{ secrets.IMAGE_COOKIE }}
+      run: python ./test/unit_test/test_generate_image/test_generation.py
\ No newline at end of file
diff --git a/.github/workflows/re_edge_gpt_stable_python3_11.yml b/.github/workflows/re_edge_gpt_stable_python3_11.yml
index fb859ee..a957bdb 100644
--- a/.github/workflows/re_edge_gpt_stable_python3_11.yml
+++ b/.github/workflows/re_edge_gpt_stable_python3_11.yml
@@ -28,4 +28,8 @@ jobs:
     - name: Test Bot
       env:
         EDGE_COOKIES: ${{ secrets.EDGE_COOKIES }}
-      run: python ./test/unit_test/test_bot.py
\ No newline at end of file
+      run: python ./test/unit_test/test_bot/test_bot.py
+    - name: Test Generation Image
+      env:
+        IMAGE_COOKIE: ${{ secrets.IMAGE_COOKIE }}
+      run: python ./test/unit_test/test_generate_image/test_generation.py
\ No newline at end of file
diff --git a/.github/workflows/re_edge_gpt_stable_python3_8.yml b/.github/workflows/re_edge_gpt_stable_python3_8.yml
deleted file mode 100644
index ce59ef7..0000000
--- a/.github/workflows/re_edge_gpt_stable_python3_8.yml
+++ /dev/null
@@ -1,31 +0,0 @@
-name: ReEdgeGPT Stable Python3.8
-
-on:
-  push:
-    branches: [ "main" ]
-  pull_request:
-    branches: [ "main" ]
-  schedule:
-    - cron: "0 8 * * *"
-
-permissions:
-  contents: read
-
-jobs:
-  build_dev_version:
-    runs-on: windows-2019
-
-    steps:
-    - uses: actions/checkout@v3
-    - name: Set up Python 3.8
-      uses: actions/setup-python@v3
-      with:
-        python-version: "3.8"
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip wheel
-        pip install -r requirements.txt
-    - name: Test Bot
-      env:
-        EDGE_COOKIES: ${{ secrets.EDGE_COOKIES }}
-      run: python ./test/unit_test/test_bot.py
\ No newline at end of file
diff --git a/.github/workflows/re_edge_gpt_stable_python3_9.yml b/.github/workflows/re_edge_gpt_stable_python3_9.yml
index e468c68..15f3b27 100644
--- a/.github/workflows/re_edge_gpt_stable_python3_9.yml
+++ b/.github/workflows/re_edge_gpt_stable_python3_9.yml
@@ -28,4 +28,8 @@ jobs:
     - name: Test Bot
       env:
         EDGE_COOKIES: ${{ secrets.EDGE_COOKIES }}
-      run: python ./test/unit_test/test_bot.py
\ No newline at end of file
+      run: python ./test/unit_test/test_bot/test_bot.py
+    - name: Test Generation Image
+      env:
+        IMAGE_COOKIE: ${{ secrets.IMAGE_COOKIE }}
+      run: python ./test/unit_test/test_generate_image/test_generation.py
\ No newline at end of file
diff --git a/README.md b/README.md
index 3000912..eeecdbc 100644
--- a/README.md
+++ b/README.md
@@ -18,7 +18,7 @@ python3 -m pip install re_edge_gpt --upgrade
 
 ## Requirements
 
-- python 3.8+
+- python 3.9+
 - A Microsoft Account with access to <https://bing.com/chat> (Optional, depending on your region)
 - Required in a supported country or region with New Bing (Chinese mainland VPN required)
 
@@ -55,8 +55,12 @@ If you receive the following error, you can try **providing a cookie** and see i
 import json
 from re_edge_gpt import Chatbot
 
-cookies = json.loads(open("./path/to/bing_cookies.json", encoding="utf-8").read())
-bot = await Chatbot.create(cookies=cookies)
+
+async def create_bot():
+    cookies = json.loads(open("./path/to/bing_cookies.json", encoding="utf-8").read())
+    bot = await Chatbot.create(cookies=cookies)
+    return bot
+
 ```
 
 </details>
@@ -134,11 +138,84 @@ async def main():
     await bot.close()
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(main())
 ```
 
 </details>
 
+<summary>
+
+# How to generate image
+
+## Getting authentication
+> ### Chromium based browsers (Edge, Opera, Vivaldi, Brave)
+> * Go to https://bing.com/.
+> * F12 to open console
+> * In the JavaScript console, type cookieStore.get("_U").then(result => console.log(result.value)) and press enter
+> * Copy the output. This is used in --U or auth_cookie.
+
+> ### Firefox
+> * Go to https://bing.com/.
+> * F12 to open developer tools
+> * navigate to the storage tab
+> * expand the cookies tab
+> * click on the https://bing.com cookie
+> * copy the value from the _U
+
+```python
+import asyncio
+import os
+import shutil
+from pathlib import Path
+
+from re_edge_gpt import ImageGen, ImageGenAsync
+
+# create a temporary output directory for testing purposes
+test_output_dir = "test_output"
+# download a test image
+test_image_url = "https://picsum.photos/200"
+auth_cooker = open("bing_cookies.txt", "r+").read()
+sync_gen = ImageGen(auth_cookie=auth_cooker)
+async_gen = ImageGenAsync(auth_cookie=auth_cooker)
+
+
+def test_save_images_sync():
+    sync_gen.save_images([test_image_url], test_output_dir)
+    sync_gen.save_images([test_image_url], test_output_dir, file_name="test_image")
+    # check if the image was downloaded and saved correctly
+    assert os.path.exists(os.path.join(test_output_dir, "test_image_0.jpeg"))
+    assert os.path.exists(os.path.join(test_output_dir, "0.jpeg"))
+
+
+# Generate image list sync
+def test_generate_image_sync():
+    image_list = sync_gen.get_images("tree")
+    print(image_list)
+
+
+# Generate image list async
+async def test_generate_image_async():
+    image_list = await async_gen.get_images("tree")
+    print(image_list)
+
+
+if __name__ == "__main__":
+    # Make dir to save image
+    Path("test_output").mkdir(exist_ok=True)
+    # Save image
+    test_save_images_sync()
+    # Generate image sync
+    test_generate_image_sync()
+    # Generate image async
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(test_generate_image_async())
+    # Remove dir
+    shutil.rmtree(test_output_dir)
+```
+
+</summary>
+
 # Q&A
 
 <summary>
diff --git a/stable.toml b/dev.toml
similarity index 94%
rename from stable.toml
rename to dev.toml
index 1cbcbed..1a29e83 100644
--- a/stable.toml
+++ b/dev.toml
@@ -5,13 +5,13 @@ requires = ["setuptools>=61.0"]
 build-backend = "setuptools.build_meta"
 
 [project]
-name = "re_edge_gpt"
-version = "0.0.04"
+name = "re_edge_gpt_dev"
+version = "0.0.06"
 authors = [
     { name = "JE-Chen", email = "jechenmailman@gmail.com" },
 ]
 dependencies = [
-    "aiohttp", "certifi", "httpx", "prompt_toolkit", "requests", "rich"
+    "aiohttp", "certifi", "httpx", "prompt_toolkit", "requests", "rich", "regex"
 ]
 description = "Microsoft's Bing Chat AI"
 requires-python = ">=3.9"
diff --git a/dev_requirements.txt b/dev_requirements.txt
index 3ea7c3d..36fbbfa 100644
--- a/dev_requirements.txt
+++ b/dev_requirements.txt
@@ -4,4 +4,4 @@ httpx
 prompt_toolkit
 requests
 rich
-re_edge_gpt
+re_edge_gpt_dev
diff --git a/pyproject.toml b/pyproject.toml
index a4bcbc2..f68492c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,13 +5,13 @@ requires = ["setuptools>=61.0"]
 build-backend = "setuptools.build_meta"
 
 [project]
-name = "re_edge_gpt_dev"
-version = "0.0.04"
+name = "re_edge_gpt"
+version = "0.0.05"
 authors = [
     { name = "JE-Chen", email = "jechenmailman@gmail.com" },
 ]
 dependencies = [
-    "aiohttp", "certifi", "httpx", "prompt_toolkit", "requests", "rich"
+    "aiohttp", "certifi", "httpx", "prompt_toolkit", "requests", "rich", "regex"
 ]
 description = "Microsoft's Bing Chat AI"
 requires-python = ">=3.9"
diff --git a/re_edge_gpt/__init__.py b/re_edge_gpt/__init__.py
index e6b4d53..2dfdf59 100644
--- a/re_edge_gpt/__init__.py
+++ b/re_edge_gpt/__init__.py
@@ -1,6 +1,8 @@
 from re_edge_gpt.re_edge_gpt import Chatbot
 from re_edge_gpt.re_edge_gpt import ConversationStyle
+from re_edge_gpt.image_genearation import ImageGen
+from re_edge_gpt.image_genearation import ImageGenAsync
 
 __all__ = [
-    "Chatbot", "ConversationStyle"
+    "Chatbot", "ConversationStyle", "ImageGen", "ImageGenAsync"
 ]
diff --git a/re_edge_gpt/image_genearation.py b/re_edge_gpt/image_genearation.py
new file mode 100644
index 0000000..6c43063
--- /dev/null
+++ b/re_edge_gpt/image_genearation.py
@@ -0,0 +1,475 @@
+import argparse
+import asyncio
+import contextlib
+import json
+import os
+import socket
+import sys
+import time
+from functools import partial
+from typing import Dict
+from typing import List
+from typing import Union
+
+import httpx
+import pkg_resources
+import regex
+import requests
+
+take_ip_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+take_ip_socket.connect(("8.8.8.8", 80))
+FORWARDED_IP: str = take_ip_socket.getsockname()[0]
+take_ip_socket.close()
+
+BING_URL = os.getenv("BING_URL", "https://www.bing.com")
+
+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-language": "en-US,en;q=0.9",
+    "cache-control": "max-age=0",
+    "content-type": "application/x-www-form-urlencoded",
+    "referrer": "https://www.bing.com/images/create/",
+    "origin": "https://www.bing.com",
+    "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+                  "AppleWebKit/537.36 (KHTML, like Gecko) "
+                  "Chrome/115.0.0.0 "
+                  "Safari/537.36 "
+                  "Edg/115.0.1901.188",
+    "x-forwarded-for": FORWARDED_IP,
+}
+
+# Error messages
+error_timeout = "Your request has timed out."
+error_redirect = "Redirect failed"
+error_blocked_prompt = (
+    "Your prompt has been blocked by Bing. Try to change any bad words and try again."
+)
+error_being_reviewed_prompt = "Your prompt is being reviewed by Bing. Try to change any sensitive words and try again."
+error_noresults = "Could not get results"
+error_unsupported_lang = "\nthis language is currently not supported by bing"
+error_bad_images = "Bad images"
+error_no_images = "No images"
+# Action messages
+sending_message = "Sending request..."
+wait_message = "Waiting for results..."
+download_message = "\nDownloading images..."
+
+
+def debug(debug_file, text_var):
+    """helper function for debug"""
+    with open(f"{debug_file}", "a", encoding="utf-8") as f:
+        f.write(str(text_var))
+        f.write("\n")
+
+
+class ImageGen:
+    """
+    Image generation by Microsoft Bing
+    Parameters:
+        auth_cookie: str
+    Optional Parameters:
+        debug_file: str
+        quiet: bool
+        all_cookies: List[Dict]
+    """
+
+    def __init__(
+            self,
+            auth_cookie: str,
+            debug_file: Union[str, None] = None,
+            quiet: bool = False,
+            all_cookies: List[Dict] = None,
+    ) -> None:
+        self.session: requests.Session = requests.Session()
+        self.session.headers = HEADERS
+        self.session.cookies.set("_U", auth_cookie)
+        if all_cookies:
+            for cookie in all_cookies:
+                self.session.cookies.set(cookie["name"], cookie["value"])
+        self.quiet = quiet
+        self.debug_file = debug_file
+        if self.debug_file:
+            self.debug = partial(debug, self.debug_file)
+
+    def get_images(self, prompt: str) -> list:
+        """
+        Fetches image links from Bing
+        Parameters:
+            prompt: str
+        """
+        if not self.quiet:
+            print(sending_message)
+        if self.debug_file:
+            self.debug(sending_message)
+        url_encoded_prompt = requests.utils.quote(prompt)
+        payload = f"q={url_encoded_prompt}&qs=ds"
+        # https://www.bing.com/images/create?q=<PROMPT>&rt=3&FORM=GENCRE
+        url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=4&FORM=GENCRE"
+        response = self.session.post(
+            url,
+            allow_redirects=False,
+            data=payload,
+            timeout=200,
+        )
+        # check for content waring message
+        if "this prompt is being reviewed" in response.text.lower():
+            if self.debug_file:
+                self.debug(f"ERROR: {error_being_reviewed_prompt}")
+            raise Exception(
+                error_being_reviewed_prompt,
+            )
+        if "this prompt has been blocked" in response.text.lower():
+            if self.debug_file:
+                self.debug(f"ERROR: {error_blocked_prompt}")
+            raise Exception(
+                error_blocked_prompt,
+            )
+        if (
+                "we're working hard to offer image creator in more languages"
+                in response.text.lower()
+        ):
+            if self.debug_file:
+                self.debug(f"ERROR: {error_unsupported_lang}")
+            raise Exception(error_unsupported_lang)
+        if response.status_code != 302:
+            # if rt4 fails, try rt3
+            url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=3&FORM=GENCRE"
+            response = self.session.post(url, allow_redirects=False, timeout=200)
+            if response.status_code != 302:
+                if self.debug_file:
+                    self.debug(f"ERROR: {error_redirect}")
+                print(f"ERROR: {response.text}")
+                raise Exception(error_redirect)
+        # Get redirect URL
+        redirect_url = response.headers["Location"].replace("&nfy=1", "")
+        request_id = redirect_url.split("id=")[-1]
+        self.session.get(f"{BING_URL}{redirect_url}")
+        # https://www.bing.com/images/create/async/results/{ID}?q={PROMPT}
+        polling_url = f"{BING_URL}/images/create/async/results/{request_id}?q={url_encoded_prompt}"
+        # Poll for results
+        if self.debug_file:
+            self.debug("Polling and waiting for result")
+        if not self.quiet:
+            print("Waiting for results...")
+        start_wait = time.time()
+        while True:
+            if int(time.time() - start_wait) > 200:
+                if self.debug_file:
+                    self.debug(f"ERROR: {error_timeout}")
+                raise Exception(error_timeout)
+            if not self.quiet:
+                print(".", end="", flush=True)
+            response = self.session.get(polling_url)
+            if response.status_code != 200:
+                if self.debug_file:
+                    self.debug(f"ERROR: {error_noresults}")
+                raise Exception(error_noresults)
+            if not response.text or response.text.find("errorMessage") != -1:
+                time.sleep(1)
+                continue
+            else:
+                break
+        # Use regex to search for src=""
+        image_links = regex.findall(r'src="([^"]+)"', response.text)
+        # Remove size limit
+        normal_image_links = [link.split("?w=")[0] for link in image_links]
+        # Remove duplicates
+        normal_image_links = list(set(normal_image_links))
+
+        # Bad images
+        bad_images = [
+            "https://r.bing.com/rp/in-2zU3AJUdkgFe7ZKv19yPBHVs.png",
+            "https://r.bing.com/rp/TX9QuO3WzcCJz1uaaSwQAz39Kb0.jpg",
+        ]
+        for img in normal_image_links:
+            if img in bad_images:
+                raise Exception("Bad images")
+        # No images
+        if not normal_image_links:
+            raise Exception(error_no_images)
+        return normal_image_links
+
+    def save_images(
+            self,
+            links: list,
+            output_dir: str,
+            file_name: str = None,
+            download_count: int = None,
+    ) -> None:
+        """
+        Saves images to output directory
+        Parameters:
+            links: list[str]
+            output_dir: str
+            file_name: str
+            download_count: int
+        """
+        if self.debug_file:
+            self.debug(download_message)
+        if not self.quiet:
+            print(download_message)
+        with contextlib.suppress(FileExistsError):
+            os.mkdir(output_dir)
+        try:
+            fn = f"{file_name}_" if file_name else ""
+            jpeg_index = 0
+
+            if download_count:
+                links = links[:download_count]
+
+            for link in links:
+                while os.path.exists(
+                        os.path.join(output_dir, f"{fn}{jpeg_index}.jpeg")
+                ):
+                    jpeg_index += 1
+                response = self.session.get(link)
+                if response.status_code != 200:
+                    raise Exception(f"Could not download image response code {response.status_code}")
+                # save response to file
+                with open(
+                        os.path.join(output_dir, f"{fn}{jpeg_index}.jpeg"), "wb"
+                ) as output_file:
+                    output_file.write(response.content)
+                jpeg_index += 1
+
+        except requests.exceptions.MissingSchema as url_exception:
+            raise Exception(
+                "Inappropriate contents found in the generated images. Please try again or try another prompt.",
+            ) from url_exception
+
+
+class ImageGenAsync:
+    """
+    Image generation by Microsoft Bing
+    Parameters:
+        auth_cookie: str
+    Optional Parameters:
+        debug_file: str
+        quiet: bool
+        all_cookies: list[dict]
+    """
+
+    def __init__(
+            self,
+            auth_cookie: str = None,
+            debug_file: Union[str, None] = None,
+            quiet: bool = False,
+            all_cookies: List[Dict] = None,
+    ) -> None:
+        if auth_cookie is None and not all_cookies:
+            raise Exception("No auth cookie provided")
+        self.session = httpx.AsyncClient(
+            headers=HEADERS,
+            trust_env=True,
+        )
+        if auth_cookie:
+            self.session.cookies.update({"_U": auth_cookie})
+        if all_cookies:
+            for cookie in all_cookies:
+                self.session.cookies.update(
+                    {cookie["name"]: cookie["value"]},
+                )
+        self.quiet = quiet
+        self.debug_file = debug_file
+        if self.debug_file:
+            self.debug = partial(debug, self.debug_file)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *excinfo) -> None:
+        await self.session.aclose()
+
+    async def get_images(self, prompt: str) -> list:
+        """
+        Fetches image links from Bing
+        Parameters:
+            prompt: str
+        """
+        if not self.quiet:
+            print("Sending request...")
+        url_encoded_prompt = requests.utils.quote(prompt)
+        # https://www.bing.com/images/create?q=<PROMPT>&rt=3&FORM=GENCRE
+        url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=3&FORM=GENCRE"
+        payload = f"q={url_encoded_prompt}&qs=ds"
+        response = await self.session.post(
+            url,
+            follow_redirects=False,
+            data=payload,
+        )
+        content = response.text
+        if "this prompt has been blocked" in content.lower():
+            raise Exception(
+                "Your prompt has been blocked by Bing. Try to change any bad words and try again.",
+            )
+        if response.status_code != 302:
+            # if rt4 fails, try rt3
+            url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=4&FORM=GENCRE"
+            response = await self.session.post(
+                url,
+                follow_redirects=False,
+                timeout=200,
+            )
+            if response.status_code != 302:
+                print(f"ERROR: {response.text}")
+                raise Exception("Redirect failed")
+        # Get redirect URL
+        redirect_url = response.headers["Location"].replace("&nfy=1", "")
+        request_id = redirect_url.split("id=")[-1]
+        await self.session.get(f"{BING_URL}{redirect_url}")
+        # https://www.bing.com/images/create/async/results/{ID}?q={PROMPT}
+        polling_url = f"{BING_URL}/images/create/async/results/{request_id}?q={url_encoded_prompt}"
+        # Poll for results
+        if not self.quiet:
+            print("Waiting for results...")
+        while True:
+            if not self.quiet:
+                print(".", end="", flush=True)
+            # By default, timeout is 300s, change as needed
+            response = await self.session.get(polling_url)
+            if response.status_code != 200:
+                raise Exception("Could not get results")
+            content = response.text
+            if content and content.find("errorMessage") == -1:
+                break
+
+            await asyncio.sleep(1)
+            continue
+        # Use regex to search for src=""
+        image_links = regex.findall(r'src="([^"]+)"', content)
+        # Remove size limit
+        normal_image_links = [link.split("?w=")[0] for link in image_links]
+        # Remove duplicates
+        normal_image_links = list(set(normal_image_links))
+
+        # Bad images
+        bad_images = [
+            "https://r.bing.com/rp/in-2zU3AJUdkgFe7ZKv19yPBHVs.png",
+            "https://r.bing.com/rp/TX9QuO3WzcCJz1uaaSwQAz39Kb0.jpg",
+        ]
+        for im in normal_image_links:
+            if im in bad_images:
+                raise Exception("Bad images")
+        # No images
+        if not normal_image_links:
+            raise Exception("No images")
+        return normal_image_links
+
+
+async def async_image_gen(
+        prompt: str,
+        download_count: int,
+        output_dir: str,
+        u_cookie=None,
+        debug_file=None,
+        quiet=False,
+        all_cookies=None,
+):
+    async with ImageGenAsync(
+            u_cookie,
+            debug_file=debug_file,
+            quiet=quiet,
+            all_cookies=all_cookies,
+    ) as image_generator:
+        images = await image_generator.get_images(prompt)
+        await image_generator.save_images(
+            images, output_dir=output_dir, download_count=download_count
+        )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-U", help="Auth cookie from browser", type=str)
+    parser.add_argument("--cookie-file", help="File containing auth cookie", type=str)
+    parser.add_argument(
+        "--prompt",
+        help="Prompt to generate images for",
+        type=str,
+        required=True,
+    )
+
+    parser.add_argument(
+        "--output-dir",
+        help="Output directory",
+        type=str,
+        default="./output",
+    )
+
+    parser.add_argument(
+        "--download-count",
+        help="Number of images to download, value must be less than five",
+        type=int,
+        default=4,
+    )
+
+    parser.add_argument(
+        "--debug-file",
+        help="Path to the file where debug information will be written.",
+        type=str,
+    )
+
+    parser.add_argument(
+        "--quiet",
+        help="Disable pipeline messages",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--asyncio",
+        help="Run ImageGen using asyncio",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--version",
+        action="store_true",
+        help="Print the version number",
+    )
+
+    args = parser.parse_args()
+
+    if args.version:
+        print(pkg_resources.get_distribution("BingImageCreator").version)
+        sys.exit()
+
+    # Load auth cookie
+    cookie_json = None
+    if args.cookie_file is not None:
+        with contextlib.suppress(Exception):
+            with open(args.cookie_file, encoding="utf-8") as file:
+                cookie_json = json.load(file)
+
+    if args.U is None and args.cookie_file is None:
+        raise Exception("Could not find auth cookie")
+
+    if args.download_count > 4:
+        raise Exception("The number of downloads must be less than five")
+
+    if not args.asyncio:
+        # Create image generator
+        image_generator = ImageGen(
+            args.U,
+            args.debug_file,
+            args.quiet,
+            all_cookies=cookie_json,
+        )
+        image_generator.save_images(
+            image_generator.get_images(args.prompt),
+            output_dir=args.output_dir,
+            download_count=args.download_count,
+        )
+    else:
+        asyncio.run(
+            async_image_gen(
+                args.prompt,
+                args.download_count,
+                args.output_dir,
+                args.U,
+                args.debug_file,
+                args.quiet,
+                all_cookies=cookie_json,
+            ),
+        )
+
+
+if __name__ == "__main__":
+    main()


Test patch:
diff --git a/test/unit_test/test_bot.py b/test/unit_test/test_bot/test_bot.py
similarity index 100%
rename from test/unit_test/test_bot.py
rename to test/unit_test/test_bot/test_bot.py
diff --git a/test/unit_test/test_bot_manual.py b/test/unit_test/test_bot/test_bot_manual.py
similarity index 100%
rename from test/unit_test/test_bot_manual.py
rename to test/unit_test/test_bot/test_bot_manual.py
diff --git a/test/unit_test/test_generate_image/test_generation.py b/test/unit_test/test_generate_image/test_generation.py
new file mode 100644
index 0000000..02ac0e7
--- /dev/null
+++ b/test/unit_test/test_generate_image/test_generation.py
@@ -0,0 +1,41 @@
+import asyncio
+import os
+import shutil
+from pathlib import Path
+
+from re_edge_gpt import ImageGen, ImageGenAsync
+
+# create a temporary output directory for testing purposes
+test_output_dir = "test_output"
+# download a test image
+test_image_url = "https://picsum.photos/200"
+auth_cooker = str(os.getenv("IMAGE_COOKIE"))
+sync_gen = ImageGen(auth_cookie=auth_cooker)
+async_gen = ImageGenAsync(auth_cookie=auth_cooker)
+
+
+def test_save_images_sync():
+    sync_gen.save_images([test_image_url], test_output_dir)
+    sync_gen.save_images([test_image_url], test_output_dir, file_name="test_image")
+    # check if the image was downloaded and saved correctly
+    assert os.path.exists(os.path.join(test_output_dir, "test_image_0.jpeg"))
+    assert os.path.exists(os.path.join(test_output_dir, "0.jpeg"))
+
+
+def test_generate_image_sync():
+    image_list = sync_gen.get_images("tree")
+    print(image_list)
+
+
+async def test_generate_image_async():
+    image_list = await async_gen.get_images("tree")
+    print(image_list)
+
+
+if __name__ == "__main__":
+    Path("test_output").mkdir(exist_ok=True)
+    test_save_images_sync()
+    test_generate_image_sync()
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(test_generate_image_async())
+    shutil.rmtree(test_output_dir)
diff --git a/test/unit_test/test_generate_image/test_generation_manual.py b/test/unit_test/test_generate_image/test_generation_manual.py
new file mode 100644
index 0000000..84fc77f
--- /dev/null
+++ b/test/unit_test/test_generate_image/test_generation_manual.py
@@ -0,0 +1,48 @@
+import asyncio
+import os
+import shutil
+from pathlib import Path
+
+from re_edge_gpt import ImageGen, ImageGenAsync
+
+# create a temporary output directory for testing purposes
+test_output_dir = "test_output"
+# download a test image
+test_image_url = "https://picsum.photos/200"
+auth_cooker = open("bing_cookies.txt", "r+").read()
+sync_gen = ImageGen(auth_cookie=auth_cooker)
+async_gen = ImageGenAsync(auth_cookie=auth_cooker)
+
+
+def test_save_images_sync():
+    sync_gen.save_images([test_image_url], test_output_dir)
+    sync_gen.save_images([test_image_url], test_output_dir, file_name="test_image")
+    # check if the image was downloaded and saved correctly
+    assert os.path.exists(os.path.join(test_output_dir, "test_image_0.jpeg"))
+    assert os.path.exists(os.path.join(test_output_dir, "0.jpeg"))
+
+
+# Generate image list sync
+def test_generate_image_sync():
+    image_list = sync_gen.get_images("tree")
+    print(image_list)
+
+
+# Generate image list async
+async def test_generate_image_async():
+    image_list = await async_gen.get_images("tree")
+    print(image_list)
+
+
+if __name__ == "__main__":
+    # Make dir to save image
+    Path("test_output").mkdir(exist_ok=True)
+    # Save image
+    test_save_images_sync()
+    # Generate image sync
+    test_generate_image_sync()
+    # Generate image async
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(test_generate_image_async())
+    # Remove dir
+    shutil.rmtree(test_output_dir)