Update app.py

app.py

@@ -1,73 +1,47 @@
-import os
-import torch
+# app.py
+
 import gradio as gr
-from huggingface_hub import login
-from transformers import (
-    AutoTokenizer,
-    AutoModelForSeq2SeqLM,
-    GPT2LMHeadModel,
-    GPT2Tokenizer
-)
+from transformers import pipeline, MarianMTModel, MarianTokenizer
 from diffusers import StableDiffusionPipeline
+import torch
+import os
 
-# Authenticate via token
+# Get Hugging Face token from environment (in Hugging Face Spaces this is auto-populated from secrets)
 hf_token = os.getenv("HUGGINGFACE_TOKEN")
-if hf_token:
-    login(token=hf_token)
 
-# 📚 Tamil ↔ English Translation (Multilingual M2M100 model)
-trans_checkpoint = "Hemanth-thunder/english-tamil-mt"
-trans_tokenizer = AutoTokenizer.from_pretrained(trans_checkpoint)
-trans_model = AutoModelForSeq2SeqLM.from_pretrained(trans_checkpoint)
+# Load Translation Pipeline: Tamil → English using MarianMT
+translation_model_name = "Helsinki-NLP/opus-mt-ta-en"
+translator = pipeline("translation", model=translation_model_name)
 
-# 🧠 GPT-2 English Text Generation
-gpt_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-gpt_model = GPT2LMHeadModel.from_pretrained("gpt2")
-gpt_model.eval()
-
-# 🎨 Stable Diffusion
+# Load Stable Diffusion Pipeline: English → Image
 device = "cuda" if torch.cuda.is_available() else "cpu"
-sd_pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
+
+image_pipe = StableDiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4",
     use_auth_token=hf_token,
     torch_dtype=torch.float16 if device == "cuda" else torch.float32
 ).to(device)
 
-# Pipeline: Tamil → English → GPT-2 → Image
-def tam_to_image_pipeline(tamil_text):
-    # Translate Tamil → English
-    inputs = trans_tokenizer(tamil_text, return_tensors="pt", truncation=True)
-    translated_ids = trans_model.generate(**inputs, max_length=128)
-    english_text = trans_tokenizer.decode(translated_ids[0], skip_special_tokens=True)
-
-    # Generate additional English Description via GPT-2
-    input_ids = gpt_tokenizer.encode(english_text, return_tensors="pt")
-    with torch.no_grad():
-        gpt_output = gpt_model.generate(
-            input_ids,
-            max_length=60,
-            num_return_sequences=1,
-            no_repeat_ngram_size=2,
-            pad_token_id=gpt_tokenizer.eos_token_id
-        )
-    generated_text = gpt_tokenizer.decode(gpt_output[0], skip_special_tokens=True)
-
-    # Generate image from description
-    image = sd_pipe(generated_text).images[0]
-
-    return english_text, generated_text, image
+# Core function
+def translate_and_generate(tamil_text):
+    # Step 1: Translate Tamil to English
+    english_output = translator(tamil_text)[0]['translation_text']
+    
+    # Step 2: Generate Image from English Text
+    image = image_pipe(english_output).images[0]
+    
+    return english_output, image
 
 # Gradio UI
-interface = gr.Interface(
-    fn=tam_to_image_pipeline,
-    inputs=gr.Textbox(label="Enter Tamil Text"),
+iface = gr.Interface(
+    fn=translate_and_generate,
+    inputs=gr.Textbox(lines=2, label="Enter Tamil Text"),
     outputs=[
         gr.Textbox(label="Translated English Text"),
-        gr.Textbox(label="Generated Description"),
         gr.Image(label="Generated Image")
     ],
-    title="Tamil → Image Generator",
-    description="📘 Tamil to English (M2M100) → GPT‑2 → Image via Stable Diffusion"
+    title="Tamil-to-Image Generator 🌸",
+    description="Enter Tamil text. It will be translated to English and visualized using Stable Diffusion."
 )
 
-interface.launch()
+iface.launch()