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://.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://.hf.space/gradio_api/call/generate/` " "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://.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://.hf.space/gradio_api/call/edit/` " "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()