Update app.py

app.py

@@ -1,57 +1,54 @@
 import streamlit as st
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from diffusers import StableDiffusionPipeline
 import torch
 
+# Load models only once
 @st.cache_resource
 def load_all_models():
-    # Load IndicTrans2 Tamil-to-English model
+    # Load translation model
     trans_model_id = "ai4bharat/indictrans2-indic-en-dist-200M"
-    trans_tokenizer = AutoTokenizer.from_pretrained(trans_model_id, trust_remote_code=True)
-    trans_model = AutoModelForSeq2SeqLM.from_pretrained(trans_model_id, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(trans_model_id, trust_remote_code=True)
+    model = AutoModelForSeq2SeqLM.from_pretrained(trans_model_id, trust_remote_code=True)
+    translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer)
 
-    # Load English text generation model (you can use GPT2 or Falcon, etc.)
-    text_gen = pipeline("text-generation", model="gpt2")
-
-    # Load Stable Diffusion for image generation
+    # Load image generation model (Stable Diffusion 2.1)
     img_pipe = StableDiffusionPipeline.from_pretrained(
-        "runwayml/stable-diffusion-v1-5",
-        torch_dtype=torch.float16,
-        revision="fp16"
-    ).to("cuda" if torch.cuda.is_available() else "cpu")
-
-    return trans_tokenizer, trans_model, text_gen, img_pipe
+        "stabilityai/stable-diffusion-2-1",
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        revision="fp16" if torch.cuda.is_available() else None,
+    )
+    img_pipe = img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-def translate_text(text, tokenizer, model):
-    input_text = f"translate Tamil to English: {text}"
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True)
-    outputs = model.generate(**inputs, max_length=128)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return tokenizer, model, translation_pipeline, img_pipe
 
+# Streamlit UI
 def main():
-    st.title("Multimodal Tamil to Image Generator 🚀")
-    st.markdown("Enter Tamil text, we translate it to English, continue the sentence, and generate an image!")
-
-    user_input = st.text_area("Enter Tamil text:", "")
-
-    if st.button("Generate"):
-        with st.spinner("Loading models..."):
-            tokenizer, model, text_gen, img_pipe = load_all_models()
-
+    st.set_page_config(page_title="Tamil to English to Image Generator", layout="centered")
+    st.title("📸 Tamil → English → AI Image Generator")
+    st.markdown("Translate Tamil text to English and generate an image from it!")
+
+    # Load models
+    with st.spinner("Loading models..."):
+        tokenizer, model, translation_pipeline, img_pipe = load_all_models()
+
+    # Input
+    tamil_text = st.text_area("Enter Tamil text here:", height=150)
+    
+    if st.button("Generate Image"):
+        if tamil_text.strip() == "":
+            st.warning("Please enter some Tamil text.")
+            return
+
+        # Step 1: Translate Tamil to English
         with st.spinner("Translating to English..."):
-            english_text = translate_text(user_input, tokenizer, model)
-            st.subheader("Translated English:")
-            st.write(english_text)
-
-        with st.spinner("Generating continuation..."):
-            continuation = text_gen(english_text, max_length=50, do_sample=True)[0]['generated_text']
-            st.subheader("Generated Text:")
-            st.write(continuation)
-
-        with st.spinner("Generating Image..."):
-            image = img_pipe(continuation).images[0]
-            st.subheader("Generated Image:")
-            st.image(image)
+            translated = translation_pipeline(tamil_text, src_lang="ta", tgt_lang="en")[0]["translation_text"]
+            st.success(f"🔤 English Translation: `{translated}`")
+
+        # Step 2: Generate image
+        with st.spinner("Generating image..."):
+            image = img_pipe(prompt=translated).images[0]
+            st.image(image, caption="Generated Image", use_column_width=True)
 
 if __name__ == "__main__":
     main()