Update app.py

app.py

@@ -14,98 +14,80 @@ if hf_token is None:
 # Login to Hugging Face
 login(token=hf_token)
 
-# Define language codes for around 10 languages
-language_codes = {
-    "French": "fr",
-    "Spanish": "es",
-    "German": "de",
-    "Tamil": "ta",
-    "Hindi": "hi",
-    "Chinese": "zh",
-    "Russian": "ru",
-    "Japanese": "ja",
-    "Korean": "ko",
-    "Arabic": "ar",
-    "Portuguese": "pt",
-    "Italian": "it"
+# Define models for specific languages
+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
 }
 
-model_name = "Helsinki-NLP/opus-mt-mul-en"
-tokenizer = MarianTokenizer.from_pretrained(model_name)
-model = MarianMTModel.from_pretrained(model_name)
-translator = pipeline("translation", model=model, tokenizer=tokenizer)
+# 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 for translation
-def translate_text(input_text, src_lang):
-    try:
-        src_prefix = f">>{src_lang}<< " + input_text
-        translation = translator(src_prefix, max_length=40)
-        translated_text = translation[0]['translation_text']
-        return translated_text
-    except Exception as e:
-        return f"An error occurred: {str(e)}"
-
-# API credentials and endpoint for FLUX
+# FLUX model API settings
 flux_API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
 flux_headers = {"Authorization": f"Bearer {hf_token}"}
 
-# Function to generate image based on prompt
-def generate_image(prompt):
+# Function for translation, creative text generation, and image creation
+def translate_generate_image_and_text(input_text, src_lang_code):
     try:
-        response = requests.post(flux_API_URL, headers=flux_headers, json={"inputs": prompt})
-        if response.status_code == 200:
-            image_bytes = response.content
+        # Step 1: Get translator and translate text
+        translator = get_translator(src_lang_code)
+        translation = translator(input_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
             image = Image.open(io.BytesIO(image_bytes))
-            return image
         else:
-            print(f"Failed to get image: Status code {response.status_code}")
-            return None
-    except Exception as e:
-        print(f"An error occurred: {e}")
-        return None
+            image = None
 
-# API setup for Mistral model
-mistral_API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-v0.1"
-mistral_headers = {"Authorization": f"Bearer {hf_token}"}
+        return translated_text, creative_text, image
 
-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)}"
-
-# Main function to handle full workflow
-def translate_generate_image_and_text(input_text, src_lang):
-    # Step 1: Translate input text
-    translated_text = translate_text(input_text, language_codes[src_lang])
-
-    # Step 2: Generate an image
-    image = generate_image(translated_text)
-
-    # Step 3: Generate creative text based on the translation
-    creative_text = generate_creative_text(translated_text)
-
-    return translated_text, creative_text, image
+        return f"An error occurred: {str(e)}", None, None
 
-# Gradio interface
+# Gradio interface setup
 interface = gr.Interface(
     fn=translate_generate_image_and_text,
     inputs=[
         gr.Textbox(label="Enter text for translation"),
-        gr.Dropdown(choices=list(language_codes.keys()), label="Source Language")
+        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")
     ],
-    title="Multilingual Translation, Image, and Creative Text Generator",
-    description="Translates text from multiple languages to English, generates images, and creates creative text."
+    title="Multilingual Translation, Creative Content, and Image Generator",
+    description="Select a language and translate text to English, generate creative content, and produce an image."
 )
 
 # Launch the Gradio app