Update app.py

app.py

@@ -1,58 +1,61 @@
+# app.py
 import streamlit as st
 from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer, pipeline
-from diffusers import StableDiffusionPipeline
+from diffusers import DiffusionPipeline
 import torch
 
-@st.cache_resource
-def load_models():
-    # Translation model
-    model_name = "alirezamsh/small100"
+@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)
 
-    # Creative text generation (GPT-2)
-    text_generator = pipeline("text-generation", model="gpt2-medium")
+    # Load creative text model (smaller GPT-2)
+    textgen = pipeline("text-generation", model="gpt2", device=-1)
 
-    # Stable diffusion image generation (CPU-friendly)
-    image_pipe = StableDiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-2-base", torch_dtype=torch.float32
+    # Load lightweight image generation pipeline
+    img_pipe = DiffusionPipeline.from_pretrained(
+        "stabilityai/sdxl-lite", torch_dtype=torch.float32
     ).to("cpu")
 
-    return tokenizer, model, text_generator, image_pipe
+    return tokenizer, model, textgen, img_pipe
 
-def translate_tamil_to_english(text, tokenizer, model):
+def translate(text, tokenizer, model):
     tokenizer.src_lang = "ta"
-    inputs = tokenizer(text, return_tensors="pt", truncation=True)
-    generated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.get_lang_id("en"))
-    return tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+    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_creative_text(prompt, generator):
-    response = generator(f"Describe creatively: {prompt}", max_length=60, num_return_sequences=1)
-    return response[0]["generated_text"]
+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 → Image & Text", layout="centered")
-    st.title("🪔 தமிழ் உரையை படமாக மாற்றவும்")
-    
-    user_input = st.text_area("தமிழ் உரை உள்ளிடவும்:", height=150)
-    if st.button("உருவாக்கவும்"):
-        if not user_input.strip():
-            st.warning("தயவுசெய்து ஒரு உரை உள்ளிடவும்.")
+    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("மொழிபெயர்ப்பு நடக்கிறது..."):
-            tokenizer, model, textgen, pipe = load_models()
-            english_text = translate_tamil_to_english(user_input, tokenizer, model)
-            st.success(f"Translated to English: {english_text}")
+        with st.spinner("மொழிபெயர்ப்பு..."):
+            english_text = translate(tamil_text, tokenizer, model)
+        st.success(f"🔁 Translated: {english_text}")
 
-        with st.spinner("சிறந்த உரையை உருவாக்குகிறது..."):
-            creative = generate_creative_text(english_text, textgen)
-            st.info("✍️ Creative Text:")
-            st.write(creative)
+        with st.spinner("உரையாக்கம்..."):
+            creative_text = generate_text(english_text, textgen)
+        st.info("📝 Creative Output:")
+        st.write(creative_text)
 
         with st.spinner("படம் உருவாக்கப்படுகிறது..."):
-            image = pipe(english_text, num_inference_steps=25).images[0]
-            st.image(image, caption="🎨 Generated Image", use_column_width=True)
+            image = img_pipe(english_text).images[0]
+        st.image(image, caption="🎨 Generated Image", use_column_width=True)
 
 if __name__ == "__main__":
     main()