Spaces:
Sleeping
Sleeping
| 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-2 model and tokenizer (smaller and faster than GPT-Neo) | |
| gpt_model_name = "gpt2" | |
| 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_2(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 for speed | |
| 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-2 | |
| creative_output = query_gpt_2(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() | |