transart / app.py
pravin0077's picture
Update app.py
3a6b0f7 verified
raw
history blame
3.25 kB
import requests
import io
from PIL import Image
import gradio as gr
from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
import os
# Load the translation model
model_name = "Helsinki-NLP/opus-mt-mul-en"
translation_model = MarianMTModel.from_pretrained(model_name)
translation_tokenizer = MarianTokenizer.from_pretrained(model_name)
# Load GPT-Neo model and tokenizer
gpt_model_name = "EleutherAI/gpt-neo-1.3B" # You can also use gpt-neo-2.7B if needed
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, max_words):
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=max_words, num_return_sequences=1)
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, max_words):
try:
# Translate the input text
translated_output = translate_text(tamil_input)
# Generate creative text using GPT-Neo
creative_output = query_gpt_neo(translated_output, max_words)
# 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 with interactive elements
interface = gr.Interface(
fn=process_input,
inputs=[
gr.Textbox(label="Input Tamil Text", placeholder="Enter your Tamil text here..."),
gr.Slider(label="Max Words for Creative Text", minimum=50, maximum=200, step=10, value=100)
],
outputs=[
gr.Textbox(label="Translated Text"),
gr.Textbox(label="Creative Text"),
gr.Image(label="Generated Image")
],
title="TRANSART - Multimodal AI App",
description="Enter Tamil text to translate to English, generate creative text, and produce an image based on the translated text.",
theme="compact", # Use the 'compact' theme for a cleaner app look
layout="vertical" # Arrange components vertically for better readability
)
interface.launch()