import spaces  # первым, до всего остального

import gc
import os
import threading
import traceback

import gradio as gr
import torch
from diffusers import (
    FluxKontextPipeline,
    StableDiffusionImg2ImgPipeline,
    StableDiffusionInpaintPipeline,
    StableDiffusionPipeline,
    StableDiffusionXLImg2ImgPipeline,
    StableDiffusionXLInpaintPipeline,
    StableDiffusionXLPipeline,
)
from huggingface_hub import hf_hub_download
from PIL import Image, ImageFilter
from transformers import pipeline as hf_pipeline


# =========================
# CONFIGURATION
# =========================

# Секрет задаётся в Settings -> Repository secrets на HF, не в коде.
API_KEY = os.environ.get("API_KEY", "")

# Нужен для скачивания гейтед-весов FLUX.1 Kontext [dev] (нужно принять
# лицензию на странице модели на HF своим аккаунтом). Обычный read-токен,
# вызовы инференса идут локально на CPU и сами по себе не тарифицируются.
HF_TOKEN = os.environ.get("HF_TOKEN", "")

DEVICE = "cpu"

NSFW_REPO = "Falconsai/nsfw_image_detection"

# =========================
# NSFW FILTER MODE (настраивается в коде, не через UI/API)
#
# "on"  -> NSFW-картинки реально замазываются (gaussian blur) перед
#          отдачей пользователю; is_nsfw всё равно возвращается, чтобы
#          было видно, что сработал фильтр.
# "off" -> фильтр отключён полностью, is_nsfw работает только как
#          информационный флаг, ничего не блокируется.
# =========================
NSFW_FILTER_MODE = "off"  # "on" | "off"

NSFW_THRESHOLD = 0.50
NSFW_BLUR_RADIUS = 40

MODELS = {
    "Counterfeit V3.0": {
        "kind": "sd15_single_file",
        "repo": "gsdf/Counterfeit-V3.0",
        "file": "Counterfeit-V3.0_fp32.safetensors",
    },
    "Pony Diffusion V6 XL": {
        "kind": "sdxl_single_file",
        "repo": "Magamanny/Pony-Diffusion-V6-XL",
        "file": "ponyDiffusionV6XL_v6StartWithThisOne.safetensors",
    },
    "URPM v2.3 Final": {
        "kind": "sd15_pretrained",
        "repo": "TheImposterImposters/URPM-v2.3Final",
    },
}

MODEL_NAMES = list(MODELS.keys())

# Отдельная модель для редактирования существующих изображений по
# текстовой инструкции (image-in -> image-out), в отличие от обычных
# text-to-image моделей выше.
EDIT_MODEL_REPO = "black-forest-labs/FLUX.1-Kontext-dev"

# =========================
# IMG2IMG MODELS (новые, отдельные от text-to-image)
# =========================
IMG2IMG_MODELS = {
    "Counterfeit V3.0 (img2img)": {
        "kind": "sd15_single_file",
        "repo": "gsdf/Counterfeit-V3.0",
        "file": "Counterfeit-V3.0_fp32.safetensors",
    },
    "SDXL Base 1.0 (img2img)": {
        "kind": "sdxl_pretrained",
        "repo": "stabilityai/stable-diffusion-xl-base-1.0",
    },
    "URPM v2.3 Final (img2img)": {
        "kind": "sd15_pretrained",
        "repo": "TheImposterImposters/URPM-v2.3Final",
    },
}
IMG2IMG_MODEL_NAMES = list(IMG2IMG_MODELS.keys())

# =========================
# INPAINTING MODELS (новые)
# =========================
INPAINT_MODELS = {
    "SD 1.5 Inpainting": {
        "kind": "sd15_pretrained",
        "repo": "runwayml/stable-diffusion-inpainting",
    },
    "SDXL Inpainting 0.1": {
        "kind": "sdxl_pretrained",
        "repo": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
    },
}
INPAINT_MODEL_NAMES = list(INPAINT_MODELS.keys())

current_pipe = None
current_model = None
nsfw_detector = None

edit_pipe = None  # ленивая загрузка FLUX.1 Kontext [dev]

img2img_pipe = None
img2img_model = None

inpaint_pipe = None
inpaint_model = None

generation_lock = threading.Lock()
edit_lock = threading.Lock()
img2img_lock = threading.Lock()
inpaint_lock = threading.Lock()


