Update app.py

app.py

@@ -1,9 +1,11 @@
-import gradio as gr
+from flask import Flask, request, send_file
 from selenium import webdriver
 from selenium.common.exceptions import WebDriverException
 from PIL import Image
 from io import BytesIO
 
+app = Flask(__name__)
+
 def take_screenshot(url):
     options = webdriver.ChromeOptions()
     options.add_argument('--headless')
@@ -17,6 +19,7 @@ def take_screenshot(url):
         wd.implicitly_wait(10)
         screenshot = wd.get_screenshot_as_png()
     except WebDriverException as e:
+        # If there's an error, return a blank image as a placeholder
         return Image.new('RGB', (1, 1))
     finally:
         if wd:
@@ -24,12 +27,21 @@ def take_screenshot(url):
 
     return Image.open(BytesIO(screenshot))
 
-iface = gr.Interface(
-    fn=take_screenshot,
-    inputs=gr.Textbox(label="Website URL", value="https://kargaranamir.github.io"),
-    outputs=gr.Image(type="pil", label="Screenshot", height=360, width=540),  # Adjust the image size here
-    title="Website Screenshot",
-    description="Take a screenshot of a website."
-)
+@app.route('/screenshot', methods=['GET'])
+def screenshot():
+    url = request.args.get('url')
+    if not url:
+        return "URL parameter is required.", 400
+
+    # Take the screenshot of the provided URL
+    image = take_screenshot(url)
+
+    # Save the screenshot to an in-memory file and return as response
+    img_io = BytesIO()
+    image.save(img_io, 'PNG')
+    img_io.seek(0)
+
+    return send_file(img_io, mimetype='image/png')
 
-iface.launch()
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860)