Update app.py

app.py

@@ -1,74 +1,77 @@
+from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
 import requests
-import io
 from PIL import Image
+import io
 import gradio as gr
-from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
-import os
 
-# Load the translation model
+# 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-2 model and tokenizer (smaller and faster than GPT-Neo)
-gpt_model_name = "gpt2"
-gpt_tokenizer = AutoTokenizer.from_pretrained(gpt_model_name)
-gpt_model = AutoModelForCausalLM.from_pretrained(gpt_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": "HUGGINGFACE_API_KEY"}  # Replace with your API key
+
+# Query Hugging Face API to generate image
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.content
 
+# Translate Tamil text to English
 def translate_text(tamil_text):
-    inputs = translation_tokenizer(tamil_text, return_tensors="pt")
+    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
 
-def query_gpt_2(translated_text):
-    prompt = f"Continue the story based on the following text: {translated_text}"
-    inputs = gpt_tokenizer(prompt, return_tensors="pt")
-    outputs = gpt_model.generate(inputs['input_ids'], max_length=50, num_return_sequences=1)  # Reduced max_length for speed
-    creative_text = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return creative_text
+# Generate an image based on the translated text
+def generate_image(prompt):
+    image_bytes = query({"inputs": prompt})
+    image = Image.open(io.BytesIO(image_bytes))
+    return image
 
-def query_image(payload):
-    huggingface_api_key = os.getenv('HUGGINGFACE_API_KEY')
-    if not huggingface_api_key:
-        return "Error: Hugging Face API key not set."
+# 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
 
-    API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
-    headers = {"Authorization": f"Bearer {huggingface_api_key}"}
-    response = requests.post(API_URL, headers=headers, json=payload)
+# 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)
     
-    if response.status_code == 200:
-        return response.content
-    else:
-        return f"Error: {response.status_code} - {response.text}"
-
-def process_input(tamil_input):
-    try:
-        # Translate the input text
-        translated_output = translate_text(tamil_input)
-        
-        # Generate creative text using GPT-2
-        creative_output = query_gpt_2(translated_output)
-        
-        # Generate an image using Hugging Face's FLUX model
-        image_bytes = query_image({"inputs": translated_output})
-        image = Image.open(io.BytesIO(image_bytes))
-        
-        return translated_output, creative_output, image
-    except Exception as e:
-        return f"Error occurred: {str(e)}", "", None
+    # 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)
+    
+    return translated_text, creative_text, image
 
-# Create a Gradio interface
+# Create Gradio interface
 interface = gr.Interface(
-    fn=process_input,
-    inputs=[gr.Textbox(label="Input Tamil Text")],
+    fn=translate_generate_image_and_text,
+    inputs=gr.Textbox(label="Enter Tamil Text"),  # Input for Tamil text
     outputs=[
-        gr.Textbox(label="Translated Text"),
-        gr.Textbox(label="Creative Text"),
-        gr.Image(label="Generated Image")
+        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
     ],
-    title="TRANSART",
-    description="Enter Tamil text to translate to English and generate an image based on the translated text."
+    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."
 )
 
-interface.launch()
+# Launch Gradio app
+interface.launch()