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

st.set_page_config(
    page_title="Tamil Creative Studio",
    page_icon="🇮🇳",
    layout="centered",
)

def load_css():
    st.markdown(
        """<style>
        .header {
            text-align: center;
            padding: 20px;
            background: #f9f9f9;
            border-radius: 10px;
            margin-bottom: 20px;
        }
        .header h1 { color: #cc0000; }
        .header p { color: #333; font-style: italic; }
        </style>""",
        unsafe_allow_html=True,
    )
    
@st.cache_resource
def load_all_models():
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
    from diffusers import StableDiffusionPipeline
    import torch

    # Translation model
    model_id = "ai4bharat/indictrans2-indic-en-dist-200M"
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_id, trust_remote_code=True)

    # Text generation model (simple GPT-2)
    text_gen = pipeline("text-generation", model="gpt2")

    # Image generation model (Stable Diffusion)
    img_pipe = StableDiffusionPipeline.from_pretrained(
        "CompVis/stable-diffusion-v1-4",
        revision="fp16",
        torch_dtype=torch.float16
    ).to("cuda")

    return tokenizer, model, text_gen, img_pipe


def translate_tamil(text, tokenizer, model):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
    outs = model.generate(**inputs, max_length=150, num_beams=5, early_stopping=True)
    return tokenizer.decode(outs[0], skip_special_tokens=True)

def main():
    load_css()
    st.markdown(
        '<div class="header"><h1>🌐 தமிழ் → English → Creative Studio</h1>'
        '<p>Translate Tamil text and generate creative content</p></div>',
        unsafe_allow_html=True
    )
    tokenizer, model, text_gen, img_pipe = load_all_models()
    tamil_text = st.text_area("**தமிழ் உரை:**", height=150, placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்...")

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

        with st.spinner("மொழிபெயர்க்கிறது..."):
            eng = translate_tamil(tamil_text, tokenizer, model)
        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)

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

if __name__ == "__main__":
    main()