# =========================
# ZEROGPU STARTUP COMPATIBILITY
# =========================
# Нужна, чтобы ZeroGPU-сканер при старте Space нашёл @spaces.GPU
# функцию. Подключена к скрытой кнопке в Gradio ниже, чтобы
# сканер видел её как часть интерфейса. Реальный инференс работает
# на CPU (DEVICE = "cpu").

@spaces.GPU(duration=1)
def _zerogpu_compatibility():
    return "ok"


# =========================
# NSFW DETECTOR + FILTER
# =========================

def get_nsfw_detector():
    global nsfw_detector

    if nsfw_detector is None:
        print("Loading NSFW classifier on CPU...")
        nsfw_detector = hf_pipeline(
            "image-classification",
            model=NSFW_REPO,
            device=-1,
        )
        print("NSFW classifier loaded.")

    return nsfw_detector


def check_nsfw(image: Image.Image) -> bool:
    """Возвращает True/False по результату классификатора."""
    predictions = get_nsfw_detector()(image)

    scores = [
        float(item["score"])
        for item in predictions
        if "nsfw" in item["label"].lower()
    ]

    return bool(scores and max(scores) >= NSFW_THRESHOLD)


def apply_nsfw_filter(image: Image.Image):
    """Прогоняет картинку через детектор и, если включён режим "on"
    и картинка признана NSFW, реально замазывает её перед отдачей.
    Возвращает (итоговая_картинка, is_nsfw)."""

    is_nsfw = check_nsfw(image)

    if is_nsfw and NSFW_FILTER_MODE == "on":
        blurred = image.filter(ImageFilter.GaussianBlur(radius=NSFW_BLUR_RADIUS))
        return blurred, is_nsfw

    return image, is_nsfw


# =========================
# MODEL LOADING (text-to-image)
# =========================

def unload_model():
    global current_pipe, current_model

    if current_pipe is not None:
        del current_pipe
        current_pipe = None
        current_model = None

    gc.collect()


def load_model(model_name: str):
    global current_pipe, current_model

    if model_name not in MODELS:
        raise ValueError(f"Unknown model: {model_name}")

    if current_pipe is not None and current_model == model_name:
        return current_pipe

    unload_model()

    spec = MODELS[model_name]
    kind = spec["kind"]

    print(f"Loading {model_name} on CPU ({kind})...")

    if kind == "sd15_single_file":
        model_path = hf_hub_download(repo_id=spec["repo"], filename=spec["file"])
        pipe = StableDiffusionPipeline.from_single_file(
            model_path,
            torch_dtype=torch.float32,
            use_safetensors=True,
        )

    elif kind == "sdxl_single_file":
        model_path = hf_hub_download(repo_id=spec["repo"], filename=spec["file"])
        pipe = StableDiffusionXLPipeline.from_single_file(
            model_path,
            config="stabilityai/stable-diffusion-xl-base-1.0",
            torch_dtype=torch.float32,
            use_safetensors=True,
        )

    elif kind == "sd15_pretrained":
        pipe = StableDiffusionPipeline.from_pretrained(
            spec["repo"],
            torch_dtype=torch.float32,
            use_safetensors=True,
        )

    elif kind == "sdxl_pretrained":
        pipe = StableDiffusionXLPipeline.from_pretrained(
            spec["repo"],
            torch_dtype=torch.float32,
            use_safetensors=True,
        )

    else:
        raise ValueError(f"Unsupported model kind: {kind}")

    pipe = pipe.to(DEVICE)
    pipe.enable_attention_slicing()

    current_pipe = pipe
    current_model = model_name

    print(f"{model_name} loaded on CPU.")
    return current_pipe


# =========================
# MODEL LOADING (image editing — FLUX.1 Kontext [dev])
# =========================

def load_edit_model():
    """Грузит FLUX.1 Kontext [dev] один раз и держит в памяти.
    Репозиторий гейтед — нужен HF_TOKEN с принятой лицензией модели."""
    global edit_pipe

    if edit_pipe is not None:
        return edit_pipe

    if not HF_TOKEN:
        raise ValueError(
            "HF_TOKEN is not set. Add it in Settings -> Repository secrets "
            "(a normal read token is enough) and accept the license on the "
            f"{EDIT_MODEL_REPO} model page with the same account."
        )

    print(f"Loading {EDIT_MODEL_REPO} on CPU (this is a large model, be patient)...")

    # bfloat16 вместо float32, чтобы не раздувать RAM ещё сильнее —
    # на CPU это по-прежнему работает, просто медленнее float32 на
    # некоторых операциях.
    pipe = FluxKontextPipeline.from_pretrained(
        EDIT_MODEL_REPO,
        torch_dtype=torch.bfloat16,
        token=HF_TOKEN,
    )
    pipe = pipe.to(DEVICE)
    pipe.enable_attention_slicing()

    edit_pipe = pipe
    print(f"{EDIT_MODEL_REPO} loaded on CPU.")
    return edit_pipe


