Update app.py

app.py

@@ -1,5 +1,5 @@
-import torch  # Explicitly import torch to avoid import issues
-from transformers import MBartForConditionalGeneration, MBart50Tokenizer, AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
+from transformers import MBartForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM, pipeline
 import gradio as gr
 import requests
 import io
@@ -15,18 +15,18 @@ headers = {"Authorization": f"Bearer {hf_api_key}"}
 # Define the text-to-image model URL
 API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
 
-# Load the translation model and tokenizer
+# Use AutoTokenizer to avoid tokenizer mismatch warnings
 translation_model_name = "facebook/mbart-large-50-many-to-one-mmt"
-tokenizer = MBart50Tokenizer.from_pretrained(translation_model_name)
+tokenizer = AutoTokenizer.from_pretrained(translation_model_name)  # Use AutoTokenizer to avoid warnings
 translation_model = MBartForConditionalGeneration.from_pretrained(translation_model_name)
 
 # Load a text generation model from Hugging Face using accelerate for memory optimization
-text_generation_model_name = "EleutherAI/gpt-neo-2.7B"  # You can switch to "EleutherAI/gpt-j-6B" if available
+text_generation_model_name = "EleutherAI/gpt-neo-2.7B"
 text_tokenizer = AutoTokenizer.from_pretrained(text_generation_model_name)
 text_model = AutoModelForCausalLM.from_pretrained(
     text_generation_model_name, 
-    device_map="auto",  # Automatically allocate model layers to devices (requires accelerate)
-    torch_dtype=torch.float32  # Specify dtype to optimize memory usage
+    device_map="auto", 
+    torch_dtype=torch.float32
 )
 
 # Create a pipeline for text generation
@@ -35,14 +35,10 @@ text_generator = pipeline("text-generation", model=text_model, tokenizer=text_to
 # Function to generate an image using Hugging Face's text-to-image model
 def generate_image_from_text(translated_text):
     try:
-        # Send the translated text to the text-to-image model
         response = requests.post(API_URL, headers=headers, json={"inputs": translated_text})
-        
-        # Check if the response is successful
         if response.status_code != 200:
             return None, f"Error generating image: {response.text}"
         
-        # Read and return the generated image
         image_bytes = response.content
         image = Image.open(io.BytesIO(image_bytes))
         return image, None
@@ -52,7 +48,6 @@ def generate_image_from_text(translated_text):
 # Define the function to translate Tamil text, generate an image, and create a descriptive text
 def translate_generate_image_and_text(tamil_text):
     try:
-        # Step 1: Translate Tamil text to English
         tokenizer.src_lang = "ta_IN"
         inputs = tokenizer(tamil_text, return_tensors="pt")
         translated_tokens = translation_model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id["en_XX"])
@@ -61,7 +56,6 @@ def translate_generate_image_and_text(tamil_text):
         return f"Error during translation: {e}", None, None
 
     try:
-        # Step 2: Use the translated English text to generate an image directly
         image, error_message = generate_image_from_text(translated_text)
         if error_message:
             return translated_text, None, error_message
@@ -69,7 +63,6 @@ def translate_generate_image_and_text(tamil_text):
         return translated_text, None, f"Error during image generation: {e}"
 
     try:
-        # Step 3: Generate a descriptive English text using GPT-Neo based on the translated text
         descriptive_text = text_generator(translated_text, max_length=100, num_return_sequences=1, temperature=0.7, top_p=0.9)[0]['generated_text']
     except Exception as e:
         return translated_text, image, f"Error during text generation: {e}"