import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import torch

# Set Hugging Face token if using private models or rate limits apply
# from huggingface_hub import login
# login(token="your_huggingface_token")

# Load Tamil to English translation model (use M2M100 for better support)
from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
translator_tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
translator_model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")

# Load text generation model
text_generator = pipeline("text-generation", model="gpt2")

# Load image generation model (e.g., SD 1.5)
from diffusers import StableDiffusionPipeline
import torch

image_pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", 
    torch_dtype=torch.float16
)
image_pipe.to("cuda" if torch.cuda.is_available() else "cpu")

def process_input(tamil_text):
    try:
        # Step 1: Translate Tamil to English
        translator_tokenizer.src_lang = "ta"
        encoded = translator_tokenizer(tamil_text, return_tensors="pt")
        generated_tokens = translator_model.generate(**encoded, forced_bos_token_id=translator_tokenizer.get_lang_id("en"))
        english_text = translator_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]

        # Step 2: Generate additional text
        generated_output = text_generator(english_text, max_length=50, num_return_sequences=1)[0]["generated_text"]

        # Step 3: Generate image from the final English text
        image = image_pipe(generated_output).images[0]

        return english_text, generated_output, image
    except Exception as e:
        return str(e), "", None

# Gradio interface
demo = gr.Interface(
    fn=process_input,
    inputs=gr.Textbox(label="Enter Tamil Text"),
    outputs=[
        gr.Textbox(label="Translated English Text"),
        gr.Textbox(label="Generated Description"),
        gr.Image(label="Generated Image")
    ],
    title="Tamil to English → Text → Image Generator",
    description="This app takes Tamil input, translates it to English, generates detailed text, and creates an image."
)

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