Update app.py

app.py

@@ -10,6 +10,8 @@ from selenium.common.exceptions import WebDriverException
 from PIL import Image
 from io import BytesIO
 
+from playwright.sync_api import sync_playwright
+
 HF_TOKEN = os.getenv("HF_TOKEN")
 hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus-08-2024", token=HF_TOKEN)
 
@@ -102,30 +104,19 @@ def on_select(space):
         print(f"Error in on_select: {str(e)}")
         print(traceback.format_exc())
         return f"오류가 발생했습니다: {str(e)}", "", Image.new('RGB', (1, 1))
-        
-def take_screenshot(url):
-    options = webdriver.ChromeOptions()
-    options.add_argument('--headless')
-    options.add_argument('--no-sandbox')
-    options.add_argument('--disable-dev-shm-usage')
 
-    wd = None
+def take_screenshot(url):
     try:
-        wd = webdriver.Chrome(options=options)
-        wd.set_window_size(1080, 720)
-        wd.get(url)
-        wd.implicitly_wait(10)
-        screenshot = wd.get_screenshot_as_png()
+        with sync_playwright() as p:
+            browser = p.chromium.launch()
+            page = browser.new_page()
+            page.goto(url)
+            screenshot = page.screenshot()
+            browser.close()
         return Image.open(BytesIO(screenshot))
-    except WebDriverException as e:
-        print(f"WebDriver error: {str(e)}")
-        return Image.new('RGB', (1, 1))
     except Exception as e:
-        print(f"Unexpected error in take_screenshot: {str(e)}")
+        print(f"Screenshot error: {str(e)}")
         return Image.new('RGB', (1, 1))
-    finally:
-        if wd:
-            wd.quit()
 
 
 def create_ui():