Update app.py

app.py

@@ -7,12 +7,12 @@ import io
 import os
 from typing import Tuple
 
-# Load HF token
-HF_API_KEY = os.getenv("HF_API_KEY") or "your_hf_token_here"  # Replace this with your token if local
+# Load Hugging Face token
+HF_API_KEY = os.getenv("HF_API_KEY") or "your_hf_token_here"
 if not HF_API_KEY:
     raise ValueError("HF_API_KEY is not set.")
 
-# Hugging Face image model
+# Hugging Face inference API endpoint
 IMAGE_GEN_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
 HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}
 
@@ -20,51 +20,54 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Translation model (Tamil to English)
 translator_model = "Helsinki-NLP/opus-mt-mul-en"
-translator = MarianMTModel.from_pretrained(translator_model).to(device)
 translator_tokenizer = MarianTokenizer.from_pretrained(translator_model)
+translator = MarianMTModel.from_pretrained(translator_model).to(device)
 
 # Text generation model
-generator_model = "EleutherAI/gpt-neo-1.3B"
-generator = AutoModelForCausalLM.from_pretrained(generator_model).to(device)
-generator_tokenizer = AutoTokenizer.from_pretrained(generator_model)
-generator_tokenizer.pad_token = generator_tokenizer.eos_token
+text_model = "EleutherAI/gpt-neo-1.3B"
+text_tokenizer = AutoTokenizer.from_pretrained(text_model)
+text_generator = AutoModelForCausalLM.from_pretrained(text_model).to(device)
+text_tokenizer.pad_token = text_tokenizer.eos_token
 
+# Step 1: Tamil to English translation
 def translate_tamil_to_english(text: str) -> str:
     inputs = translator_tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
-    output = translator.generate(**inputs)
-    return translator_tokenizer.decode(output[0], skip_special_tokens=True)
+    outputs = translator.generate(**inputs)
+    return translator_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+# Step 2: Generate creative text
 def generate_text(prompt: str) -> str:
-    inputs = generator_tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(device)
-    output = generator.generate(**inputs, max_length=100, num_return_sequences=1)
-    return generator_tokenizer.decode(output[0], skip_special_tokens=True)
+    inputs = text_tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(device)
+    outputs = text_generator.generate(**inputs, max_length=100, num_return_sequences=1)
+    return text_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+# Step 3: Generate image
 def generate_image(prompt: str) -> Image.Image:
     response = requests.post(IMAGE_GEN_URL, headers=HEADERS, json={"inputs": prompt})
-    try:
-        if response.status_code == 200 and response.headers["content-type"].startswith("image"):
-            return Image.open(io.BytesIO(response.content))
-    except Exception as e:
-        print("Image generation failed:", e)
-    return Image.new("RGB", (300, 300), color="gray")
+    if response.status_code == 200 and response.headers.get("content-type", "").startswith("image"):
+        return Image.open(io.BytesIO(response.content))
+    else:
+        return Image.new("RGB", (512, 512), color="gray")
 
+# Master function
 def process_input(tamil_text: str) -> Tuple[str, str, Image.Image]:
-    english_text = translate_tamil_to_english(tamil_text)
-    creative_text = generate_text(english_text)
-    image = generate_image(english_text)
-    return english_text, creative_text, image
+    english = translate_tamil_to_english(tamil_text)
+    creative = generate_text(english)
+    image = generate_image(english)
+    return english, creative, image
 
-# Gradio app
+# Gradio UI using Blocks API
 with gr.Blocks() as demo:
-    gr.Markdown("## Tamil to English Translator with Text and Image Generator")
+    gr.Markdown("## 🌍 Tamil to English | Text & Image Generator")
 
-    tamil_input = gr.Textbox(label="Enter Tamil Text")
-    translate_btn = gr.Button("Translate & Generate")
+    with gr.Row():
+        tamil_input = gr.Textbox(label="📝 Enter Tamil Text", placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்...", lines=2)
+        generate_btn = gr.Button("Translate & Generate")
 
-    english_output = gr.Textbox(label="Translated English")
-    creative_output = gr.Textbox(label="Creative Text")
-    image_output = gr.Image(label="Generated Image")
+    english_output = gr.Textbox(label="🇬🇧 Translated English")
+    creative_output = gr.Textbox(label="✨ Generated Text")
+    image_output = gr.Image(label="🖼️ Generated Image", type="pil")
 
-    translate_btn.click(fn=process_input, inputs=tamil_input, outputs=[english_output, creative_output, image_output])
+    generate_btn.click(fn=process_input, inputs=tamil_input, outputs=[english_output, creative_output, image_output])
 
 demo.launch()