Update app.py

app.py

@@ -4,7 +4,7 @@ import requests
 from PIL import Image
 from io import BytesIO
 from selenium import webdriver
-from selenium.webdriver.chrome.service import Service
+from selenium.webdriver.chrome.service import Service, Options
 from webdriver_manager.chrome import ChromeDriverManager
 from diffusers import StableDiffusionPipeline
 import torch
@@ -13,7 +13,12 @@ import gradio as gr
 # ---------- Step 1: Scrape Celebrity Images ----------
 def scrape_images(celebrity_name, num_images=20):
     search_url = f"https://www.google.com/search?q={celebrity_name}+portrait&tbm=isch"
-    driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()))
+    chrome_options = Options()
+    chrome_options.add_argument("--headless")
+    chrome_options.add_argument("--no-sandbox")
+    chrome_options.add_argument("--disable-dev-shm-usage")
+
+    driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=chrome_options)
     driver.get(search_url)
     os.makedirs(f"data/{celebrity_name}", exist_ok=True)