Update app.py

app.py

@@ -6,7 +6,11 @@ from typing import List, Dict, Union
 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
 
 
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -89,34 +93,42 @@ def get_app_py_content(space_id: str) -> str:
     except requests.RequestException:
         return f"Error fetching app.py content for space: {space_id}"
         
-def on_select(space):
+async def on_select(space):
     try:
         summary = summarize_space(space)
         app_content = get_app_py_content(space['id'])
-        thumbnail = get_space_thumbnail(space['id'])
+        screenshot = await 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"
         info += f"URL: {space['url']}\n\n"
         info += f"요약:\n{summary}"
-        return info, app_content, thumbnail
+        return info, app_content, screenshot
     except Exception as e:
         print(f"Error in on_select: {str(e)}")
         print(traceback.format_exc())
-        return f"오류가 발생했습니다: {str(e)}", "", Image.new('RGB', (400, 300), color='lightgray')
+        return f"오류가 발생했습니다: {str(e)}", "", Image.new('RGB', (1080, 720), color='lightgray')
+
+async def take_screenshot(url):
+    options = webdriver.ChromeOptions()
+    options.add_argument('--headless')
+    options.add_argument('--no-sandbox')
+    options.add_argument('--disable-dev-shm-usage')
 
-def get_space_thumbnail(space_id: str) -> Image.Image:
     try:
-        thumbnail_url = f"https://huggingface.co/spaces/{space_id}/resolve/main/thumbnail.png"
-        response = requests.get(thumbnail_url)
-        if response.status_code == 200:
-            return Image.open(BytesIO(response.content))
-        else:
-            print(f"Thumbnail not found for space: {space_id}")
-            return Image.new('RGB', (400, 300), color='lightgray')
+        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))
     except Exception as e:
-        print(f"Error fetching thumbnail: {str(e)}")
-        return Image.new('RGB', (400, 300), color='lightgray')
+        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:
@@ -146,18 +158,16 @@ def create_ui():
 
                 with gr.Column(scale=1):
                     info_output = gr.Textbox(label="Space 정보 및 요약", lines=20)
-                    thumbnail_output = gr.Image(type="pil", label="Space 썸네일", height=300, width=400)
+                    screenshot_output = gr.Image(type="pil", label="Space 스크린샷", height=360, width=540)
                     app_py_content = gr.Code(language="python", label="app.py 내용")
 
                 for _, button, space in space_rows:
                     button.click(
-                        lambda s=space: on_select(s),
+                        lambda s=space: asyncio.run(on_select(s)),
                         inputs=[],
-                        outputs=[info_output, app_py_content, thumbnail_output]
+                        outputs=[info_output, app_py_content, screenshot_output]
                     )
 
-
-
         return demo
 
     except Exception as e: