Update app.py

app.py

@@ -1,60 +1,67 @@
 # app.py
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 from diffusers import StableDiffusionPipeline
 import torch
 import re
 
-# Initialize models (load once at startup)
-# Translation pipeline (Tamil to English)
-translator = pipeline(
-    "translation", 
-    model="Helsinki-NLP/opus-mt-ta-en",
-    device=0 if torch.cuda.is_available() else -1
-)
-
-# Text generation pipeline (English creative content)
-text_generator = pipeline(
-    "text-generation", 
-    model="gpt2-medium",
-    device=0 if torch.cuda.is_available() else -1
-)
+# Initialize models
+def load_models():
+    # Load translation model (NLLB - best for Tamil)
+    translator = pipeline(
+        "translation",
+        model="facebook/nllb-200-distilled-600M",
+        src_lang="tam_Taml",
+        tgt_lang="eng_Latn",
+        device=0 if torch.cuda.is_available() else -1
+    )
+    
+    # Text generation pipeline
+    text_generator = pipeline(
+        "text-generation", 
+        model="gpt2-medium",
+        device=0 if torch.cuda.is_available() else -1
+    )
+    
+    # Image generation pipeline
+    if torch.cuda.is_available():
+        image_pipe = StableDiffusionPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5",
+            torch_dtype=torch.float16,
+            revision="fp16"
+        ).to("cuda")
+    else:
+        image_pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+    
+    return translator, text_generator, image_pipe
 
-# Image generation pipeline (English text to image)
-if torch.cuda.is_available():
-    image_pipe = StableDiffusionPipeline.from_pretrained(
-        "runwayml/stable-diffusion-v1-5",
-        torch_dtype=torch.float16,
-        revision="fp16"
-    ).to("cuda")
-else:
-    image_pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+# Load models at startup
+translator, text_generator, image_pipe = load_models()
 
-# Clean and format generated text
 def clean_text(text):
     """Remove special characters and truncate to complete sentences"""
-    cleaned = re.split(r'(?<=[.!?])\s+', text)[0]  # Take first complete sentence
+    cleaned = re.split(r'(?<=[.!?])\s+', text)[0]
     return re.sub(r'[^a-zA-Z0-9,.!?\'"\- ]+', '', cleaned).strip()
 
 def process_content(tamil_input, creativity_level):
-    """Main processing pipeline: Translate → Generate Image → Create Text"""
     outputs = {}
     error = ""
     
     try:
         # Translate Tamil to English
-        translation = translator(tamil_input)
-        english_text = translation[0]['translation_text']
+        translation_result = translator(tamil_input)
+        english_text = translation_result[0]['translation_text']
         outputs["translation"] = english_text
         
-        # Generate image from translated text
+        # Generate image
         image = image_pipe(
             english_text, 
-            guidance_scale=creativity_level
+            guidance_scale=creativity_level,
+            num_inference_steps=30
         ).images[0]
         outputs["image"] = image
         
-        # Generate creative text from translation
+        # Generate creative text
         creative_output = text_generator(
             f"Create creative content about: {english_text}",
             max_length=150,
@@ -70,42 +77,54 @@ def process_content(tamil_input, creativity_level):
 
 # Gradio UI
 with gr.Blocks(theme=gr.themes.Soft()) as app:
-    gr.Markdown("# 🌐 Tamil Creative Content Generator")
-    gr.Markdown("Enter Tamil text → Get English translation + AI-generated image + creative English text")
+    gr.Markdown("# 🌐 தமிழ் உரை முதல் பட உருவாக்கம் (Tamil to Image Generator)")
+    gr.Markdown("தமிழில் உள்ளீடு செய்து → ஆங்கில மொழிபெயர்ப்பு + AI உருவம் + படைப்பு உரை பெறவும்")
     
     with gr.Row():
         with gr.Column():
             tamil_input = gr.Textbox(
-                label="தமிழ் உள்ளீடு (Tamil Input)", 
+                label="தமிழ் உள்ளீடு", 
                 placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்...",
                 lines=3
             )
             creativity = gr.Slider(
-                label="Creativity Level", 
+                label="படைப்பாற்றல் நிலை", 
                 minimum=1, maximum=10, value=7, step=1
             )
-            submit_btn = gr.Button("Generate Content")
+            submit_btn = gr.Button("உருவாக்கு")
             
         with gr.Column():
-            translation_box = gr.Textbox(label="English Translation")
-            creative_output = gr.Textbox(label="Creative English Content")
-            image_output = gr.Image(label="Generated Image")
-            error_output = gr.Textbox(label="System Messages", visible=True)
+            translation_box = gr.Textbox(label="ஆங்கில மொழிபெயர்ப்பு")
+            creative_output = gr.Textbox(label="படைப்பு உரை", lines=3)
+            image_output = gr.Image(label="உருவாக்கப்பட்ட படம்")
+            error_output = gr.Textbox(label="பிழை செய்திகள்", visible=True)
     
     examples = gr.Examples(
         examples=[
-            ["ஒரு மலையின் மீது உள்ள அழகிய கோட்டை", 8],
-            ["விண்மீன்கள் நிறைந்த இரவு வானத்தில் பறக்கும் யானை", 9],
-            ["தாவரங்கள் வளரும் புதிரான கோளம்", 10]
+            ["கடலின் அடியில் மறைந்திருக்கும் பழைய நகரம்", 8],
+            ["பனி படர்ந்த குளிர்காலத்தில் வெப்பமான காபி குடிக்கும் பழங்குடி பெண்", 7],
+            ["வேறு கிரகத்தில் இருந்து வந்த அறிவார்ந்த இயந்திரங்கள்", 9]
         ],
         inputs=[tamil_input, creativity]
     )
     
+    # Clear inputs button
+    clear_btn = gr.Button("துடைத்து துவக்கவும்")
+    
+    def clear_all():
+        return {"": ""}, "", "", None, ""
+    
     submit_btn.click(
         fn=process_content,
         inputs=[tamil_input, creativity],
         outputs=[{"translation": translation_box, "creative_text": creative_output, "image": image_output}, error_output]
     )
+    
+    clear_btn.click(
+        fn=clear_all,
+        inputs=[],
+        outputs=[tamil_input, translation_box, creative_output, image_output, error_output]
+    )
 
 if __name__ == "__main__":
-    app.launch(share=True)
+    app.launch()