Update app.py

app.py

@@ -4,7 +4,7 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Load NLLB translation model
+# Load NLLB-200 translation model
 @st.cache_resource
 def load_translation_model():
     model_name = "facebook/nllb-200-distilled-600M"
@@ -12,20 +12,20 @@ def load_translation_model():
     model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
     return tokenizer, model
 
-# ✅ Load Stable Diffusion model for CPU
+# Load Stable Diffusion model
 @st.cache_resource
 def load_diffusion_model():
     pipe = StableDiffusionPipeline.from_pretrained(
         "CompVis/stable-diffusion-v1-4",
-        torch_dtype=torch.float32  # ✅ Important for CPU
+        torch_dtype=torch.float32  # CPU only
     )
-    pipe = pipe.to("cpu")  # ✅ Force to CPU
+    pipe = pipe.to("cpu")
     return pipe
 
-# Translate Tamil to English using NLLB
+# Translate Tamil to English using prompt-style method
 def translate_text(text, tokenizer, model):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    inputs["forced_bos_token_id"] = tokenizer.lang_code_to_id["eng_Latn"]
+    prompt = f"Translate this from Tamil to English: {text}"
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
     with torch.no_grad():
         translated_tokens = model.generate(**inputs, max_length=512)
     return tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
@@ -47,7 +47,6 @@ def main():
 
         with st.spinner("மொழிபெயர்ப்பு மற்றும் பட உருவாக்கம் நடைபெறுகிறது..."):
             tokenizer, model = load_translation_model()
-            tokenizer.src_lang = "tam_Taml"  # Tamil
             translated_text = translate_text(user_input, tokenizer, model)
             st.success(f"மொழிபெயர்ப்பு (ஆங்கிலம்): {translated_text}")