# app.py

import gradio as gr
from transformers import pipeline, MarianMTModel, MarianTokenizer
from diffusers import StableDiffusionPipeline
import torch
import os

# Get Hugging Face token from environment (in Hugging Face Spaces this is auto-populated from secrets)
hf_token = os.getenv("HUGGINGFACE_TOKEN")

# Load Translation Pipeline: Tamil → English using MarianMT
translation_model_name = "Helsinki-NLP/opus-mt-ta-en"
translator = pipeline("translation", model=translation_model_name)

# Load Stable Diffusion Pipeline: English → Image
device = "cuda" if torch.cuda.is_available() else "cpu"

image_pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    use_auth_token=hf_token,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
).to(device)

# Core function
def translate_and_generate(tamil_text):
    # Step 1: Translate Tamil to English
    english_output = translator(tamil_text)[0]['translation_text']
    
    # Step 2: Generate Image from English Text
    image = image_pipe(english_output).images[0]
    
    return english_output, image

# Gradio UI
iface = gr.Interface(
    fn=translate_and_generate,
    inputs=gr.Textbox(lines=2, label="Enter Tamil Text"),
    outputs=[
        gr.Textbox(label="Translated English Text"),
        gr.Image(label="Generated Image")
    ],
    title="Tamil-to-Image Generator 🌸",
    description="Enter Tamil text. It will be translated to English and visualized using Stable Diffusion."
)

iface.launch()