Update app.py

app.py

@@ -4,14 +4,9 @@ import os
 import requests
 from typing import List, Dict, Union
 import traceback
-
-from selenium import webdriver
-from selenium.common.exceptions import WebDriverException
 from PIL import Image
 from io import BytesIO
 
-from playwright.sync_api import sync_playwright
-from selenium.webdriver.chrome.options import Options
 
 
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -98,37 +93,30 @@ def on_select(space):
     try:
         summary = summarize_space(space)
         app_content = get_app_py_content(space['id'])
-        screenshot = take_screenshot(space['url'])
+        thumbnail = get_space_thumbnail(space['id'])
         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, screenshot
+        return info, app_content, thumbnail
     except Exception as e:
         print(f"Error in on_select: {str(e)}")
         print(traceback.format_exc())
-        return f"오류가 발생했습니다: {str(e)}", "", Image.new('RGB', (1, 1))
+        return f"오류가 발생했습니다: {str(e)}", "", Image.new('RGB', (400, 300), color='lightgray')
 
-def take_screenshot(url):
+def get_space_thumbnail(space_id: str) -> Image.Image:
     try:
-        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=1080,720")
-        
-        driver = webdriver.Chrome(options=chrome_options)
-        
-        driver.get(url)
-        screenshot = driver.get_screenshot_as_png()
-        driver.quit()
-        
-        return Image.open(BytesIO(screenshot))
+        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')
     except Exception as e:
-        print(f"Screenshot error: {str(e)}")
-        return Image.new('RGB', (1, 1))
+        print(f"Error fetching thumbnail: {str(e)}")
+        return Image.new('RGB', (400, 300), color='lightgray')
 
 def create_ui():
     try:
@@ -158,14 +146,14 @@ def create_ui():
 
                 with gr.Column(scale=1):
                     info_output = gr.Textbox(label="Space 정보 및 요약", lines=20)
-                    screenshot_output = gr.Image(type="pil", label="Live 화면", height=360, width=540)
+                    thumbnail_output = gr.Image(type="pil", label="Space 썸네일", height=300, width=400)
                     app_py_content = gr.Code(language="python", label="app.py 내용")
 
                 for _, button, space in space_rows:
                     button.click(
                         lambda s=space: on_select(s),
                         inputs=[],
-                        outputs=[info_output, app_py_content, screenshot_output]
+                        outputs=[info_output, app_py_content, thumbnail_output]
                     )