Update app.py

app.py

@@ -1,6 +1,8 @@
 from huggingface_hub import InferenceClient
 import gradio as gr
 import os
+import re
+import requests
 import http.client
 import typing
 import urllib.request
@@ -11,6 +13,19 @@ vertexai.init(project=os.getenv('project_id'))
 model = GenerativeModel("gemini-1.0-pro-vision")
 client = InferenceClient("google/gemma-7b-it")
 
+def extract_image_urls(text):
+    url_regex = r"(https?:\/\/.*\.(?:png|jpg|jpeg|gif|webp|svg))"
+    image_urls = re.findall(url_regex, text, flags=re.IGNORECASE)
+    valid_image_url = ""
+    for url in image_urls:
+        try:
+            response = requests.head(url)  # Use HEAD request for efficiency
+            if response.status_code in range(200, 300) and 'image' in response.headers.get('content-type', ''):
+                valid_image_url = url
+        except requests.exceptions.RequestException:
+            pass  # Ignore inaccessible URLs
+    return valid_image_url
+
 def load_image_from_url(image_url: str) -> Image:
     with urllib.request.urlopen(image_url) as response:
         response = typing.cast(http.client.HTTPResponse, response)
@@ -47,6 +62,9 @@ def generate(
         seed=42,
     )
 
+    image = extract_image_urls(prompt)
+    if image:
+        prompt = prompt.replace(image, search(image))
     formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
     output = ""