Update app.py

app.py

@@ -5,16 +5,32 @@ import gradio as gr
 from transformers import MarianMTModel, MarianTokenizer
 import os
 
+# Model and language settings
 model_name = "Helsinki-NLP/opus-mt-mul-en"
 model = MarianMTModel.from_pretrained(model_name)
 tokenizer = MarianTokenizer.from_pretrained(model_name)
 
-def translate_text(tamil_text):
-    inputs = tokenizer(tamil_text, return_tensors="pt")
+# Define a function for language translation
+def translate_text(input_text, language):
+    # Dictionary mapping the selected language to the respective language tag
+    language_map = {
+        "Tamil": "ta",
+        "French": "fr",
+        "Hindi": "hi",
+        "German": "de"
+    }
+    
+    # Prepare input with the appropriate language code for the model
+    lang_prefix = f">>{language_map[language]}<< "  # Add the prefix needed for the model
+    text_with_lang = lang_prefix + input_text
+    inputs = tokenizer(text_with_lang, return_tensors="pt", padding=True)
+    
+    # Perform translation
     translated_tokens = model.generate(**inputs)
     translation = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
     return translation
 
+# Function to query the Gemini API for creative text generation
 def query_gemini_api(translated_text, gemini_api_key):
     url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent"
     headers = {"Content-Type": "application/json"}
@@ -31,6 +47,7 @@ def query_gemini_api(translated_text, gemini_api_key):
     else:
         return f"Error: {response.status_code} - {response.text}"
 
+# Function to query the image generation API
 def query_image(payload):
     huggingface_api_key = os.getenv('HUGGINGFACE_API_KEY')
     API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
@@ -38,25 +55,28 @@ def query_image(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
     return response.content
 
-def process_input(tamil_input):
-    gemini_api_key = os.getenv('GEMINI_API_KEY')  
-    translated_output = translate_text(tamil_input)
-    creative_output = query_gemini_api(translated_output, gemini_api_key)
-    image_bytes = query_image({"inputs": translated_output})
-    image = Image.open(io.BytesIO(image_bytes))
+# Function to process the entire input
+def process_input(input_text, language):
+    gemini_api_key = os.getenv('GEMINI_API_KEY')  # Use your Gemini API key
+    translated_output = translate_text(input_text, language)  # Translate text from selected language to English
+    creative_output = query_gemini_api(translated_output, gemini_api_key)  # Generate creative text
+    image_bytes = query_image({"inputs": translated_output})  # Generate an image based on the translated text
+    image = Image.open(io.BytesIO(image_bytes))  # Convert bytes to image
+    
     return translated_output, creative_output, image
 
-
+# Create Gradio Interface
 iface = gr.Interface(
     fn=process_input,
-    inputs=[gr.Textbox(label="Input Tamil Text")],
+    inputs=[gr.Textbox(label="Enter text in the selected language"),
+            gr.Dropdown(choices=["Tamil", "French", "Hindi", "German"], label="Select Source Language")],
     outputs=[
-        gr.Textbox(label="Translated Text"),
+        gr.Textbox(label="Translated Text (English)"),
         gr.Textbox(label="Creative Text"),
         gr.Image(label="Generated Image")
     ],
-    title="TRANSART🎨 BY Sakthi",
-    description="Enter Tamil text to translate to English and generate an image based on the translated text."
+    title="TRANSART 🎨",
+    description="Select a language, enter text, and get the English translation along with creative content and image."
 )
 
 iface.launch()