app.py

# app.py
import streamlit as st
from transformers import pipeline
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch
import re
import os
from huggingface_hub import login
import time
import io

# Set page config
st.set_page_config(
    page_title="தமிழ் உரை முதல் பட உருவாக்கம்",
    page_icon="🖼️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Get Hugging Face token from environment variable
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    st.error("Hugging Face token not found in environment variables!")
    st.stop()
else:
    login(token=HF_TOKEN)

# Global flag for safety checker
SAFETY_CHECK_ENABLED = True

# Cache model loading
@st.cache_resource(show_spinner=False)
def load_models():
    with st.spinner("மாதிரிகள் ஏற்றப்படுகின்றன... இது சில நிமிடங்கள் எடுக்கலாம்"):
        # Translation model: Tamil → English
        translator = pipeline(
            "translation",
            model="facebook/nllb-200-distilled-600M",
            src_lang="tam_Taml",
            tgt_lang="eng_Latn",
            device=0 if torch.cuda.is_available() else -1,
            token=HF_TOKEN
        )

        # Text generation model
        text_generator = pipeline(
            "text-generation",
            model="gpt2-medium",
            device=0 if torch.cuda.is_available() else -1,
            token=HF_TOKEN
        )

        # Stable Diffusion for image generation
        if torch.cuda.is_available():
            image_pipe = StableDiffusionPipeline.from_pretrained(
                "runwayml/stable-diffusion-v1-5",
                torch_dtype=torch.float16,
                revision="fp16",
                token=HF_TOKEN,
                safety_checker=None if not SAFETY_CHECK_ENABLED else None
            )
            image_pipe.scheduler = DPMSolverMultistepScheduler.from_config(image_pipe.scheduler.config)
            image_pipe = image_pipe.to("cuda")
            image_pipe.enable_attention_slicing()
        else:
            image_pipe = StableDiffusionPipeline.from_pretrained(
                "runwayml/stable-diffusion-v1-5",
                token=HF_TOKEN
            )
            image_pipe.enable_attention_slicing()

        return translator, text_generator, image_pipe

# Load models
try:
    translator, text_generator, image_pipe = load_models()
except Exception as e:
    st.error(f"மாதிரிகள் ஏற்றுவதில் தோல்வி: {str(e)}")
    st.stop()

# Clean generated text
def clean_text(text):
    cleaned = re.sub(r'[^a-zA-Z0-9,.!?\'"\- ]+', '', text).strip()
    sentences = re.split(r'(?<=[.!?])\s+', cleaned)
    return ' '.join(sentences[:2])  # return first 2 sentences

# Main processing function
def process_content(tamil_input, creativity_level):
    try:
        # Track processing time
        start_time = time.time()
        
        # Translation
        with st.spinner("மொழிபெயர்ப்பு... (Translating)"):
            translation_result = translator(tamil_input)
            english_text = translation_result[0]['translation_text']
        
        # Image generation
        with st.spinner("படம் உருவாக்கப்படுகிறது... (Generating Image)"):
            image = image_pipe(
                english_text,
                guidance_scale=7 + creativity_level,  # Range 8-17
                num_inference_steps=30 + creativity_level * 2,
                height=512,
                width=512
            ).images[0]
        
        # Text generation
        with st.spinner("உரை உருவாக்கப்படுகிறது... (Generating Text)"):
            creative_output = text_generator(
                f"Create creative content about: {english_text}",
                max_length=150,
                temperature=creativity_level / 10,
                num_return_sequences=1
            )
            cleaned_text = clean_text(creative_output[0]['generated_text'])
        
        # Calculate processing time
        proc_time = time.time() - start_time
        
        return english_text, cleaned_text, image, f"⏱️ செயலாக்க நேரம்: {proc_time:.1f} வினாடிகள்", ""
    
    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        return "", "", None, "", "⚠️ GPU மெமரி நிரம்பிவிட்டது! தயவுசெய்து உள்ளீட்டை குறைக்கவும் அல்லது படைப்புத்திறன் அளவை குறைக்கவும்"
    except Exception as e:
        return "", "", None, "", f"⚠️ பிழை: {str(e)}"

# Initialize session state
if 'tamil_input' not in st.session_state:
    st.session_state.tamil_input = ""
if 'creativity' not in st.session_state:
    st.session_state.creativity = 7
if 'outputs' not in st.session_state:
    st.session_state.outputs = []

# Custom CSS for Tamil font support
def local_css(file_name):
    try:
        with open(file_name) as f:
            st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
    except:
        st.markdown("""
        <style>
        @import url('https://fonts.googleapis.com/css2?family=Catamaran:wght@400;700&family=Hind+Madurai:wght@400;700&display=swap');
        
        /* Apply Tamil font to specific elements */
        body, .stTextArea>label, .stSlider>label, .stButton>button, .stSelectbox>label {
            font-family: 'Catamaran', 'Hind Madurai', sans-serif !important;
        }
        
        /* Custom styling */
        .stTextInput input, .stTextArea textarea {
            border: 2px solid #4CAF50 !important;
        }
        
        .stButton>button {
            background-color: #4CAF50 !important;
            color: white !important;
            font-weight: bold;
        }
        
        .stSlider>div>div>div>div {
            background-color: #4CAF50 !important;
        }
        </style>
        """, unsafe_allow_html=True)

# Apply CSS
local_css("style.css")

st.title("🌐 தமிழ் உரை முதல் பட உருவாக்கம்")
st.markdown("தமிழில் உள்ளீடு செய்து → ஆங்கில மொழிபெயர்ப்பு + AI உருவம் + படைப்பு உரை பெறவும்")

# Sidebar with examples and info
with st.sidebar:
    st.header("உதாரணங்கள்")
    examples = [
        ("கடலின் அடியில் மறைந்திருக்கும் பழைய நகரம்", 8),
        ("பனி படர்ந்த குளிர்காலத்தில் வெப்பமான காபி குடிக்கும் பழங்குடி பெண்", 7),
        ("வேறு கிரகத்தில் இருந்து வந்த அறிவார்ந்த இயந்திரங்கள்", 9)
    ]
    
    for text, creativity in examples:
        if st.button(text, use_container_width=True):
            st.session_state.tamil_input = text
            st.session_state.creativity = creativity
    
    st.divider()
    st.header("விவரங்கள்")
    st.markdown("""
    - **மொழிபெயர்ப்பு மாதிரி**: Facebook NLLB-200 (Tamil → English)
    - **உரை உருவாக்கம்**: GPT-2 Medium
    - **பட உருவாக்கம்**: Stable Diffusion v1.5
    """)
    st.divider()
    st.markdown("""
    **அறிவுறுத்தல்கள்**:
    1. தமிழில் உங்கள் யோசனையை உள்ளிடவும்
    2. படைப்புத்திறன் அளவை சரிசெய்யவும் (1-10)
    3. "உருவாக்கு" பொத்தானை அழுத்தவும்
    4. உங்கள் படம் மற்றும் உரை விளைவுகளைப் பாருங்கள்!
    """)

# Main content
with st.form("input_form"):
    col1, col2 = st.columns([3, 1])
    
    with col1:
        tamil_input = st.text_area(
            "தமிழ் உள்ளீடு",
            value=st.session_state.tamil_input,
            placeholder="உதாரணம்: பனி படர்ந்த குளிர்காலத்தில் வெப்பமான காபி குடிக்கும் பழங்குடி பெண்",
            height=150
        )
    
    with col2:
        creativity = st.slider(
            "படைப்பாற்றல் நிலை",
            min_value=1, max_value=10, value=st.session_state.creativity, step=1,
            help="அதிக எண் = அதிக புதுமை ஆனால் குறைந்த துல்லியம்"
        )
        submit_btn = st.form_submit_button("உருவாக்கு", use_container_width=True)
        clear_btn = st.form_submit_button("துடைத்து துவக்கவும்", use_container_width=True)

# Process inputs
if submit_btn and tamil_input:
    with st.spinner("உருவாக்கம் நடந்து கொண்டிருக்கிறது..."):
        english_text, creative_text, image, proc_time, error = process_content(tamil_input, creativity)
        
        # Save outputs to session state
        st.session_state.outputs.append({
            "tamil_input": tamil_input,
            "english_text": english_text,
            "creative_text": creative_text,
            "image": image,
            "proc_time": proc_time
        })
        
        # Display results
        st.subheader("முடிவுகள்")
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.text_area("ஆங்கில மொழிபெயர்ப்பு", value=english_text, height=100, disabled=True)
            st.text_area("படைப்பு உரை", value=creative_text, height=150, disabled=True)
            if proc_time:
                st.info(proc_time)
        
        with col2:
            if image:
                st.image(image, caption="உருவாக்கப்பட்ட படம்", use_column_width=True)
                
                # Add download button
                buf = io.BytesIO()
                image.save(buf, format="PNG")
                byte_im = buf.getvalue()
                st.download_button(
                    label="படத்தை பதிவிறக்குக",
                    data=byte_im,
                    file_name="tamil_ai_image.png",
                    mime="image/png",
                    use_container_width=True
                )
        
        if error:
            st.error(error)

# Clear button functionality
if clear_btn:
    st.session_state.tamil_input = ""
    st.session_state.creativity = 7
    st.session_state.outputs = []
    st.experimental_rerun()

# Display history
if st.session_state.outputs:
    st.divider()
    st.subheader("முந்தைய உருவாக்கங்கள்")
    
    for i, output in enumerate(reversed(st.session_state.outputs)):
        with st.expander(f"உருவாக்கம் #{len(st.session_state.outputs)-i}: {output['tamil_input'][:50]}..."):
            col1, col2 = st.columns(2)
            
            with col1:
                st.text_area(f"மொழிபெயர்ப்பு #{len(st.session_state.outputs)-i}", 
                            value=output['english_text'], height=100, disabled=True, key=f"eng_{i}")
                st.text_area(f"படைப்பு உரை #{len(st.session_state.outputs)-i}", 
                            value=output['creative_text'], height=150, disabled=True, key=f"text_{i}")
                if output.get('proc_time'):
                    st.info(output['proc_time'])
            
            with col2:
                if output['image']:
                    st.image(output['image'], caption="உருவாக்கப்பட்ட படம்", use_column_width=True)
                    buf = io.BytesIO()
                    output['image'].save(buf, format="PNG")
                    byte_im = buf.getvalue()
                    st.download_button(
                        label=f"படத்தை பதிவிறக்குக #{len(st.session_state.outputs)-i}",
                        data=byte_im,
                        file_name=f"tamil_ai_image_{len(st.session_state.outputs)-i}.png",
                        mime="image/png",
                        use_container_width=True,
                        key=f"download_{i}"
                    )