import gradio as gr
import numpy as np
import random

# import spaces #[uncomment to use ZeroGPU]
from diffusers import DiffusionPipeline, AutoPipelineForText2Image
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "stabilityai/stable-diffusion-xl-base-1.0"

if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

# Load pipeline with safety checker disabled
pipe = AutoPipelineForText2Image.from_pretrained(
    model_repo_id, 
    torch_dtype=torch_dtype,
    safety_checker=None,
    requires_safety_checker=False
)
pipe = pipe.to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024


# @spaces.GPU #[uncomment to use ZeroGPU]
def infer(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)

    # SDXL-Turbo works best with these settings:
    # - guidance_scale should be 0.0 (it's trained without guidance)
    # - num_inference_steps should be 1-4 for best results
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=4,  # SDXL-Turbo requires 0.0
        num_inference_steps=max(1, min(4, num_inference_steps)),  # Clamp to 1-4
        width=width,
        height=height,
        generator=generator,
    ).images[0]

    return image, seed


examples = [
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    "An astronaut riding a green horse",
    "A delicious ceviche cheesecake slice",
    "Beautiful landscape with mountains and lake at sunset",
]

css = """
/* ChatGPT-inspired Dark Theme */
:root {
    --bg-primary: #0D0D0D;
    --bg-secondary: #171717;
    --bg-tertiary: #212121;
    --text-primary: #ECECEC;
    --text-secondary: #9B9B9B;
    --accent: #10A37F;
    --accent-hover: #0E8C6F;
    --border-color: #2D2D2D;
    --input-bg: #2D2D2D;
}

/* Global container */
.gradio-container {
    background: var(--bg-primary) !important;
    font-family: 'Söhne', 'Segoe UI', 'Helvetica Neue', sans-serif;
    color: var(--text-primary) !important;
}

/* Main content area */
#col-container {
    margin: 0 auto;
    max-width: 900px;
    padding: 40px 20px;
}

/* Title styling */
.gradio-container h1 {
    color: var(--text-primary) !important;
    font-size: 2rem;
    font-weight: 600;
    margin-bottom: 30px;
    text-align: center;
    letter-spacing: -0.5px;
}

/* All text and labels */
.gradio-container label,
.gradio-container .label,
.gradio-container p {
    color: var(--text-primary) !important;
    font-size: 14px;
    font-weight: 500;
    margin-bottom: 8px;
}

/* Input fields */
.gradio-container .gradio-textbox input,
.gradio-container .gradio-textbox textarea,
.gradio-container input[type="text"] {
    background: var(--input-bg) !important;
    border: 1px solid var(--border-color) !important;
    border-radius: 8px !important;
    color: var(--text-primary) !important;
    padding: 12px 16px !important;
    font-size: 15px;
    transition: all 0.2s ease;
}

.gradio-container .gradio-textbox input:focus,
.gradio-container .gradio-textbox textarea:focus {
    outline: none !important;
    border-color: var(--accent) !important;
    box-shadow: 0 0 0 3px rgba(16, 163, 127, 0.1) !important;
}

.gradio-container .gradio-textbox input::placeholder,
.gradio-container .gradio-textbox textarea::placeholder {
    color: var(--text-secondary) !important;
}

/* Buttons */
.gradio-container button,
.gradio-container .gradio-button {
    background: var(--accent) !important;
    border: none !important;
    border-radius: 8px !important;
    color: white !important;
    font-weight: 500 !important;
    font-size: 14px !important;
    padding: 12px 24px !important;
    cursor: pointer;
    transition: all 0.2s ease;
    box-shadow: none !important;
}

.gradio-container button:hover,
.gradio-container .gradio-button:hover {
    background: var(--accent-hover) !important;
    transform: translateY(-1px);
}

.gradio-container button:active {
    transform: translateY(0);
}

/* Secondary buttons (examples) */
.gradio-container .gradio-examples button {
    background: var(--bg-tertiary) !important;
    border: 1px solid var(--border-color) !important;
    color: var(--text-primary) !important;
}

.gradio-container .gradio-examples button:hover {
    background: var(--input-bg) !important;
    border-color: var(--accent) !important;
}

/* Image container */
.gradio-container .gradio-image {
    background: var(--bg-secondary) !important;
    border: 1px solid var(--border-color) !important;
    border-radius: 12px !important;
    overflow: hidden;
    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.4) !important;
}

/* Accordion */
.gradio-container .gradio-accordion {
    background: var(--bg-secondary) !important;
    border: 1px solid var(--border-color) !important;
    border-radius: 10px !important;
    margin: 20px 0;
}

.gradio-container .gradio-accordion summary {
    background: transparent !important;
    color: var(--text-primary) !important;
    padding: 16px 20px !important;
    font-weight: 500;
    cursor: pointer;
    border: none !important;
}

.gradio-container .gradio-accordion[open] summary {
    border-bottom: 1px solid var(--border-color) !important;
}

/* Sliders */
.gradio-container .gradio-slider {
    background: transparent !important;
}

.gradio-container .gradio-slider input[type="range"] {
    -webkit-appearance: none;
    appearance: none;
    height: 4px;
    border-radius: 2px;
    background: var(--input-bg) !important;
    outline: none;
}

.gradio-container .gradio-slider input[type="range"]::-webkit-slider-thumb {
    -webkit-appearance: none;
    appearance: none;
    width: 18px;
    height: 18px;
    border-radius: 50%;
    background: var(--accent) !important;
    cursor: pointer;
    transition: all 0.2s ease;
}

.gradio-container .gradio-slider input[type="range"]::-webkit-slider-thumb:hover {
    transform: scale(1.1);
    box-shadow: 0 0 0 4px rgba(16, 163, 127, 0.2);
}

.gradio-container .gradio-slider input[type="range"]::-moz-range-thumb {
    width: 18px;
    height: 18px;
    border-radius: 50%;
    background: var(--accent) !important;
    border: none;
    cursor: pointer;
}

/* Slider number display */
.gradio-container .gradio-slider input[type="number"] {
    background: var(--input-bg) !important;
    border: 1px solid var(--border-color) !important;
    color: var(--text-primary) !important;
    border-radius: 6px !important;
}

/* Checkbox */
.gradio-container .gradio-checkbox input[type="checkbox"] {
    width: 18px;
    height: 18px;
    border-radius: 4px;
    border: 2px solid var(--border-color) !important;
    background: var(--input-bg) !important;
    cursor: pointer;
}

.gradio-container .gradio-checkbox input[type="checkbox"]:checked {
    background: var(--accent) !important;
    border-color: var(--accent) !important;
}

.gradio-container .gradio-checkbox label {
    color: var(--text-primary) !important;
}

/* Examples section */
.gradio-container .gradio-examples {
    background: var(--bg-secondary) !important;
    border: 1px solid var(--border-color) !important;
    border-radius: 10px;
    padding: 20px;
    margin: 20px;
}

.gradio-container .gradio-examples .label {
    color: var(--text-primary) !important;
    font-weight: 500;
    margin-bottom: 12px;
}

/* Rows */
.gradio-container .gradio-row {
    gap: 16px;
}

/* Progress bar */
.gradio-container .gradio-progress {
    background: var(--bg-secondary) !important;
    border-radius: 8px;
    overflow: hidden;
}

.gradio-container .gradio-progress-bar {
    background: var(--accent) !important;
}

/* Remove default Gradio backgrounds */
.gradio-container .gradio-box,
.gradio-container .gradio-form,
.gradio-container .gradio-group {
    background: transparent !important;
    border: none !important;
}

/* Scrollbar styling */
::-webkit-scrollbar {
    width: 8px;
    height: 8px;
}

::-webkit-scrollbar-track {
    background: var(--bg-secondary);
}

::-webkit-scrollbar-thumb {
    background: var(--input-bg);
    border-radius: 4px;
}

::-webkit-scrollbar-thumb:hover {
    background: var(--border-color);
}

/* Responsive design */
@media (max-width: 768px) {
    #col-container {
        padding: 20px 15px;
    }
    
    .gradio-container h1 {
        font-size: 1.5rem;
    }
    
    .gradio-container .gradio-row {
        flex-direction: column;
    }
}

/* Smooth animations */
* {
    transition: background-color 0.2s ease, border-color 0.2s ease;
}
"""

with gr.Blocks(css=css, title="AI Image Generator", theme=gr.themes.Base()) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# AI Image Generator")
        gr.Markdown("*Powered by SDXL-Turbo - Optimized for fast generation (1-4 steps)*")

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=True,
                max_lines=1,
                placeholder="Describe your image...",
                container=True,
                scale=4
            )

            run_button = gr.Button("Generate", scale=1, variant="primary")

        result = gr.Image(label="Generated Image", show_label=True, height=400)

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(
                label="Negative Prompt",
                max_lines=1,
                placeholder="What to exclude from the image...",
                visible=True,
            )

            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )

            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=512,
                )

                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=512,
                )

            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance Scale (Ignored - SDXL-Turbo uses 0.0)",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=0.0,
                    interactive=False,
                )

                num_inference_steps = gr.Slider(
                    label="Inference Steps (1-4 recommended)",
                    minimum=1,
                    maximum=10,
                    step=1,
                    value=2,
                )

        gr.Examples(
            examples=examples, 
            inputs=[prompt],
            label="Example Prompts"
        )
        
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed],
    )

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