"""
Depth ControlNet Logo Generator
--------------------------------
- Nimmt ein hochgeladenes Logo/Referenzbild
- Berechnet daraus eine Depth Map (Midas)
- Generiert per Stable Diffusion 1.5 + ControlNet-Depth ein neues,
  strukturell ähnliches Bild nach Text-Prompt
- Läuft auf HuggingFace ZeroGPU (dynamisch zugewiesene GPU pro Request)
"""

import spaces
import gradio as gr
import numpy as np
import torch
from PIL import Image
from diffusers import (
    StableDiffusionControlNetPipeline,
    ControlNetModel,
    UniPCMultistepScheduler,
)
from transformers import pipeline as hf_pipeline

# ---------------------------------------------------------------------------
# Konfiguration
# ---------------------------------------------------------------------------
SD_MODEL_ID = "runwayml/stable-diffusion-v1-5"
CONTROLNET_ID = "lllyasviel/sd-controlnet-depth"
DEPTH_MODEL_ID = "Intel/dpt-hybrid-midas"  # gleicher Preprocessor wie beim ControlNet-Training
RESOLUTION = 512  # native SD1.5 Auflösung

device = "cuda"
dtype = torch.float16

# ---------------------------------------------------------------------------
# Modelle laden (einmalig beim Space-Start)
# Bei ZeroGPU: .to("cuda") hier ist ok, tatsächliche GPU wird erst bei
# @spaces.GPU-Aufrufen zugewiesen.
# ---------------------------------------------------------------------------
print("Lade Depth-Estimator ...")
depth_estimator = hf_pipeline("depth-estimation", model=DEPTH_MODEL_ID)

print("Lade ControlNet + Stable Diffusion Pipeline ...")
controlnet = ControlNetModel.from_pretrained(CONTROLNET_ID, torch_dtype=dtype)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    SD_MODEL_ID,
    controlnet=controlnet,
    torch_dtype=dtype,
    safety_checker=None,
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to(device)

try:
    pipe.enable_xformers_memory_efficient_attention()
except Exception:
    pass


# ---------------------------------------------------------------------------
# Hilfsfunktionen
# ---------------------------------------------------------------------------
def preprocess_image(image: Image.Image, resolution: int = RESOLUTION) -> Image.Image:
    """Logo auf quadratische Zielauflösung bringen."""
    image = image.convert("RGB")
    image = image.resize((resolution, resolution), Image.LANCZOS)
    return image


def get_depth_map(image: Image.Image) -> Image.Image:
    """Depth Map aus dem Logo berechnen (das ist das ControlNet-Kontrollbild)."""
    depth = depth_estimator(image)["depth"]
    depth = np.array(depth)
    depth = depth[:, :, None]
    depth = np.concatenate([depth, depth, depth], axis=2)
    return Image.fromarray(depth)


# ---------------------------------------------------------------------------
# Haupt-Generierungsfunktion (läuft auf der ZeroGPU-Instanz)
# ---------------------------------------------------------------------------
@spaces.GPU(duration=60)
def generate(
    logo_image,
    prompt,
    negative_prompt,
    controlnet_scale,
    num_steps,
    guidance_scale,
    seed,
    progress=gr.Progress(track_tqdm=True),
):
    if logo_image is None:
        raise gr.Error("Bitte zuerst ein Logo-/Referenzbild hochladen.")
    if not prompt or prompt.strip() == "":
        raise gr.Error("Bitte einen Prompt eingeben.")

    logo_image = preprocess_image(logo_image)
    depth_image = get_depth_map(logo_image)

    seed = int(seed)
    if seed < 0:
        generator = None  # zufälliger Seed
    else:
        generator = torch.Generator(device=device).manual_seed(seed)

    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=depth_image,
        num_inference_steps=int(num_steps),
        guidance_scale=float(guidance_scale),
        controlnet_conditioning_scale=float(controlnet_scale),
        generator=generator,
    ).images[0]

    return result, depth_image


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Depth ControlNet Logo Generator") as demo:
    gr.Markdown(
        """
        # 🎨 Depth ControlNet — Logo Generator
        Lade ein bestehendes Logo hoch, gib einen neuen Stil-Prompt ein.
        Die **Tiefenstruktur / Silhouette** des Logos bleibt erhalten,
        während Stil, Farben und Textur komplett neu generiert werden.
        """
    )

    with gr.Row():
        with gr.Column():
            logo_input = gr.Image(
                label="1️⃣ Logo / Referenzbild hochladen",
                type="pil",
                height=300,
            )
            prompt = gr.Textbox(
                label="2️⃣ Prompt",
                placeholder="z.B. vintage japanese emblem logo, ink brush style, minimal, black and red, flat vector, white background",
                lines=3,
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                value="blurry, low quality, watermark, text, extra elements, photo, 3d render",
                lines=2,
            )

            with gr.Accordion("⚙️ Erweiterte Einstellungen", open=False):
                controlnet_scale = gr.Slider(
                    0.0, 2.0, value=1.0, step=0.05,
                    label="ControlNet Conditioning Scale (Struktur-Treue)",
                )
                num_steps = gr.Slider(10, 50, value=25, step=1, label="Inference Steps")
                guidance_scale = gr.Slider(1.0, 20.0, value=7.5, step=0.5, label="Guidance Scale")
                seed = gr.Slider(-1, 999999, value=-1, step=1, label="Seed (-1 = zufällig)")

            run_button = gr.Button("🚀 Generieren", variant="primary")

        with gr.Column():
            output_image = gr.Image(label="Ergebnis", height=400)
            depth_preview = gr.Image(label="Erkannte Depth Map (Kontrollbild)", height=200)

    run_button.click(
        fn=generate,
        inputs=[
            logo_input,
            prompt,
            negative_prompt,
            controlnet_scale,
            num_steps,
            guidance_scale,
            seed,
        ],
        outputs=[output_image, depth_preview],
    )

    gr.Markdown(
        """
        ---
        💡 **Tipps:**
        - Höherer *ControlNet Conditioning Scale* = Form/Struktur des Original-Logos wird strenger befolgt.
        - Niedrigerer Wert = mehr kreative Freiheit, weniger Ähnlichkeit zur Vorlage.
        - Für saubere Vektor-Optik: Begriffe wie "flat vector", "clean lines", "white background" im Prompt verwenden.
        """
    )

demo.queue(max_size=20)
demo.launch()