Update app.py

app.py

@@ -10,21 +10,22 @@ from transformers import (
 )
 from diffusers import StableDiffusionPipeline
 
-# Authenticate with Hugging Face Token
+# Authenticate via token
 hf_token = os.getenv("HUGGINGFACE_TOKEN")
 if hf_token:
     login(token=hf_token)
 
-# Load Tamil to English Translation Model
-trans_tokenizer = AutoTokenizer.from_pretrained("nandhinivaradharajan14/tam-eng-translator")
-trans_model = AutoModelForSeq2SeqLM.from_pretrained("nandhinivaradharajan14/tam-eng-translator")
+# 📚 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 GPT-2 for English Text Generation
+# 🧠 GPT-2 English Text Generation
 gpt_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 gpt_model = GPT2LMHeadModel.from_pretrained("gpt2")
 gpt_model.eval()
 
-# Load Stable Diffusion
+# 🎨 Stable Diffusion
 device = "cuda" if torch.cuda.is_available() else "cpu"
 sd_pipe = StableDiffusionPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
@@ -32,31 +33,31 @@ sd_pipe = StableDiffusionPipeline.from_pretrained(
     torch_dtype=torch.float16 if device == "cuda" else torch.float32
 ).to(device)
 
-# Main function
+# Pipeline: Tamil → English → GPT-2 → Image
 def tam_to_image_pipeline(tamil_text):
-    # 1. Tamil to English Translation
-    inputs = trans_tokenizer(tamil_text, return_tensors="pt")
-    translated = trans_model.generate(**inputs)
-    english_text = trans_tokenizer.decode(translated[0], skip_special_tokens=True)
+    # 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)
 
-    # 2. Generate Descriptive Text using GPT-2
-    gpt_input = gpt_tokenizer.encode(english_text, return_tensors="pt")
+    # 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(
-            gpt_input,
-            max_length=50,
+            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)
 
-    # 3. Generate Image using Stable Diffusion
+    # Generate image from description
     image = sd_pipe(generated_text).images[0]
 
     return english_text, generated_text, image
 
-# Gradio Interface
+# Gradio UI
 interface = gr.Interface(
     fn=tam_to_image_pipeline,
     inputs=gr.Textbox(label="Enter Tamil Text"),
@@ -65,9 +66,8 @@ interface = gr.Interface(
         gr.Textbox(label="Generated Description"),
         gr.Image(label="Generated Image")
     ],
-    title="Tamil to Image Generator",
-    description="🔤 Tamil → English → GPT-2 Description → 🎨 Stable Diffusion Image Generator"
+    title="Tamil → Image Generator",
+    description="📘 Tamil to English (M2M100) → GPT‑2 → Image via Stable Diffusion"
 )
 
-# Launch app
 interface.launch()