Update app.py

app.py

@@ -1,73 +1,61 @@
 import gradio as gr
-import requests
-from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
+from transformers import MarianMTModel, MarianTokenizer, BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
 import torch
-import io
-import os
-from typing import Tuple
 
-# 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 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}"}
-
-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_tokenizer = MarianTokenizer.from_pretrained(translator_model)
-translator = MarianMTModel.from_pretrained(translator_model).to(device)
-
-# Text generation model
-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)
-    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 = 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})
-    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 = translate_tamil_to_english(tamil_text)
-    creative = generate_text(english)
-    image = generate_image(english)
-    return english, creative, image
-
-# Gradio UI using Blocks API
+# Load the Tamil-to-English translation model
+model_name = "Helsinki-NLP/opus-mt-ta-en"
+tokenizer = MarianTokenizer.from_pretrained(model_name)
+translation_model = MarianMTModel.from_pretrained(model_name)
+
+# Load the BLIP model for image captioning
+caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+
+def translate_tamil_to_english(tamil_text):
+    inputs = tokenizer(tamil_text, return_tensors="pt", padding=True)
+    translated = translation_model.generate(**inputs)
+    english_text = tokenizer.decode(translated[0], skip_special_tokens=True)
+    return english_text
+
+# Generate image using text (stub – replace with actual model if needed)
+def generate_image_from_text(text_prompt):
+    # Instead of using Stable Diffusion, just show a sample image
+    img = Image.new('RGB', (512, 512), color='lightblue')
+    return img
+
+def describe_image(image):
+    inputs = caption_processor(images=image, return_tensors="pt")
+    out = caption_model.generate(**inputs)
+    caption = caption_processor.decode(out[0], skip_special_tokens=True)
+    return caption
+
+def full_pipeline(tamil_text):
+    english_text = translate_tamil_to_english(tamil_text)
+    generated_image = generate_image_from_text(english_text)
+    description = describe_image(generated_image)
+    return english_text, generated_image, description
+
+# Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("## 🌍 Tamil to English | Text & Image Generator")
+    gr.Markdown("## Tamil to English → Image → Description")
+
+    with gr.Row():
+        tamil_input = gr.Textbox(label="Enter Tamil Text", lines=2, placeholder="உதாரணம்: ஒரு பூந்தோட்டத்தில் செருப்புகள் இருக்கின்றன")
+    
+    with gr.Row():
+        translate_btn = gr.Button("Translate and 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 Text")
+        description_output = gr.Textbox(label="Image Description")
 
-    english_output = gr.Textbox(label="🇬🇧 Translated English")
-    creative_output = gr.Textbox(label="✨ Generated Text")
-    image_output = gr.Image(label="🖼️ Generated Image", type="pil")
+    image_output = gr.Image(label="Generated Image")
 
-    generate_btn.click(fn=process_input, inputs=tamil_input, outputs=[english_output, creative_output, image_output])
+    translate_btn.click(
+        fn=full_pipeline,
+        inputs=tamil_input,
+        outputs=[english_output, image_output, description_output]
+    )
 
 demo.launch()