1 files changed, 7 insertions, 14 deletions

diff --git a/g4f/Provider/TalkAi.py b/g4f/Provider/TalkAi.py
index 5b03b91e..20ba65b5 100644
--- a/g4f/Provider/TalkAi.py
+++ b/g4f/Provider/TalkAi.py
@@ -4,7 +4,7 @@ import time, json, time
 
 from ..typing import CreateResult, Messages
 from .base_provider import BaseProvider
-from .helper import WebDriver, get_browser
+from .helper import WebDriver, WebDriverSession
 
 class TalkAi(BaseProvider):
     url = "https://talkai.info"
@@ -19,16 +19,14 @@ class TalkAi(BaseProvider):
         messages: Messages,
         stream: bool,
         proxy: str = None,
-        browser: WebDriver = None,
+        web_driver: WebDriver = None,
         **kwargs
     ) -> CreateResult:
-        driver = browser if browser else get_browser("", False, proxy)
+        with WebDriverSession(web_driver, "", virtual_display=True, proxy=proxy) as driver:
+            from selenium.webdriver.common.by import By
+            from selenium.webdriver.support.ui import WebDriverWait
+            from selenium.webdriver.support import expected_conditions as EC
 
-        from selenium.webdriver.common.by import By
-        from selenium.webdriver.support.ui import WebDriverWait
-        from selenium.webdriver.support import expected_conditions as EC
-
-        try:
             driver.get(f"{cls.url}/chat/")
     
             # Wait for page load
@@ -86,9 +84,4 @@ return content;
                 elif chunk != "":
                     break
                 else:
-                    time.sleep(0.1)
-        finally:
-            if not browser:
-                driver.close()
-                time.sleep(0.1)
-                driver.quit()
\ No newline at end of file
+                    time.sleep(0.1)
\ No newline at end of file