File size: 3,006 Bytes
b88f708 |
1 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import os
import torch
from transformers import MBartForConditionalGeneration, MBart50Tokenizer, AutoTokenizer, AutoModelForCausalLM
from diffusers import StableDiffusionPipeline
from PIL import Image
import tempfile
import time
import streamlit as st
# Use CPU (Hugging Face Spaces free tier)
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load translation model
translator_tokenizer = MBart50Tokenizer.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
translator_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt").to(device)
translator_tokenizer.src_lang = "ta_IN"
# Load text generation model
gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
gen_model = AutoModelForCausalLM.from_pretrained("gpt2").to(device)
# Load image generation model
pipe = StableDiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1-base",
torch_dtype=torch.float32,
safety_checker=None
).to(device)
def translate_tamil_to_english(text):
inputs = translator_tokenizer(text, return_tensors="pt").to(device)
output = translator_model.generate(
**inputs,
forced_bos_token_id=translator_tokenizer.lang_code_to_id["en_XX"]
)
translated = translator_tokenizer.batch_decode(output, skip_special_tokens=True)[0]
return translated
def generate_creative_text(prompt, max_length=100):
input_ids = gen_tokenizer.encode(prompt, return_tensors="pt").to(device)
output = gen_model.generate(
input_ids, max_length=max_length, do_sample=True, top_k=50, temperature=0.9
)
return gen_tokenizer.decode(output[0], skip_special_tokens=True)
def generate_image(prompt):
image = pipe(prompt).images[0]
temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
image.save(temp_file.name)
return temp_file.name
# Streamlit UI
st.set_page_config(page_title="Tamil → English + AI", layout="centered")
st.title("🌐 Tamil to English + AI Image Generator")
tamil_input = st.text_area("✍️ Enter Tamil Text", height=150)
if st.button("🚀 Generate"):
if not tamil_input.strip():
st.warning("Please enter Tamil text.")
else:
with st.spinner("Translating..."):
translated = translate_tamil_to_english(tamil_input)
st.success("✅ Translated!")
st.markdown(f"**English:** `{translated}`")
with st.spinner("Generating creative text..."):
creative_text = generate_creative_text(translated)
st.success("✅ Creative text generated!")
st.markdown(f"**Creative Prompt:** `{creative_text}`")
with st.spinner("Generating image..."):
image_path = generate_image(translated)
st.success("✅ Image generated!")
st.image(Image.open(image_path), caption="🖼️ AI Generated Image", use_column_width=True)
st.markdown("---")
st.markdown("🔧 Powered by MBart, GPT2 & Stable Diffusion - Deployed on Hugging Face 🤗")
|