Update app.py

app.py

@@ -1,52 +1,73 @@
 import os
+import torch
 import gradio as gr
-from transformers import pipeline
+from huggingface_hub import login
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    GPT2LMHeadModel,
+    GPT2Tokenizer
+)
 from diffusers import StableDiffusionPipeline
-import torch
 
+# Authenticate with Hugging Face Token
+hf_token = os.getenv("HUGGINGFACE_TOKEN")
+if hf_token:
+    login(token=hf_token)
 
-HF_TOKEN = os.getenv("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")
 
-# 1. Tamil to English translator (public model, no token required)
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-mul-en")
+# Load GPT-2 for English Text Generation
+gpt_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+gpt_model = GPT2LMHeadModel.from_pretrained("gpt2")
+gpt_model.eval()
 
-# 2. English text generator (GPT-2, public model, no token required)
-generator = pipeline("text-generation", model="gpt2")
-
-# 3. Stable Diffusion image generator (needs token)
+# Load Stable Diffusion
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
-image_pipe = StableDiffusionPipeline.from_pretrained(
-    "CompVis/stable-diffusion-v1-4",
-    use_auth_token=HF_TOKEN,
+sd_pipe = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    use_auth_token=hf_token,
     torch_dtype=torch.float16 if device == "cuda" else torch.float32
-)
-image_pipe = image_pipe.to(device)
-
-def generate_image_from_tamil(tamil_text):
-    # Translate Tamil → English
-    translated = translator(tamil_text, max_length=100)[0]['translation_text']
-    
-    # Generate English text from translated sentence
-    generated = generator(translated, max_length=50, num_return_sequences=1)[0]['generated_text']
-    
-    # Generate image from generated English text
-    image = image_pipe(generated).images[0]
-    
-    return translated, generated, image
-
-# Create Gradio interface
-iface = gr.Interface(
-    fn=generate_image_from_tamil,
-    inputs=gr.Textbox(lines=2, label="Enter Tamil Text"),
+).to(device)
+
+# Main function
+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)
+
+    # 2. Generate Descriptive Text using GPT-2
+    gpt_input = gpt_tokenizer.encode(english_text, return_tensors="pt")
+    with torch.no_grad():
+        gpt_output = gpt_model.generate(
+            gpt_input,
+            max_length=50,
+            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
+    image = sd_pipe(generated_text).images[0]
+
+    return english_text, generated_text, image
+
+# Gradio Interface
+interface = gr.Interface(
+    fn=tam_to_image_pipeline,
+    inputs=gr.Textbox(label="Enter Tamil Text"),
     outputs=[
         gr.Textbox(label="Translated English Text"),
-        gr.Textbox(label="Generated English Prompt"),
+        gr.Textbox(label="Generated Description"),
         gr.Image(label="Generated Image")
     ],
-    title="Tamil Text to English and Image Generator",
-    description="Translate Tamil to English, generate English text, and create image using Stable Diffusion."
+    title="Tamil to Image Generator",
+    description="🔤 Tamil → English → GPT-2 Description → 🎨 Stable Diffusion Image Generator"
 )
 
-# Launch Gradio app with public link
-iface.launch(share=True)
+# Launch app
+interface.launch()