import streamlit as st
from transformers import pipeline
from diffusers import StableDiffusionPipeline
import torch

@st.cache_resource
def load_all_models():
    # Load Tamil ➝ English translation model
    translation_pipeline = pipeline(
        "translation",
        model="ai4bharat/indictrans2-indic-en-dist-200M",
        tokenizer="ai4bharat/indictrans2-indic-en-dist-200M",
        src_lang="ta", tgt_lang="en",
        device=0 if torch.cuda.is_available() else -1,
    )

    # Load image generation model
    img_pipe = StableDiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-2-1",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        revision="fp16" if torch.cuda.is_available() else None,
    )
    img_pipe = img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")

    return translation_pipeline, img_pipe

def main():
    st.set_page_config(page_title="Tamil ➝ English ➝ Image", layout="centered")
    st.title("🌐 Tamil to English to Image Generator")

    tamil_text = st.text_area("Enter Tamil text:", height=150)

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

        with st.spinner("Loading models..."):
            translation_pipeline, img_pipe = load_all_models()

        with st.spinner("Translating Tamil to English..."):
            translated = translation_pipeline(tamil_text)[0]["translation_text"]
            st.success(f"🔤 Translated English: `{translated}`")

        with st.spinner("Generating image..."):
            image = img_pipe(prompt=translated).images[0]
            st.image(image, caption="🎨 AI-generated Image", use_column_width=True)

if __name__ == "__main__":
    main()