Update app.py

app.py

@@ -1,83 +1,115 @@
-from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
+# Import necessary libraries
 import requests
-from PIL import Image
 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
-
-# Load the MarianMT model and tokenizer for translation (Tamil to English)
-model_name = "Helsinki-NLP/opus-mt-mul-en"
-translation_model = MarianMTModel.from_pretrained(model_name)
-translation_tokenizer = MarianTokenizer.from_pretrained(model_name)
-
-# Load GPT-Neo for creative text generation
-text_generation_model_name = "EleutherAI/gpt-neo-1.3B"
-text_generation_model = AutoModelForCausalLM.from_pretrained(text_generation_model_name)
-text_generation_tokenizer = AutoTokenizer.from_pretrained(text_generation_model_name)
-
-# Add padding token to GPT-Neo tokenizer if not present
-if text_generation_tokenizer.pad_token is None:
-    text_generation_tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-
-# Hugging Face API for FLUX.1 image generation
-API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
-headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_API_KEY')}"}
-
-# Query Hugging Face API to generate image
-def query(payload):
-    response = requests.post(API_URL, headers=headers, json=payload)
-    if response.status_code == 200:
-        return response.content
+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('HF_API_TOKEN')}"}
+
+# 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)}")
+            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:
-        return f"Error: {response.status_code} - {response.text}"
-
-# Translate Tamil text to English
-def translate_text(tamil_text):
-    inputs = translation_tokenizer(tamil_text, return_tensors="pt", padding=True, truncation=True)
-    translated_tokens = translation_model.generate(**inputs)
-    translation = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
-    return translation
-
-# Generate an image based on the translated text
-def generate_image(prompt):
-    image_bytes = query({"inputs": prompt})
-    if isinstance(image_bytes, str) and "Error" in image_bytes:
-        return image_bytes  # Return the error message if there's a problem with the API call
-    image = Image.open(io.BytesIO(image_bytes))
-    return image
-
-# Generate creative text based on the translated English text
-def generate_creative_text(translated_text):
-    inputs = text_generation_tokenizer(translated_text, return_tensors="pt", padding=True, truncation=True)
-    generated_tokens = text_generation_model.generate(**inputs, max_length=100)
-    creative_text = text_generation_tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
-    return creative_text
-
-# Function to handle the full workflow
-def translate_generate_image_and_text(tamil_text):
-    # Step 1: Translate Tamil to English
-    translated_text = translate_text(tamil_text)
-    
-    # Step 2: Generate an image from the translated text
-    image = generate_image(translated_text)
-    
-    # Step 3: Generate creative text from the translated text
-    creative_text = generate_creative_text(translated_text)
-    
+        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)
     return translated_text, creative_text, image
 
+
 # Create Gradio interface
 interface = gr.Interface(
-    fn=translate_generate_image_and_text,
-    inputs=gr.Textbox(label="Enter Tamil Text"),  # Input for Tamil text
-    outputs=[
-        gr.Textbox(label="Translated Text"),        # Output for translated text
-        gr.Textbox(label="Creative Generated Text"),# Output for creative text
-        gr.Image(label="Generated Image")           # Output for generated image
-    ],
+    fn=gradio_interface,
+    inputs="text",
+    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."
 )
 
 # Launch Gradio app
-interface.launch()
+if __name__ == "__main__":
+    interface.launch()