app.py

import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
from diffusers import StableDiffusionPipeline
import torch

# Cache models for faster loading
@st.cache_resource
def load_all_models():
    # Translation model
    translation_model = AutoModelForSeq2SeqLM.from_pretrained(
        "ai4bharat/indictrans2-indic-en-dist-200M", trust_remote_code=True
    )
    translation_tokenizer = AutoTokenizer.from_pretrained(
        "ai4bharat/indictrans2-indic-en-dist-200M", trust_remote_code=True
    )
    translation_pipeline = pipeline(
        "text2text-generation", model=translation_model, tokenizer=translation_tokenizer
    )

    # Image generation model (Stable Diffusion)
    img_pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
    )
    img_pipe = img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
    
    return translation_pipeline, img_pipe

def main():
    st.title("📘 Tamil to English Translator & Image Generator")

    tamil_text = st.text_area("📝 Enter Tamil text (word or sentence)", height=100)

    if st.button("🔄 Translate & Generate Image"):
        if not tamil_text.strip():
            st.warning("Please enter some Tamil text.")
            return

        try:
            translation_pipeline, img_pipe = load_all_models()

            # Prepare translation input
            formatted_input = "<2en><|ta|>" + tamil_text.strip()
            translated = translation_pipeline(formatted_input, max_length=256)[0]["generated_text"]

            st.success("✅ English Translation:")
            st.write(translated)

            with st.spinner("🖼️ Generating image..."):
                image = img_pipe(translated).images[0]
                st.image(image, caption="🖼️ Generated from English text")

        except Exception as e:
            st.error(f"❌ Error: {str(e)}")

if __name__ == "__main__":
    main()