Update app.py

app.py

@@ -4,24 +4,23 @@ from PIL import Image
 import gradio as gr
 from transformers import MarianMTModel, MarianTokenizer
 
-# Load the translation model
 model_name = "Helsinki-NLP/opus-mt-mul-en"
 model = MarianMTModel.from_pretrained(model_name)
 tokenizer = MarianTokenizer.from_pretrained(model_name)
 
-# Function to translate Tamil text to English
+
 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)
     return translation
 
-# Function to query Gemini API for creative text
 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": translated_text}]}]
+        "contents": [{"parts": [{"text": prompt}]}]
     }
     response = requests.post(f"{url}?key={gemini_api_key}", headers=headers, json=payload)
 
@@ -32,14 +31,15 @@ def query_gemini_api(translated_text, gemini_api_key):
     else:
         return f"Error: {response.status_code} - {response.text}"
 
-# Function to query the Hugging Face API for image generation
+
+
 def query_image(payload, huggingface_api_key):
     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
 
-# Main function to process input
+
 def process_input(tamil_input, gemini_api_key, huggingface_api_key):
     translated_output = translate_text(tamil_input)
     creative_output = query_gemini_api(translated_output, gemini_api_key)
@@ -47,13 +47,13 @@ def process_input(tamil_input, gemini_api_key, huggingface_api_key):
     image = Image.open(io.BytesIO(image_bytes))
     return translated_output, creative_output, image
 
-# Gradio interface
+
 iface = gr.Interface(
     fn=process_input,
     inputs=[
         gr.Textbox(label="Input Tamil Text"),
         gr.Textbox(label="Enter your Gemini API Key", type="password"),
-        gr.Textbox(label="Enter your Hugging Face API Key", type="password")
+        gr.Textbox(label="Enter your Hugging Face API Key", type="password"),
     ],
     outputs=[
         gr.Textbox(label="Translated Text"),
@@ -63,5 +63,4 @@ iface = gr.Interface(
     title="TRANSART🎨 BY Sakthi",
     description="Enter Tamil text to translate to English and generate an image based on the translated text."
 )
-
 iface.launch()