Update app.py

app.py

@@ -7,10 +7,7 @@ import traceback
 from PIL import Image
 from io import BytesIO
 import asyncio
-from selenium import webdriver
-from selenium.webdriver.chrome.service import Service
-from webdriver_manager.chrome import ChromeDriverManager
-from selenium.webdriver.chrome.options import Options
+from gradio_client import Client
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus-08-2024", token=HF_TOKEN)
@@ -92,11 +89,11 @@ def get_app_py_content(space_id: str) -> str:
     except requests.RequestException:
         return f"Error fetching app.py content for space: {space_id}"
         
-async def on_select(space):
+def on_select(space):
     try:
         summary = summarize_space(space)
         app_content = get_app_py_content(space['id'])
-        screenshot = await take_screenshot(space['url'])
+        screenshot = take_screenshot(space['url'])
         info = f"선택된 Space: {space['name']} (ID: {space['id']})\n"
         info += f"Author: {space['author']}\n"
         info += f"Likes: {space['likes']}\n"
@@ -107,27 +104,15 @@ async def on_select(space):
         print(f"Error in on_select: {str(e)}")
         print(traceback.format_exc())
         return f"오류가 발생했습니다: {str(e)}", "", Image.new('RGB', (1080, 720), color='lightgray')
-
-async def take_screenshot(url):
-    options = Options()
-    options.add_argument('--headless')
-    options.add_argument('--no-sandbox')
-    options.add_argument('--disable-dev-shm-usage')
-
+        
+def take_screenshot(url):
     try:
-        service = Service(ChromeDriverManager().install())
-        driver = webdriver.Chrome(service=service, options=options)
-        driver.set_window_size(1080, 720)
-        driver.get(url)
-        await asyncio.sleep(5)  # Wait for the page to load
-        screenshot = driver.get_screenshot_as_png()
-        return Image.open(BytesIO(screenshot))
+        client = Client("ginipick/selenium-screenshot-gradio")
+        result = client.predict(url=url, api_name="/predict")
+        return Image.open(result)
     except Exception as e:
         print(f"Screenshot error: {str(e)}")
         return Image.new('RGB', (1080, 720), color='lightgray')
-    finally:
-        if 'driver' in locals():
-            driver.quit()
             
 def create_ui():
     try:
@@ -162,11 +147,12 @@ def create_ui():
 
                 for _, button, space in space_rows:
                     button.click(
-                        lambda s=space: asyncio.run(on_select(s)),
-                        inputs=[],
-                        outputs=[info_output, app_py_content, screenshot_output]
+                        lambda s=space: on_select(s),
+                            inputs=[],
+                            outputs=[info_output, app_py_content, screenshot_output]
                     )
 
+
         return demo
 
     except Exception as e: