from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer, pipeline
import gradio as gr

# Load the multilingual translation model (supports Tamil <-> English)
model_name = "facebook/m2m100_418M"
tokenizer = M2M100Tokenizer.from_pretrained(model_name)
model = M2M100ForConditionalGeneration.from_pretrained(model_name)

# Load text generation pipeline (you can replace "gpt2" with your preferred model)
text_generator = pipeline("text-generation", model="gpt2")

def translate_tamil_to_english(text):
    if not text.strip():
        return ""
    tokenizer.src_lang = "ta"
    encoded = tokenizer(text, return_tensors="pt")
    generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id("en"))
    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]

def process_text(tamil_text):
    if not tamil_text.strip():
        return "Please enter Tamil text.", None, ""

    try:
        # Step 1: Translate Tamil to English
        translation = translate_tamil_to_english(tamil_text)
    except Exception as e:
        return f"Translation error: {str(e)}", None, ""

    try:
        # Step 2: Generate English text based on translation
        generated = text_generator(translation, max_length=50, num_return_sequences=1)
        generated_text = generated[0]['generated_text']
    except Exception as e:
        return translation, None, f"Text generation error: {str(e)}"

    # Step 3: Placeholder for image generation
    # Replace with your own image generation model/function
    image_url = "https://via.placeholder.com/512.png?text=Generated+Image"

    # Step 4: Description based on generated text (you can use an image captioning model here)
    description = f"Generated description: {generated_text}"

    return translation, image_url, description

iface = gr.Interface(
    fn=process_text,
    inputs=gr.Textbox(lines=2, label="Enter Tamil text"),
    outputs=[
        gr.Textbox(label="Translated English Text"),
        gr.Image(label="Generated Image"),
        gr.Textbox(label="Image Description"),
    ],
    title="Tamil to English Translation + Text & Image Generation",
    description="Enter Tamil text, get English translation, generated text, image, and description."
)

if __name__ == "__main__":
    iface.launch()