# =========================
# MODEL LOADING (img2img)
# =========================

def unload_img2img_model():
    global img2img_pipe, img2img_model

    if img2img_pipe is not None:
        del img2img_pipe
        img2img_pipe = None
        img2img_model = None

    gc.collect()


def load_img2img_model(model_name: str):
    global img2img_pipe, img2img_model

    if model_name not in IMG2IMG_MODELS:
        raise ValueError(f"Unknown img2img model: {model_name}")

    if img2img_pipe is not None and img2img_model == model_name:
        return img2img_pipe

    unload_img2img_model()

    spec = IMG2IMG_MODELS[model_name]
    kind = spec["kind"]

    print(f"Loading {model_name} on CPU ({kind}, img2img)...")

    if kind == "sd15_single_file":
        model_path = hf_hub_download(repo_id=spec["repo"], filename=spec["file"])
        pipe = StableDiffusionImg2ImgPipeline.from_single_file(
            model_path,
            torch_dtype=torch.float32,
            use_safetensors=True,
        )

    elif kind == "sdxl_single_file":
        model_path = hf_hub_download(repo_id=spec["repo"], filename=spec["file"])
        pipe = StableDiffusionXLImg2ImgPipeline.from_single_file(
            model_path,
            config="stabilityai/stable-diffusion-xl-base-1.0",
            torch_dtype=torch.float32,
            use_safetensors=True,
        )

    elif kind == "sd15_pretrained":
        pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
            spec["repo"],
            torch_dtype=torch.float32,
            use_safetensors=True,
        )

    elif kind == "sdxl_pretrained":
        pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
            spec["repo"],
            torch_dtype=torch.float32,
            use_safetensors=True,
        )

    else:
        raise ValueError(f"Unsupported img2img model kind: {kind}")

    pipe = pipe.to(DEVICE)
    pipe.enable_attention_slicing()

    img2img_pipe = pipe
    img2img_model = model_name

    print(f"{model_name} loaded on CPU.")
    return img2img_pipe


# =========================
# MODEL LOADING (inpainting)
# =========================

def unload_inpaint_model():
    global inpaint_pipe, inpaint_model

    if inpaint_pipe is not None:
        del inpaint_pipe
        inpaint_pipe = None
        inpaint_model = None

    gc.collect()


def load_inpaint_model(model_name: str):
    global inpaint_pipe, inpaint_model

    if model_name not in INPAINT_MODELS:
        raise ValueError(f"Unknown inpaint model: {model_name}")

    if inpaint_pipe is not None and inpaint_model == model_name:
        return inpaint_pipe

    unload_inpaint_model()

    spec = INPAINT_MODELS[model_name]
    kind = spec["kind"]

    print(f"Loading {model_name} on CPU ({kind}, inpaint)...")

    if kind == "sd15_pretrained":
        pipe = StableDiffusionInpaintPipeline.from_pretrained(
            spec["repo"],
            torch_dtype=torch.float32,
            use_safetensors=True,
        )

    elif kind == "sdxl_pretrained":
        pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
            spec["repo"],
            torch_dtype=torch.float32,
            use_safetensors=True,
        )

    else:
        raise ValueError(f"Unsupported inpaint model kind: {kind}")

    pipe = pipe.to(DEVICE)
    pipe.enable_attention_slicing()

    inpaint_pipe = pipe
    inpaint_model = model_name

    print(f"{model_name} loaded on CPU.")
    return inpaint_pipe


# =========================
# IMAGE GENERATION (общее ядро, text-to-image)
# =========================

