Update app.py

app.py

@@ -6,7 +6,14 @@ Original file is located at
 """
 
 import os
+import asyncio
 from huggingface_hub import login
+from transformers import MarianMTModel, MarianTokenizer, pipeline, AutoTokenizer, AutoModelForCausalLM
+import aiohttp
+import io
+from PIL import Image
+import matplotlib.pyplot as plt
+import gradio as gr
 
 # Retrieve the actual token from the environment variable
 hf_token = os.getenv("HF_TOKEN")
@@ -18,23 +25,23 @@ if hf_token:
 else:
     raise ValueError("Hugging Face token not found in environment variables.")
 
-# Import necessary libraries
-from transformers import MarianMTModel, MarianTokenizer, pipeline
-import requests
-import io
-from PIL import Image
-import matplotlib.pyplot as plt
-import gradio as gr
-
-# Load the translation model and tokenizer
+# Load the translation model and tokenizer (cached for faster loading)
 model_name = "Helsinki-NLP/opus-mt-mul-en"
-tokenizer = MarianTokenizer.from_pretrained(model_name)
-model = MarianMTModel.from_pretrained(model_name)
+tokenizer = MarianTokenizer.from_pretrained(model_name, cache_dir="./cache")
+model = MarianMTModel.from_pretrained(model_name, cache_dir="./cache")
 
 # Create a translation pipeline
 translator = pipeline("translation", model=model, tokenizer=tokenizer)
 
-# Function for translation
+# Load GPT-Neo model for creative text generation (cached)
+gpt_neo_tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M", cache_dir="./cache")
+gpt_neo_model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M", cache_dir="./cache")
+
+# API credentials and endpoint for image generation
+API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
+headers = {"Authorization": f"Bearer {hf_token}"}
+
+# Function for translation (batch translation for multiple inputs)
 def translate_text(tamil_text):
     try:
         translation = translator(tamil_text, max_length=40)
@@ -43,78 +50,68 @@ def translate_text(tamil_text):
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
-# API credentials and endpoint
-API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
-headers = {"Authorization": f"Bearer {hf_token}"}
-
-# Function to send payload and generate image
-def generate_image(prompt):
+# Asynchronous function to send payload and generate image
+async def generate_image_async(prompt):
     try:
-        response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
-
-        # Check if the response is successful
-        if response.status_code == 200:
-            print("API call successful, generating image...")
-            image_bytes = response.content
-
-            # Try opening the image
-            try:
-                image = Image.open(io.BytesIO(image_bytes))
-                return image
-            except Exception as e:
-                print(f"Error opening image: {e}")
-                return None
-        else:
-            print(f"Failed to get image: Status code {response.status_code}")
-            print("Response content:", response.text)  # Print response for debugging
-            return None
-
+        async with aiohttp.ClientSession() as session:
+            async with session.post(API_URL, headers=headers, json={"inputs": prompt}) as response:
+                if response.status == 200:
+                    print("API call successful, generating image...")
+                    image_bytes = await response.read()
+
+                    # Try opening the image
+                    try:
+                        image = Image.open(io.BytesIO(image_bytes))
+                        return image
+                    except Exception as e:
+                        print(f"Error opening image: {e}")
+                        return None
+                else:
+                    print(f"Failed to get image: Status code {response.status}")
+                    return None
     except Exception as e:
         print(f"An error occurred: {e}")
         return None
 
-# Display image
-def show_image(image):
-    if image:
-        plt.imshow(image)
-        plt.axis('off')  # Hide axes
-        plt.show()
-    else:
-        print("No image to display")
-
-# Load GPT-Neo model for creative text generation
-from transformers import AutoTokenizer, AutoModelForCausalLM
-gpt_neo_tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
-gpt_neo_model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
-
-# Function to generate creative text based on translated text
-def generate_creative_text(translated_text):
+# Generate creative text based on the translated text (with optimization for generation)
+def generate_creative_text(translated_text, max_length=50):
     input_ids = gpt_neo_tokenizer(translated_text, return_tensors='pt').input_ids
-    generated_text_ids = gpt_neo_model.generate(input_ids, max_length=100)
+    generated_text_ids = gpt_neo_model.generate(input_ids, max_length=max_length, num_return_sequences=1, do_sample=True, top_k=50)
     creative_text = gpt_neo_tokenizer.decode(generated_text_ids[0], skip_special_tokens=True)
     return creative_text
 
-# Function to handle the full workflow
-def translate_generate_image_and_text(tamil_text):
+# Handle the full workflow: translate, generate image, generate creative text
+async def translate_generate_image_and_text(tamil_text):
     # Step 1: Translate Tamil text to English
     translated_text = translate_text(tamil_text)
 
-    # Step 2: Generate an image based on the translated text
-    image = generate_image(translated_text)
+    # Step 2: Generate an image based on the translated text asynchronously
+    image = await generate_image_async(translated_text)
 
     # Step 3: Generate creative text based on the translated text
     creative_text = generate_creative_text(translated_text)
 
     return translated_text, creative_text, image
 
-# Create Gradio interface
+# Display image
+def show_image(image):
+    if image:
+        plt.imshow(image)
+        plt.axis('off')  # Hide axes
+        plt.show()
+    else:
+        print("No image to display")
+
+# Create Gradio interface with live updates for faster feedback
 interface = gr.Interface(
-    fn=translate_generate_image_and_text,
+    fn=lambda tamil_text: asyncio.run(translate_generate_image_and_text(tamil_text)),
     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."
+    title="Optimized 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.",
+    live=True  # Enables real-time outputs for faster feedback
 )
 
 # Launch Gradio app
 interface.launch()
+