Update app.py

app.py

@@ -1,85 +1,57 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
 from diffusers import StableDiffusionPipeline
 import torch
 
-st.set_page_config(
-    page_title="Tamil Creative Studio",
-    page_icon="🇮🇳",
-    layout="centered",
-)
-
-def load_css():
-    st.markdown(
-        """<style>
-        .header {
-            text-align: center;
-            padding: 20px;
-            background: #f9f9f9;
-            border-radius: 10px;
-            margin-bottom: 20px;
-        }
-        .header h1 { color: #cc0000; }
-        .header p { color: #333; font-style: italic; }
-        </style>""",
-        unsafe_allow_html=True,
-    )
-    
 @st.cache_resource
 def load_all_models():
-    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-    from diffusers import StableDiffusionPipeline
-    import torch
-
-    # Translation model
-    model_id = "ai4bharat/indictrans2-indic-en-dist-200M"
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_id, trust_remote_code=True)
+    # Load IndicTrans2 Tamil-to-English 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)
 
-    # Text generation model (simple GPT-2)
+    # Load English text generation model (you can use GPT2 or Falcon, etc.)
     text_gen = pipeline("text-generation", model="gpt2")
 
-    # Image generation model (Stable Diffusion)
+    # Load Stable Diffusion for image generation
     img_pipe = StableDiffusionPipeline.from_pretrained(
-        "CompVis/stable-diffusion-v1-4",
-        revision="fp16",
-        torch_dtype=torch.float16
-    ).to("cuda")
+        "runwayml/stable-diffusion-v1-5",
+        torch_dtype=torch.float16,
+        revision="fp16"
+    ).to("cuda" if torch.cuda.is_available() else "cpu")
 
-    return tokenizer, model, text_gen, img_pipe
+    return trans_tokenizer, trans_model, text_gen, img_pipe
 
-
-def translate_tamil(text, tokenizer, model):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
-    outs = model.generate(**inputs, max_length=150, num_beams=5, early_stopping=True)
-    return tokenizer.decode(outs[0], skip_special_tokens=True)
+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)
 
 def main():
-    load_css()
-    st.markdown(
-        '<div class="header"><h1>🌐 தமிழ் → English → Creative Studio</h1>'
-        '<p>Translate Tamil text and generate creative content</p></div>',
-        unsafe_allow_html=True
-    )
-    tokenizer, model, text_gen, img_pipe = load_all_models()
-    tamil_text = st.text_area("**தமிழ் உரை:**", height=150, placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்...")
+    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("உருவாக்கு"):
-        if not tamil_text.strip():
-            st.warning("உரையை உள்ளிடவும்.")
-            return
+    if st.button("Generate"):
+        with st.spinner("Loading models..."):
+            tokenizer, model, text_gen, img_pipe = load_all_models()
 
-        with st.spinner("மொழிபெயர்க்கிறது..."):
-            eng = translate_tamil(tamil_text, tokenizer, model)
-        st.success(eng)
+        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("உரை உருவாக்குதல்..."):
-            creative = text_gen(f"Create a creative description about: {eng}", max_length=80, num_return_sequences=1)[0]["generated_text"]
-        st.info(creative)
+        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("படத்தை உருவாக்குதல்..."):
-            img = img_pipe(eng, num_inference_steps=40, guidance_scale=8.5).images[0]
-        st.image(img, caption="Generated Image", use_column_width=True)
+        with st.spinner("Generating Image..."):
+            image = img_pipe(continuation).images[0]
+            st.subheader("Generated Image:")
+            st.image(image)
 
 if __name__ == "__main__":
     main()