import gradio as gr
import random

# Load the model
model = gr.load("models/Purz/face-projection")

def generate_image(text, seed):
    if seed is not None:
        random.seed(seed)

    if text in [example[0] for example in examples]:
        print(f"Using example: {text}")

    # Generate four images by calling the model four times
    images = [model(text) for _ in range(4)]
    return images

# Examples for user guidance
examples = [
    ["Humanoid Cat Warrior, Full View", None],
    ["Warhammer Sisterhood", None],
    ["Future Robots war", None],
    ["Fantasy dragon", None]
]

# Custom CSS for gradient buttons and capsule shapes
css = """
/* General button styles */
button {
    border-radius: 25px; /* Capsule shape */
    padding: 10px 20px;
    font-size: 16px;
    font-weight: bold;
    color: white;
    background: linear-gradient(135deg, #6D5BBA, #8A2387); /* Gradient color */
    border: none;
    cursor: pointer;
    transition: background 0.3s ease;
}

button:hover {
    background: linear-gradient(135deg, #8A2387, #6D5BBA); /* Gradient reverse on hover */
}

/* Slider handle style */
input[type=range] {
    accent-color: #8A2387; /* Gradient color accent */
}

/* Textbox styling */
.gr-textbox {
    border-radius: 8px;
    border: 1px solid #DDD;
    padding: 10px;
}

/* Output image styling */
.output_image {
    border: 2px solid #6D5BBA;
    border-radius: 10px;
    padding: 10px;
    background-color: #f9f9f9;
}
"""

# Gradio interface with custom styles
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Type your imagination:", placeholder="Type or select an example..."),
        gr.Slider(minimum=0, maximum=10000, step=1, label="Seed (optional)")
    ],
    # Specify a list of 4 Image outputs for displaying multiple images
    outputs=[gr.Image(label=f"Generated Image {i+1}") for i in range(4)],
    examples=examples,
    css=css,  # Apply custom CSS styling
    description="Note: The model is running on CPU; processing might take a bit longer. Thank you for your patience!"
)

# Launch the interface
interface.launch()