7 files changed, 32 insertions, 97 deletions

diff --git a/g4f/Provider/AItianhuSpace.py b/g4f/Provider/AItianhuSpace.py
index a9a824cf..545944c0 100644
--- a/g4f/Provider/AItianhuSpace.py
+++ b/g4f/Provider/AItianhuSpace.py
@@ -24,7 +24,7 @@ class AItianhuSpace(BaseProvider):
         proxy: str = None,
         timeout: int = 120,
         browser: WebDriver = None,
-        hidden_display: bool = True,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
         if not model:
@@ -38,13 +38,7 @@ class AItianhuSpace(BaseProvider):
             print(f"AItianhuSpace | using domain: {domain}")
         url = f"https://{domain}"
         prompt = format_prompt(messages)
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser("", True, proxy)
-            else:
-                driver = get_browser("", False, proxy)
+        driver = browser if browser else get_browser("", headless, proxy)
 
         from selenium.webdriver.common.by import By
         from selenium.webdriver.support.ui import WebDriverWait
@@ -66,6 +60,7 @@ document.getElementById('sheet').addEventListener('click', () => {{
         original_window = driver.current_window_handle
         for window_handle in driver.window_handles:
             if window_handle != original_window:
+                driver.close()
                 driver.switch_to.window(window_handle)
                 break
 
@@ -120,9 +115,7 @@ return "";
                 else:
                     time.sleep(0.1)
         finally:
-            driver.close()
             if not browser:
+                driver.close()
                 time.sleep(0.1)
-                driver.quit()
-            if hidden_display:
-                display.stop()
\ No newline at end of file
+                driver.quit()
\ No newline at end of file
diff --git a/g4f/Provider/MyShell.py b/g4f/Provider/MyShell.py
index 0efeb0e8..5580ff6e 100644
--- a/g4f/Provider/MyShell.py
+++ b/g4f/Provider/MyShell.py
@@ -21,16 +21,10 @@ class MyShell(BaseProvider):
         proxy: str = None,
         timeout: int = 120,
         browser: WebDriver = None,
-        hidden_display: bool = True,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser("", True, proxy)
-            else:
-                driver = get_browser("", False, proxy)
+        driver = browser if browser else get_browser("", headless, proxy)
 
         from selenium.webdriver.common.by import By
         from selenium.webdriver.support.ui import WebDriverWait
@@ -87,9 +81,7 @@ return content;
                 elif chunk != "":
                     break
         finally:
-            driver.close()
             if not browser:
+                driver.close()
                 time.sleep(0.1)
-                driver.quit()
-            if hidden_display:
-                display.stop()
\ No newline at end of file
+                driver.quit()
\ No newline at end of file
diff --git a/g4f/Provider/PerplexityAi.py b/g4f/Provider/PerplexityAi.py
index b9de429e..1159840a 100644
--- a/g4f/Provider/PerplexityAi.py
+++ b/g4f/Provider/PerplexityAi.py
@@ -22,16 +22,10 @@ class PerplexityAi(BaseProvider):
         timeout: int = 120,
         browser: WebDriver = None,
         copilot: bool = False,
-        hidden_display: bool = True,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser("", True, proxy)
-            else:
-                driver = get_browser("", False, proxy)
+        driver = browser if browser else get_browser("", headless, proxy)
 
         from selenium.webdriver.common.by import By
         from selenium.webdriver.support.ui import WebDriverWait
@@ -112,9 +106,7 @@ if(window._message && window._message != window._last_message) {
                 else:
                     time.sleep(0.1)
         finally:
-            driver.close()
             if not browser:
+                driver.close()
                 time.sleep(0.1)
-                driver.quit()
-            if hidden_display:
-                display.stop()
\ No newline at end of file
+                driver.quit()
\ No newline at end of file
diff --git a/g4f/Provider/Phind.py b/g4f/Provider/Phind.py
index 2612a92d..0cb17ddb 100644
--- a/g4f/Provider/Phind.py
+++ b/g4f/Provider/Phind.py
@@ -23,16 +23,10 @@ class Phind(BaseProvider):
         timeout: int = 120,
         browser: WebDriver = None,
         creative_mode: bool = None,
-        hidden_display: bool = True,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser("", True, proxy)
-            else:
-                driver = get_browser("", False, proxy)
+        driver = browser if browser else get_browser("", headless, proxy)
 
         from selenium.webdriver.common.by import By
         from selenium.webdriver.support.ui import WebDriverWait
@@ -111,9 +105,7 @@ if(window.reader) {
                 else:
                     time.sleep(0.1)
         finally:
-            driver.close()
             if not browser:
+                driver.close()
                 time.sleep(0.1)
-                driver.quit()
-            if hidden_display:
-                display.stop()
\ No newline at end of file
+                driver.quit()
\ No newline at end of file
diff --git a/g4f/Provider/helper.py b/g4f/Provider/helper.py
index b10c912a..f879816a 100644
--- a/g4f/Provider/helper.py
+++ b/g4f/Provider/helper.py
@@ -31,14 +31,6 @@ except ImportError:
     class ChromeOptions():
         def add_argument():
             pass
-try:
-    from pyvirtualdisplay import Display
-except ImportError:
-    class Display():
-        def start():
-            pass
-        def stop():
-            pass
 
 from ..typing import Dict, Messages, Union, Tuple
 from .. import debug
@@ -139,24 +131,16 @@ def format_prompt(messages: Messages, add_special_tokens=False) -> str:
 
 def get_browser(
     user_data_dir: str = None,
-    hidden_display: bool = False,
+    headless: bool = False,
     proxy: str = None,
     options: ChromeOptions = None
-) -> Union[Chrome, Tuple[Chrome, Display]] :
+) -> Chrome:
     if user_data_dir == None:
         user_data_dir = user_config_dir("g4f")
 
-    if hidden_display:
-        display = Display(visible=0, size=(1920, 1080))
-        display.start()
-
     if proxy:
         if not options:
             options = ChromeOptions()
         options.add_argument(f'--proxy-server={proxy}')
 
-    browser = Chrome(user_data_dir=user_data_dir, options=options)
-    if hidden_display:
-        return browser, display
-
-    return browser
\ No newline at end of file
+    return Chrome(user_data_dir=user_data_dir, options=options, headless=headless)
\ No newline at end of file
diff --git a/g4f/Provider/needs_auth/Bard.py b/g4f/Provider/needs_auth/Bard.py
index 6cb40c90..7f73f1b3 100644
--- a/g4f/Provider/needs_auth/Bard.py
+++ b/g4f/Provider/needs_auth/Bard.py
@@ -19,17 +19,12 @@ class Bard(BaseProvider):
         stream: bool,
         proxy: str = None,
         browser: WebDriver = None,
-        hidden_display: bool = True,
+        user_data_dir: str = None,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
         prompt = format_prompt(messages)
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser(None, True, proxy)
-            else:
-                driver = get_browser(None, False, proxy)
+        driver = browser if browser else get_browser(user_data_dir, headless, proxy)
 
         from selenium.webdriver.common.by import By
         from selenium.webdriver.support.ui import WebDriverWait
@@ -43,9 +38,6 @@ class Bard(BaseProvider):
             # Reopen browser for login
             if not browser:
                 driver.quit()
-                # New browser should be visible
-                if hidden_display:
-                    display.stop()
                 driver = get_browser(None, False, proxy)
                 driver.get(f"{cls.url}/chat")
                 wait = WebDriverWait(driver, 240)
@@ -83,9 +75,7 @@ XMLHttpRequest.prototype.open = function(method, url) {
                 else:
                     time.sleep(0.1)
         finally:
-            driver.close()
             if not browser:
+                driver.close()
                 time.sleep(0.1)
-                driver.quit()
-            if hidden_display:
-                display.stop()
\ No newline at end of file
+                driver.quit()
\ No newline at end of file
diff --git a/g4f/Provider/needs_auth/Poe.py b/g4f/Provider/needs_auth/Poe.py
index 6fbf7fd4..a894bcb1 100644
--- a/g4f/Provider/needs_auth/Poe.py
+++ b/g4f/Provider/needs_auth/Poe.py
@@ -22,6 +22,7 @@ models = {
 class Poe(BaseProvider):
     url = "https://poe.com"
     working = True
+    needs_auth = True
     supports_gpt_35_turbo = True
     supports_stream = True
 
@@ -33,7 +34,8 @@ class Poe(BaseProvider):
         stream: bool,
         proxy: str = None,
         browser: WebDriver = None,
-        hidden_display: bool = True,
+        user_data_dir: str = None,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
         if not model:
@@ -41,13 +43,7 @@ class Poe(BaseProvider):
         elif model not in models:
             raise ValueError(f"Model are not supported: {model}")
         prompt = format_prompt(messages)
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser(None, True, proxy)
-            else:
-                driver = get_browser(None, False, proxy)
+        driver = browser if browser else get_browser(user_data_dir, headless, proxy)
 
         script = """
 window._message = window._last_message = "";
@@ -80,14 +76,12 @@ window.WebSocket = ProxiedWebSocket;
 
         try:
             driver.get(f"{cls.url}/{models[model]['name']}")
-            wait = WebDriverWait(driver, 10 if hidden_display else 240)
+            wait = WebDriverWait(driver, 10 if headless else 240)
             wait.until(EC.visibility_of_element_located((By.CSS_SELECTOR, "textarea[class^='GrowingTextArea']")))
         except:
             # Reopen browser for login
             if not browser:
                 driver.quit()
-                if hidden_display:
-                    display.stop()
                 driver = get_browser(None, False, proxy)
                 driver.get(f"{cls.url}/{models[model]['name']}")
                 wait = WebDriverWait(driver, 240)
@@ -121,9 +115,7 @@ if(window._message && window._message != window._last_message) {
                 else:
                     time.sleep(0.1)
         finally:
-            driver.close()
             if not browser:
+                driver.close()
                 time.sleep(0.1)
-                driver.quit()
-            if hidden_display:
-                display.stop()
\ No newline at end of file
+                driver.quit()
\ No newline at end of file