24Sureshkumar's picture
Update app.py
a558492 verified
raw
history blame
2.28 kB
import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import torch
# Set Hugging Face token if using private models or rate limits apply
# from huggingface_hub import login
# login(token="your_huggingface_token")
# Load Tamil to English translation model (use M2M100 for better support)
from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
translator_tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
translator_model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
# Load text generation model
text_generator = pipeline("text-generation", model="gpt2")
# Load image generation model (e.g., SD 1.5)
from diffusers import StableDiffusionPipeline
import torch
image_pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16
)
image_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
def process_input(tamil_text):
try:
# Step 1: Translate Tamil to English
translator_tokenizer.src_lang = "ta"
encoded = translator_tokenizer(tamil_text, return_tensors="pt")
generated_tokens = translator_model.generate(**encoded, forced_bos_token_id=translator_tokenizer.get_lang_id("en"))
english_text = translator_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
# Step 2: Generate additional text
generated_output = text_generator(english_text, max_length=50, num_return_sequences=1)[0]["generated_text"]
# Step 3: Generate image from the final English text
image = image_pipe(generated_output).images[0]
return english_text, generated_output, image
except Exception as e:
return str(e), "", None
# Gradio interface
demo = gr.Interface(
fn=process_input,
inputs=gr.Textbox(label="Enter Tamil Text"),
outputs=[
gr.Textbox(label="Translated English Text"),
gr.Textbox(label="Generated Description"),
gr.Image(label="Generated Image")
],
title="Tamil to English β†’ Text β†’ Image Generator",
description="This app takes Tamil input, translates it to English, generates detailed text, and creates an image."
)
if __name__ == "__main__":
demo.launch()