def generate_core(
    model_name,
    prompt,
    negative_prompt,
    steps,
    guidance_scale,
    width,
    height,
    seed,
):
    if not prompt or not prompt.strip():
        raise ValueError("prompt cannot be empty")

    steps = int(steps)
    width = int(width)
    height = int(height)
    guidance_scale = float(guidance_scale)
    seed = int(seed)

    if not 1 <= steps <= 40:
        raise ValueError("steps must be between 1 and 40")

    if not 256 <= width <= 1024 or width % 8:
        raise ValueError("width must be 256-1024 and divisible by 8")

    if not 256 <= height <= 1024 or height % 8:
        raise ValueError("height must be 256-1024 and divisible by 8")

    with generation_lock:
        pipe = load_model(model_name)

        generator = (
            torch.Generator(device="cpu").manual_seed(seed)
            if seed >= 0
            else None
        )

        with torch.inference_mode():
            result = pipe(
                prompt=prompt.strip(),
                negative_prompt=negative_prompt or "",
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                width=width,
                height=height,
                generator=generator,
            )

        image = result.images[0]
        image, is_nsfw = apply_nsfw_filter(image)

        return image, is_nsfw


# =========================
# IMAGE EDITING (ядро, FLUX.1 Kontext [dev])
# =========================

def edit_core(
    input_image,
    prompt,
    steps,
    guidance_scale,
    seed,
):
    if input_image is None:
        raise ValueError("input image is required")

    if not prompt or not prompt.strip():
        raise ValueError("prompt cannot be empty")

    steps = int(steps)
    guidance_scale = float(guidance_scale)
    seed = int(seed)

    if not 1 <= steps <= 50:
        raise ValueError("steps must be between 1 and 50")

    if not isinstance(input_image, Image.Image):
        input_image = Image.fromarray(input_image)

    input_image = input_image.convert("RGB")

    with edit_lock:
        pipe = load_edit_model()

        generator = (
            torch.Generator(device="cpu").manual_seed(seed)
            if seed >= 0
            else None
        )

        with torch.inference_mode():
            result = pipe(
                image=input_image,
                prompt=prompt.strip(),
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                generator=generator,
            )

        image = result.images[0]
        image, is_nsfw = apply_nsfw_filter(image)

        return image, is_nsfw


# =========================
# IMG2IMG (ядро)
# =========================

def img2img_core(
    model_name,
    input_image,
    prompt,
    negative_prompt,
    strength,
    steps,
    guidance_scale,
    width,
    height,
    seed,
):
    if input_image is None:
        raise ValueError("input image is required")

    if not prompt or not prompt.strip():
        raise ValueError("prompt cannot be empty")

    steps = int(steps)
    guidance_scale = float(guidance_scale)
    strength = float(strength)
    seed = int(seed)
    width = int(width)
    height = int(height)

    if not 1 <= steps <= 50:
        raise ValueError("steps must be between 1 and 50")

    if not 0.0 < strength <= 1.0:
        raise ValueError("strength must be between 0 and 1")

    if not 256 <= width <= 1024 or width % 8:
        raise ValueError("width must be 256-1024 and divisible by 8")

    if not 256 <= height <= 1024 or height % 8:
        raise ValueError("height must be 256-1024 and divisible by 8")

    if not isinstance(input_image, Image.Image):
        input_image = Image.fromarray(input_image)

    input_image = input_image.convert("RGB").resize((width, height))

    with img2img_lock:
        pipe = load_img2img_model(model_name)

        generator = (
            torch.Generator(device="cpu").manual_seed(seed)
            if seed >= 0
            else None
        )

        with torch.inference_mode():
            result = pipe(
                prompt=prompt.strip(),
                negative_prompt=negative_prompt or "",
                image=input_image,
                strength=strength,
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                generator=generator,
            )

        image = result.images[0]
        image, is_nsfw = apply_nsfw_filter(image)

        return image, is_nsfw


# =========================
# INPAINTING (ядро)
# =========================

def inpaint_core(
    model_name,
    input_image,
    mask_image,
    prompt,
    negative_prompt,
    steps,
    guidance_scale,
    seed,
):
    if input_image is None:
        raise ValueError("input image is required")

    if mask_image is None:
        raise ValueError("mask image is required")

    if not prompt or not prompt.strip():
        raise ValueError("prompt cannot be empty")

    steps = int(steps)
    guidance_scale = float(guidance_scale)
    seed = int(seed)

    if not 1 <= steps <= 50:
        raise ValueError("steps must be between 1 and 50")

    if not isinstance(input_image, Image.Image):
        input_image = Image.fromarray(input_image)

    if not isinstance(mask_image, Image.Image):
        mask_image = Image.fromarray(mask_image)

    input_image = input_image.convert("RGB")
    mask_image = mask_image.convert("L")

    if mask_image.size != input_image.size:
        mask_image = mask_image.resize(input_image.size)

    with inpaint_lock:
        pipe = load_inpaint_model(model_name)

        generator = (
            torch.Generator(device="cpu").manual_seed(seed)
            if seed >= 0
            else None
        )

        with torch.inference_mode():
            result = pipe(
                prompt=prompt.strip(),
                negative_prompt=negative_prompt or "",
                image=input_image,
                mask_image=mask_image,
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                generator=generator,
            )

        image = result.images[0]
        image, is_nsfw = apply_nsfw_filter(image)

        return image, is_nsfw


