Update app.py

app.py

@@ -1,307 +1,60 @@
 # app.py
 import streamlit as st
-from transformers import pipeline
-from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from diffusers import StableDiffusionPipeline
 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}"
-                    )
+# Load NLLB translation model
+@st.cache_resource
+def load_translation_model():
+    model_name = "facebook/nllb-200-distilled-600M"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    return tokenizer, model
+
+# Load Stable Diffusion model
+@st.cache_resource
+def load_diffusion_model():
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "CompVis/stable-diffusion-v1-4",
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    )
+    if torch.cuda.is_available():
+        pipe = pipe.to("cuda")
+    return pipe
+
+# Translate Tamil to English using NLLB
+def translate_text(text, tokenizer, model):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    inputs["forced_bos_token_id"] = tokenizer.lang_code_to_id["eng_Latn"]
+    with torch.no_grad():
+        translated_tokens = model.generate(**inputs, max_length=512)
+    return tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+
+# Streamlit App
+def main():
+    st.set_page_config(page_title="Tamil to Image Generator", layout="centered")
+    with open("style.css") as f:
+        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
+
+    st.title("🪔 தமிழ் உரை -> பட உருவாக்கம்")
+
+    user_input = st.text_area("உங்கள் உரையை உள்ளிடவும் (தமிழில்)", height=150)
+
+    if st.button("படம் உருவாக்கவும்"):
+        if not user_input.strip():
+            st.warning("தயவு செய்து தமிழ் உரையை உள்ளிடவும்.")
+            return
+
+        with st.spinner("மொழிபெயர்ப்பு மற்றும் பட உருவாக்��ம் நடைபெறுகிறது..."):
+            tokenizer, model = load_translation_model()
+            tokenizer.src_lang = "tam_Taml"  # Tamil
+            translated_text = translate_text(user_input, tokenizer, model)
+            st.success(f"மொழிபெயர்ப்பு (ஆங்கிலம்): {translated_text}")
+
+            pipe = load_diffusion_model()
+            image = pipe(translated_text).images[0]
+            st.image(image, caption="உங்கள் உருவாக்கப்பட்ட படம்", use_column_width=True)
+
+if __name__ == "__main__":
+    main()