app.py

# -*- coding: utf-8 -*-
"""gen ai project f.ipynb
Automatically generated by Colab.
Original file is located at
    https://colab.research.google.com/drive/1iF7hdOjWNeFUtGvUYdaFsBErJGnY1h5J
"""

import os
from transformers import MarianMTModel, MarianTokenizer
import gradio as gr
from PIL import Image, UnidentifiedImageError
import requests
import io

# Load translation models
model_name = "Helsinki-NLP/opus-mt-mul-en"
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)

# Define language map
language_map = {
    "Tamil": "ta",
    "Russian": "rus"
}

def translate_text(input_text, selected_languages):
    if not selected_languages:
        return "Please select at least one language."
    
    selected_language = selected_languages[0]  # Pick the first selected language
    lang_code = language_map[selected_language]
    lang_prefix = f">>{lang_code}<< "
    text_with_lang = lang_prefix + input_text
    inputs = tokenizer(text_with_lang, return_tensors="pt", padding=True)
    translated_tokens = model.generate(**inputs)
    translation = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
    return translation

def generate_image(prompt):
    API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
    hf_token = os.getenv("HF_TOKEN")
    headers = {"Authorization": f"Bearer {hf_token}"}
    
    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})

    if response.status_code == 200:
        image_bytes = response.content
        try:
            image = Image.open(io.BytesIO(image_bytes))
            return image
        except UnidentifiedImageError:
            return None
    else:
        return None

def process_input(text_input, selected_languages):
    translated_output = translate_text(text_input, selected_languages)
    image = generate_image(translated_output)
    return translated_output, image

# Gradio interface
interface = gr.Interface(
    fn=process_input,
    inputs=[gr.Textbox(label="Input Text"), gr.CheckboxGroup(choices=["Tamil", "Russian"], label="Select Language")],
    outputs=[gr.Textbox(label="Translated Text"), gr.Image(label="Generated Image")],
    title="Multilingual Translation and Image Generation",
    description="Translate Tamil or Russian text to English and generate an image."
)

interface.launch()