Update app.py

app.py

@@ -4,13 +4,14 @@ from selenium.webdriver.common.by import By
 from selenium.webdriver.common.keys import Keys
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
+from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.chrome.service import Service
 from webdriver_manager.chrome import ChromeDriverManager
-from selenium.webdriver.chrome.options import Options
 import time
 
 # Function to generate image URL
 def generate_image(prompt):
+    # Configure Chrome options for headless mode
     chrome_options = Options()
     chrome_options.add_argument("--headless")
     chrome_options.add_argument("--no-sandbox")
@@ -36,7 +37,7 @@ def generate_image(prompt):
         # Wait for the textarea input field
         input_field = WebDriverWait(driver, 10).until(
             EC.presence_of_element_located(
-                (By.CLASS_NAME, "wp-content-visualizer-chat-input")
+                (By.CSS_SELECTOR, "textarea.wp-content-visualizer-chat-input")
             )
         )
         # Enter the prompt
@@ -45,22 +46,22 @@ def generate_image(prompt):
         # Wait for the "Generate Image" button and click it
         generate_button = WebDriverWait(driver, 10).until(
             EC.element_to_be_clickable(
-                (By.CLASS_NAME, "wp-content-visualizer-chat-send")
+                (By.CSS_SELECTOR, "button.wp-content-visualizer-chat-send")
             )
         )
         generate_button.click()
 
         # Wait for the generated image
-        generated_image = WebDriverWait(driver, 520).until(
+        generated_image = WebDriverWait(driver, 20).until(
             EC.presence_of_element_located(
-                (By.CLASS_NAME, "wp-content-visualizer-generated-image")
+                (By.CSS_SELECTOR, "img.wp-content-visualizer-generated-image")
             )
         )
-        
+
         # Extract the image URL
         image_url = generated_image.get_attribute("src")
         return image_url
-        
+
     except Exception as e:
         return f"Error: {str(e)}"
     finally:
@@ -71,6 +72,8 @@ def generate_image(prompt):
 def gradio_interface(prompt):
     try:
         image_url = generate_image(prompt)
+        if image_url.startswith("Error"):
+            return image_url, None
         return f"Generated Image URL: {image_url}", image_url
     except Exception as e:
         return f"Error: {str(e)}", None