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

@st.cache_resource(show_spinner=False)
def load_all_models():
    # Load IndicTrans2 model for Tamil→English
    indic_tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indictrans2-ta-en-dist-200M")
    indic_model = AutoModelForSeq2SeqLM.from_pretrained("ai4bharat/indictrans2-ta-en-dist-200M")
    
    # Load other models
    text_gen = pipeline("text-generation", model="gpt2", device=-1)
    img_pipe = StableDiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-2-base", 
        torch_dtype=torch.float32
    ).to("cpu")
    
    return indic_tokenizer, indic_model, text_gen, img_pipe

def translate_tamil(text, tokenizer, model):
    # Tokenize with batch processing
    inputs = tokenizer(
        text,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=128
    )
    
    # Generate translation
    generated = model.generate(
        **inputs,
        max_length=150,
        num_beams=5,
        early_stopping=True
    )
    
    # Decode with normalization
    return tokenizer.batch_decode(
        generated, 
        skip_special_tokens=True,
        clean_up_tokenization_spaces=True
    )[0]

def main():
    st.set_page_config(page_title="Tamil→EN→Image/Text", layout="centered")
    st.title("🌐 தமிழ் → English → Creative Text & Image")
    
    # Load models once
    tokenizer, model, text_gen, img_pipe = load_all_models()
    
    # Input with Tamil placeholder
    tamil_text = st.text_area(
        "தமிழ் உரை:", 
        height=150,
        placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்..."
    )

    if st.button("உருவாக்கு", type="primary"):
        if not tamil_text.strip():
            st.warning("தயவு செய்து உரையை உள்ளிடவும்.")
            return

        with st.spinner("மொழிபெயர்க்கிறது..."):
            eng = translate_tamil(tamil_text, tokenizer, model)
        
        with st.expander("🔤 Translation", expanded=True):
            st.success(eng)

        with st.spinner("உரை உருவாக்குதல்..."):
            creative = text_gen(
                f"Create a creative description about: {eng}",
                max_length=80,
                num_return_sequences=1
            )[0]["generated_text"]
        
        st.info("📝 Creative Text:")
        st.write(creative)

        with st.spinner("படத்தை உருவாக்குதல்..."):
            img = img_pipe(
                eng,
                num_inference_steps=35,
                guidance_scale=8.0
            ).images[0]
        
        st.image(img, caption="🎨 Generated Image")

if __name__ == "__main__":
    main()