# app.py
import streamlit as st
from transformers import T5Tokenizer, T5ForConditionalGeneration
from diffusers import StableDiffusionPipeline
import torch

# Load translation model
@st.cache_resource
def load_translation_model():
    model_name = "nithalis/tamil_translation_t5"
    tokenizer = T5Tokenizer.from_pretrained(model_name)
    model = T5ForConditionalGeneration.from_pretrained(model_name)
    return tokenizer, model

# Load faster image model
@st.cache_resource
def load_diffusion_model():
    pipe = StableDiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-2-base",
        torch_dtype=torch.float32
    ).to("cpu")
    return pipe

# Translate Tamil to English
def translate_tamil_to_english(text, tokenizer, model):
    prompt = f"translate Tamil to English: {text}"
    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        output = model.generate(**inputs, max_length=128)
    return tokenizer.decode(output[0], skip_special_tokens=True)

# Main Streamlit app
def main():
    st.set_page_config(page_title="தமிழ்→Image", layout="centered")
    st.title("🌐 தமிழ் → ஆங்கிலம் → படம்")

    user_input = st.text_area("உங்கள் உரையை தமிழ் எழுத்துப்படியாக உள்ளிடவும்:", height=150)

    if st.button("உருவாக்கு"):
        if not user_input.strip():
            st.warning("உரை பதிவுசெய்யவும்")
            return
        with st.spinner("மொழிபெயர்ப்பு மற்றும் பட உருவாக்கம்..."):
            tok, mdl = load_translation_model()
            translation = translate_tamil_to_english(user_input, tok, mdl)
            st.success(f"ஆங்கில மொழிபெயர்ப்பு: {translation}")

            pipe = load_diffusion_model()
            img = pipe(translation, num_inference_steps=25, guidance_scale=7.5).images[0]
            st.image(img, caption="உருவாக்கப்பட்ட படம்", use_column_width=True)

if __name__ == "__main__":
    main()