Update app.py

app.py

@@ -1,10 +1,11 @@
+import streamlit as st
+import base64
 import os
 import requests
 from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
 from PIL import Image, ImageDraw
 import io
 import torch
-import streamlit as st
 
 # Detect if GPU is available
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -13,64 +14,40 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model_name = "Helsinki-NLP/opus-mt-mul-en"
 translation_model = MarianMTModel.from_pretrained(model_name).to(device)
 translation_tokenizer = MarianTokenizer.from_pretrained(model_name, clean_up_tokenization_spaces=True)
+
+# Set padding token for the translation tokenizer
 translation_tokenizer.pad_token = translation_tokenizer.eos_token
 
 # Load GPT-Neo for creative text generation
 text_generation_model_name = "EleutherAI/gpt-neo-1.3B"
 text_generation_model = AutoModelForCausalLM.from_pretrained(text_generation_model_name).to(device)
 text_generation_tokenizer = AutoTokenizer.from_pretrained(text_generation_model_name, clean_up_tokenization_spaces=True)
-text_generation_tokenizer.pad_token = text_generation_tokenizer.eos_token
 
-# Access API keys from environment variables
-hf_api_key = os.getenv('HF_API_KEY')
-unsplash_api_key = os.getenv('UNSPLASH_API_KEY')
+# Set padding token for the text generation tokenizer
+text_generation_tokenizer.pad_token = text_generation_tokenizer.eos_token
 
-# Ensure API keys are set
-if hf_api_key is None:
-    raise ValueError("Hugging Face API key is not set. Please set it as an environment variable in Hugging Face Spaces.")
-if unsplash_api_key is None:
-    raise ValueError("Unsplash API key is not set. Please set it as an environment variable in Hugging Face Spaces.")
+# Retrieve Hugging Face API key from environment variable
+api_key = os.getenv('HF_API_KEY')
+if api_key is None:
+    raise ValueError("Hugging Face API key is not set. Please set it in your environment.")
 
-headers = {"Authorization": f"Bearer {hf_api_key}"}
+headers = {"Authorization": f"Bearer {api_key}"}
 
 # Define the API URL for image generation
 API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
 
-# Query Hugging Face API to generate image with error handling
-def query(payload):
-    try:
-        response = requests.post(API_URL, headers=headers, json=payload)
-        response.raise_for_status()  # Raise an error for bad status codes
-        return response.content
-    except requests.exceptions.RequestException as e:
-        st.error(f"API call failed: {e}")
-        return None
-
-# Function to get wallpaper from Unsplash using the API
-def get_random_wallpaper():
-    unsplash_url = "https://api.unsplash.com/photos/random"
-    params = {
-        "query": "wallpaper",
-        "client_id": unsplash_api_key  # Use Unsplash API key
-    }
-    try:
-        response = requests.get(unsplash_url, params=params)
-        response.raise_for_status()
-        image_url = response.json()["urls"]["full"]
-        return image_url
-    except requests.exceptions.RequestException as e:
-        st.error(f"Failed to fetch wallpaper: {e}")
-        return None
-
-# Function to set background image from Unsplash
-def set_background_image_from_unsplash():
-    image_url = get_random_wallpaper()
-    if image_url:
+# Function to set background image
+def set_background_image(uploaded_file):
+    if uploaded_file is not None:
+        # Read the uploaded image file and encode it in base64
+        image_bytes = uploaded_file.read()
+        b64 = base64.b64encode(image_bytes).decode()
+        # Set the CSS style to set the background image
         st.markdown(
             f"""
             <style>
             .stApp {{
-                background-image: url({image_url});
+                background-image: url(data:image/png;base64,{b64});
                 background-size: cover;
                 background-position: center;
                 height: 100vh;
@@ -81,6 +58,16 @@ def set_background_image_from_unsplash():
             unsafe_allow_html=True
         )
 
+# Query Hugging Face API to generate image with error handling
+def query(payload):
+    try:
+        response = requests.post(API_URL, headers=headers, json=payload)
+        response.raise_for_status()  # Raise an error for bad status codes
+        return response.content
+    except requests.exceptions.RequestException as e:
+        st.error(f"API call failed: {e}")
+        return None
+
 # Translate Tamil text to English
 def translate_text(tamil_text):
     inputs = translation_tokenizer(tamil_text, return_tensors="pt", padding=True, truncation=True).to(device)
@@ -88,18 +75,12 @@ def translate_text(tamil_text):
     translation = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
     return translation
 
-# Generate creative text based on the translated English text
-def generate_creative_text(translated_text):
-    inputs = text_generation_tokenizer(translated_text, return_tensors="pt", padding=True, truncation=True).to(device)
-    generated_tokens = text_generation_model.generate(**inputs, max_length=100)
-    creative_text = text_generation_tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
-    return creative_text
-
 # Generate an image based on the translated text with error handling
 def generate_image(prompt):
     image_bytes = query({"inputs": prompt})
 
     if image_bytes is None:
+        # Return a blank image with an error message
         error_img = Image.new('RGB', (300, 300), color=(255, 0, 0))
         d = ImageDraw.Draw(error_img)
         d.text((10, 150), "Image Generation Failed", fill=(255, 255, 255))
@@ -110,24 +91,41 @@ def generate_image(prompt):
         return image
     except Exception as e:
         st.error(f"Error loading image: {e}")
+        # Return an error image in case of failure
         error_img = Image.new('RGB', (300, 300), color=(255, 0, 0))
         d = ImageDraw.Draw(error_img)
         d.text((10, 150), "Invalid Image Data", fill=(255, 255, 255))
         return error_img
 
+# Generate creative text based on the translated English text
+def generate_creative_text(translated_text):
+    inputs = text_generation_tokenizer(translated_text, return_tensors="pt", padding=True, truncation=True).to(device)
+    generated_tokens = text_generation_model.generate(**inputs, max_length=100)
+    creative_text = text_generation_tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+    return creative_text
+
 # Function to handle the full workflow
 def translate_generate_image_and_text(tamil_text):
+    # Step 1: Translate Tamil to English
     translated_text = translate_text(tamil_text)
+
+    # Step 2: Generate an image from the translated text
     image = generate_image(translated_text)
+
+    # Step 3: Generate creative text from the translated text
     creative_text = generate_creative_text(translated_text)
+
     return translated_text, creative_text, image
 
 # Streamlit app title and subtitle
 st.markdown("<h1 style='text-align: center;'>A Multimodal Application</h1>", unsafe_allow_html=True)
 st.markdown("<h3 style='text-align: center;'>Language Translation and Image Synthesis</h3>", unsafe_allow_html=True)
 
-# Automatically set background from Unsplash
-set_background_image_from_unsplash()
+# Upload image for background
+uploaded_file = st.file_uploader("Upload a JPEG image for the background", type=["jpg", "jpeg"])
+
+# Set the background image if uploaded
+set_background_image(uploaded_file)
 
 # User input for Tamil text
 tamil_input = st.text_area("Enter Tamil Text", height=150)
@@ -137,13 +135,13 @@ if st.button("Generate"):
     if tamil_input:
         with st.spinner("Generating..."):
             translated_output, creative_text_output, generated_image_output = translate_generate_image_and_text(tamil_input)
-
+        
         st.subheader("Translated Text")
         st.write(translated_output)
-
+        
         st.subheader("Creative Generated Text")
         st.write(creative_text_output)
-
+        
         st.subheader("Generated Image")
         st.image(generated_image_output, caption="Generated Image", use_column_width=True)
     else: