Update app.py

app.py

@@ -6,89 +6,95 @@ import io
 from PIL import Image
 import gradio as gr
 
-# Set Hugging Face API key
+# Retrieve the token from the environment variable
 hf_token = os.getenv("HUGGINGFACE_API_KEY")
-if hf_token is None:
-    raise ValueError("Hugging Face API key not found in environment variables.")
+if not hf_token:
+    raise ValueError("Hugging Face token not found in environment variables.")
 
-# Login to Hugging Face
-login(token=hf_token)
+login(token=hf_token, add_to_git_credential=True)
 
-# Define models for specific languages
+# Define available languages with their respective Helsinki model names
 language_models = {
-    "fra": "Helsinki-NLP/opus-mt-fr-en",
-    "spa": "Helsinki-NLP/opus-mt-es-en",
-    "tam": "Helsinki-NLP/opus-mt-tam-en",
-    "deu": "Helsinki-NLP/opus-mt-de-en",
-    "jpn": "Helsinki-NLP/opus-mt-ja-en",
-    "rus": "Helsinki-NLP/opus-mt-ru-en",
-    "kor": "Helsinki-NLP/opus-mt-ko-en",
-    "hin": "Helsinki-NLP/opus-mt-hi-en",
-    "ita": "Helsinki-NLP/opus-mt-it-en",
-    "por": "Helsinki-NLP/opus-mt-pt-en"
-    # Add more language models as needed
+    "Abkhaz": "Helsinki-NLP/opus-mt-abk-en",
+    "Arabic": "Helsinki-NLP/opus-mt-ar-en",
+    "Azerbaijani": "Helsinki-NLP/opus-mt-az-en",
+    "Bengali": "Helsinki-NLP/opus-mt-bn-en",
+    "Chinese": "Helsinki-NLP/opus-mt-zh-en",
+    "Danish": "Helsinki-NLP/opus-mt-da-en",
+    "Finnish": "Helsinki-NLP/opus-mt-fi-en",
+    "French": "Helsinki-NLP/opus-mt-fr-en",
+    "Hindi": "Helsinki-NLP/opus-mt-hi-en",
+    "Tamil": "Helsinki-NLP/opus-mt-mul-en"  # Using multilingual model for Tamil
 }
 
-# Function to get translator pipeline for specific language
-def get_translator(language_code):
-    model_name = language_models.get(language_code)
-    if model_name:
-        tokenizer = MarianTokenizer.from_pretrained(model_name)
-        model = MarianMTModel.from_pretrained(model_name)
-        return pipeline("translation", model=model, tokenizer=tokenizer)
-    else:
-        raise ValueError(f"No translation model found for language code '{language_code}'.")
+# Function to load a translation model dynamically
+def load_translation_pipeline(language):
+    model_name = language_models[language]
+    tokenizer = MarianTokenizer.from_pretrained(model_name)
+    model = MarianMTModel.from_pretrained(model_name)
+    return pipeline("translation", model=model, tokenizer=tokenizer)
 
-# FLUX model API settings
+# API credentials and endpoint for FLUX (Image generation)
 flux_API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
 flux_headers = {"Authorization": f"Bearer {hf_token}"}
 
-# Function for translation, creative text generation, and image creation
-def translate_generate_image_and_text(input_text, src_lang_code):
+# Function for translation
+def translate_text(tamil_text, language):
+    translator = load_translation_pipeline(language)
     try:
-        # Step 1: Get translator and translate text
-        translator = get_translator(src_lang_code)
-        translation = translator(input_text, max_length=40)
+        translation = translator(tamil_text, max_length=40)
         translated_text = translation[0]['translation_text']
-        
-        # Step 2: Generate creative text with Mistral model
-        mistral_API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-v0.1"
-        mistral_headers = {"Authorization": f"Bearer {hf_token}"}
-        mistral_response = requests.post(mistral_API_URL, headers=mistral_headers, json={"inputs": translated_text})
-        
-        if mistral_response.status_code == 200:
-            creative_text = mistral_response.json()[0]['generated_text']
-        else:
-            creative_text = "Error generating creative text"
-        
-        # Step 3: Generate an image with FLUX model
-        flux_response = requests.post(flux_API_URL, headers=flux_headers, json={"inputs": creative_text})
-        if flux_response.status_code == 200:
-            image_bytes = flux_response.content
+        return translated_text
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
+
+# Function to send payload and generate an image
+def generate_image(prompt):
+    try:
+        response = requests.post(flux_API_URL, headers=flux_headers, json={"inputs": prompt})
+        if response.status_code == 200:
+            image_bytes = response.content
             image = Image.open(io.BytesIO(image_bytes))
+            return image
         else:
-            image = None
+            return None
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return None
 
-        return translated_text, creative_text, image
+# Function for Mistral API call to generate creative text
+mistral_API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-v0.1"
+mistral_headers = {"Authorization": f"Bearer {hf_token}"}
 
+def generate_creative_text(translated_text):
+    try:
+        response = requests.post(mistral_API_URL, headers=mistral_headers, json={"inputs": translated_text})
+        if response.status_code == 200:
+            creative_text = response.json()[0]['generated_text']
+            return creative_text
+        else:
+            return "Error generating creative text"
     except Exception as e:
-        return f"An error occurred: {str(e)}", None, None
+        return None
 
-# Gradio interface setup
+# Function to handle the full workflow
+def translate_generate_image_and_text(tamil_text, language):
+    translated_text = translate_text(tamil_text, language)
+    image = generate_image(translated_text)
+    creative_text = generate_creative_text(translated_text)
+    return translated_text, creative_text, image
+
+# Create Gradio interface with language selection
 interface = gr.Interface(
     fn=translate_generate_image_and_text,
     inputs=[
-        gr.Textbox(label="Enter text for translation"),
-        gr.Textbox(label="Source Language Code", placeholder="e.g., 'fra' for French, 'spa' for Spanish")
-    ],
-    outputs=[
-        gr.Textbox(label="Translated Text"),
-        gr.Textbox(label="Creative Text"),
-        gr.Image(label="Generated Image")
+        gr.Textbox(label="Input Text in Source Language"),
+        gr.Dropdown(choices=list(language_models.keys()), label="Source Language")
     ],
-    title="Multilingual Translation, Creative Content, and Image Generator",
-    description="Select a language and translate text to English, generate creative content, and produce an image."
+    outputs=["text", "text", "image"],
+    title="Multilingual Translation, Image Generation & Creative Text",
+    description="Enter text to translate to English, generate an image, and create creative content based on the translation."
 )
 
-# Launch the Gradio app
+# Launch Gradio app
 interface.launch()