Update app.py

app.py

@@ -3,6 +3,7 @@ import requests
 from PIL import Image
 from io import BytesIO
 import os
+from bs4 import BeautifulSoup
 
 # Hugging Face API endpoint and model (Stable Diffusion in this case)
 HF_MODEL_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2"
@@ -10,7 +11,25 @@ HF_MODEL_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-d
 # Get API key securely from environment variable
 HF_API_KEY = os.getenv("HF_API_KEY")
 
-# Function to generate image based on topic
+# Function to perform image search using DuckDuckGo and return image URL
+def search_images(topic):
+    search_url = f"https://duckduckgo.com/html/?q={topic}+images"
+    
+    # Send the request to DuckDuckGo and retrieve search results
+    response = requests.get(search_url)
+    soup = BeautifulSoup(response.text, "html.parser")
+    
+    # Find all image links in the search results
+    image_urls = []
+    for img in soup.find_all("img", {"class": "tile--img__img"}):
+        img_url = img.get("src")
+        if img_url and img_url.startswith("http"):
+            image_urls.append(img_url)
+    
+    # Return the first 3 image URLs as potential suggestions
+    return image_urls[:3]
+
+# Function to generate image based on topic using Hugging Face API (Stable Diffusion)
 def generate_image(topic):
     # Prepare the payload for the Hugging Face API
     payload = {
@@ -31,16 +50,35 @@ def generate_image(topic):
     else:
         return "Error generating image. Please try again."
 
+# Function to combine image search and generation
+def generate_and_search_images(topic):
+    # Generate an image using Hugging Face (Stable Diffusion)
+    generated_image = generate_image(topic)
+    
+    # Perform image search using DuckDuckGo
+    search_results = search_images(topic)
+    
+    # Display the results
+    images = []
+    for url in search_results:
+        # Download the image from the URL
+        img_response = requests.get(url)
+        img = Image.open(BytesIO(img_response.content))
+        images.append(img)
+    
+    return generated_image, images
+
 # Build the Gradio interface
 interface = gr.Interface(
-    fn=generate_image,
+    fn=generate_and_search_images,
     inputs=gr.Textbox(label="Enter Topic (e.g., Dog, Space, Nature)"),
-    outputs=gr.Image(),
+    outputs=[gr.Image(label="Generated Image"), gr.Gallery(label="Suggested Images")],
     live=True,
-    title="Children's School Project Image Generator",
-    description="Enter a topic and generate an image related to it for school projects. Example topics: Dog, Space, Nature."
+    title="Children's School Project Image Generator and Image Search",
+    description="Enter a topic and generate an image related to it for school projects. Additionally, get suggested images based on the topic."
 )
 
 # Launch the app
 interface.launch(share=True)
 
+