1067 lines
42 KiB
Python
1067 lines
42 KiB
Python
"""
|
||
SAM Smart Inpainter — ComfyUI Custom Node v0.3 (Fixed latent_image API + decode)
|
||
All‑in‑one text‑prompted segmentation + diffusion inpainting node.
|
||
Supports SAM3 models.
|
||
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import math
|
||
import traceback
|
||
import torch
|
||
import numpy as np
|
||
import safetensors.torch
|
||
import folder_paths
|
||
import comfy.utils
|
||
import comfy.sample
|
||
import comfy.samplers
|
||
import comfy.model_management
|
||
from PIL import Image
|
||
|
||
# ============================================================================
|
||
# 1. Model folder registration
|
||
# ============================================================================
|
||
_NODE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
|
||
def _ensure_model_folder(key, subfolder):
|
||
try:
|
||
base = folder_paths.models_dir
|
||
except AttributeError:
|
||
base = os.path.join(os.path.dirname(os.path.dirname(_NODE_DIR)), "models")
|
||
path = os.path.join(base, subfolder)
|
||
os.makedirs(path, exist_ok=True)
|
||
exts = getattr(folder_paths, "supported_pt_extensions",
|
||
{".pt", ".pth", ".bin", ".safetensors", ".ckpt"})
|
||
if key not in folder_paths.folder_names_and_paths:
|
||
folder_paths.folder_names_and_paths[key] = ([path], exts)
|
||
elif path not in folder_paths.folder_names_and_paths[key][0]:
|
||
folder_paths.folder_names_and_paths[key][0].append(path)
|
||
|
||
|
||
_ensure_model_folder("sam3", "sam3")
|
||
_ensure_model_folder("sam2", "sam2")
|
||
_ensure_model_folder("sams", "sams")
|
||
_SAM_SEARCH_KEYS = ["sam3", "sams", "sam2", "checkpoints"]
|
||
|
||
|
||
def _list_all_sam_models():
|
||
seen, out = set(), []
|
||
for key in _SAM_SEARCH_KEYS:
|
||
try:
|
||
for name in folder_paths.get_filename_list(key):
|
||
if name not in seen:
|
||
seen.add(name)
|
||
out.append(name)
|
||
except Exception:
|
||
pass
|
||
return out or ["<no models found>"]
|
||
|
||
|
||
def _resolve_path(filename):
|
||
for key in _SAM_SEARCH_KEYS:
|
||
try:
|
||
p = folder_paths.get_full_path(key, filename)
|
||
if p and os.path.isfile(p):
|
||
return p
|
||
except Exception:
|
||
pass
|
||
return None
|
||
|
||
|
||
# ============================================================================
|
||
# 2. Model‑family detection
|
||
# ============================================================================
|
||
def _is_sam3_model(path):
|
||
name = os.path.basename(path).lower()
|
||
return "sam3" in name or "multiplex" in name
|
||
|
||
|
||
# ============================================================================
|
||
# 3. Safe weight loading
|
||
# ============================================================================
|
||
def _load_state_dict_safe(path):
|
||
try:
|
||
return safetensors.torch.load_file(path)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
return torch.load(path, map_location="cpu", weights_only=False)
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
f"Failed to load checkpoint: {path}\n{type(e).__name__}: {e}"
|
||
)
|
||
|
||
|
||
# ============================================================================
|
||
# 4. SAM3 package discovery
|
||
# ============================================================================
|
||
def _find_sam3_package():
|
||
try:
|
||
import sam3
|
||
return None
|
||
except ImportError:
|
||
pass
|
||
comfy_root = os.path.dirname(os.path.dirname(_NODE_DIR))
|
||
nodes_root = os.path.join(comfy_root, "custom_nodes")
|
||
search_roots = [nodes_root]
|
||
custom = folder_paths.folder_names_and_paths.get("custom_nodes")
|
||
if custom:
|
||
search_roots.extend(custom[0])
|
||
|
||
for root in search_roots:
|
||
if not os.path.isdir(root):
|
||
continue
|
||
for pack in os.listdir(root):
|
||
sam3_dir = os.path.join(root, pack, "sam3")
|
||
init_py = os.path.join(sam3_dir, "__init__.py")
|
||
if os.path.isdir(sam3_dir) and os.path.isfile(init_py):
|
||
parent = os.path.join(root, pack)
|
||
if parent not in sys.path:
|
||
sys.path.insert(0, parent)
|
||
try:
|
||
import sam3
|
||
return parent
|
||
except ImportError:
|
||
if parent in sys.path:
|
||
sys.path.remove(parent)
|
||
return None
|
||
|
||
|
||
_sam3_pkg_added = _find_sam3_package()
|
||
try:
|
||
from sam3.model_builder import build_sam3_image_model as _build_sam3_image_model
|
||
from sam3.model.sam3_image_processor import Sam3Processor as _Sam3Processor
|
||
_SAM3_AVAILABLE = True
|
||
except ImportError:
|
||
_SAM3_AVAILABLE = False
|
||
_build_sam3_image_model = None
|
||
_Sam3Processor = None
|
||
|
||
|
||
# ============================================================================
|
||
# 5. BPE tokenizer locator
|
||
# ============================================================================
|
||
def _find_sam3_bpe_path():
|
||
target = "bpe_simple_vocab_16e6.txt.gz"
|
||
try:
|
||
import sam3 as _s
|
||
pkg_dir = os.path.dirname(_s.__file__)
|
||
candidates = [
|
||
os.path.join(pkg_dir, "assets", target),
|
||
os.path.join(pkg_dir, "..", "assets", target),
|
||
]
|
||
for c in candidates:
|
||
c = os.path.normpath(c)
|
||
if os.path.isfile(c):
|
||
return c
|
||
except Exception:
|
||
pass
|
||
comfy_root = os.path.dirname(os.path.dirname(_NODE_DIR))
|
||
nodes_root = os.path.join(comfy_root, "custom_nodes")
|
||
if os.path.isdir(nodes_root):
|
||
for pack in os.listdir(nodes_root):
|
||
for sub in (
|
||
os.path.join(nodes_root, pack, "sam3", "assets", target),
|
||
os.path.join(nodes_root, pack, "models", "sam3", "assets", target),
|
||
os.path.join(nodes_root, pack, "assets", target),
|
||
):
|
||
if os.path.isfile(sub):
|
||
return sub
|
||
return None
|
||
|
||
|
||
# ============================================================================
|
||
# 6. SAM3 loader
|
||
# ============================================================================
|
||
def _load_sam3_predictor(path, device):
|
||
if not _SAM3_AVAILABLE:
|
||
raise RuntimeError(
|
||
"sam3 package not found. Install ComfyUI‑Easy‑Sam3 or ComfyUI‑RMBG."
|
||
)
|
||
bpe_path = _find_sam3_bpe_path()
|
||
device_str = str(device)
|
||
|
||
err_a = "not attempted"
|
||
try:
|
||
kwargs = dict(eval_mode=True, checkpoint_path=path, load_from_HF=False)
|
||
if bpe_path:
|
||
kwargs["bpe_path"] = bpe_path
|
||
try:
|
||
model = _build_sam3_image_model(device=device_str, **kwargs)
|
||
except TypeError:
|
||
model = _build_sam3_image_model(**kwargs)
|
||
model = model.to(device).float()
|
||
model.eval()
|
||
processor = _make_sam3_processor(model, device_str)
|
||
return _SAM3Wrapper(processor, device)
|
||
except Exception as e:
|
||
err_a = f"{type(e).__name__}: {e}"
|
||
|
||
err_b = "not attempted"
|
||
try:
|
||
build_kwargs = dict(load_from_HF=False, eval_mode=True)
|
||
if bpe_path:
|
||
build_kwargs["bpe_path"] = bpe_path
|
||
model = _build_sam3_image_model(**build_kwargs)
|
||
sd = _load_state_dict_safe(path)
|
||
|
||
def _remap(state_dict):
|
||
out = {}
|
||
for k, v in state_dict.items():
|
||
if k.startswith("detector."):
|
||
out[k[len("detector."):]] = v
|
||
elif not k.startswith("tracker."):
|
||
out[k] = v
|
||
return out
|
||
|
||
loaded = False
|
||
for sd_try in [sd, _remap(sd)]:
|
||
try:
|
||
model.load_state_dict(sd_try, strict=False)
|
||
loaded = True
|
||
break
|
||
except Exception:
|
||
pass
|
||
if not loaded:
|
||
raise RuntimeError("load_state_dict failed for all remapping strategies")
|
||
|
||
model = model.to(device).float()
|
||
model.eval()
|
||
processor = _make_sam3_processor(model, device_str)
|
||
return _SAM3Wrapper(processor, device)
|
||
except Exception as e:
|
||
err_b = f"{type(e).__name__}: {e}"
|
||
|
||
err_c = "not attempted"
|
||
try:
|
||
model = _build_sam3_image_model()
|
||
model = model.to(device).float()
|
||
model.eval()
|
||
processor = _make_sam3_processor(model, device_str)
|
||
return _SAM3Wrapper(processor, device)
|
||
except Exception as e:
|
||
err_c = f"{type(e).__name__}: {e}"
|
||
|
||
raise RuntimeError(
|
||
f"All three SAM3 load strategies failed for: {os.path.basename(path)}\n"
|
||
f" A (explicit params): {err_a}\n"
|
||
f" B (manual sd load): {err_b}\n"
|
||
f" C (default / HF): {err_c}\n"
|
||
f"BPE file found: {bpe_path or 'NOT FOUND'}"
|
||
)
|
||
|
||
|
||
def _make_sam3_processor(model, device_str):
|
||
try:
|
||
return _Sam3Processor(model, device=device_str)
|
||
except TypeError:
|
||
return _Sam3Processor(model)
|
||
|
||
|
||
# ============================================================================
|
||
# 7. SAM2 loader
|
||
# ============================================================================
|
||
def _sam2_yaml_name(model_path):
|
||
name = os.path.basename(model_path).lower()
|
||
v21 = "2.1" in name or "sam2.1" in name
|
||
if "large" in name or "_l." in name or "hiera_l" in name:
|
||
sz = "l"
|
||
elif "base_p" in name or "_b+" in name or "hiera_b" in name:
|
||
sz = "b+"
|
||
elif "small" in name or "_s." in name or "hiera_s" in name:
|
||
sz = "s"
|
||
elif "tiny" in name or "_t." in name or "hiera_t" in name:
|
||
sz = "t"
|
||
else:
|
||
sz = "l"
|
||
return ("sam2.1", f"sam2.1_hiera_{sz}.yaml") if v21 else ("sam2", f"sam2_hiera_{sz}.yaml")
|
||
|
||
|
||
def _find_sam2_yaml(model_path):
|
||
sub, fname = _sam2_yaml_name(model_path)
|
||
candidates = []
|
||
try:
|
||
import sam2 as _s2
|
||
d = os.path.dirname(_s2.__file__)
|
||
candidates += [
|
||
os.path.join(d, "configs", sub, fname),
|
||
os.path.join(d, "configs", fname),
|
||
os.path.join(d, fname),
|
||
]
|
||
except ImportError:
|
||
pass
|
||
comfy_root = os.path.dirname(os.path.dirname(_NODE_DIR))
|
||
nodes_root = os.path.join(comfy_root, "custom_nodes")
|
||
if os.path.isdir(nodes_root):
|
||
for pack in os.listdir(nodes_root):
|
||
for cfgdir in ("sam2_configs", "configs"):
|
||
candidates += [
|
||
os.path.join(nodes_root, pack, cfgdir, fname),
|
||
os.path.join(nodes_root, pack, cfgdir, sub, fname),
|
||
]
|
||
for p in candidates:
|
||
if p and os.path.isfile(p):
|
||
return p
|
||
return None
|
||
|
||
|
||
def _load_sam2_predictor(path, device):
|
||
yaml_path = _find_sam2_yaml(path)
|
||
if yaml_path is None:
|
||
return _load_sam2_via_build_sam2(path, device)
|
||
try:
|
||
from omegaconf import OmegaConf
|
||
from hydra.utils import instantiate
|
||
cfg = OmegaConf.load(yaml_path)
|
||
model = instantiate(cfg.model, _recursive_=True)
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
f"Failed to instantiate SAM2 model from YAML '{yaml_path}': {e}"
|
||
)
|
||
|
||
sd = _load_state_dict_safe(path)
|
||
|
||
def _try_load(state_dict, strict):
|
||
try:
|
||
missing, unexpected = model.load_state_dict(state_dict, strict=strict)
|
||
return not (strict and (missing or unexpected))
|
||
except Exception:
|
||
return False
|
||
|
||
loaded = (
|
||
_try_load(sd, strict=True)
|
||
or _try_load({k.lstrip("model."): v for k, v in sd.items()}, strict=False)
|
||
or _try_load(sd, strict=False)
|
||
)
|
||
if not loaded:
|
||
raise RuntimeError("SAM2 state_dict load failed.")
|
||
|
||
model.to(device=device, dtype=torch.float32)
|
||
model.eval()
|
||
|
||
try:
|
||
from sam2.sam2_image_predictor import SAM2ImagePredictor
|
||
predictor = SAM2ImagePredictor(model)
|
||
except ImportError:
|
||
raise RuntimeError("sam2 package not found. Install: pip install sam2")
|
||
return _SAM2Wrapper(predictor, device)
|
||
|
||
|
||
def _load_sam2_via_build_sam2(path, device):
|
||
try:
|
||
from sam2.build_sam import build_sam2
|
||
from sam2.sam2_image_predictor import SAM2ImagePredictor
|
||
except ImportError:
|
||
raise RuntimeError("sam2 package not found.")
|
||
sub, fname = _sam2_yaml_name(path)
|
||
fname_noext = fname.replace(".yaml", "")
|
||
candidates = [f"configs/{sub}/{fname}", f"configs/{fname}", fname, fname_noext]
|
||
err_last = None
|
||
for cfg_str in candidates:
|
||
try:
|
||
sam = build_sam2(cfg_str, path, device=device)
|
||
return _SAM2Wrapper(SAM2ImagePredictor(sam), device)
|
||
except Exception as e:
|
||
err_last = e
|
||
raise RuntimeError(f"build_sam2 failed. Last error: {err_last}")
|
||
|
||
|
||
# ============================================================================
|
||
# 8. Unified predictor builder
|
||
# ============================================================================
|
||
def _build_predictor(path, device):
|
||
return (
|
||
_load_sam3_predictor(path, device)
|
||
if _is_sam3_model(path)
|
||
else _load_sam2_predictor(path, device)
|
||
)
|
||
|
||
|
||
# ============================================================================
|
||
# 9. Predictor wrappers
|
||
# ============================================================================
|
||
class _SAM3Wrapper:
|
||
family = "sam3"
|
||
|
||
def __init__(self, processor, device):
|
||
self.processor = processor
|
||
self.device = device
|
||
self._state = None
|
||
|
||
def set_image(self, img_np_rgb):
|
||
pil = Image.fromarray(img_np_rgb.astype(np.uint8))
|
||
self._state = self.processor.set_image(pil)
|
||
|
||
def predict_with_text(self, prompt, threshold=0.35):
|
||
if self._state is None:
|
||
return []
|
||
try:
|
||
output = self.processor.set_text_prompt(state=self._state, prompt=prompt)
|
||
masks = output.get("masks", [])
|
||
scores = output.get("scores", [])
|
||
result = []
|
||
for i, m in enumerate(masks):
|
||
m = m.cpu().numpy() if isinstance(m, torch.Tensor) else np.asarray(m, dtype=bool)
|
||
if m.ndim == 3:
|
||
m = m[0]
|
||
score = float(scores[i]) if i < len(scores) else 1.0
|
||
if score >= threshold:
|
||
result.append((m, score))
|
||
return result
|
||
except Exception as e:
|
||
print(f"[SAM_SmartInpainter] SAM3 text predict error: {e}")
|
||
return []
|
||
|
||
def predict_with_point(self, point_xy):
|
||
if self._state is None:
|
||
return []
|
||
try:
|
||
output = self.processor.set_point_prompt(
|
||
state=self._state, points=[list(point_xy)], labels=[1]
|
||
)
|
||
masks = output.get("masks", [])
|
||
scores = output.get("scores", [])
|
||
return [
|
||
(
|
||
m.cpu().numpy()[0] if isinstance(m, torch.Tensor) else np.asarray(m, dtype=bool),
|
||
float(scores[i]) if i < len(scores) else 1.0,
|
||
)
|
||
for i, m in enumerate(masks)
|
||
]
|
||
except Exception:
|
||
return []
|
||
|
||
|
||
class _SAM2Wrapper:
|
||
family = "sam2"
|
||
|
||
def __init__(self, predictor, device):
|
||
self.predictor = predictor
|
||
self.device = device
|
||
|
||
def set_image(self, img_np_rgb):
|
||
if hasattr(self.predictor, "set_image"):
|
||
self.predictor.set_image(img_np_rgb)
|
||
elif hasattr(self.predictor, "set_torch_image"):
|
||
t = torch.from_numpy(img_np_rgb).permute(2, 0, 1).float() / 255.0
|
||
self.predictor.set_torch_image(t, img_np_rgb.shape[:2])
|
||
|
||
def predict_with_boxes(self, boxes_px):
|
||
results = []
|
||
for box in boxes_px:
|
||
box_arr = np.array(box, dtype=np.float32)
|
||
try:
|
||
masks, scores, _ = self.predictor.predict(
|
||
box=box_arr, multimask_output=False, return_logits=False
|
||
)
|
||
except TypeError:
|
||
masks, scores, _ = self.predictor.predict(
|
||
box=box_arr, multimask_output=False
|
||
)
|
||
if len(masks):
|
||
results.append((masks[0].astype(bool), float(scores[0])))
|
||
return results
|
||
|
||
def predict_with_point(self, point_xy):
|
||
pt = np.array([list(point_xy)], dtype=np.float32)
|
||
lbl = np.array([1], dtype=np.int32)
|
||
try:
|
||
masks, scores, _ = self.predictor.predict(
|
||
point_coords=pt, point_labels=lbl,
|
||
multimask_output=False, return_logits=False
|
||
)
|
||
except TypeError:
|
||
masks, scores, _ = self.predictor.predict(
|
||
point_coords=pt, point_labels=lbl, multimask_output=False
|
||
)
|
||
return [(masks[0].astype(bool), float(scores[0]))] if len(masks) else []
|
||
|
||
|
||
# ============================================================================
|
||
# 10. Optional Grounding‑DINO
|
||
# ============================================================================
|
||
_GDINO_AVAILABLE = False
|
||
try:
|
||
from groundingdino.util.inference import load_model as _gdino_load
|
||
from groundingdino.util.inference import predict as _gdino_predict
|
||
_GDINO_AVAILABLE = True
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _gdino_boxes(image_np_rgb, text, threshold=0.30, device="cuda"):
|
||
if not _GDINO_AVAILABLE:
|
||
return []
|
||
try:
|
||
from groundingdino.util.inference import load_image
|
||
pil = Image.fromarray(image_np_rgb)
|
||
img_t, _ = load_image(pil)
|
||
ckpt = _find_gdino_ckpt()
|
||
if ckpt is None:
|
||
return []
|
||
model = _gdino_load(
|
||
"groundingdino/config/GroundingDINO_SwinT_OGC.py", ckpt
|
||
).to(device)
|
||
boxes, logits, _ = _gdino_predict(model, img_t, text, threshold, threshold)
|
||
h, w = image_np_rgb.shape[:2]
|
||
out = []
|
||
for b in boxes:
|
||
cx, cy, bw, bh = b.tolist()
|
||
out.append([
|
||
max(0, int((cx - bw / 2) * w)),
|
||
max(0, int((cy - bh / 2) * h)),
|
||
min(w, int((cx + bw / 2) * w)),
|
||
min(h, int((cy + bh / 2) * h)),
|
||
])
|
||
return out
|
||
except Exception as e:
|
||
print(f"[SAM_SmartInpainter] GDino error: {e}")
|
||
return []
|
||
|
||
|
||
def _find_gdino_ckpt():
|
||
for key in ("checkpoints", "models"):
|
||
try:
|
||
folders = folder_paths.get_folder_paths(key)
|
||
except Exception:
|
||
continue
|
||
for folder in folders:
|
||
if not os.path.isdir(folder):
|
||
continue
|
||
for f in os.listdir(folder):
|
||
if "groundingdino" in f.lower() and f.endswith((".pth", ".pt", ".bin")):
|
||
return os.path.join(folder, f)
|
||
return None
|
||
|
||
|
||
# ============================================================================
|
||
# 11. Image / mask utilities & Conditioning Normalizer
|
||
# ============================================================================
|
||
def _gaussian_blur_mask(mask_t, blur_px):
|
||
if blur_px <= 0:
|
||
return mask_t
|
||
k = blur_px * 2 + 1
|
||
try:
|
||
import torchvision.transforms.functional as TF
|
||
return TF.gaussian_blur(
|
||
mask_t.unsqueeze(0).unsqueeze(0), kernel_size=k, sigma=float(blur_px)
|
||
).squeeze()
|
||
except Exception:
|
||
return mask_t
|
||
|
||
|
||
def _erode_dilate(mask_np, erosion, dilation):
|
||
try:
|
||
import cv2
|
||
m = (mask_np * 255).astype(np.uint8)
|
||
if erosion > 0:
|
||
k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (erosion * 2 + 1,) * 2)
|
||
m = cv2.erode(m, k)
|
||
if dilation > 0:
|
||
k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation * 2 + 1,) * 2)
|
||
m = cv2.dilate(m, k)
|
||
return m.astype(np.float32) / 255.0
|
||
except Exception:
|
||
return mask_np
|
||
|
||
|
||
def _snap8(v):
|
||
return math.ceil(v / 8) * 8
|
||
|
||
|
||
def _normalize_conditioning(cond):
|
||
"""
|
||
Ensure conditioning is in ComfyUI's canonical format: list of (tensor, dict) tuples.
|
||
Handles all known formats produced by CLIP encode nodes across ComfyUI versions.
|
||
"""
|
||
if isinstance(cond, list) and len(cond) > 0:
|
||
first = cond[0]
|
||
# Already canonical: list of (tensor, dict) tuples
|
||
if isinstance(first, (list, tuple)) and len(first) == 2:
|
||
tensor_part, extra_part = first
|
||
if isinstance(tensor_part, torch.Tensor) and isinstance(extra_part, dict):
|
||
return cond
|
||
# Old dict-list format
|
||
result = []
|
||
for item in cond:
|
||
if isinstance(item, dict):
|
||
tensor = item.get("cond") or item.get("samples")
|
||
extra = {}
|
||
if "pooled_output" in item:
|
||
extra["pooled_output"] = item["pooled_output"]
|
||
result.append((tensor, extra))
|
||
elif isinstance(item, (list, tuple)) and len(item) == 2:
|
||
result.append(tuple(item))
|
||
else:
|
||
result.append((item, {}))
|
||
return result
|
||
elif isinstance(cond, dict):
|
||
tensor = cond.get("cond") or cond.get("samples")
|
||
extra = {}
|
||
if "pooled_output" in cond:
|
||
extra["pooled_output"] = cond["pooled_output"]
|
||
return [(tensor, extra)]
|
||
elif isinstance(cond, torch.Tensor):
|
||
return [(cond, {})]
|
||
return cond
|
||
|
||
|
||
def _prepare_noise_mask_for_latent(mask_pixel, latent_tensor, feather_px):
|
||
"""
|
||
Take a pixel-space mask and produce a latent-space noise_mask tensor
|
||
with shape [1, H_latent, W_latent] ready for comfy.sample.sample.
|
||
"""
|
||
# mask_pixel: [H_pixel, W_pixel]
|
||
nm = mask_pixel.unsqueeze(0) # [1, H, W]
|
||
if feather_px > 0:
|
||
nm = _gaussian_blur_mask(nm.squeeze(0), feather_px).unsqueeze(0)
|
||
|
||
lh, lw = latent_tensor.shape[2], latent_tensor.shape[3]
|
||
# Interpolate to latent spatial size: [1, 1, lh, lw] → squeeze → [1, lh, lw]
|
||
noise_mask = torch.nn.functional.interpolate(
|
||
nm.unsqueeze(0).float(),
|
||
size=(lh, lw),
|
||
mode="bilinear",
|
||
align_corners=False,
|
||
).squeeze(0) # [1, lh, lw]
|
||
return noise_mask.clamp(0.0, 1.0)
|
||
|
||
|
||
# ============================================================================
|
||
# 12. Main node: SAM Smart Inpainter
|
||
# ============================================================================
|
||
class SAM_SmartInpainter:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
models = _list_all_sam_models()
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"sam_model": (
|
||
models,
|
||
{"tooltip": "SAM3 → models/sam3/ | SAM2 → models/sam2/ or models/sams/"},
|
||
),
|
||
"prompt": ("STRING", {"default": "face", "multiline": True}),
|
||
},
|
||
"optional": {
|
||
"text_threshold": ("FLOAT", {"default": 0.35, "min": 0.05, "max": 0.95, "step": 0.05}),
|
||
"fallback_point_mode": (["center", "face_region", "custom"], {"default": "face_region"}),
|
||
"custom_point_x": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"custom_point_y": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"multi_mask_merge": (["union", "largest", "highest_score"], {"default": "union"}),
|
||
"min_mask_area": ("FLOAT", {"default": 0.005, "min": 0.0001, "max": 0.5, "step": 0.001}),
|
||
"max_mask_area": ("FLOAT", {"default": 0.95, "min": 0.1, "max": 1.0, "step": 0.01}),
|
||
"mask_blur": ("INT", {"default": 6, "min": 0, "max": 64}),
|
||
"mask_erosion": ("INT", {"default": 0, "min": 0, "max": 32}),
|
||
"mask_dilation": ("INT", {"default": 4, "min": 0, "max": 64}),
|
||
"mask_threshold": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 0.99, "step": 0.01}),
|
||
"invert_mask": ("BOOLEAN", {"default": False}),
|
||
"crop_padding": ("INT", {"default": 32, "min": 0, "max": 512, "step": 8}),
|
||
"crop_padding_mode": (["aspect", "square"], {"default": "aspect"}),
|
||
"inpaint_resolution": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 64}),
|
||
"upscale_method": (["bicubic", "bilinear", "lanczos", "nearest"], {"default": "bicubic"}),
|
||
"blend_mode": (["feather", "mask", "hard"], {"default": "feather"}),
|
||
"blend_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}),
|
||
"model": ("MODEL",),
|
||
"clip": ("CLIP",),
|
||
"vae": ("VAE",),
|
||
"positive": ("CONDITIONING",),
|
||
"negative": ("CONDITIONING",),
|
||
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "dpmpp_2m"}),
|
||
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "karras"}),
|
||
"steps": ("INT", {"default": 20, "min": 1, "max": 150}),
|
||
"cfg": ("FLOAT", {"default": 7.0, "min": 1.0, "max": 30.0, "step": 0.5}),
|
||
"denoise": ("FLOAT", {"default": 0.55, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
||
"noise_mask_feather": ("INT", {"default": 8, "min": 0, "max": 64}),
|
||
"device": (["auto", "cuda", "cpu", "mps"], {"default": "auto"}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE", "MASK", "STRING")
|
||
RETURN_NAMES = ("output_image", "inpainted_patch", "segmentation_mask", "debug_info")
|
||
FUNCTION = "execute"
|
||
CATEGORY = "image/inpainting"
|
||
|
||
# ------------------------------------------------------------------
|
||
def execute(self, image, sam_model, prompt, **kwargs):
|
||
log = []
|
||
device = (
|
||
comfy.model_management.get_torch_device()
|
||
if kwargs.get("device", "auto") == "auto"
|
||
else torch.device(kwargs["device"])
|
||
)
|
||
log.append(f"Device: {device}")
|
||
|
||
img_t = image[0]
|
||
img_np = (img_t.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
||
if img_np.shape[-1] == 4:
|
||
img_np = img_np[..., :3]
|
||
H, W = img_np.shape[:2]
|
||
|
||
model_path = _resolve_path(sam_model)
|
||
if model_path is None:
|
||
return self._passthrough(image, f"Model file not found: {sam_model}")
|
||
|
||
family = "sam3" if _is_sam3_model(model_path) else "sam2"
|
||
log.append(f"Family: {family.upper()}")
|
||
|
||
try:
|
||
predictor = _build_predictor(model_path, device)
|
||
predictor.set_image(img_np)
|
||
except Exception as e:
|
||
return self._passthrough(
|
||
image, f"Failed to load model: {e}\n{traceback.format_exc()[:500]}"
|
||
)
|
||
log.append("Model loaded OK")
|
||
|
||
masks_and_scores = self._get_masks(
|
||
predictor, img_np, prompt.strip(), H, W,
|
||
kwargs.get("text_threshold", 0.35),
|
||
kwargs.get("fallback_point_mode", "face_region"),
|
||
kwargs.get("custom_point_x", 0.5),
|
||
kwargs.get("custom_point_y", 0.25),
|
||
device, log,
|
||
)
|
||
|
||
if not masks_and_scores:
|
||
return self._passthrough(image, f"No masks produced for: '{prompt}'")
|
||
|
||
masks_only = [m for m, _ in masks_and_scores]
|
||
scores = [s for _, s in masks_and_scores]
|
||
merge = kwargs.get("multi_mask_merge", "union")
|
||
if merge == "largest":
|
||
idx = max(range(len(masks_only)), key=lambda i: masks_only[i].sum())
|
||
merged = masks_only[idx].astype(np.float32)
|
||
elif merge == "highest_score":
|
||
idx = max(range(len(scores)), key=lambda i: scores[i])
|
||
merged = masks_only[idx].astype(np.float32)
|
||
else:
|
||
merged = np.zeros((H, W), dtype=np.float32)
|
||
for m in masks_only:
|
||
merged = np.maximum(merged, m.astype(np.float32))
|
||
|
||
log.append(f"Raw mask coverage: {merged.mean() * 100:.1f}%")
|
||
merged = _erode_dilate(
|
||
merged, kwargs.get("mask_erosion", 0), kwargs.get("mask_dilation", 4)
|
||
)
|
||
if kwargs.get("invert_mask", False):
|
||
merged = 1.0 - merged
|
||
|
||
mask_t = torch.from_numpy(merged).float()
|
||
mask_t = _gaussian_blur_mask(mask_t, kwargs.get("mask_blur", 6))
|
||
mask_bin = (mask_t > kwargs.get("mask_threshold", 0.5)).float()
|
||
frac = mask_bin.mean().item()
|
||
log.append(f"Final mask: {frac * 100:.1f}%")
|
||
|
||
min_area = kwargs.get("min_mask_area", 0.005)
|
||
max_area = kwargs.get("max_mask_area", 0.95)
|
||
if frac < min_area:
|
||
return self._passthrough(
|
||
image, f"Mask too small ({frac * 100:.2f}% < {min_area * 100:.2f}%)"
|
||
)
|
||
if frac > max_area:
|
||
return self._passthrough(
|
||
image, f"Mask too large ({frac * 100:.2f}% > {max_area * 100:.2f}%)"
|
||
)
|
||
|
||
y_idx, x_idx = torch.where(mask_bin > 0.5)
|
||
if len(x_idx) == 0:
|
||
return self._passthrough(image, "No valid mask pixels")
|
||
|
||
# ------------------------------------------------------------------
|
||
# Crop region around the mask
|
||
# ------------------------------------------------------------------
|
||
pad = kwargs.get("crop_padding", 32)
|
||
x1r = max(0, x_idx.min().item() - pad)
|
||
x2r = min(W, x_idx.max().item() + pad)
|
||
y1r = max(0, y_idx.min().item() - pad)
|
||
y2r = min(H, y_idx.max().item() + pad)
|
||
|
||
if kwargs.get("crop_padding_mode", "aspect") == "square":
|
||
side = max(x2r - x1r, y2r - y1r)
|
||
cx, cy = (x1r + x2r) // 2, (y1r + y2r) // 2
|
||
x1 = max(0, cx - side // 2)
|
||
x2 = min(W, cx + side // 2)
|
||
y1 = max(0, cy - side // 2)
|
||
y2 = min(H, cy + side // 2)
|
||
else:
|
||
x1, x2, y1, y2 = x1r, x2r, y1r, y2r
|
||
|
||
x2 = min(W, x1 + _snap8(x2 - x1))
|
||
y2 = min(H, y1 + _snap8(y2 - y1))
|
||
log.append(f"Crop: ({x1},{y1})→({x2},{y2})")
|
||
|
||
crop_img = image[:, y1:y2, x1:x2, :] # [1, ch, cw, 3]
|
||
crop_mask = mask_t[y1:y2, x1:x2] # [ch, cw]
|
||
crop_h, crop_w = y2 - y1, x2 - x1
|
||
|
||
# ------------------------------------------------------------------
|
||
# Upscale crop to inpaint resolution
|
||
# ------------------------------------------------------------------
|
||
iw_base = _snap8(crop_w)
|
||
ih_base = _snap8(crop_h)
|
||
target = kwargs.get("inpaint_resolution", 768)
|
||
scale = target / max(iw_base, ih_base)
|
||
iw = _snap8(int(iw_base * scale))
|
||
ih = _snap8(int(ih_base * scale))
|
||
|
||
up_method = kwargs.get("upscale_method", "bicubic")
|
||
resized_img = comfy.utils.common_upscale(
|
||
crop_img.movedim(-1, 1), iw, ih, up_method, "disabled"
|
||
).movedim(1, -1) # [1, ih, iw, 3]
|
||
resized_mask = comfy.utils.common_upscale(
|
||
crop_mask.unsqueeze(0).unsqueeze(0), iw, ih, up_method, "disabled"
|
||
).squeeze(0).squeeze(0) # [ih, iw]
|
||
|
||
# Default result = the (possibly upscaled) crop (no inpainting)
|
||
result_patch = resized_img
|
||
|
||
# ------------------------------------------------------------------
|
||
# Diffusion inpainting
|
||
# ------------------------------------------------------------------
|
||
diff_model = kwargs.get("model")
|
||
clip = kwargs.get("clip")
|
||
vae = kwargs.get("vae")
|
||
positive = kwargs.get("positive")
|
||
negative = kwargs.get("negative")
|
||
|
||
if all(x is not None for x in [diff_model, vae, clip, positive, negative]):
|
||
try:
|
||
# Step 1 — normalise conditioning
|
||
positive = _normalize_conditioning(positive)
|
||
negative = _normalize_conditioning(negative)
|
||
log.append("Conditioning normalised OK")
|
||
|
||
# Step 2 — VAE encode the cropped+upscaled image
|
||
encode_input = resized_img[..., :3] # drop alpha if present
|
||
latent_img = vae.encode(encode_input) # returns Tensor [B, C, H, W]
|
||
log.append(f"VAE encode OK — latent shape: {list(latent_img.shape)}")
|
||
|
||
# Step 3 — optional fix_empty_latent_channels (newer ComfyUI)
|
||
if hasattr(comfy.sample, "fix_empty_latent_channels"):
|
||
latent_img = comfy.sample.fix_empty_latent_channels(diff_model, latent_img)
|
||
|
||
# Step 4 — build noise mask in latent space [1, H_lat, W_lat]
|
||
noise_mask = _prepare_noise_mask_for_latent(
|
||
resized_mask, latent_img, kwargs.get("noise_mask_feather", 8)
|
||
)
|
||
log.append(f"Noise mask shape: {list(noise_mask.shape)}")
|
||
|
||
# Step 5 — prepare noise
|
||
seed = kwargs.get("seed", 0)
|
||
actual_seed = (
|
||
seed if seed != 0 else torch.randint(0, 2**32, (1,)).item()
|
||
)
|
||
noise = comfy.sample.prepare_noise(latent_img, actual_seed)
|
||
log.append(f"Noise prepared (seed {actual_seed})")
|
||
|
||
# ----------------------------------------------------------
|
||
# Step 6 — run the sampler
|
||
#
|
||
# CRITICAL FIX (v5.3):
|
||
# comfy.sample.sample() expects:
|
||
# latent_image → a raw TENSOR (NOT {"samples": tensor})
|
||
# and returns:
|
||
# samples → a raw TENSOR (NOT a dict)
|
||
#
|
||
# Passing the dict caused:
|
||
# 'dict' object has no attribute 'is_nested'
|
||
# because PyTorch / ComfyUI internally calls .is_nested on
|
||
# what it expects to be a Tensor.
|
||
# ----------------------------------------------------------
|
||
samples = comfy.sample.sample(
|
||
model=diff_model,
|
||
noise=noise,
|
||
steps=kwargs.get("steps", 20),
|
||
cfg=kwargs.get("cfg", 7.0),
|
||
sampler_name=kwargs.get("sampler_name", "dpmpp_2m"),
|
||
scheduler=kwargs.get("scheduler", "karras"),
|
||
positive=positive,
|
||
negative=negative,
|
||
latent_image=latent_img, # ← RAW TENSOR (not dict)
|
||
denoise=kwargs.get("denoise", 0.55),
|
||
disable_noise=False,
|
||
start_step=None,
|
||
last_step=None,
|
||
force_full_denoise=False,
|
||
noise_mask=noise_mask, # ← separate kwarg, shape [1,H,W]
|
||
)
|
||
log.append(f"Sampling OK — output shape: {list(samples.shape)}")
|
||
|
||
# Step 7 — VAE decode: samples is a Tensor, decode it directly
|
||
result_patch = vae.decode(samples) # ← NOT samples["samples"]
|
||
log.append(f"VAE decode OK — patch shape: {list(result_patch.shape)}")
|
||
|
||
except Exception as e:
|
||
log.append(f"Inpainting failed: {e}")
|
||
log.append(traceback.format_exc()[:600])
|
||
else:
|
||
log.append("Inpainting skipped — diffusion inputs not connected")
|
||
|
||
# ------------------------------------------------------------------
|
||
# Downscale result_patch back to original crop size, blend into image
|
||
# ------------------------------------------------------------------
|
||
restored = comfy.utils.common_upscale(
|
||
result_patch.movedim(-1, 1), crop_w, crop_h, up_method, "disabled"
|
||
).movedim(1, -1) # [1, crop_h, crop_w, 3]
|
||
|
||
output = image.clone()
|
||
blend_strength = kwargs.get("blend_strength", 1.0)
|
||
blend_mode = kwargs.get("blend_mode", "feather")
|
||
|
||
if blend_mode == "hard":
|
||
blend_mask = (crop_mask > kwargs.get("mask_threshold", 0.5)).float()
|
||
else:
|
||
blend_mask = crop_mask.clamp(0, 1)
|
||
|
||
bm = (blend_mask * blend_strength).unsqueeze(-1).unsqueeze(0) # [1,H,W,1]
|
||
output[:, y1:y2, x1:x2, :] = (
|
||
output[:, y1:y2, x1:x2, :] * (1.0 - bm) + restored * bm
|
||
)
|
||
|
||
mask_out = mask_t.unsqueeze(0) # [1, H, W]
|
||
return (output, result_patch, mask_out, "\n".join(log))
|
||
|
||
# ------------------------------------------------------------------
|
||
def _get_masks(
|
||
self, predictor, img_np, prompt, H, W,
|
||
threshold, fallback_point_mode, cx_norm, cy_norm, device, log
|
||
):
|
||
if predictor.family == "sam3":
|
||
log.append("Using SAM3 native text prompt")
|
||
results = predictor.predict_with_text(prompt, threshold=threshold)
|
||
if not results:
|
||
log.append("SAM3 text returned no masks — trying point fallback")
|
||
results = predictor.predict_with_point((W // 2, H // 2))
|
||
return results
|
||
else:
|
||
log.append(f"GDino available: {_GDINO_AVAILABLE}")
|
||
boxes = (
|
||
_gdino_boxes(img_np, prompt, threshold, str(device))
|
||
if _GDINO_AVAILABLE
|
||
else []
|
||
)
|
||
log.append(f"GDino detected {len(boxes)} box(es)")
|
||
if boxes:
|
||
return predictor.predict_with_boxes(boxes)
|
||
if fallback_point_mode == "custom":
|
||
px, py = int(cx_norm * W), int(cy_norm * H)
|
||
elif fallback_point_mode == "face_region":
|
||
px, py = W // 2, int(H * 0.25)
|
||
else:
|
||
px, py = W // 2, H // 2
|
||
log.append(f"SAM2 point fallback: ({px},{py})")
|
||
return predictor.predict_with_point((px, py))
|
||
|
||
def _passthrough(self, image, msg):
|
||
print(f"[SAM_SmartInpainter] {msg}")
|
||
return (
|
||
image,
|
||
image[:1],
|
||
torch.zeros(1, image.shape[1], image.shape[2]),
|
||
msg,
|
||
)
|
||
|
||
|
||
# ============================================================================
|
||
# 13. Mask‑only node
|
||
# ============================================================================
|
||
class SAM_MaskOnly:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"sam_model": (_list_all_sam_models(), {}),
|
||
"prompt": ("STRING", {"default": "face", "multiline": True}),
|
||
},
|
||
"optional": {
|
||
"text_threshold": ("FLOAT", {"default": 0.35, "min": 0.05, "max": 0.95, "step": 0.05}),
|
||
"mask_blur": ("INT", {"default": 6, "min": 0, "max": 64}),
|
||
"mask_dilation": ("INT", {"default": 4, "min": 0, "max": 64}),
|
||
"mask_erosion": ("INT", {"default": 0, "min": 0, "max": 32}),
|
||
"invert_mask": ("BOOLEAN", {"default": False}),
|
||
"device": (["auto", "cuda", "cpu", "mps"], {"default": "auto"}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("MASK", "IMAGE", "STRING")
|
||
RETURN_NAMES = ("mask", "mask_preview", "debug_info")
|
||
FUNCTION = "execute"
|
||
CATEGORY = "image/inpainting"
|
||
|
||
def execute(
|
||
self, image, sam_model, prompt,
|
||
text_threshold=0.35, mask_blur=6, mask_dilation=4,
|
||
mask_erosion=0, invert_mask=False, device="auto"
|
||
):
|
||
log = []
|
||
dev = (
|
||
comfy.model_management.get_torch_device()
|
||
if device == "auto"
|
||
else torch.device(device)
|
||
)
|
||
img_np = (image[0].cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
||
if img_np.shape[-1] == 4:
|
||
img_np = img_np[..., :3]
|
||
H, W = img_np.shape[:2]
|
||
|
||
model_path = _resolve_path(sam_model)
|
||
if model_path is None:
|
||
return (torch.zeros(1, H, W), image, f"Model not found: {sam_model}")
|
||
try:
|
||
predictor = _build_predictor(model_path, dev)
|
||
predictor.set_image(img_np)
|
||
except Exception as e:
|
||
return (torch.zeros(1, H, W), image, f"Model load error: {e}")
|
||
|
||
if predictor.family == "sam3":
|
||
results = predictor.predict_with_text(prompt.strip(), threshold=text_threshold)
|
||
if not results:
|
||
results = predictor.predict_with_point((W // 2, H // 2))
|
||
else:
|
||
boxes = (
|
||
_gdino_boxes(img_np, prompt.strip(), text_threshold, str(dev))
|
||
if _GDINO_AVAILABLE
|
||
else []
|
||
)
|
||
results = (
|
||
predictor.predict_with_boxes(boxes)
|
||
if boxes
|
||
else predictor.predict_with_point((W // 2, H // 2))
|
||
)
|
||
|
||
if not results:
|
||
return (torch.zeros(1, H, W), image, f"No masks found for: '{prompt}'")
|
||
|
||
merged = np.zeros((H, W), dtype=np.float32)
|
||
for m, _ in results:
|
||
merged = np.maximum(merged, m.astype(np.float32))
|
||
if mask_erosion > 0 or mask_dilation > 0:
|
||
merged = _erode_dilate(merged, mask_erosion, mask_dilation)
|
||
if invert_mask:
|
||
merged = 1.0 - merged
|
||
|
||
mask_t = torch.from_numpy(merged).float()
|
||
mask_t = _gaussian_blur_mask(mask_t, mask_blur)
|
||
log.append(f"Coverage: {mask_t.mean().item() * 100:.1f}%")
|
||
|
||
mask_preview = mask_t.unsqueeze(-1).expand(H, W, 3).unsqueeze(0)
|
||
return (mask_t.unsqueeze(0), mask_preview, "\n".join(log))
|
||
|
||
|
||
# ============================================================================
|
||
# 14. Node registration
|
||
# ============================================================================
|
||
NODE_CLASS_MAPPINGS = {
|
||
"SAM_SmartInpainter": SAM_SmartInpainter,
|
||
"SAM_MaskOnly": SAM_MaskOnly,
|
||
}
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"SAM_SmartInpainter": "SAM Smart Inpainter (SAM3 v0.3)",
|
||
"SAM_MaskOnly": "SAM Mask Generator (SAM3 v0.3)",
|
||
}
|