app.py

import os
import requests
import streamlit as st
from PIL import Image
from io import BytesIO
from transformers import MBartForConditionalGeneration, MBart50Tokenizer
import time

# Fetch the API keys from Hugging Face Secrets
HUGGINGFACE_TOKEN = os.getenv("Hugging_face_token")
GROQ_API_KEY = os.getenv("Groq_api")

# Hugging Face API endpoint
HF_API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
hf_headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"}

# Groq API endpoint
groq_url = "https://api.groq.com/openai/v1/chat/completions"
groq_headers = {
    "Authorization": f"Bearer {GROQ_API_KEY}",
    "Content-Type": "application/json"
}

# Function to query Hugging Face model for image generation
def query_huggingface(payload):
    response = requests.post(HF_API_URL, headers=hf_headers, json=payload)
    if response.status_code != 200:
        st.error(f"Error: {response.status_code} - {response.text}")
        return None
    return response.content

# Function to generate text using Groq API
def generate_response(prompt):
    payload = {
        "model": "mixtral-8x7b-32768",
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ],
        "max_tokens": 100,
        "temperature": 0.7
    }
    response = requests.post(groq_url, json=payload, headers=groq_headers)
    if response.status_code == 200:
        result = response.json()
        return result['choices'][0]['message']['content']
    else:
        st.error(f"Error: {response.status_code} - {response.text}")
        return None

# Function to translate Tamil to English using MBart model
def translate_tamil_to_english(tamil_text):
    model_name = "facebook/mbart-large-50-many-to-one-mmt"
    model = MBartForConditionalGeneration.from_pretrained(model_name)
    tokenizer = MBart50Tokenizer.from_pretrained(model_name, src_lang="ta_IN")

    inputs = tokenizer(tamil_text, return_tensors="pt")
    translated = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id["en_XX"])
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    return translated_text

# Main function to generate text and image
def generate_image_and_text(user_input):
    with st.spinner("Generating results..."):
        time.sleep(2)  # Simulate some processing time

        # Translate Tamil to English
        english_input = translate_tamil_to_english(user_input)
        if not english_input:
            st.error("Error in translation.")
            return

        # Generate text description (100 tokens) and image prompt (30 tokens) using Groq API
        full_text_description = generate_response(english_input)
        if not full_text_description:
            st.error("Error in text generation.")
            return

        # Create image prompt based on the full text description
        image_prompt = generate_response(f"Create a concise image prompt from the following text: {full_text_description}")
        if not image_prompt:
            st.error("Error in generating image prompt.")
            return

        # Request an image based on the generated image prompt
        image_data = query_huggingface({"inputs": image_prompt})
        if not image_data:
            st.error("Error in image generation.")
            return

        # Display the results
        st.markdown("### Translated English Text:")
        st.write(english_input)

        st.markdown("### Generated Text Response:")
        st.write(full_text_description)

        try:
            # Load and display the image
            image = Image.open(BytesIO(image_data))
            st.image(image, caption="Generated Image", use_column_width=True)
        except Exception as e:
            st.error(f"Failed to display image: {e}")

# Streamlit interface
st.title("Multi-Modal Generator (Tamil to English)")
st.write("Enter a prompt in Tamil to generate both text and an image.")

# Input field for Tamil text
user_input = st.text_input("Enter Tamil text here:")

# Generate results when button is clicked
if st.button("Generate"):
    if user_input:
        generate_image_and_text(user_input)
    else:
        st.error("Please enter a Tamil text.")