fix: use PIL for image rendering

app.py

@@ -2,12 +2,15 @@ from __future__ import annotations
 from selenium import webdriver
 from typing import Iterable
 
-import re
-import os
 import gradio as Gradio
 from gradio.themes.base import Base
 from gradio.themes.utils import colors, fonts, sizes
 
+import re
+from PIL import Image
+from io import BytesIO
+from selenium.common.exceptions import WebDriverException
+
 theme = Gradio.themes.Monochrome(
     primary_hue="purple",
     secondary_hue="purple",
@@ -73,19 +76,31 @@ def run_imaginor(text: str):
     is_url = re.search(regex, text)
 
     if is_url:
-        driver = webdriver.Chrome()
-        driver.get(text)
-        screenshot = driver.save_screenshot('./site_screenshot.png')
-        driver.quit()
+        options = webdriver.ChromeOptions()
+        options.add_argument('--headless')
+        options.add_argument('--no-sandbox')
+        options.add_argument('--disable-dev-shm-usage')
+
+        try:
+            driver = webdriver.Chrome(options=options)
+            driver.get(text)
+            driver.implicitly_wait(10)
+            driver.set_window_size(1920, 1080)
+            screenshot = driver.get_screenshot_as_png()
+        except WebDriverException as e:
+            return [Image.new('RGB', (1, 1)), 'Website imagined by Imaginor, operation failed.']
+        finally:
+            if driver:
+                driver.quit()
 
-        return ['./site_screenshot.png', 'Website imagined by Imaginor, operation success.']
+        return [Image.open(BytesIO(screenshot)), 'Website imagined by Imaginor, operation success.']
     else:
         return [None, 'Please enter a valid URL of a website/host.']
 
 Gradio.Interface(
     run_imaginor,
     Gradio.Textbox(placeholder="Enter your URL here", label="Website URL / Endpoint"),
-    [Gradio.Image(type="filepath", label="Screenshot"),
+    [Gradio.Image(type="pil", label="Screenshot"),
     Gradio.Textbox(label="Text result", interactive=False)],
     title=title,
     description=description,