Update app.py

app.py

@@ -9,14 +9,19 @@ 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")
-    # chrome_options.add_argument("--disable-dev-shm-usage")
+    chrome_options = Options()
+    chrome_options.add_argument("--headless")
+    chrome_options.add_argument("--no-sandbox")
+    chrome_options.add_argument("--disable-dev-shm-usage")
+    chrome_options.add_argument("--disable-gpu")
+    chrome_options.add_argument("window-size=1920x1080")
 
     # Initialize the WebDriver
-    driver = webdriver.Chrome() # options=chrome_options
+    try:
+        driver = webdriver.Chrome(options=chrome_options)
+    except Exception as e:
+        print(f"Failed to initialize ChromeDriver: {e}. Falling back to webdriver-manager.")
+        driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=chrome_options)
 
     try:
         # Open the website
@@ -49,7 +54,9 @@ def generate_image(prompt):
         # Extract the image URL
         image_url = generated_image.get_attribute("src")
         return image_url
-
+        
+    except Exception as e:
+        return f"Error: {str(e)}"
     finally:
         # Close the browser
         driver.quit()