import spaces
import gradio as gr
import torch
import numpy as np
from PIL import Image
import cv2
from diffusers import (
    StableDiffusionXLControlNetPipeline,
    ControlNetModel,
    AutoencoderKL,
    EulerAncestralDiscreteScheduler,
)

DTYPE = torch.float16

# ---------------------------------------------------------------------------
# Model loading (runs once on startup, stays on GPU via ZeroGPU)
# ---------------------------------------------------------------------------
controlnet = ControlNetModel.from_pretrained(
    "xinsir/controlnet-canny-sdxl-1.0",
    torch_dtype=DTYPE,
)
vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix",
    torch_dtype=DTYPE,
)
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    controlnet=controlnet,
    vae=vae,
    torch_dtype=DTYPE,
    safety_checker=None,
)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

# ---------------------------------------------------------------------------
# Helper: extract Canny edges (resized to ~1024 for best SDXL performance)
# ---------------------------------------------------------------------------
def extract_canny(image: Image.Image, low: int, high: int):
    img = np.array(image.convert("RGB"))
    h, w, _ = img.shape
    ratio = np.sqrt(1024.0 * 1024.0 / (w * h))
    new_w, new_h = int(w * ratio), int(h * ratio)
    img = cv2.resize(img, (new_w, new_h))
    edges = cv2.Canny(img, low, high)
    edges = np.concatenate([edges[:, :, None]] * 3, axis=2)
    return Image.fromarray(edges), new_w, new_h

# ---------------------------------------------------------------------------
# Main generation function — @spaces.GPU activates ZeroGPU during the call
# ---------------------------------------------------------------------------
@spaces.GPU(duration=90)
def generate(input_image, prompt, negative_prompt, canny_low, canny_high,
             guidance_scale, steps, cn_scale, seed):
    if input_image is None:
        raise gr.Error("Bitte lade ein Bild hoch.")
    if not prompt.strip():
        raise gr.Error("Bitte gib einen Prompt ein.")

    pil_image = Image.fromarray(input_image)
    control_image, new_w, new_h = extract_canny(pil_image, int(canny_low), int(canny_high))

    generator = torch.manual_seed(int(seed)) if seed >= 0 else None

    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt or None,
        image=control_image,
        controlnet_conditioning_scale=float(cn_scale),
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        width=new_w,
        height=new_h,
        generator=generator,
    ).images[0]

    return control_image, result

# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
css = """
body { font-family: 'Inter', sans-serif; background: #0f0f11; color: #e8e8f0; }
.gradio-container { max-width: 1100px; margin: 0 auto; }
#title { text-align: center; padding: 2rem 0 0.5rem; }
#title h1 { font-size: 2rem; font-weight: 700; letter-spacing: -0.5px;
             background: linear-gradient(90deg, #f59e0b, #ef4444);
             -webkit-background-clip: text; -webkit-text-fill-color: transparent; }
#title p  { color: #9090a8; font-size: 0.95rem; margin-top: 0.25rem; }
.panel    { background: #1a1a22; border: 1px solid #2a2a38; border-radius: 12px; padding: 1.25rem; }
.generate-btn { background: linear-gradient(135deg, #f59e0b, #ef4444) !important;
                color: white !important; border: none !important;
                font-weight: 600 !important; font-size: 1rem !important;
                border-radius: 8px !important; height: 48px !important; }
.generate-btn:hover { opacity: 0.9 !important; }
"""

with gr.Blocks(css=css, title="ControlNet SDXL Canny") as demo:
    gr.HTML("""
        <div id="title">
            <h1>🔥 ControlNet · SDXL Canny</h1>
            <p>Hochwertige Bildgenerierung mit SDXL auf ZeroGPU. Lade ein Bild hoch, schreib einen Prompt – die Struktur deines Originals bleibt erhalten.</p>
        </div>
    """)

    with gr.Row():
        with gr.Column(scale=1, elem_classes="panel"):
            gr.Markdown("### 📥 Eingabe")
            input_image = gr.Image(label="Referenzbild", type="numpy", height=300)
            prompt = gr.Textbox(label="Prompt", lines=3,
                placeholder="a rugged pirate on a wooden ship, photorealistic, cinematic, 8k")
            negative_prompt = gr.Textbox(label="Negative Prompt (optional)", lines=2,
                placeholder="blurry, low quality, deformed, extra limbs")

            with gr.Accordion("⚙️ Erweiterte Einstellungen", open=False):
                with gr.Row():
                    canny_low  = gr.Slider(0, 255, value=100, step=1, label="Canny Low")
                    canny_high = gr.Slider(0, 255, value=200, step=1, label="Canny High")
                with gr.Row():
                    guidance_scale = gr.Slider(1, 15, value=6.0, step=0.5, label="Guidance Scale")
                    steps          = gr.Slider(15, 50, value=30, step=1, label="Inference Steps")
                cn_scale = gr.Slider(0.1, 2.0, value=0.8, step=0.05,
                    label="ControlNet Stärke (niedriger = mehr Freiheit)")
                seed = gr.Number(value=42, label="Seed (-1 = zufällig)", precision=0)

            run_btn = gr.Button("🎨  Generieren", elem_classes="generate-btn")

        with gr.Column(scale=1, elem_classes="panel"):
            gr.Markdown("### 📤 Ergebnis")
            canny_out  = gr.Image(label="Canny-Kantenbild", height=250)
            result_out = gr.Image(label="Generiertes Bild", height=400)

    run_btn.click(
        fn=generate,
        inputs=[input_image, prompt, negative_prompt, canny_low, canny_high,
                guidance_scale, steps, cn_scale, seed],
        outputs=[canny_out, result_out],
    )

demo.queue().launch()