Update app.py

app.py

@@ -1,115 +1,93 @@
-# Import necessary libraries
+import os
+import concurrent.futures
+from huggingface_hub import login
+from transformers import MarianMTModel, MarianTokenizer, pipeline
 import requests
 import io
 from PIL import Image
-import matplotlib.pyplot as plt
-from transformers import MarianMTModel, MarianTokenizer, pipeline
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-import os  # For accessing environment variables
-
-# Constants for model names and API URLs
-class Constants:
-    TRANSLATION_MODEL_NAME = "Helsinki-NLP/opus-mt-mul-en"
-    IMAGE_GENERATION_API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
-    GPT_NEO_MODEL_NAME = "EleutherAI/gpt-neo-125M"
-    # Get the Hugging Face API token from environment variables
-    HEADERS = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_API_KEY')}"}
-
-# Translation Class
-class Translator:
-    def __init__(self):
-        self.tokenizer = MarianTokenizer.from_pretrained(Constants.TRANSLATION_MODEL_NAME)
-        self.model = MarianMTModel.from_pretrained(Constants.TRANSLATION_MODEL_NAME)
-        self.pipeline = pipeline("translation", model=self.model, tokenizer=self.tokenizer)
-
-    def translate(self, tamil_text):
-        """Translate Tamil text to English."""
-        try:
-            translation = self.pipeline(tamil_text, max_length=40)
-            return translation[0]['translation_text']
-        except Exception as e:
-            return f"Translation error: {str(e)}"
 
-
-# Image Generation Class
-class ImageGenerator:
-    def __init__(self):
-        self.api_url = Constants.IMAGE_GENERATION_API_URL
-
-    def generate(self, prompt):
-        """Generate an image based on the given prompt."""
-        try:
-            response = requests.post(self.api_url, headers=Constants.HEADERS, json={"inputs": prompt})
-            if response.status_code == 200:
-                image_bytes = response.content
-                return Image.open(io.BytesIO(image_bytes))
-            else:
-                print(f"Image generation failed: Status code {response.status_code}")
-                return None
-        except Exception as e:
-            print(f"Image generation error: {str(e)}")
+# Login with Hugging Face token
+hf_token = os.getenv("HF_TOKEN")
+if hf_token:
+    login(token=hf_token, add_to_git_credential=True)
+else:
+    raise ValueError("Hugging Face token 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
+
+# Translation function with reduced max_length
+def translate_text(text, src_lang, tgt_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']
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
+
+# Image generation with reduced resolution
+flux_API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
+flux_headers = {"Authorization": f"Bearer {hf_token}"}
+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
+            return image
+        else:
             return None
-
-
-# Creative Text Generation Class
-class CreativeTextGenerator:
-    def __init__(self):
-        self.tokenizer = AutoTokenizer.from_pretrained(Constants.GPT_NEO_MODEL_NAME)
-        self.model = AutoModelForCausalLM.from_pretrained(Constants.GPT_NEO_MODEL_NAME)
-
-    def generate(self, translated_text):
-        """Generate creative text based on translated text."""
-        input_ids = self.tokenizer(translated_text, return_tensors='pt').input_ids
-        generated_text_ids = self.model.generate(input_ids, max_length=100)
-        return self.tokenizer.decode(generated_text_ids[0], skip_special_tokens=True)
-
-
-# Main Application Class
-class TransArtApp:
-    def __init__(self):
-        self.translator = Translator()
-        self.image_generator = ImageGenerator()
-        self.creative_text_generator = CreativeTextGenerator()
-
-    def process(self, tamil_text):
-        """Handle the full workflow: translate, generate image, and creative text."""
-        translated_text = self.translator.translate(tamil_text)
-        image = self.image_generator.generate(translated_text)
-        creative_text = self.creative_text_generator.generate(translated_text)
-        return translated_text, creative_text, image
-
-
-# Function to display images
-def show_image(image):
-    """Display an image using matplotlib."""
-    if image:
-        plt.imshow(image)
-        plt.axis('off')  # Hide axes
-        plt.show()
-    else:
-        print("No image to display.")
-
-
-# Create an instance of the TransArt app
-app = TransArtApp()
-
-# Gradio interface function
-def gradio_interface(tamil_text):
-    """Interface function for Gradio."""
-    translated_text, creative_text, image = app.process(tamil_text)
+    except Exception as e:
+        print(f"Error in image generation: {e}")
+        return None
+
+# Creative text generation with reduced length
+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})
+        if response.status_code == 200:
+            return response.json()[0]['generated_text']
+        else:
+            return "Error generating creative text"
+    except Exception as e:
+        print(f"Error in creative text generation: {e}")
+        return None
+
+# 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
 
+# Language options for Gradio dropdown
+language_codes = {
+    "Tamil": "ta", "English": "en", "French": "fr", "Spanish": "es", "German": "de"
+}
 
-# Create Gradio interface
+# Gradio Interface
 interface = gr.Interface(
-    fn=gradio_interface,
-    inputs="text",
+    fn=translate_generate_image_and_text,
+    inputs=[
+        gr.Textbox(label="Enter text"),
+        gr.Dropdown(choices=list(language_codes.keys()), label="Source Language", default="Tamil"),
+        gr.Dropdown(choices=list(language_codes.keys()), label="Target Language", default="English"),
+    ],
     outputs=["text", "text", "image"],
-    title="Tamil to English Translation, Image Generation & Creative Text",
-    description="Enter Tamil text to translate to English, generate an image, and create creative text based on the translation."
+    title="Multilingual Translation, Image Generation & Creative Text",
+    description="Translate text between languages, generate images based on translation, and create creative text.",
 )
 
 # Launch Gradio app
-if __name__ == "__main__":
-    interface.launch()
+interface.launch()