File size: 5,045 Bytes
c6141f4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
import logging
import requests
import tempfile
import dotenv
import os
from datetime import datetime
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.keys import Keys
from selenium.common.exceptions import TimeoutException
from selenium.common.exceptions import NoSuchElementException
dotenv.load_dotenv(".env_d3")
class StellaD3():
def __init__(self, local_save_dir:str="outputs", local_save:bool=True, show_debug_log:bool=False):
try:
if show_debug_log: logging.basicConfig(level=logging.DEBUG)
self.local_save = local_save
if self.local_save:
self.local_save_dir = local_save_dir
os.makedirs(self.local_save_dir, exist_ok=True)
else: self.local_save_dir = None
options = Options()
options.add_argument("--headless=new")
options.add_experimental_option("excludeSwitches", ["ignore-certificate-errors"])
options.add_argument('--disable-gpu')
options.add_argument('--disable-blink-features=AutomationControlled')
self.driver = webdriver.Chrome(options=options)
self.driver.get(os.getenv("CONNECTOR_URL"))
email_input = WebDriverWait(self.driver, 2).until(EC.visibility_of_element_located((By.ID, "i0116")))
email_input.send_keys(os.getenv("CONNECTOR_ID"))
next_button = WebDriverWait(self.driver, 2).until(EC.element_to_be_clickable((By.ID, "idSIButton9")))
next_button.click()
password_input = WebDriverWait(self.driver, 2).until(EC.visibility_of_element_located((By.ID, "i0118")))
password_input.send_keys(os.getenv("CONNECTOR_KEY"))
next_button = WebDriverWait(self.driver, 2).until(EC.element_to_be_clickable((By.ID, "idSIButton9")))
next_button.click()
accept_button = WebDriverWait(self.driver, 2).until(EC.element_to_be_clickable((By.ID, "acceptButton")))
accept_button.click()
print('D3 initialized!')
except Exception as e:
print(f'Error: {e}')
quit()
def image_generator_d3(self, prompt):
self.driver.get(os.getenv("CREATOR_URL"))
self.driver.refresh()
self.driver.find_element(By.ID, "sb_form_q").send_keys(prompt)
self.driver.find_element(By.ID, "create_btn_c").click()
output_files = []
if self.local_save:
temp_dir = str(self.local_save_dir)
temp_dir = os.path.join(temp_dir, datetime.now().strftime("%Y-%m-%d"))
os.makedirs(temp_dir, exist_ok=True)
else: temp_dir = None
try:
WebDriverWait(self.driver, 15).until(EC.presence_of_element_located((By.CLASS_NAME, "gil_err_tc")))
raise Exception('GIL_ERR_TC (Prompt Blocked)')
except TimeoutException:
pass
while True:
print("GI_REFRESH (Refresh Initiated)")
self.driver.refresh()
try:
WebDriverWait(self.driver, 15).until(EC.presence_of_element_located((By.CLASS_NAME, "img_cont")))
divs = self.driver.find_elements(By.CLASS_NAME, "img_cont")
urls = [div.find_element(By.TAG_NAME, "img").get_attribute("src").split("?")[0] for div in divs]
print('IMG_CONT (Complete!)')
for url in urls:
response = requests.get(url)
with tempfile.NamedTemporaryFile(delete=False, prefix=f"{self.__class__.__name__}_", suffix=".png", dir=temp_dir) as output:
output.write(response.content)
output_files.append(output.name)
return output_files
except TimeoutException:
try:
img = self.driver.find_element(By.CLASS_NAME, "gir_mmimg")
src = img.get_attribute("src").split("?")[0]
print('GIR_MMIMG (Complete!)')
response = requests.get(src)
with tempfile.NamedTemporaryFile(delete=False, prefix=f"{self.__class__.__name__}_", suffix=".png", dir=temp_dir) as output:
output.write(response.content)
output_files.append(output.name)
return output_files
except NoSuchElementException:
raise Exception('GIL_NOT_FOUND (Element Not Found)')
if __name__ == "__main__":
print("Import StellaD3 in another file to start using it!") |