Update app.py

app.py

@@ -6,14 +6,7 @@ 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")
@@ -25,23 +18,23 @@ if hf_token:
 else:
     raise ValueError("Hugging Face token not found in environment variables.")
 
-# Load the translation model and tokenizer (cached for faster loading)
+# 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
 model_name = "Helsinki-NLP/opus-mt-mul-en"
-tokenizer = MarianTokenizer.from_pretrained(model_name, cache_dir="./cache")
-model = MarianMTModel.from_pretrained(model_name, cache_dir="./cache")
+tokenizer = MarianTokenizer.from_pretrained(model_name)
+model = MarianMTModel.from_pretrained(model_name)
 
 # Create a translation pipeline
 translator = pipeline("translation", model=model, tokenizer=tokenizer)
 
-# 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)
+# Function for translation
 def translate_text(tamil_text):
     try:
         translation = translator(tamil_text, max_length=40)
@@ -50,68 +43,78 @@ def translate_text(tamil_text):
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
-# Asynchronous function to send payload and generate image
-async def generate_image_async(prompt):
+# 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):
     try:
-        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
+        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
+
     except Exception as e:
         print(f"An error occurred: {e}")
         return None
 
-# Generate creative text based on the translated text (with optimization for generation)
-def generate_creative_text(translated_text, max_length=50):
+# 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):
     input_ids = gpt_neo_tokenizer(translated_text, return_tensors='pt').input_ids
-    generated_text_ids = gpt_neo_model.generate(input_ids, max_length=max_length, num_return_sequences=1, do_sample=True, top_k=50)
+    generated_text_ids = gpt_neo_model.generate(input_ids, max_length=100)
     creative_text = gpt_neo_tokenizer.decode(generated_text_ids[0], skip_special_tokens=True)
     return creative_text
 
-# Handle the full workflow: translate, generate image, generate creative text
-async def translate_generate_image_and_text(tamil_text):
+# Function to handle the full workflow
+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 asynchronously
-    image = await generate_image_async(translated_text)
+    # Step 2: Generate an image based on the translated text
+    image = generate_image(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
 
-# 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
+# Create Gradio interface
 interface = gr.Interface(
-    fn=lambda tamil_text: asyncio.run(translate_generate_image_and_text(tamil_text)),
+    fn=translate_generate_image_and_text,
     inputs="text",
     outputs=["text", "text", "image"],
-    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
+    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()
-