transart / app.py
pravin0077's picture
Update app.py
3f10588 verified
raw
history blame
2.83 kB
import requests
import io
from PIL import Image
import gradio as gr
from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
import os
# Load models and tokenizers globally to avoid reloading them for every request
model_name = "Helsinki-NLP/opus-mt-mul-en"
translation_model = MarianMTModel.from_pretrained(model_name)
translation_tokenizer = MarianTokenizer.from_pretrained(model_name)
gpt_model_name = "EleutherAI/gpt-neo-1.3B"
gpt_tokenizer = AutoTokenizer.from_pretrained(gpt_model_name)
gpt_model = AutoModelForCausalLM.from_pretrained(gpt_model_name)
def translate_text(tamil_text):
inputs = translation_tokenizer(tamil_text, return_tensors="pt")
translated_tokens = translation_model.generate(**inputs)
translation = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
return translation
def query_gpt_neo(translated_text):
prompt = f"Continue the story based on the following text: {translated_text}"
inputs = gpt_tokenizer(prompt, return_tensors="pt")
outputs = gpt_model.generate(inputs['input_ids'], max_length=50, num_return_sequences=1) # Reduced max_length
creative_text = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
return creative_text
def query_image(payload):
huggingface_api_key = os.getenv('HUGGINGFACE_API_KEY')
if not huggingface_api_key:
return "Error: Hugging Face API key not set."
API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
headers = {"Authorization": f"Bearer {huggingface_api_key}"}
response = requests.post(API_URL, headers=headers, json=payload)
if response.status_code == 200:
return response.content
else:
return f"Error: {response.status_code} - {response.text}"
def process_input(tamil_input):
try:
# Translate the input text
translated_output = translate_text(tamil_input)
# Generate creative text using GPT-Neo
creative_output = query_gpt_neo(translated_output)
# Generate an image using Hugging Face's FLUX model
image_bytes = query_image({"inputs": translated_output})
image = Image.open(io.BytesIO(image_bytes))
return translated_output, creative_output, image
except Exception as e:
return f"Error occurred: {str(e)}", "", None
# Create a Gradio interface
interface = gr.Interface(
fn=process_input,
inputs=[gr.Textbox(label="Input Tamil Text")],
outputs=[
gr.Textbox(label="Translated Text"),
gr.Textbox(label="Creative Text"),
gr.Image(label="Generated Image")
],
title="TRANSART",
description="Enter Tamil text to translate to English and generate an image based on the translated text."
)
interface.launch()