Spaces:

AstrodiR
/

YESNS

Running on Zero

YESNS / app.py

Update app.py

1bda36b verified 2 days ago

18.7 kB

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

import gc
import os
import threading
import traceback

import gradio as gr
import torch
from diffusers import (
    FluxKontextPipeline,
    StableDiffusionPipeline,
    StableDiffusionXLPipeline,
)
from huggingface_hub import hf_hub_download
from PIL import Image
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"

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"

current_pipe = None
current_model = None
nsfw_detector = None

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

generation_lock = threading.Lock()
edit_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 (информативный флаг, без блокировки)
# =========================

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) >= 0.50)


# =========================
# 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


# =========================
# 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]
        is_nsfw = check_nsfw(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]
        is_nsfw = check_nsfw(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


# =========================
# 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. `is_nsfw` is an informational flag only "
            "(detector result), it does not block output.\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",
        )

    # Скрытая кнопка только для 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()

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

	import gc
	import os
	import threading
	import traceback

	import gradio as gr
	import torch
	from diffusers import (
	FluxKontextPipeline,
	StableDiffusionPipeline,
	StableDiffusionXLPipeline,
	)
	from huggingface_hub import hf_hub_download
	from PIL import Image
	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"

	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"

	current_pipe = None
	current_model = None
	nsfw_detector = None

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

	generation_lock = threading.Lock()
	edit_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 (информативный флаг, без блокировки)
	# =========================

	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) >= 0.50)


	# =========================
	# 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


	# =========================
	# 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]
	is_nsfw = check_nsfw(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]
	is_nsfw = check_nsfw(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


	# =========================
	# 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. `is_nsfw` is an informational flag only "
	"(detector result), it does not block output.\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",
	)

	# Скрытая кнопка только для 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()