import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image
import io
import os

API_KEY = os.getenv("HF_API_TOKEN")
MODEL_NAME = "black-forest-labs/FLUX.1-dev"

client = InferenceClient(api_key=API_KEY)

def generate_image(prompt):
    """Generate an image from a text prompt using Hugging Face API."""
    try:
        image = client.text_to_image(prompt, model=MODEL_NAME)
        img_byte_array = io.BytesIO()
        image.save(img_byte_array, format="PNG")
        img_byte_array.seek(0)
        return image
    except Exception as e:
        return f"Error generating image: {e}"

css = """
body {
    background: linear-gradient(135deg, #121212, #1e1e2e);
    color: white;
    font-family: 'Poppins', sans-serif;
    text-align: center;
}
.gradio-container {
    background: rgba(30, 30, 46, 0.9);
    padding: 25px;
    border-radius: 12px;
    box-shadow: 0px 4px 10px rgba(0, 0, 0, 0.3);
    max-width: 600px;
    margin: auto;
}
button {
    background: linear-gradient(135deg, #ff416c, #ff4b4b);
    color: white;
    border: none;
    border-radius: 10px;
    font-size: 18px;
    padding: 12px;
    transition: 0.3s;
}
button:hover {
    background: linear-gradient(135deg, #ff4b4b, #ff416c);
    transform: scale(1.05);
}
h1, h2, h3, p, label {
    color: white !important;
}
input, textarea {
    background: #252537;
    color: white;
    border: 1px solid #444;
    padding: 10px;
    border-radius: 8px;
}
input:focus, textarea:focus {
    background: #252537;
    outline: 2px solid #ff416c;
    color: white;
}
"""

with gr.Blocks(css=css) as demo:
    title="🚀 Agents Valley AI Image Generator",
    gr.Markdown(
        "<h2 style='color: white; font-size: 35px; font-weight: 600; margin-bottom: 100px;'>AI image Generator powered by Agents Valley 🤖</h2>"
    )
    gr.Markdown(
        "<h2 style='color: white; font-size: 15px; font-weight: 600; margin-bottom: 100px;'>FLUX.1-dev</h2>"
    )
    prompt_input = gr.Textbox(label="Enter your prompt", value="A robot on a stallion.")
    output_image = gr.Image(label="Generated Image by Agents Valley")
    generate_button = gr.Button("Generate Image")
    generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image)

demo.launch()