# =========================
# GRADIO FUNCTIONS
# =========================

def generate_ui(
    model_name,
    prompt,
    negative_prompt,
    steps,
    guidance_scale,
    width,
    height,
    seed,
):
    """Для кнопки в веб-интерфейсе — без проверки ключа."""
    try:
        image, is_nsfw = generate_core(
            model_name, prompt, negative_prompt,
            steps, guidance_scale, width, height, seed,
        )
        return image, str(is_nsfw).lower()

    except Exception as exc:
        print(f"[generate_ui] ERROR: {exc}")
        traceback.print_exc()
        raise gr.Error(f"Generation error: {exc}") from exc


def generate_api(
    model_name,
    prompt,
    negative_prompt,
    steps,
    guidance_scale,
    width,
    height,
    seed,
    api_key,
):
    """Для вызова через Gradio API (POST /gradio_api/call/generate).
    Ключ передаётся последним параметром в data, а не в заголовке —
    у Gradio API нет кастомных заголовков."""

    if API_KEY and api_key != API_KEY:
        raise gr.Error("Unauthorized: invalid api_key")

    try:
        image, is_nsfw = generate_core(
            model_name, prompt, negative_prompt,
            steps, guidance_scale, width, height, seed,
        )
        # image (PIL) отдаём как есть — Gradio сам сохранит во временный
        # файл и в ответе API пришлёт на него ссылку (url/path).
        return image, str(is_nsfw).lower()

    except gr.Error:
        raise
    except Exception as exc:
        print(f"[generate_api] ERROR: {exc}")
        traceback.print_exc()
        raise gr.Error(f"Generation error: {exc}") from exc


def edit_ui(input_image, prompt, steps, guidance_scale, seed):
    """Для кнопки в веб-интерфейсе — без проверки ключа."""
    try:
        image, is_nsfw = edit_core(input_image, prompt, steps, guidance_scale, seed)
        return image, str(is_nsfw).lower()

    except Exception as exc:
        print(f"[edit_ui] ERROR: {exc}")
        traceback.print_exc()
        raise gr.Error(f"Edit error: {exc}") from exc


def edit_api(input_image, prompt, steps, guidance_scale, seed, api_key):
    """Для вызова через Gradio API (POST /gradio_api/call/edit)."""

    if API_KEY and api_key != API_KEY:
        raise gr.Error("Unauthorized: invalid api_key")

    try:
        image, is_nsfw = edit_core(input_image, prompt, steps, guidance_scale, seed)
        return image, str(is_nsfw).lower()

    except gr.Error:
        raise
    except Exception as exc:
        print(f"[edit_api] ERROR: {exc}")
        traceback.print_exc()
        raise gr.Error(f"Edit error: {exc}") from exc


def img2img_ui(model_name, input_image, prompt, negative_prompt, strength, steps, guidance_scale, width, height, seed):
    """Для кнопки в веб-интерфейсе — без проверки ключа."""
    try:
        image, is_nsfw = img2img_core(
            model_name, input_image, prompt, negative_prompt,
            strength, steps, guidance_scale, width, height, seed,
        )
        return image, str(is_nsfw).lower()

    except Exception as exc:
        print(f"[img2img_ui] ERROR: {exc}")
        traceback.print_exc()
        raise gr.Error(f"Img2img error: {exc}") from exc


def img2img_api(model_name, input_image, prompt, negative_prompt, strength, steps, guidance_scale, width, height, seed, api_key):
    """Для вызова через Gradio API (POST /gradio_api/call/img2img)."""

    if API_KEY and api_key != API_KEY:
        raise gr.Error("Unauthorized: invalid api_key")

    try:
        image, is_nsfw = img2img_core(
            model_name, input_image, prompt, negative_prompt,
            strength, steps, guidance_scale, width, height, seed,
        )
        return image, str(is_nsfw).lower()

    except gr.Error:
        raise
    except Exception as exc:
        print(f"[img2img_api] ERROR: {exc}")
        traceback.print_exc()
        raise gr.Error(f"Img2img error: {exc}") from exc


