Update app.py
Browse files
app.py
CHANGED
@@ -1,44 +1,77 @@
|
|
1 |
-
import
|
2 |
-
|
3 |
-
|
4 |
-
|
5 |
-
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import torch
|
3 |
+
from transformers import MBartForConditionalGeneration, MBart50Tokenizer, AutoTokenizer, AutoModelForCausalLM
|
4 |
+
from diffusers import StableDiffusionPipeline
|
5 |
+
from PIL import Image
|
6 |
+
import tempfile
|
7 |
+
import time
|
8 |
+
import streamlit as st
|
9 |
+
|
10 |
+
# Use CPU (Hugging Face Spaces free tier)
|
11 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
12 |
+
|
13 |
+
# Load translation model
|
14 |
+
translator_tokenizer = MBart50Tokenizer.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
|
15 |
+
translator_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt").to(device)
|
16 |
+
translator_tokenizer.src_lang = "ta_IN"
|
17 |
+
|
18 |
+
# Load text generation model
|
19 |
+
gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
|
20 |
+
gen_model = AutoModelForCausalLM.from_pretrained("gpt2").to(device)
|
21 |
+
|
22 |
+
# Load image generation model
|
23 |
+
pipe = StableDiffusionPipeline.from_pretrained(
|
24 |
+
"stabilityai/stable-diffusion-2-1-base",
|
25 |
+
torch_dtype=torch.float32,
|
26 |
+
safety_checker=None
|
27 |
+
).to(device)
|
28 |
+
|
29 |
+
def translate_tamil_to_english(text):
|
30 |
+
inputs = translator_tokenizer(text, return_tensors="pt").to(device)
|
31 |
+
output = translator_model.generate(
|
32 |
+
**inputs,
|
33 |
+
forced_bos_token_id=translator_tokenizer.lang_code_to_id["en_XX"]
|
34 |
+
)
|
35 |
+
translated = translator_tokenizer.batch_decode(output, skip_special_tokens=True)[0]
|
36 |
+
return translated
|
37 |
+
|
38 |
+
def generate_creative_text(prompt, max_length=100):
|
39 |
+
input_ids = gen_tokenizer.encode(prompt, return_tensors="pt").to(device)
|
40 |
+
output = gen_model.generate(
|
41 |
+
input_ids, max_length=max_length, do_sample=True, top_k=50, temperature=0.9
|
42 |
+
)
|
43 |
+
return gen_tokenizer.decode(output[0], skip_special_tokens=True)
|
44 |
+
|
45 |
+
def generate_image(prompt):
|
46 |
+
image = pipe(prompt).images[0]
|
47 |
+
temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
|
48 |
+
image.save(temp_file.name)
|
49 |
+
return temp_file.name
|
50 |
+
|
51 |
+
# Streamlit UI
|
52 |
+
st.set_page_config(page_title="Tamil → English + AI", layout="centered")
|
53 |
+
st.title("🌐 Tamil to English + AI Image Generator")
|
54 |
+
|
55 |
+
tamil_input = st.text_area("✍️ Enter Tamil Text", height=150)
|
56 |
+
|
57 |
+
if st.button("🚀 Generate"):
|
58 |
+
if not tamil_input.strip():
|
59 |
+
st.warning("Please enter Tamil text.")
|
60 |
+
else:
|
61 |
+
with st.spinner("Translating..."):
|
62 |
+
translated = translate_tamil_to_english(tamil_input)
|
63 |
+
st.success("✅ Translated!")
|
64 |
+
st.markdown(f"**English:** `{translated}`")
|
65 |
+
|
66 |
+
with st.spinner("Generating creative text..."):
|
67 |
+
creative_text = generate_creative_text(translated)
|
68 |
+
st.success("✅ Creative text generated!")
|
69 |
+
st.markdown(f"**Creative Prompt:** `{creative_text}`")
|
70 |
+
|
71 |
+
with st.spinner("Generating image..."):
|
72 |
+
image_path = generate_image(translated)
|
73 |
+
st.success("✅ Image generated!")
|
74 |
+
st.image(Image.open(image_path), caption="🖼️ AI Generated Image", use_column_width=True)
|
75 |
+
|
76 |
+
st.markdown("---")
|
77 |
+
st.markdown("🔧 Powered by MBart, GPT2 & Stable Diffusion - Deployed on Hugging Face 🤗")
|