Update app.py

app.py

@@ -2,60 +2,78 @@ import requests
 import io
 from PIL import Image
 import gradio as gr
-from transformers import MarianMTModel, MarianTokenizer
+from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
 import os
 
+# Load the translation model
 model_name = "Helsinki-NLP/opus-mt-mul-en"
-model = MarianMTModel.from_pretrained(model_name)
-tokenizer = MarianTokenizer.from_pretrained(model_name)
+translation_model = MarianMTModel.from_pretrained(model_name)
+translation_tokenizer = MarianTokenizer.from_pretrained(model_name)
+
+# Load GPT-Neo model and tokenizer
+gpt_model_name = "EleutherAI/gpt-neo-1.3B"  # You can also use gpt-neo-2.7B if needed
+gpt_tokenizer = AutoTokenizer.from_pretrained(gpt_model_name)
+gpt_model = AutoModelForCausalLM.from_pretrained(gpt_model_name)
 
 def translate_text(tamil_text):
-    inputs = tokenizer(tamil_text, return_tensors="pt")
-    translated_tokens = model.generate(**inputs)
-    translation = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+    inputs = translation_tokenizer(tamil_text, return_tensors="pt")
+    translated_tokens = translation_model.generate(**inputs)
+    translation = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
     return translation
 
-def query_gemini_api(translated_text, gemini_api_key):
-    url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent"
-    headers = {"Content-Type": "application/json"}
-    prompt = f"Based on the following sentence, continue the story: {translated_text}"
-    payload = {
-        "contents": [{"parts": [{"text": prompt}]}]
-    }
-    response = requests.post(f"{url}?key={gemini_api_key}", headers=headers, json=payload)
-
-    if response.status_code == 200:
-        result = response.json()
-        creative_text = result['candidates'][0]['content']['parts'][0]['text']
-        return creative_text
-    else:
-        return f"Error: {response.status_code} - {response.text}"
+def query_gpt_neo(translated_text, max_words):
+    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=max_words, num_return_sequences=1)
+    creative_text = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return creative_text
 
 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."
+
     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)
-    return response.content
-
-def process_input(tamil_input):
-    gemini_api_key = os.getenv('GEMINI_API_KEY')  
-    translated_output = translate_text(tamil_input)
-    creative_output = query_gemini_api(translated_output, gemini_api_key)
-    image_bytes = query_image({"inputs": translated_output})
-    image = Image.open(io.BytesIO(image_bytes))
-    return translated_output, creative_output, image
+    
+    if response.status_code == 200:
+        return response.content
+    else:
+        return f"Error: {response.status_code} - {response.text}"
 
+def process_input(tamil_input, max_words):
+    try:
+        # Translate the input text
+        translated_output = translate_text(tamil_input)
+        
+        # Generate creative text using GPT-Neo
+        creative_output = query_gpt_neo(translated_output, max_words)
+        
+        # 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
 
+# Create a Gradio interface with interactive elements
 interface = gr.Interface(
     fn=process_input,
-    inputs=[gr.Textbox(label="Input Tamil Text")],
+    inputs=[
+        gr.Textbox(label="Input Tamil Text", placeholder="Enter your Tamil text here..."),
+        gr.Slider(label="Max Words for Creative Text", minimum=50, maximum=200, step=10, value=100)
+    ],
     outputs=[
         gr.Textbox(label="Translated Text"),
         gr.Textbox(label="Creative Text"),
         gr.Image(label="Generated Image")
     ],
-    title="TRANSART",
-    description="Enter Tamil text to translate to English and generate an image based on the translated text."
+    title="TRANSART - Multimodal AI App",
+    description="Enter Tamil text to translate to English, generate creative text, and produce an image based on the translated text.",
+    theme="compact",  # Use the 'compact' theme for a cleaner app look
+    layout="vertical"  # Arrange components vertically for better readability
 )
+
 interface.launch()