Update app.py

app.py

@@ -1,68 +1,59 @@
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer, pipeline
 import gradio as gr
-from transformers import pipeline
-from diffusers import StableDiffusionPipeline
-import torch
-import os
 
-# 1. Use Hugging Face token securely
-HF_TOKEN = os.getenv("HF_TOKEN", None)
+# Load the multilingual translation model (supports Tamil <-> English)
+model_name = "facebook/m2m100_418M"
+tokenizer = M2M100Tokenizer.from_pretrained(model_name)
+model = M2M100ForConditionalGeneration.from_pretrained(model_name)
 
-# 2. Set device
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Load text generation pipeline (you can replace "gpt2" with your preferred model)
+text_generator = pipeline("text-generation", model="gpt2")
 
-# 3. Load translator (Tamil → English using multilingual model)
-try:
-    translator = pipeline(
-        "translation",
-        model="Helsinki-NLP/opus-mt-mul-en",
-        use_auth_token=HF_TOKEN
-    )
-except Exception as e:
-    translator = None
-    print(f"Error loading translator: {e}")
+def translate_tamil_to_english(text):
+    if not text.strip():
+        return ""
+    tokenizer.src_lang = "ta"
+    encoded = tokenizer(text, return_tensors="pt")
+    generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id("en"))
+    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
-# 4. Load GPT2 for English text generation
-try:
-    generator = pipeline("text-generation", model="gpt2")
-except Exception as e:
-    generator = None
-    print(f"Error loading GPT2: {e}")
+def process_text(tamil_text):
+    if not tamil_text.strip():
+        return "Please enter Tamil text.", None, ""
 
-# 5. Load Stable Diffusion for image generation
-try:
-    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)
-except Exception as e:
-    image_pipe = None
-    print(f"Error loading Stable Diffusion: {e}")
-
-# 6. Full pipeline function with safe error handling
-def generate_image_from_tamil(tamil_text):
-    if not translator or not generator or not image_pipe:
-        return "Model load error", "Model load error", None
+    try:
+        # Step 1: Translate Tamil to English
+        translation = translate_tamil_to_english(tamil_text)
+    except Exception as e:
+        return f"Translation error: {str(e)}", None, ""
 
     try:
-        translated = translator(tamil_text, max_length=100)[0]['translation_text']
-        prompt = generator(translated, max_length=50, num_return_sequences=1)[0]['generated_text']
-        image = image_pipe(prompt).images[0]
-        return translated, prompt, image
+        # Step 2: Generate English text based on translation
+        generated = text_generator(translation, max_length=50, num_return_sequences=1)
+        generated_text = generated[0]['generated_text']
     except Exception as e:
-        return f"Translation/Image generation error: {str(e)}", "", None
+        return translation, None, f"Text generation error: {str(e)}"
+
+    # Step 3: Placeholder for image generation
+    # Replace with your own image generation model/function
+    image_url = "https://via.placeholder.com/512.png?text=Generated+Image"
+
+    # Step 4: Description based on generated text (you can use an image captioning model here)
+    description = f"Generated description: {generated_text}"
+
+    return translation, image_url, description
 
-# 7. Gradio UI
 iface = gr.Interface(
-    fn=generate_image_from_tamil,
-    inputs=gr.Textbox(lines=2, label="Enter Tamil Text"),
+    fn=process_text,
+    inputs=gr.Textbox(lines=2, label="Enter Tamil text"),
     outputs=[
         gr.Textbox(label="Translated English Text"),
-        gr.Textbox(label="Generated Prompt"),
-        gr.Image(label="Generated Image")
+        gr.Image(label="Generated Image"),
+        gr.Textbox(label="Image Description"),
     ],
-    title="Tamil to Image Generator",
-    description="Translate Tamil ➜ Generate English Text ➜ Create Image"
+    title="Tamil to English Translation + Text & Image Generation",
+    description="Enter Tamil text, get English translation, generated text, image, and description."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()