mirror of
https://github.com/SillyTavern/SillyTavern-Extras.git
synced 2026-04-29 10:51:19 +00:00
Add colab detection util
This commit is contained in:
@@ -4,20 +4,26 @@ from selenium.webdriver.chrome.options import Options as ChromeOptions
|
||||
from selenium.webdriver.firefox.options import Options as FirefoxOptions
|
||||
from selenium.webdriver.chrome.service import Service as ChromeService
|
||||
from selenium.webdriver.firefox.service import Service as FirefoxService
|
||||
from modules.utils import is_colab
|
||||
import atexit
|
||||
|
||||
|
||||
def get_driver():
|
||||
try:
|
||||
print("Initializing Chrome driver...")
|
||||
chromeService = ChromeService()
|
||||
options = ChromeOptions()
|
||||
options.add_argument('--disable-infobars')
|
||||
options.add_argument("--headless")
|
||||
options.add_argument("--disable-gpu")
|
||||
options.add_argument("--no-sandbox")
|
||||
options.add_argument('--disable-dev-shm-usage')
|
||||
options.add_argument("--lang=en-GB")
|
||||
return webdriver.Chrome(service=chromeService, options=options)
|
||||
|
||||
if is_colab():
|
||||
return webdriver.Chrome('chromedriver', options=options)
|
||||
else:
|
||||
chromeService = ChromeService()
|
||||
return webdriver.Chrome(service=chromeService, options=options)
|
||||
except:
|
||||
print("Chrome not found, using Firefox instead.")
|
||||
firefoxService = FirefoxService()
|
||||
|
||||
Reference in New Issue
Block a user