387 lines
14 KiB
Python
387 lines
14 KiB
Python
# AI selection for the layered editor.
|
|
#
|
|
# "Select Subject": BiRefNet (the current salient-object/DIS SOTA,
|
|
# MIT-licensed — the same model class behind Photoshop's one-click Select
|
|
# Subject) segments the posted composite and returns the raw sigmoid output
|
|
# as a grayscale PNG. The matte is SOFT on purpose: hair/fabric edges carry
|
|
# partial alpha, so the editor applies it directly as selection-mask alpha
|
|
# without thresholding. Weights (~0.9GB) download into the HuggingFace cache
|
|
# on the first request; the model then stays resident (~0.5GB FP16).
|
|
#
|
|
# "Select Object" (click-to-select): SAM2 point prompt via ultralytics
|
|
# (already installed for RegionalDetailer; sam2.1_b.pt ~154MB auto-download).
|
|
# The encoder result is cached per image hash, so successive clicks on the
|
|
# same composite only run the prompt decoder (~25ms).
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import io as _io
|
|
import logging
|
|
import threading
|
|
|
|
import server
|
|
from aiohttp import web
|
|
from PIL import Image
|
|
|
|
from .api_utils import error_response
|
|
|
|
routes = server.PromptServer.instance.routes
|
|
log = logging.getLogger("promptchain")
|
|
|
|
BIREFNET_ID = "ZhengPeng7/BiRefNet"
|
|
INFER_SIZE = 1024 # BiRefNet's native training resolution
|
|
|
|
_model = None
|
|
_device = None
|
|
# One segmentation at a time: the lock serializes both the lazy first-load
|
|
# (two concurrent loads would download/initialize twice) and GPU inference.
|
|
_lock = threading.Lock()
|
|
|
|
|
|
def _get_model():
|
|
global _model, _device
|
|
if _model is None:
|
|
import torch
|
|
from transformers import AutoModelForImageSegmentation
|
|
|
|
log.info("[SelectSubject] loading %s (first call downloads ~0.9GB)", BIREFNET_ID)
|
|
try:
|
|
model = AutoModelForImageSegmentation.from_pretrained(
|
|
BIREFNET_ID, trust_remote_code=True)
|
|
except Exception as exc:
|
|
raise RuntimeError(
|
|
f"Couldn't load the BiRefNet selection model. On first use it "
|
|
f"downloads ~0.9GB — check your internet connection, or pre-install "
|
|
f"it from Settings → PromptChain → Features (AI Selection). "
|
|
f"[{type(exc).__name__}: {exc}]"
|
|
) from exc
|
|
_device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
model = model.to(_device).eval()
|
|
if _device == "cuda":
|
|
model = model.half()
|
|
_model = model
|
|
log.info("[SelectSubject] model ready on %s", _device)
|
|
return _model
|
|
|
|
|
|
def _subject_mask(img: Image.Image) -> Image.Image:
|
|
import torch
|
|
from torchvision import transforms
|
|
|
|
with _lock:
|
|
model = _get_model()
|
|
prep = transforms.Compose([
|
|
transforms.Resize((INFER_SIZE, INFER_SIZE)),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
|
])
|
|
inp = prep(img.convert("RGB")).unsqueeze(0).to(_device)
|
|
if _device == "cuda":
|
|
inp = inp.half()
|
|
with torch.no_grad():
|
|
pred = model(inp)[-1].sigmoid().cpu()
|
|
mask = transforms.ToPILImage()(pred[0].squeeze().float())
|
|
return mask.resize(img.size, Image.BILINEAR)
|
|
|
|
|
|
@routes.post("/promptchain/select-subject")
|
|
async def select_subject(request):
|
|
form = await request.post()
|
|
upload = form.get("image")
|
|
if upload is None or not getattr(upload, "file", None):
|
|
return error_response("missing image upload")
|
|
raw = upload.file.read()
|
|
if not raw:
|
|
return error_response("empty image upload")
|
|
try:
|
|
img = Image.open(_io.BytesIO(raw))
|
|
img.load()
|
|
except OSError:
|
|
return error_response("could not decode image")
|
|
|
|
try:
|
|
mask = await asyncio.to_thread(_subject_mask, img)
|
|
except Exception as e: # model load/download/inference — surface the cause
|
|
log.exception("[SelectSubject] segmentation failed")
|
|
return error_response(f"subject segmentation failed: {e}", 500)
|
|
|
|
buf = _io.BytesIO()
|
|
mask.save(buf, format="PNG")
|
|
return web.Response(body=buf.getvalue(), content_type="image/png")
|
|
|
|
|
|
# ── click-to-select (SAM2 point prompt) ──
|
|
|
|
SAM_MODEL = "sam2.1_b.pt"
|
|
_sam = None # SAM2Predictor, loaded once
|
|
_sam_image_key = None # sha1 of the image the cached encoder features belong to
|
|
_sam_lock = threading.Lock()
|
|
|
|
|
|
def _sam_weights_path() -> str:
|
|
"""models/sam/<weights> — ultralytics otherwise drops the download into the
|
|
process cwd (the ComfyUI root)."""
|
|
import os
|
|
import shutil
|
|
|
|
import folder_paths
|
|
|
|
sam_dir = os.path.join(folder_paths.models_dir, "sam")
|
|
os.makedirs(sam_dir, exist_ok=True)
|
|
path = os.path.join(sam_dir, SAM_MODEL)
|
|
if not os.path.isfile(path):
|
|
from ultralytics.utils.downloads import attempt_download_asset
|
|
|
|
log.info("[SelectObject] downloading %s (~154MB)", SAM_MODEL)
|
|
try:
|
|
downloaded = attempt_download_asset(SAM_MODEL) # lands relative to cwd
|
|
shutil.move(str(downloaded), path)
|
|
except Exception as exc:
|
|
# Clean a partial drop in cwd so a retry starts fresh, then surface a
|
|
# clear message instead of a raw ultralytics/network traceback.
|
|
stray = os.path.join(os.getcwd(), SAM_MODEL)
|
|
if os.path.isfile(stray):
|
|
try:
|
|
os.remove(stray)
|
|
except OSError:
|
|
pass
|
|
raise RuntimeError(
|
|
f"Couldn't download the SAM2 selection model ({SAM_MODEL}, ~154MB). "
|
|
f"Check your internet connection and try again. "
|
|
f"[{type(exc).__name__}: {exc}]"
|
|
) from exc
|
|
return path
|
|
|
|
|
|
def _get_sam():
|
|
global _sam
|
|
if _sam is None:
|
|
from ultralytics.models.sam import SAM2Predictor
|
|
|
|
_sam = SAM2Predictor(overrides=dict(
|
|
conf=0.25, task="segment", mode="predict", model=_sam_weights_path(),
|
|
imgsz=1024, save=False, verbose=False))
|
|
log.info("[SelectObject] predictor ready")
|
|
return _sam
|
|
|
|
|
|
def _object_mask(img: Image.Image, x: int, y: int) -> Image.Image:
|
|
"""One positive point → the WHOLE object under it.
|
|
|
|
SAM2's three multimask granularities (sub-part / part / whole) come with an
|
|
unreliable score head (the whole-figure mask measured 0.012 vs 0.766 for a
|
|
foot), so the pick is by AREA — Photoshop's Object Selection bias — then
|
|
cleaned to the connected component containing the click.
|
|
"""
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from ultralytics.utils import ops
|
|
|
|
rgb = np.asarray(img.convert("RGB"))
|
|
h, w = rgb.shape[:2]
|
|
x = max(0, min(w - 1, int(x)))
|
|
y = max(0, min(h - 1, int(y)))
|
|
|
|
with _sam_lock:
|
|
global _sam_image_key
|
|
sam = _get_sam()
|
|
key = hashlib.sha1(rgb.tobytes()).hexdigest()
|
|
if key != _sam_image_key:
|
|
sam.set_image(rgb)
|
|
_sam_image_key = key
|
|
with torch.no_grad():
|
|
points, labels, _ = sam._prepare_prompts((1024, 1024), (h, w), None, [[x, y]], [1], None)
|
|
pred_masks, _scores = sam._inference_features(
|
|
sam.features, points, labels, None, multimask_output=True)
|
|
full = ops.scale_masks(pred_masks[None].float(), (h, w), padding=False)[0]
|
|
binaries = (full > 0).cpu().numpy().astype(np.uint8)
|
|
|
|
best = max(range(len(binaries)), key=lambda i: int(binaries[i].sum()))
|
|
mask = binaries[best]
|
|
# Drop disconnected speckle: keep the component under the click (fall back
|
|
# to the largest one when the click sits in a hole of the winning mask).
|
|
n, lab = cv2.connectedComponents(mask, connectivity=8)
|
|
if n > 2:
|
|
target = lab[y, x]
|
|
if target == 0:
|
|
counts = np.bincount(lab.ravel())
|
|
counts[0] = 0
|
|
target = int(counts.argmax())
|
|
mask = (lab == target).astype(np.uint8)
|
|
elif mask[y, x] == 0 and n == 2:
|
|
mask = (lab == 1).astype(np.uint8)
|
|
return Image.fromarray(mask * 255, "L")
|
|
|
|
|
|
@routes.post("/promptchain/select-object")
|
|
async def select_object(request):
|
|
form = await request.post()
|
|
upload = form.get("image")
|
|
if upload is None or not getattr(upload, "file", None):
|
|
return error_response("missing image upload")
|
|
raw = upload.file.read()
|
|
if not raw:
|
|
return error_response("empty image upload")
|
|
try:
|
|
x = int(float(form.get("x", "")))
|
|
y = int(float(form.get("y", "")))
|
|
except ValueError:
|
|
return error_response("missing click coordinates")
|
|
try:
|
|
img = Image.open(_io.BytesIO(raw))
|
|
img.load()
|
|
except OSError:
|
|
return error_response("could not decode image")
|
|
|
|
try:
|
|
mask = await asyncio.to_thread(_object_mask, img, x, y)
|
|
except Exception as e:
|
|
log.exception("[SelectObject] segmentation failed")
|
|
return error_response(f"object segmentation failed: {e}", 500)
|
|
|
|
buf = _io.BytesIO()
|
|
mask.save(buf, format="PNG")
|
|
return web.Response(body=buf.getvalue(), content_type="image/png")
|
|
|
|
|
|
# ── install / status ──────────────────────────────────────────────
|
|
# BiRefNet/SAM2 aren't ComfyUI nodes, so they can't use node_install_api's
|
|
# proof_nodes gate — they're our own routes backed by pip libs (transformers /
|
|
# ultralytics) + model files. These routes give selection the same first-use
|
|
# installer (with progress) and Full-Install coverage every other 3rd-party
|
|
# feature has, instead of a silent multi-GB download or a raw import error.
|
|
import importlib.util as _ilu
|
|
import sys as _sys
|
|
|
|
from .node_install_api import (
|
|
_CRITICAL, _log, _open_sse, _pkg_version, _restore_critical, _run, _send,
|
|
)
|
|
|
|
|
|
def _have(mod: str) -> bool:
|
|
try:
|
|
return _ilu.find_spec(mod) is not None
|
|
except (ImportError, ValueError):
|
|
return False
|
|
|
|
|
|
def _birefnet_cached() -> bool:
|
|
"""BiRefNet downloads into the HuggingFace hub cache (not models/), so probe
|
|
the cache dir rather than a single file path."""
|
|
try:
|
|
import os
|
|
from huggingface_hub.constants import HF_HUB_CACHE
|
|
snaps = os.path.join(HF_HUB_CACHE, "models--" + BIREFNET_ID.replace("/", "--"), "snapshots")
|
|
return os.path.isdir(snaps) and any(os.scandir(snaps))
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _sam_cached() -> bool:
|
|
try:
|
|
import os
|
|
import folder_paths
|
|
return os.path.isfile(os.path.join(folder_paths.models_dir, "sam", SAM_MODEL))
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _subject_status() -> dict:
|
|
deps, model = _have("transformers") and _have("torchvision"), _birefnet_cached()
|
|
return {"ready": deps and model, "deps_ok": deps, "model_ok": model,
|
|
"label": "Select Subject (BiRefNet)", "size": "~0.9 GB"}
|
|
|
|
|
|
def _object_status() -> dict:
|
|
deps, model = _have("ultralytics"), _sam_cached()
|
|
return {"ready": deps and model, "deps_ok": deps, "model_ok": model,
|
|
"label": "Object Select (SAM2)", "size": "~154 MB"}
|
|
|
|
|
|
@routes.get("/promptchain/subject/status")
|
|
async def subject_status(request):
|
|
which = (request.query.get("which") or "").strip()
|
|
out = {"subject": _subject_status(), "object": _object_status()}
|
|
if which in out:
|
|
return web.json_response(out[which])
|
|
out["ready"] = out["subject"]["ready"] and out["object"]["ready"]
|
|
return web.json_response(out)
|
|
|
|
|
|
# (_pip_install removed — PromptChain declares its deps in requirements.txt and
|
|
# never runs pip at runtime; missing packages are surfaced to the user instead.)
|
|
|
|
|
|
def _download_birefnet():
|
|
# snapshot_download fetches the repo (weights + remote code) WITHOUT loading
|
|
# the model onto the GPU — first real use does the load.
|
|
from huggingface_hub import snapshot_download
|
|
try:
|
|
snapshot_download(BIREFNET_ID)
|
|
except Exception as exc:
|
|
raise RuntimeError(
|
|
f"Couldn't download the BiRefNet selection model (~0.9GB) from Hugging "
|
|
f"Face. Check your internet connection and try again. "
|
|
f"[{type(exc).__name__}: {exc}]"
|
|
) from exc
|
|
|
|
|
|
async def install_subject_target(resp, target: str, snapshot: dict) -> bool:
|
|
"""Install one AI-selection target (pip deps + model) on the open SSE `resp`,
|
|
skipping anything already present. Returns True on success / already-ready.
|
|
Shared by the /subject/install route and the section installer; the caller
|
|
owns the SSE lifecycle, the snapshot, and the final _restore_critical."""
|
|
if target == "subject":
|
|
miss = [m for m in ("transformers",) if not _have(m)]
|
|
if miss:
|
|
await _send(resp, {"error": f"missing package(s): {', '.join(miss)}. Run `pip install -r requirements.txt` for PromptChain (or reinstall it via Manager) and restart."})
|
|
return False
|
|
if not _birefnet_cached():
|
|
await _send(resp, {"stage": "download", "file": "BiRefNet (~0.9GB)", "pct": 0})
|
|
await _log(resp, "Downloading BiRefNet weights (~0.9GB)…")
|
|
await asyncio.to_thread(_download_birefnet)
|
|
return True
|
|
if target == "object":
|
|
miss = [m for m in ("ultralytics",) if not _have(m)]
|
|
if miss:
|
|
await _send(resp, {"error": f"missing package(s): {', '.join(miss)}. Run `pip install -r requirements.txt` for PromptChain (or reinstall it via Manager) and restart."})
|
|
return False
|
|
if not _sam_cached():
|
|
await _send(resp, {"stage": "download", "file": "SAM2 (~154MB)", "pct": 0})
|
|
await _log(resp, "Downloading SAM2 weights (~154MB)…")
|
|
await asyncio.to_thread(_sam_weights_path)
|
|
return True
|
|
await _send(resp, {"error": f"unknown selection target: {target}"})
|
|
return False
|
|
|
|
|
|
@routes.post("/promptchain/subject/install")
|
|
async def subject_install(request):
|
|
try:
|
|
body = await request.json()
|
|
except Exception:
|
|
body = {}
|
|
which = (body.get("which") or "both").strip()
|
|
targets = ["subject", "object"] if which == "both" else [which]
|
|
resp = await _open_sse(request)
|
|
snapshot = {p: _pkg_version(p) for p in _CRITICAL}
|
|
try:
|
|
for t in targets:
|
|
if not await install_subject_target(resp, t, snapshot):
|
|
return resp
|
|
await _restore_critical(resp, snapshot)
|
|
# pip libs import at call time, so the next select request picks them up
|
|
# without a restart.
|
|
await _send(resp, {"done": True, "ok": True, "needs_restart": False})
|
|
except ConnectionResetError:
|
|
pass # client navigated away; pip/download keep running
|
|
except Exception as e:
|
|
log.exception("[subject/install] failed")
|
|
await _restore_critical(resp, snapshot)
|
|
try:
|
|
await _send(resp, {"error": str(e)})
|
|
except Exception:
|
|
pass
|
|
return resp
|