Update app.py

app.py

@@ -25,7 +25,7 @@ def generate_image(prompt):
         # Wait for the textarea input field
         input_field = WebDriverWait(driver, 10).until(
             EC.presence_of_element_located(
-                (By.CSS_SELECTOR, "textarea.wp-content-visualizer-chat-input")
+                (By.CLASS_NAME, "wp-content-visualizer-chat-input")
             )
         )
         # Enter the prompt
@@ -34,18 +34,18 @@ 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.CSS_SELECTOR, "button.wp-content-visualizer-chat-send")
+                (By.CLASS_NAME, "wp-content-visualizer-chat-send")
             )
         )
         generate_button.click()
 
         # Wait for the generated image
-        generated_image = WebDriverWait(driver, 20).until(
+        generated_image = WebDriverWait(driver, 520).until(
             EC.presence_of_element_located(
-                (By.CSS_SELECTOR, "img.wp-content-visualizer-generated-image")
+                (By.CLASS_NAME, "wp-content-visualizer-generated-image")
             )
         )
-
+        
         # Extract the image URL
         image_url = generated_image.get_attribute("src")
         return image_url