24Sureshkumar's picture
Create app.py
bf7e1be verified
raw
history blame
2.95 kB
import gradio as gr
import requests
from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
from PIL import Image
import torch
import io
# Check if GPU is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load Tamil-to-English Translation Model
translator_model = "Helsinki-NLP/opus-mt-mul-en"
translator = MarianMTModel.from_pretrained(translator_model).to(device)
translator_tokenizer = MarianTokenizer.from_pretrained(translator_model)
# Load Text Generation Model
generator_model = "EleutherAI/gpt-neo-1.3B"
generator = AutoModelForCausalLM.from_pretrained(generator_model).to(device)
generator_tokenizer = AutoTokenizer.from_pretrained(generator_model)
if generator_tokenizer.pad_token is None:
generator_tokenizer.pad_token = generator_tokenizer.eos_token
# Hugging Face API for Image Generation
HF_API_KEY = "my_token" # Replace with your API key
IMAGE_GEN_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
# Get the API key from environment variables or Hugging Face secrets
HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}
def translate_tamil_to_english(text):
"""Translates Tamil text to English."""
inputs = translator_tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
output = translator.generate(**inputs)
return translator_tokenizer.decode(output[0], skip_special_tokens=True)
def generate_text(prompt):
"""Generates a creative text based on English input."""
inputs = generator_tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(device)
output = generator.generate(**inputs, max_length=100)
return generator_tokenizer.decode(output[0], skip_special_tokens=True)
def generate_image(prompt):
"""Sends request to API for image generation."""
response = requests.post(IMAGE_GEN_URL, headers=HEADERS, json={"inputs": prompt})
if response.status_code == 200:
return Image.open(io.BytesIO(response.content))
return Image.new("RGB", (300, 300), "red") # Error placeholder image
def process_input(tamil_text):
"""Complete pipeline: Translation, Text Generation, and Image Generation."""
english_text = translate_tamil_to_english(tamil_text)
creative_text = generate_text(english_text)
image = generate_image(english_text)
return english_text, creative_text, image
# Create Gradio Interface
interface = gr.Interface(
fn=process_input,
inputs=gr.Textbox(label="Enter Tamil Text"),
outputs=[
gr.Textbox(label="Translated English Text"),
gr.Textbox(label="Creative Text"),
gr.Image(label="Generated Image")
],
title="Tamil to English Translator & Image Generator",
description="Enter Tamil text, and this app will translate it, generate a creative description, and create an image based on the text."
)
# Launch the Gradio app
interface.launch()