def inpaint_ui(model_name, input_image, mask_image, prompt, negative_prompt, steps, guidance_scale, seed):
    """Для кнопки в веб-интерфейсе — без проверки ключа."""
    try:
        image, is_nsfw = inpaint_core(
            model_name, input_image, mask_image, prompt, negative_prompt,
            steps, guidance_scale, seed,
        )
        return image, str(is_nsfw).lower()

    except Exception as exc:
        print(f"[inpaint_ui] ERROR: {exc}")
        traceback.print_exc()
        raise gr.Error(f"Inpaint error: {exc}") from exc


def inpaint_api(model_name, input_image, mask_image, prompt, negative_prompt, steps, guidance_scale, seed, api_key):
    """Для вызова через Gradio API (POST /gradio_api/call/inpaint)."""

    if API_KEY and api_key != API_KEY:
        raise gr.Error("Unauthorized: invalid api_key")

    try:
        image, is_nsfw = inpaint_core(
            model_name, input_image, mask_image, prompt, negative_prompt,
            steps, guidance_scale, seed,
        )
        return image, str(is_nsfw).lower()

    except gr.Error:
        raise
    except Exception as exc:
        print(f"[inpaint_api] ERROR: {exc}")
        traceback.print_exc()
        raise gr.Error(f"Inpaint error: {exc}") from exc


# =========================
# GRADIO UI
# =========================

