Update app.py

app.py

@@ -5,21 +5,52 @@
 from urllib.parse import urljoin
 from PIL import Image
 from io import BytesIO
-from selenium import webdriver
-from selenium.webdriver.chrome.options import Options
-from selenium.webdriver.common.by import By
-from selenium.webdriver.common.keys import Keys
-from webdriver_manager.chrome import ChromeDriverManager
-from time import sleep
-from selenium.webdriver.support.ui import WebDriverWait
-from selenium.webdriver.support import expected_conditions as EC
 
 # Download NLTK data
 nltk.download('punkt')
 
 # Global variable for article text
 article_text = ""
 
+# Add Hugging Face API token (for Image Generation)
+HF_API_TOKEN = st.secrets["hug"]["NOTIMPORTANT"]  # Streamlit secrets for sensitive data
+if not HF_API_TOKEN:
+    raise ValueError("Hugging Face API token is missing!")
+
+IMAGE_GEN_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2"
+
+# Cache heavy operations like article fetching
+@st.cache_data
+def fetch_article(url):
+    response = requests.get(url)
+    response.raise_for_status()
+    soup = BeautifulSoup(response.text, 'html.parser')
+    content = soup.find("div", class_="mw-parser-output")
+    paragraphs = content.find_all("p")
+    return "\n".join([p.text for p in paragraphs])
+
+def generate_image_cached(prompt, hf_api_token):
+    headers = {
+        "Authorization": f"Bearer {hf_api_token}",
+        "Content-Type": "application/json",
+    }
+    try:
+        # Make the request to Hugging Face API
+        response = requests.post(IMAGE_GEN_URL, headers=headers, json={"inputs": prompt})
+
+        # Handle successful response
+        if response.status_code == 200:
+            image_url = response.json()[0]["url"]
+            image_response = requests.get(image_url)
+            img = Image.open(BytesIO(image_response.content))
+            return img
+        else:
+            st.sidebar.error(f"Error generating image: {response.status_code}, {response.text}")
+            return None
+    except Exception as e:
+        st.sidebar.error(f"An error occurred: {str(e)}")
+        return None
+
 # Summarize text
 def summarize_text(text, num_sentences=10):
     sentences = nltk.sent_tokenize(text)
@@ -34,58 +65,6 @@ def extract_images(url):
     image_urls = [urljoin(url, img['src']) for img in img_tags if img.has_attr('src')]
     return image_urls
 
-# Function to fetch article text
-@st.cache_data
-def fetch_article(url):
-    try:
-        response = requests.get(url)
-        response.raise_for_status()
-        soup = BeautifulSoup(response.text, 'html.parser')
-        content = soup.find("div", class_="mw-parser-output")
-        paragraphs = content.find_all("p")
-        return "\n".join([p.text for p in paragraphs])
-    except Exception as e:
-        st.error(f"Error fetching the article: {e}")
-        return ""
-
-# Function to get a driver instance with headless Chrome
-def get_driver():
-    options = Options()
-    options.add_argument("--headless")
-    driver_path = ChromeDriverManager().install()
-    return webdriver.Chrome(executable_path=driver_path, options=options)
-
-# Function to generate image using Selenium and WebDriver Manager (using Chrome)
-def generate_image_using_selenium(prompt):
-    driver = get_driver()
-
-    # Open Craiyon website
-    driver.get("https://www.craiyon.com/")
-
-    input_box = driver.find_element(By.XPATH, "//input[@name='prompt']")
-    input_box.send_keys(prompt)
-    input_box.send_keys(Keys.RETURN)
-
-    # Wait for the image to be generated
-    try:
-        image_element = WebDriverWait(driver, 30).until(
-            EC.presence_of_element_located((By.XPATH, "//img[@class='generated-img']"))
-        )
-        image_url = image_element.get_attribute("src")
-
-        # Download image
-        img_data = requests.get(image_url).content
-        img = Image.open(BytesIO(img_data))
-
-    except Exception as e:
-        print(f"Error during image generation: {e}")
-        return None
-
-    finally:
-        driver.quit()
-
-    return img
-
 # Main function
 def main():
     st.title("Wikipedia Article Analyzer")
@@ -120,12 +99,11 @@ def main():
 
     if image_prompt:
         with st.spinner("Generating image..."):
-            # Call the Selenium image generation function
-            image = generate_image_using_selenium(image_prompt)
+            image = generate_image_cached(image_prompt, HF_API_TOKEN)
             if image:
                 st.image(image, caption=f"Generated image for: {image_prompt}", use_column_width=True)
             else:
-                st.sidebar.error("Sorry, there was an error generating the image.")
+                st.sidebar.write("Sorry, there was an error generating the image.")
 
 if __name__ == "__main__":
     main()