24Sureshkumar's picture
Update app.py
52b96f9 verified
raw
history blame
3.4 kB
import gradio as gr
import requests
from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
from PIL import Image
import torch
import io
import os
from typing import Tuple
# Load Hugging Face API key securely
HF_API_KEY = os.getenv("HF_API_KEY") # You must set this as an environment variable
if not HF_API_KEY:
raise ValueError("HF_API_KEY is not set. Add it in Hugging Face 'Variables and Secrets' or local environment.")
# API Endpoint for Image Generation
IMAGE_GEN_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}
# 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
def translate_tamil_to_english(text: str) -> str:
"""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: str) -> str:
"""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: str) -> Image.Image:
"""Sends request to API for image generation."""
response = requests.post(IMAGE_GEN_URL, headers=HEADERS, json={"inputs": prompt})
try:
if response.status_code == 200 and response.headers["content-type"].startswith("image"):
return Image.open(io.BytesIO(response.content))
except Exception as e:
print("Image generation error:", e)
return Image.new("RGB", (300, 300), "red") # Fallback placeholder image
def process_input(tamil_text: str) -> Tuple[str, str, Image.Image]:
"""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
# 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.",
allow_flagging="never" # Avoids schema-related error in Spaces
)
# Launch the app
interface.launch()