# app.py
import streamlit as st
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer, pipeline
from diffusers import DiffusionPipeline
import torch

@st.cache_resource(show_spinner=False)
def load_all_models():
    # Load translation model
    model_name = "facebook/m2m100_418M"
    tokenizer = M2M100Tokenizer.from_pretrained(model_name)
    model = M2M100ForConditionalGeneration.from_pretrained(model_name)

    # Load creative text model (smaller GPT-2)
    textgen = pipeline("text-generation", model="gpt2", device=-1)

    # Load lightweight image generation pipeline
    img_pipe = DiffusionPipeline.from_pretrained(
        "stabilityai/sdxl-lite", torch_dtype=torch.float32
    ).to("cpu")

    return tokenizer, model, textgen, img_pipe

def translate(text, tokenizer, model):
    tokenizer.src_lang = "ta"
    inputs = tokenizer(text, return_tensors="pt")
    output = model.generate(inputs["input_ids"], forced_bos_token_id=tokenizer.get_lang_id("en"), max_length=100)
    return tokenizer.decode(output[0], skip_special_tokens=True)

def generate_text(prompt, pipe):
    output = pipe(prompt, max_length=60, do_sample=True)[0]
    return output["generated_text"]

def main():
    st.set_page_config(page_title="Tamil to English → Creative → Image", layout="centered")
    st.title("🌐 தமிழ் ➝ English ➝ Creative Text + Image")

    tokenizer, model, textgen, img_pipe = load_all_models()

    tamil_text = st.text_area("தமிழ் உரையை உள்ளிடவும்:", height=130)

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

        with st.spinner("மொழிபெயர்ப்பு..."):
            english_text = translate(tamil_text, tokenizer, model)
        st.success(f"🔁 Translated: {english_text}")

        with st.spinner("உரையாக்கம்..."):
            creative_text = generate_text(english_text, textgen)
        st.info("📝 Creative Output:")
        st.write(creative_text)

        with st.spinner("படம் உருவாக்கப்படுகிறது..."):
            image = img_pipe(english_text).images[0]
        st.image(image, caption="🎨 Generated Image", use_column_width=True)

if __name__ == "__main__":
    main()