import streamlit as st
from transformers import pipeline
from diffusers import DiffusionPipeline, FluxPipeline
import torch
import time

device = "cuda" if torch.cuda.is_available() else "cpu"
# Load models
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-dra-en")
summarizer = pipeline("summarization", model="Falconsai/text_summarization")
# image_pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device)
# image_pipe = fluxPipeline("image-generation", model=DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device))

image_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
image_pipe.enable_model_cpu_offload()


# Functions for each task
def translate_tamil_to_english(text):
    time.sleep(2)
    result = translator(text)
    return result[0]['translation_text']

def summarize_english_text(paragraph):
    time.sleep(2)
    summary = summarizer(paragraph, max_length=100, min_length=25, do_sample=False)
    return summary[0]['summary_text']

def english_text_to_image(text):
    image = image_pipe(text).images[0]
    return image

# Custom CSS
st.markdown("""
    <style>
        /* Background color */
        body {
            background-color: #f0f0f5;
        }

        /* Text color and font */
        .stApp {
            font-family: 'Arial', sans-serif;
            color: #333;
        }

        /* Titles and subtitles styling */
        h1 {
            color: #2E8B57;
            text-align: center;
            text-shadow: 2px 2px 5px #aaaaaa;
        }

        h2, h3 {
            color: #4682B4;
            text-shadow: 1px 1px 3px #aaaaaa;
        }

        /* Background texture */
        .stApp {
            background: linear-gradient(to bottom right, #fff7e6, #e6f7ff);
        }

        /* Button styling */
        button[kind="primary"] {
            background-color: #4682B4;
            color: white;
            border-radius: 8px;
            padding: 0.5rem 1rem;
        }

        button[kind="primary"]:hover {
            background-color: #5b9bd5;
        }

        /* Text area and input field styling */
        textarea, input {
            border-radius: 10px;
            padding: 1rem;
            border: 2px solid #ccc;
            background-color: #f9f9f9;
        }

        /* Styling the output boxes */
        .stMarkdown {
            background-color: #e6f9ff;
            padding: 1rem;
            border-radius: 10px;
            box-shadow: 2px 2px 10px #ccc;
        }
    </style>
    """, unsafe_allow_html=True)

# Streamlit app layout
st.title("🌟Multifunctional AI Application🌟")

# Row 1: Tamil to English translation
st.subheader("🌐 Translate Tamil to English")
tamil_input = st.text_area("Enter Tamil text", "")
if st.button("Translate"):
    english_output = translate_tamil_to_english(tamil_input)
    st.markdown(f"**Translated English Text**: {english_output}")

# Row 2: English paragraph summarization
st.subheader("📝 Summarize English Paragraph")
english_paragraph = st.text_area("Enter English paragraph", "")
if st.button("Summarize"):
    summary_output = summarize_english_text(english_paragraph)
    st.markdown(f"**Summary**: {summary_output}")

# Row 3: English text to image generation
st.subheader("🎨 Generate Image from English Text")
image_text = st.text_input("Enter description for image generation", "")
if st.button("Generate Image"):
    generated_image = english_text_to_image(image_text)
    st.image(generated_image, caption="Generated Image")