app.py

# app.py
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from diffusers import StableDiffusionPipeline
import torch

@st.cache_resource
def load_models():
    # Load translation model
    translation_model_name = "Helsinki-NLP/opus-mt-ta-en"
    translator_tokenizer = AutoTokenizer.from_pretrained(translation_model_name)
    translator_model = AutoModelForSeq2SeqLM.from_pretrained(translation_model_name)

    # Load text generation model
    textgen = pipeline("text-generation", model="gpt2")

    # Load stable diffusion (light version)
    image_model_name = "CompVis/stable-diffusion-v1-4"
    image_pipe = StableDiffusionPipeline.from_pretrained(
        image_model_name,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
    )
    if torch.cuda.is_available():
        image_pipe = image_pipe.to("cuda")

    return translator_tokenizer, translator_model, textgen, image_pipe

def translate(text, tokenizer, model):
    inputs = tokenizer.encode(text, return_tensors="pt", truncation=True, max_length=512)
    outputs = model.generate(inputs, max_length=512, num_beams=4)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def main():
    st.set_page_config(page_title="Tamil Text to English + Image + Text", layout="centered")
    st.title("🌐 தமிழ் உரையை ஆங்கிலம் + படம் + உரை ஆக மாற்றவும்")

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

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

        with st.spinner("மாடல்கள் ஏற்றப்படுகிறது..."):
            tokenizer, model, textgen, image_pipe = load_models()

        with st.spinner("மொழிபெயர்ப்பு நடைபெறுகிறது..."):
            english_text = translate(tamil_input, tokenizer, model)
            st.success(f"🔤 மொழிபெயர்ப்பு: {english_text}")

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

        with st.spinner("உரையை உருவாக்கப்படுகிறது..."):
            story = textgen(english_text, max_length=100, num_return_sequences=1)[0]["generated_text"]
            st.text_area("✍️ உருவாக்கப்பட்ட உரை", value=story, height=200)

if __name__ == "__main__":
    main()