Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-import concurrent.futures
 from huggingface_hub import login
 from transformers import MarianMTModel, MarianTokenizer, pipeline
 import requests
@@ -7,87 +6,107 @@ import io
 from PIL import Image
 import gradio as gr
 
-# Login with Hugging Face token
-hf_token = os.getenv("HUGGINGFACE_API_KEY")  # Updated variable name
-if hf_token:
-    login(token=hf_token, add_to_git_credential=True)
-else:
-    raise ValueError("Hugging Face token not found in environment variables.")
+# Set Hugging Face API key
+hf_token = os.getenv("HUGGINGFACE_API_KEY")
+if hf_token is None:
+    raise ValueError("Hugging Face API key not found in environment variables.")
 
-# Dynamic translation model loading
-def load_translation_model(src_lang, tgt_lang):
-    model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-    translator = pipeline("translation", model=model, tokenizer=tokenizer)
-    return translator
+# Login to Hugging Face
+login(token=hf_token)
 
-# Translation function with reduced max_length
-def translate_text(text, src_lang, tgt_lang):
+# 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"
+}
+
+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 for translation
+def translate_text(input_text, src_lang):
     try:
-        translator = load_translation_model(src_lang, tgt_lang)
-        translation = translator(text, max_length=20)  # Reduced max length for speed
-        return translation[0]['translation_text']
+        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)}"
 
-# Image generation with reduced resolution
+# API credentials and endpoint for FLUX
 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):
     try:
         response = requests.post(flux_API_URL, headers=flux_headers, json={"inputs": prompt})
         if response.status_code == 200:
-            image = Image.open(io.BytesIO(response.content))
-            image = image.resize((256, 256))  # Reduce resolution for faster processing
+            image_bytes = 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"Error in image generation: {e}")
+        print(f"An error occurred: {e}")
         return None
 
-# Creative text generation with reduced length
+# 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}"}
+
 def generate_creative_text(translated_text):
     try:
-        response = requests.post(mistral_API_URL, headers=mistral_headers, json={"inputs": translated_text, "max_length": 30})
+        response = requests.post(mistral_API_URL, headers=mistral_headers, json={"inputs": translated_text})
         if response.status_code == 200:
-            return response.json()[0]['generated_text']
+            creative_text = response.json()[0]['generated_text']
+            return creative_text
         else:
             return "Error generating creative text"
     except Exception as e:
-        print(f"Error in creative text generation: {e}")
-        return None
+        return f"An error occurred: {str(e)}"
 
-# Full workflow function with parallel processing
-def translate_generate_image_and_text(text, src_lang, tgt_lang):
-    translated_text = translate_text(text, src_lang, tgt_lang)
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        image_future = executor.submit(generate_image, translated_text)
-        creative_text_future = executor.submit(generate_creative_text, translated_text)
-        image = image_future.result()
-        creative_text = creative_text_future.result()
-    return translated_text, creative_text, image
+# 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])
 
-# Language options for Gradio dropdown
-language_codes = {
-    "Tamil": "ta", "English": "en", "French": "fr", "Spanish": "es", "German": "de"
-}
+    # 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)
 
-# Gradio Interface
+    return translated_text, creative_text, image
+
+# Gradio interface
 interface = gr.Interface(
     fn=translate_generate_image_and_text,
     inputs=[
-        gr.Textbox(label="Enter text"),
-        gr.Dropdown(choices=list(language_codes.keys()), label="Source Language", value="Tamil"),
-        gr.Dropdown(choices=list(language_codes.keys()), label="Target Language", value="English"),
+        gr.Textbox(label="Enter text for translation"),
+        gr.Dropdown(choices=list(language_codes.keys()), label="Source Language")
+    ],
+    outputs=[
+        gr.Textbox(label="Translated Text"),
+        gr.Textbox(label="Creative Text"),
+        gr.Image(label="Generated Image")
     ],
-    outputs=["text", "text", "image"],
-    title="Multilingual Translation, Image Generation & Creative Text",
-    description="Translate text between languages, generate images based on translation, and create creative text.",
+    title="Multilingual Translation, Image, and Creative Text Generator",
+    description="Translates text from multiple languages to English, generates images, and creates creative text."
 )
 
-# Launch Gradio app
+# Launch the Gradio app
 interface.launch()