shootthesound
85e034b075
Thanks to @FNGarvin (#19). SAM 3 loads via torch.load (pickle), so a checksum doubles as a tamper guard at rest. Hash is pinned to the 1038lab/sam3 file Peter's repo serves today; if upstream republishes the constant must be bumped (the verification fails closed — bad copy is deleted and we surface a clear error).
380 lines
15 KiB
Python
380 lines
15 KiB
Python
"""Angelo — SAM 3 auto-segmentation backend.
|
|
|
|
Our own lazy loader for Meta's SAM 3 (the pip-installed `sam3` package —
|
|
not the other node packs). SAM 3's headline capability is *concept /
|
|
text* prompts: give it a noun phrase ("the front wheel", "her face") and
|
|
it segments every matching instance. We run it on Angelo's current
|
|
preview image and return the detections as polygons + boxes so the JS
|
|
can overlay them, let the user click-to-confirm, and feed the chosen
|
|
shape into the refine pipeline (silhouette for Refine, bbox for Smart
|
|
Inpaint).
|
|
|
|
The model is loaded once on first use and cached. It loads to GPU on
|
|
demand and is offloaded back to CPU after each detection so it doesn't
|
|
sit on VRAM next to the diffusion model between clicks.
|
|
|
|
Routes are namespaced under /angelo/ to avoid clashing with other packs.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import threading
|
|
import traceback
|
|
|
|
_LOCK = threading.Lock()
|
|
_STATE = {
|
|
"model": None, # the SAM3 image model
|
|
"processor": None, # Sam3Processor
|
|
"device": None, # torch device the model loads onto
|
|
"import_error": None, # cached import failure message, if any
|
|
}
|
|
|
|
# Checkpoint lives in <models>/sam3/sam3.pt. Auto-downloaded from the
|
|
# public HF mirror the installed comfyui-sam3 pack uses (no HF token
|
|
# needed — the official facebook/sam3 repo is gated).
|
|
_SAM3_HF_REPO = "1038lab/sam3"
|
|
_SAM3_CKPT_NAME = "sam3.pt"
|
|
# SHA256 of the current 1038lab/sam3 sam3.pt. Verified against the file
|
|
# the public mirror serves today (and as a checksum, also a guard against
|
|
# pickle tampering at rest, since SAM 3 loads via torch.load). Update if
|
|
# the upstream repo republishes — bump it, do not silently delete this
|
|
# check. Thanks to @FNGarvin (#19) for both the suggestion and value.
|
|
_SAM3_EXPECTED_SHA256 = "9999e2341ceef5e136daa386eecb55cb414446a00ac2b55eb2dfd2f7c3cf8c9e"
|
|
|
|
|
|
def _torch_devices():
|
|
import comfy.model_management as mm
|
|
return mm.get_torch_device(), mm.unet_offload_device()
|
|
|
|
|
|
def _checkpoint_path():
|
|
"""Absolute path where sam3.pt should live (<models>/sam3/sam3.pt)."""
|
|
import folder_paths
|
|
return os.path.join(folder_paths.models_dir, "sam3", _SAM3_CKPT_NAME)
|
|
|
|
|
|
def _download_checkpoint(target_path):
|
|
"""Fetch sam3.pt into <models>/sam3/ from the public HF mirror, using
|
|
the proven download+move pattern from the installed SAM 3 node packs.
|
|
The downloaded blob is SHA256-verified before being moved into place;
|
|
a mismatch deletes the bad copy and raises so the next attempt retries
|
|
cleanly."""
|
|
import hashlib
|
|
import shutil
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
target_dir = os.path.dirname(target_path)
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
print(f"[Angelo/SAM3] sam3.pt not found — downloading from {_SAM3_HF_REPO} "
|
|
f"(this is a one-time ~GB download)...")
|
|
downloaded = hf_hub_download(
|
|
repo_id=_SAM3_HF_REPO,
|
|
filename=_SAM3_CKPT_NAME,
|
|
local_dir=target_dir,
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
|
|
h = hashlib.sha256()
|
|
with open(downloaded, "rb") as f:
|
|
for chunk in iter(lambda: f.read(1 << 20), b""):
|
|
h.update(chunk)
|
|
actual = h.hexdigest()
|
|
if actual != _SAM3_EXPECTED_SHA256:
|
|
try:
|
|
os.remove(downloaded)
|
|
except OSError:
|
|
pass
|
|
raise RuntimeError(
|
|
f"[Angelo/SAM3] Downloaded sam3.pt failed SHA256 check.\n"
|
|
f" expected: {_SAM3_EXPECTED_SHA256}\n"
|
|
f" actual: {actual}\n"
|
|
"Refusing to load a tampered or unexpected checkpoint. The "
|
|
"bad copy has been deleted; re-run to retry the download."
|
|
)
|
|
|
|
if os.path.normpath(downloaded) != os.path.normpath(target_path):
|
|
shutil.move(downloaded, target_path)
|
|
print(f"[Angelo/SAM3] Downloaded to {target_path}")
|
|
return target_path
|
|
|
|
|
|
def _find_bpe(sam3_module):
|
|
"""Locate bpe_simple_vocab_16e6.txt.gz next to the imported sam3
|
|
package. Returns the path or None (let build use its default)."""
|
|
name = "bpe_simple_vocab_16e6.txt.gz"
|
|
try:
|
|
pkg_dir = os.path.dirname(os.path.abspath(sam3_module.__file__))
|
|
except Exception:
|
|
return None
|
|
for c in (
|
|
os.path.join(pkg_dir, "assets", name), # rmbg layout
|
|
os.path.join(pkg_dir, name),
|
|
os.path.join(pkg_dir, "sam3", "assets", name), # nested (tbg layout)
|
|
os.path.join(pkg_dir, "sam3_lib", name),
|
|
):
|
|
if os.path.exists(c):
|
|
return c
|
|
return None
|
|
|
|
|
|
def _ensure_model():
|
|
"""Lazily build + cache the SAM3 image model and processor. Returns
|
|
(model, processor) or raises RuntimeError with a clear message."""
|
|
if _STATE["import_error"]:
|
|
raise RuntimeError(_STATE["import_error"])
|
|
if _STATE["model"] is not None:
|
|
return _STATE["model"], _STATE["processor"]
|
|
|
|
with _LOCK:
|
|
if _STATE["model"] is not None:
|
|
return _STATE["model"], _STATE["processor"]
|
|
try:
|
|
# SAM 3 — whatever `import sam3` resolves to in this env. This
|
|
# may be the pip package OR a copy vendored inside another
|
|
# custom node (rmbg / tbg-sam3) that's on sys.path first.
|
|
import sam3
|
|
from sam3.model_builder import build_sam3_image_model
|
|
from sam3.model.sam3_image_processor import Sam3Processor
|
|
except Exception as e: # pragma: no cover - environment dependent
|
|
_STATE["import_error"] = (
|
|
"SAM 3 Detect isn't installed (it's optional). CLOSE ComfyUI, "
|
|
"then run the installer in the ComfyUI-Angelo folder — "
|
|
"install_sam3_support.bat (Windows) or install_sam3_support.sh "
|
|
"(macOS/Linux) — and start ComfyUI again. "
|
|
f"(sam3 package not importable: {e})"
|
|
)
|
|
raise RuntimeError(_STATE["import_error"])
|
|
|
|
load_device, _ = _torch_devices()
|
|
ckpt = _checkpoint_path()
|
|
if not os.path.exists(ckpt):
|
|
try:
|
|
_download_checkpoint(ckpt)
|
|
except Exception as e:
|
|
# Not cached in import_error — a network blip should retry
|
|
# on the next Detect, not block until restart.
|
|
raise RuntimeError(
|
|
f"SAM 3 weights missing and auto-download failed: {e}\n"
|
|
f"Place sam3.pt manually at: {ckpt}"
|
|
)
|
|
# build_sam3_image_model's default bpe_path is computed relative to
|
|
# the package and is unreliable when sam3 is a vendored copy inside
|
|
# another node (it pointed at a non-existent assets/ dir). Locate
|
|
# the BPE tokenizer co-located with the loaded sam3 and pass it.
|
|
bpe_path = _find_bpe(sam3)
|
|
print(f"[Angelo/SAM3] Building image model (checkpoint={ckpt}, "
|
|
f"bpe={'auto' if not bpe_path else bpe_path})...")
|
|
# Build under a forced float32 default dtype, then hard-cast to fp32.
|
|
# Some setups leave torch's GLOBAL default dtype at bfloat16; on
|
|
# torch >= 2.10 that bakes bf16 weights into SAM 3 during the build,
|
|
# but SAM 3's image processor always feeds an fp32 image — so the
|
|
# backbone crashes with "mat1 and mat2 must have the same dtype, got
|
|
# BFloat16 and Float". fp32 is SAM 3's known-good precision (and a
|
|
# no-op when the default is already fp32). Restore the previous
|
|
# default afterwards so we don't disturb the rest of ComfyUI.
|
|
# (Root-caused by reproducing on a torch 2.10.0+cu130 RC venv.)
|
|
import torch
|
|
_prev_default_dtype = torch.get_default_dtype()
|
|
torch.set_default_dtype(torch.float32)
|
|
try:
|
|
if bpe_path:
|
|
model = build_sam3_image_model(checkpoint_path=ckpt, bpe_path=bpe_path)
|
|
else:
|
|
model = build_sam3_image_model(checkpoint_path=ckpt)
|
|
finally:
|
|
torch.set_default_dtype(_prev_default_dtype)
|
|
model.float()
|
|
processor = Sam3Processor(model)
|
|
model.processor = processor
|
|
model.eval()
|
|
|
|
_STATE["model"] = model
|
|
_STATE["processor"] = processor
|
|
_STATE["device"] = load_device
|
|
print("[Angelo/SAM3] Model ready.")
|
|
return model, processor
|
|
|
|
|
|
def _mask_to_polygons(mask_np, min_area=64, epsilon_frac=0.004):
|
|
"""Binary mask (H, W) -> list of polygons, each a flat [x0,y0,x1,y1,...]
|
|
list in image-pixel coords. Contours are simplified so the payload and
|
|
the JS overlay stay light. Tiny specks are dropped."""
|
|
import numpy as np
|
|
polys = []
|
|
m = (mask_np > 0.5).astype("uint8")
|
|
if m.sum() == 0:
|
|
return polys
|
|
try:
|
|
import cv2
|
|
except Exception:
|
|
# No OpenCV — fall back to the mask's bounding box as a 4-pt polygon.
|
|
ys, xs = np.where(m)
|
|
x1, y1, x2, y2 = int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())
|
|
return [[x1, y1, x2, y1, x2, y2, x1, y2]]
|
|
|
|
contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
for c in contours:
|
|
if cv2.contourArea(c) < min_area:
|
|
continue
|
|
peri = cv2.arcLength(c, True)
|
|
approx = cv2.approxPolyDP(c, epsilon_frac * peri, True)
|
|
pts = approx.reshape(-1, 2)
|
|
if len(pts) < 3:
|
|
continue
|
|
polys.append([int(v) for xy in pts for v in xy])
|
|
return polys
|
|
|
|
|
|
def detect_text(pil_image, text, confidence_threshold=0.3, max_detections=20):
|
|
"""Run SAM 3 text grounding on a PIL image. Returns a list of
|
|
detections: {"polygons": [[x,y,...], ...], "bbox": [x1,y1,x2,y2],
|
|
"score": float}, in image-pixel coords, sorted by score."""
|
|
import torch
|
|
import numpy as np
|
|
|
|
model, processor = _ensure_model()
|
|
load_device, offload_device = _torch_devices()
|
|
|
|
img_w, img_h = pil_image.size
|
|
results = []
|
|
try:
|
|
model.to(load_device)
|
|
if hasattr(processor, "set_confidence_threshold"):
|
|
processor.set_confidence_threshold(float(confidence_threshold))
|
|
|
|
# SAM 3 runs in fp32 (forced at build in _ensure_model) and its
|
|
# processor feeds an fp32 image, so weights and activations agree —
|
|
# no dtype juggling needed here.
|
|
state = processor.set_image(pil_image)
|
|
if text and text.strip():
|
|
state = processor.set_text_prompt(text.strip(), state)
|
|
|
|
masks = state.get("masks", None)
|
|
boxes = state.get("boxes", None)
|
|
scores = state.get("scores", None)
|
|
if masks is None or len(masks) == 0:
|
|
return []
|
|
|
|
# Sort by score (desc) and cap.
|
|
if scores is not None and len(scores) > 0:
|
|
order = torch.argsort(scores, descending=True)
|
|
masks = masks[order]
|
|
boxes = boxes[order] if boxes is not None else None
|
|
scores = scores[order]
|
|
n = len(masks)
|
|
if max_detections > 0:
|
|
n = min(n, max_detections)
|
|
|
|
for i in range(n):
|
|
mask_np = masks[i].detach().to("cpu").float().numpy()
|
|
if mask_np.ndim == 3:
|
|
mask_np = mask_np[0]
|
|
polys = _mask_to_polygons(mask_np)
|
|
if not polys:
|
|
continue
|
|
if boxes is not None:
|
|
b = boxes[i].detach().to("cpu").float().tolist()
|
|
bbox = [float(b[0]), float(b[1]), float(b[2]), float(b[3])]
|
|
else:
|
|
ys, xs = np.where(mask_np > 0.5)
|
|
bbox = [float(xs.min()), float(ys.min()), float(xs.max()), float(ys.max())]
|
|
score = float(scores[i]) if scores is not None else 0.0
|
|
results.append({"polygons": polys, "bbox": bbox, "score": score})
|
|
finally:
|
|
# Offload off the GPU so we don't squat VRAM between clicks.
|
|
try:
|
|
model.to(offload_device)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
except Exception:
|
|
pass
|
|
|
|
return {"width": img_w, "height": img_h, "detections": results}
|
|
|
|
|
|
# --- HTTP route ------------------------------------------------------
|
|
|
|
try:
|
|
from server import PromptServer
|
|
_HAS_SERVER = True
|
|
except ImportError:
|
|
_HAS_SERVER = False
|
|
|
|
|
|
def _load_preview_image(filename, subfolder, type_):
|
|
"""Resolve a ComfyUI /view-style image ref to a PIL image."""
|
|
import folder_paths
|
|
from PIL import Image
|
|
type_ = type_ or "temp"
|
|
if type_ == "output":
|
|
base = folder_paths.get_output_directory()
|
|
elif type_ == "input":
|
|
base = folder_paths.get_input_directory()
|
|
else:
|
|
base = folder_paths.get_temp_directory()
|
|
path = os.path.join(base, subfolder or "", filename or "")
|
|
path = os.path.normpath(path)
|
|
# Guard against path traversal outside the base dir.
|
|
if not os.path.normpath(path).startswith(os.path.normpath(base)):
|
|
raise ValueError("invalid image path")
|
|
return Image.open(path).convert("RGB")
|
|
|
|
|
|
if _HAS_SERVER:
|
|
from aiohttp import web
|
|
|
|
routes = PromptServer.instance.routes
|
|
|
|
@routes.post("/angelo/detect")
|
|
async def _angelo_detect(request):
|
|
try:
|
|
data = await request.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
|
|
method = (data or {}).get("method", "sam3_text")
|
|
text = (data or {}).get("text", "")
|
|
threshold = float((data or {}).get("confidence_threshold", 0.3))
|
|
max_det = int((data or {}).get("max_detections", 20))
|
|
filename = (data or {}).get("filename", "")
|
|
subfolder = (data or {}).get("subfolder", "")
|
|
type_ = (data or {}).get("type", "temp")
|
|
|
|
if not filename:
|
|
return web.json_response({"error": "no preview image (generate one first)"}, status=400)
|
|
|
|
try:
|
|
pil = _load_preview_image(filename, subfolder, type_)
|
|
except Exception as e:
|
|
return web.json_response({"error": f"could not load preview: {e}"}, status=400)
|
|
|
|
# Run the (blocking, GPU) detection off the event loop. On WSL /
|
|
# containerised CUDA, the executor thread occasionally trips a
|
|
# spurious torch.OutOfMemoryError on first allocation despite ample
|
|
# free VRAM (CUDA's per-thread context bookkeeping doesn't see what
|
|
# the main thread already reserved). Fall back to a synchronous
|
|
# main-thread run when that specific failure mode shows up — it
|
|
# blocks the event loop briefly but unblocks the user. Thanks to
|
|
# @FNGarvin (#19) for the diagnosis.
|
|
import asyncio
|
|
import torch
|
|
loop = asyncio.get_event_loop()
|
|
try:
|
|
if method == "sam3_text":
|
|
try:
|
|
result = await loop.run_in_executor(
|
|
None, detect_text, pil, text, threshold, max_det
|
|
)
|
|
except torch.cuda.OutOfMemoryError:
|
|
print("[Angelo/SAM3] Background-thread CUDA allocation failed — "
|
|
"retrying on the main thread (WSL/container quirk).")
|
|
result = detect_text(pil, text, threshold, max_det)
|
|
else:
|
|
return web.json_response({"error": f"unknown method '{method}'"}, status=400)
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
return web.json_response({"ok": True, **result})
|