with gr.Blocks(title="CPU Image Generation API") as demo:
    gr.Markdown("# Image Generation API")

    with gr.Tab("Generate"):
        gr.Markdown(
            "Generation runs on CPU. The NSFW filter mode is set in code "
            "(`NSFW_FILTER_MODE`): `\"on\"` blurs NSFW results before they "
            "are returned, `\"off\"` disables filtering and `is_nsfw` is "
            "informational only.\n\n"
            "**Programmatic access (Gradio API):**\n"
            "1. `POST https://<space>.hf.space/gradio_api/call/generate` "
            "with JSON body `{\"data\": [model_name, prompt, negative_prompt, "
            "steps, guidance_scale, width, height, seed, api_key]}`\n"
            "2. Response contains an `event_id`.\n"
            "3. `GET https://<space>.hf.space/gradio_api/call/generate/<event_id>` "
            "to fetch the result (image file URL + is_nsfw flag)."
        )

        model = gr.Dropdown(choices=MODEL_NAMES, value=MODEL_NAMES[0], label="Model")
        prompt = gr.Textbox(label="Prompt", lines=4)
        negative_prompt = gr.Textbox(
            label="Negative prompt",
            value="low quality, blurry, bad anatomy",
            lines=2,
        )

        with gr.Row():
            steps = gr.Slider(1, 40, value=20, step=1, label="Steps")
            guidance = gr.Slider(1, 12, value=7, step=0.5, label="Guidance scale")

        with gr.Row():
            width = gr.Dropdown([512, 640, 768, 1024], value=512, label="Width")
            height = gr.Dropdown([512, 640, 768, 1024], value=512, label="Height")

        seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
        generate_button = gr.Button("Generate", variant="primary")

        output = gr.Image(label="Generated image", type="pil")
        nsfw_output = gr.Textbox(label="is_nsfw", interactive=False)

        ui_inputs = [model, prompt, negative_prompt, steps, guidance, width, height, seed]

        generate_button.click(
            fn=generate_ui,
            inputs=ui_inputs,
            outputs=[output, nsfw_output],
        )

        # Скрытое поле под API-ключ + скрытая кнопка, которая публикует
        # generate_api как именованный эндпоинт Gradio API ("generate").
        api_key_box = gr.Textbox(visible=False)
        api_output_image = gr.Image(visible=False, type="pil")
        api_output_nsfw = gr.Textbox(visible=False)
        api_button = gr.Button("API endpoint", visible=False)

        api_button.click(
            fn=generate_api,
            inputs=ui_inputs + [api_key_box],
            outputs=[api_output_image, api_output_nsfw],
            api_name="generate",
        )

    with gr.Tab("Edit Image (FLUX.1 Kontext)"):
        gr.Markdown(
            f"Instruction-based image editing with `{EDIT_MODEL_REPO}`, also "
            "running on CPU. Upload a source image, describe the change in "
            "plain language (e.g. \"add a wooden chair next to the table\"), "
            "and the model edits only what you asked for, keeping the rest "
            "of the image intact.\n\n"
            "This is a large (~12B) model: first call downloads the weights "
            "(needs `HF_TOKEN` with the model's license accepted) and loads "
            "them into RAM, and each edit can take a long time on CPU.\n\n"
            "**Programmatic access (Gradio API):**\n"
            "1. `POST https://<space>.hf.space/gradio_api/call/edit` "
            "with JSON body `{\"data\": [input_image, prompt, steps, "
            "guidance_scale, seed, api_key]}` (`input_image` as a base64 "
            "data URL or a Gradio file reference)\n"
            "2. Response contains an `event_id`.\n"
            "3. `GET https://<space>.hf.space/gradio_api/call/edit/<event_id>` "
            "to fetch the result (image file URL + is_nsfw flag)."
        )

        edit_input_image = gr.Image(label="Source image", type="pil")
        edit_prompt = gr.Textbox(label="Edit instruction", lines=3, placeholder="Add a wooden chair next to the table")

        with gr.Row():
            edit_steps = gr.Slider(1, 50, value=28, step=1, label="Steps")
            edit_guidance = gr.Slider(1, 10, value=2.5, step=0.1, label="Guidance scale")

        edit_seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
        edit_button = gr.Button("Edit image", variant="primary")

        edit_output = gr.Image(label="Edited image", type="pil")
        edit_nsfw_output = gr.Textbox(label="is_nsfw", interactive=False)

        edit_ui_inputs = [edit_input_image, edit_prompt, edit_steps, edit_guidance, edit_seed]

        edit_button.click(
            fn=edit_ui,
            inputs=edit_ui_inputs,
            outputs=[edit_output, edit_nsfw_output],
        )

        edit_api_key_box = gr.Textbox(visible=False)
        edit_api_output_image = gr.Image(visible=False, type="pil")
        edit_api_output_nsfw = gr.Textbox(visible=False)
        edit_api_button = gr.Button("API endpoint", visible=False)

        edit_api_button.click(
            fn=edit_api,
            inputs=edit_ui_inputs + [edit_api_key_box],
            outputs=[edit_api_output_image, edit_api_output_nsfw],
            api_name="edit",
        )

    with gr.Tab("Img2Img"):
        gr.Markdown(
            "Генерация на основе исходного изображения + текстового промпта "
            "(`strength` — насколько сильно менять картинку: ближе к 1.0 — "
            "почти полная перегенерация, ближе к 0 — лёгкая правка). "
            "Width/Height задают итоговый размер — картинка ресайзится под "
            "них перед генерацией (есть пресеты под 9:16).\n\n"
            "**Programmatic access (Gradio API):**\n"
            "1. `POST https://<space>.hf.space/gradio_api/call/img2img` "
            "with JSON body `{\"data\": [model_name, input_image, prompt, "
            "negative_prompt, strength, steps, guidance_scale, width, height, "
            "seed, api_key]}`\n"
            "2. Response contains an `event_id`.\n"
            "3. `GET https://<space>.hf.space/gradio_api/call/img2img/<event_id>` "
            "to fetch the result (image file URL + is_nsfw flag)."
        )

        i2i_model = gr.Dropdown(choices=IMG2IMG_MODEL_NAMES, value=IMG2IMG_MODEL_NAMES[0], label="Model")
        i2i_input_image = gr.Image(label="Source image", type="pil")
        i2i_prompt = gr.Textbox(label="Prompt", lines=3)
        i2i_negative_prompt = gr.Textbox(
            label="Negative prompt",
            value="low quality, blurry, bad anatomy",
            lines=2,
        )

        with gr.Row():
            i2i_strength = gr.Slider(0.05, 1.0, value=0.6, step=0.05, label="Strength")
            i2i_steps = gr.Slider(1, 50, value=30, step=1, label="Steps")
            i2i_guidance = gr.Slider(1, 12, value=7, step=0.5, label="Guidance scale")

        with gr.Row():
            i2i_width = gr.Dropdown([512, 576, 640, 720, 768, 1024], value=576, label="Width")
            i2i_height = gr.Dropdown([512, 640, 768, 1024, 1280, 1792], value=1024, label="Height")

        i2i_seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
        i2i_button = gr.Button("Generate (img2img)", variant="primary")

        i2i_output = gr.Image(label="Result image", type="pil")
        i2i_nsfw_output = gr.Textbox(label="is_nsfw", interactive=False)
        i2i_download = gr.File(label="Экспортировать картинку")

        i2i_inputs = [
            i2i_model, i2i_input_image, i2i_prompt, i2i_negative_prompt,
            i2i_strength, i2i_steps, i2i_guidance, i2i_width, i2i_height, i2i_seed,
        ]

        i2i_button.click(
            fn=img2img_ui,
            inputs=i2i_inputs,
            outputs=[i2i_output, i2i_nsfw_output],
        ).then(
            fn=lambda img: img,
            inputs=i2i_output,
            outputs=i2i_download,
        )

        i2i_api_key_box = gr.Textbox(visible=False)
        i2i_api_output_image = gr.Image(visible=False, type="pil")
        i2i_api_output_nsfw = gr.Textbox(visible=False)
        i2i_api_button = gr.Button("API endpoint", visible=False)

        i2i_api_button.click(
            fn=img2img_api,
            inputs=i2i_inputs + [i2i_api_key_box],
            outputs=[i2i_api_output_image, i2i_api_output_nsfw],
            api_name="img2img",
        )

    with gr.Tab("Inpainting"):
        gr.Markdown(
            "Перерисовка части изображения по маске: белые области маски "
            "перерисовываются согласно промпту, чёрные остаются нетронутыми.\n\n"
            "**Programmatic access (Gradio API):**\n"
            "1. `POST https://<space>.hf.space/gradio_api/call/inpaint` "
            "with JSON body `{\"data\": [model_name, input_image, mask_image, "
            "prompt, negative_prompt, steps, guidance_scale, seed, api_key]}`\n"
            "2. Response contains an `event_id`.\n"
            "3. `GET https://<space>.hf.space/gradio_api/call/inpaint/<event_id>` "
            "to fetch the result (image file URL + is_nsfw flag)."
        )

        inp_model = gr.Dropdown(choices=INPAINT_MODEL_NAMES, value=INPAINT_MODEL_NAMES[0], label="Model")
        inp_input_image = gr.Image(label="Source image", type="pil")
        inp_mask_image = gr.Image(label="Mask (white = area to repaint)", type="pil")
        inp_prompt = gr.Textbox(label="Prompt", lines=3)
        inp_negative_prompt = gr.Textbox(
            label="Negative prompt",
            value="low quality, blurry, bad anatomy",
            lines=2,
        )

        with gr.Row():
            inp_steps = gr.Slider(1, 50, value=30, step=1, label="Steps")
            inp_guidance = gr.Slider(1, 12, value=7, step=0.5, label="Guidance scale")

        inp_seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
        inp_button = gr.Button("Generate (inpaint)", variant="primary")

        inp_output = gr.Image(label="Result image", type="pil")
        inp_nsfw_output = gr.Textbox(label="is_nsfw", interactive=False)
        inp_download = gr.File(label="Экспортировать картинку")

        inp_inputs = [
            inp_model, inp_input_image, inp_mask_image, inp_prompt,
            inp_negative_prompt, inp_steps, inp_guidance, inp_seed,
        ]

        inp_button.click(
            fn=inpaint_ui,
            inputs=inp_inputs,
            outputs=[inp_output, inp_nsfw_output],
        ).then(
            fn=lambda img: img,
            inputs=inp_output,
            outputs=inp_download,
        )

        inp_api_key_box = gr.Textbox(visible=False)
        inp_api_output_image = gr.Image(visible=False, type="pil")
        inp_api_output_nsfw = gr.Textbox(visible=False)
        inp_api_button = gr.Button("API endpoint", visible=False)

        inp_api_button.click(
            fn=inpaint_api,
            inputs=inp_inputs + [inp_api_key_box],
            outputs=[inp_api_output_image, inp_api_output_nsfw],
            api_name="inpaint",
        )

    # Скрытая кнопка только для ZeroGPU-сканера: подключает
    # _zerogpu_compatibility к интерфейсу Gradio.
    probe_btn = gr.Button("probe", visible=False)
    probe_out = gr.Textbox(visible=False)
    probe_btn.click(
        fn=_zerogpu_compatibility,
        inputs=None,
        outputs=probe_out,
        api_name="probe",
    )


# HF сам находит объект `demo` в файле и поднимает его через
# свою обвязку на порту 7860 — никакого ручного launch()/uvicorn.run()
# не требуется и не должно быть на ZeroGPU-железе. Никакого FastAPI
# и mount_gradio_app здесь больше нет — ZeroGPU-сканер с ними не
# срабатывает (проверено: со своим FastAPI-роутингом Space падал
# с "No @spaces.GPU function detected during startup").
demo.queue().launch()