app.py

# app.py

import gradio as gr
from transformers import pipeline
from diffusers import StableDiffusionPipeline
import torch
import re
import os
from huggingface_hub import login

# Get Hugging Face token from environment variable
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("Hugging Face token not found in environment variables!")
else:
    login(token=HF_TOKEN)

# Load all models
def load_models():
    # Translation model: Tamil → English
    translator = pipeline(
        "translation",
        model="facebook/nllb-200-distilled-600M",
        src_lang="tam_Taml",
        tgt_lang="eng_Latn",
        device=0 if torch.cuda.is_available() else -1,
        use_auth_token=HF_TOKEN
    )

    # Text generation model
    text_generator = pipeline(
        "text-generation",
        model="gpt2-medium",
        device=0 if torch.cuda.is_available() else -1,
        use_auth_token=HF_TOKEN
    )

    # Stable Diffusion for image generation
    if torch.cuda.is_available():
        image_pipe = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            torch_dtype=torch.float16,
            revision="fp16",
            use_auth_token=HF_TOKEN
        ).to("cuda")
    else:
        image_pipe = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            use_auth_token=HF_TOKEN
        )

    return translator, text_generator, image_pipe

# Load models once at startup
try:
    translator, text_generator, image_pipe = load_models()
except Exception as e:
    raise RuntimeError(f"Model loading failed: {str(e)}")

# Clean generated text
def clean_text(text):
    cleaned = re.sub(r'[^a-zA-Z0-9,.!?\'"\- ]+', '', text).strip()
    sentences = re.split(r'(?<=[.!?])\s+', cleaned)
    return ' '.join(sentences[:2])  # return first 2 sentences

# Main processing function
def process_content(tamil_input, creativity_level):
    try:
        # Translation
        translation_result = translator(tamil_input)
        english_text = translation_result[0]['translation_text']

        # Image generation
        image = image_pipe(
            english_text,
            guidance_scale=creativity_level,
            num_inference_steps=30
        ).images[0]

        # Text generation
        creative_output = text_generator(
            f"Create creative content about: {english_text}",
            max_length=150,
            temperature=creativity_level / 10,
            num_return_sequences=1
        )

        return english_text, clean_text(creative_output[0]['generated_text']), image, ""
    except Exception as e:
        return "", "", None, f"⚠️ Error: {str(e)}"

# Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🌐 தமிழ் உரை முதல் பட உருவாக்கம் (Tamil to Image Generator)")
    gr.Markdown("தமிழில் உள்ளீடு செய்து → ஆங்கில மொழிபெயர்ப்பு + AI உருவம் + படைப்பு உரை பெறவும்")

    with gr.Row():
        with gr.Column():
            tamil_input = gr.Textbox(
                label="தமிழ் உள்ளீடு",
                placeholder="உதாரணம்: பனி படர்ந்த குளிர்காலத்தில் வெப்பமான காபி குடிக்கும் பழங்குடி பெண்",
                lines=3
            )
            creativity = gr.Slider(
                label="படைப்பாற்றல் நிலை",
                minimum=1, maximum=10, value=7, step=1
            )
            submit_btn = gr.Button("உருவாக்கு")

        with gr.Column():
            translation_box = gr.Textbox(label="ஆங்கில மொழிபெயர்ப்பு")
            creative_output = gr.Textbox(label="படைப்பு உரை", lines=3)
            image_output = gr.Image(label="உருவாக்கப்பட்ட படம்")
            error_output = gr.Textbox(label="பிழை செய்திகள்", visible=True)

    # Example inputs
    examples = gr.Examples(
        examples=[
            ["கடலின் அடியில் மறைந்திருக்கும் பழைய நகரம்", 8],
            ["பனி படர்ந்த குளிர்காலத்தில் வெப்பமான காபி குடிக்கும் பழங்குடி பெண்", 7],
            ["வேறு கிரகத்தில் இருந்து வந்த அறிவார்ந்த இயந்திரங்கள்", 9]
        ],
        inputs=[tamil_input, creativity]
    )

    # Clear all inputs/outputs
    clear_btn = gr.Button("துடைத்து துவக்கவும்")
    def clear_all():
        return "", "", "", None, ""

    submit_btn.click(
        fn=process_content,
        inputs=[tamil_input, creativity],
        outputs=[translation_box, creative_output, image_output, error_output]
    )

    clear_btn.click(
        fn=clear_all,
        inputs=[],
        outputs=[tamil_input, translation_box, creative_output, image_output, error_output]
    )

# Launch the app
if __name__ == "__main__":
    app.queue().launch()