""" SAM3 (Segment Anything 3) nodes for detection, segmentation, and video tracking. """ from typing_extensions import override import json import os import torch import torch.nn.functional as F import comfy.model_management import comfy.utils import folder_paths from comfy_api.latest import ComfyExtension, io, ui import av from fractions import Fraction def _extract_text_prompts(conditioning, device, dtype): """Extract list of (text_embeddings, text_mask) from conditioning.""" cond_meta = conditioning[0][1] multi = cond_meta.get("sam3_multi_cond") prompts = [] if multi is not None: for entry in multi: emb = entry["cond"].to(device=device, dtype=dtype) mask = entry["attention_mask"].to(device) if entry["attention_mask"] is not None else None if mask is None: mask = torch.ones(emb.shape[0], emb.shape[1], dtype=torch.int64, device=device) prompts.append((emb, mask, entry.get("max_detections", 1))) else: emb = conditioning[0][0].to(device=device, dtype=dtype) mask = cond_meta.get("attention_mask") if mask is not None: mask = mask.to(device) else: mask = torch.ones(emb.shape[0], emb.shape[1], dtype=torch.int64, device=device) prompts.append((emb, mask, 1)) return prompts def _refine_mask(sam3_model, orig_image_hwc, coarse_mask, box_xyxy, H, W, device, dtype, iterations): """Refine a coarse detector mask via SAM decoder, cropping to the detection box. Returns: [1, H, W] binary mask """ def _coarse_fallback(): return (F.interpolate(coarse_mask.unsqueeze(0).unsqueeze(0), size=(H, W), mode="bilinear", align_corners=False)[0] > 0).float() if iterations <= 0: return _coarse_fallback() pad_frac = 0.1 x1, y1, x2, y2 = box_xyxy.tolist() bw, bh = x2 - x1, y2 - y1 cx1 = max(0, int(x1 - bw * pad_frac)) cy1 = max(0, int(y1 - bh * pad_frac)) cx2 = min(W, int(x2 + bw * pad_frac)) cy2 = min(H, int(y2 + bh * pad_frac)) if cx2 <= cx1 or cy2 <= cy1: return _coarse_fallback() crop = orig_image_hwc[cy1:cy2, cx1:cx2, :3] crop_1008 = comfy.utils.common_upscale(crop.unsqueeze(0).movedim(-1, 1), 1008, 1008, "bilinear", crop="disabled") crop_frame = crop_1008.to(device=device, dtype=dtype) crop_h, crop_w = cy2 - cy1, cx2 - cx1 # Crop coarse mask and refine via SAM on the cropped image mask_h, mask_w = coarse_mask.shape[-2:] mx1, my1 = int(cx1 / W * mask_w), int(cy1 / H * mask_h) mx2, my2 = int(cx2 / W * mask_w), int(cy2 / H * mask_h) if mx2 <= mx1 or my2 <= my1: return _coarse_fallback() mask_logit = coarse_mask[..., my1:my2, mx1:mx2].unsqueeze(0).unsqueeze(0) for _ in range(iterations): coarse_input = F.interpolate(mask_logit, size=(1008, 1008), mode="bilinear", align_corners=False) mask_logit = sam3_model.forward_segment(crop_frame, mask_inputs=coarse_input) refined_crop = F.interpolate(mask_logit, size=(crop_h, crop_w), mode="bilinear", align_corners=False) full_mask = torch.zeros(1, 1, H, W, device=device, dtype=dtype) full_mask[:, :, cy1:cy2, cx1:cx2] = refined_crop coarse_full = F.interpolate(coarse_mask.unsqueeze(0).unsqueeze(0), size=(H, W), mode="bilinear", align_corners=False) return ((full_mask[0] > 0) | (coarse_full[0] > 0)).float() class SAM3_Detect(io.ComfyNode): """Open-vocabulary detection and segmentation using text, box, or point prompts.""" @classmethod def define_schema(cls): return io.Schema( node_id="SAM3_Detect", display_name="SAM3 Detect", category="image/detection", search_aliases=["sam3", "segment anything", "open vocabulary", "text detection", "segment"], inputs=[ io.Model.Input("model", display_name="model"), io.Image.Input("image", display_name="image"), io.Conditioning.Input("conditioning", display_name="conditioning", optional=True, tooltip="Text conditioning from CLIPTextEncode"), io.BoundingBox.Input("bboxes", display_name="bboxes", force_input=True, optional=True, tooltip="Bounding boxes to segment within"), io.String.Input("positive_coords", display_name="positive_coords", force_input=True, optional=True, tooltip="Positive point prompts as JSON [{\"x\": int, \"y\": int}, ...] (pixel coords)"), io.String.Input("negative_coords", display_name="negative_coords", force_input=True, optional=True, tooltip="Negative point prompts as JSON [{\"x\": int, \"y\": int}, ...] (pixel coords)"), io.Float.Input("threshold", display_name="threshold", default=0.5, min=0.0, max=1.0, step=0.01), io.Int.Input("refine_iterations", display_name="refine_iterations", default=2, min=0, max=5, tooltip="SAM decoder refinement passes (0=use raw detector masks)"), io.Boolean.Input("individual_masks", display_name="individual_masks", default=False, tooltip="Output per-object masks instead of union"), ], outputs=[ io.Mask.Output("masks"), io.BoundingBox.Output("bboxes"), ], ) @classmethod def execute(cls, model, image, conditioning=None, bboxes=None, positive_coords=None, negative_coords=None, threshold=0.5, refine_iterations=2, individual_masks=False) -> io.NodeOutput: B, H, W, C = image.shape image_in = comfy.utils.common_upscale(image[..., :3].movedim(-1, 1), 1008, 1008, "bilinear", crop="disabled") # Convert bboxes to normalized cxcywh format, per-frame list of [1, N, 4] tensors. # Supports: single dict (all frames), list[dict] (all frames), list[list[dict]] (per-frame). def _boxes_to_tensor(box_list): coords = [] for d in box_list: cx = (d["x"] + d["width"] / 2) / W cy = (d["y"] + d["height"] / 2) / H coords.append([cx, cy, d["width"] / W, d["height"] / H]) return torch.tensor([coords], dtype=torch.float32) # [1, N, 4] per_frame_boxes = None if bboxes is not None: if isinstance(bboxes, dict): # Single box → same for all frames shared = _boxes_to_tensor([bboxes]) per_frame_boxes = [shared] * B elif isinstance(bboxes, list) and len(bboxes) > 0 and isinstance(bboxes[0], list): # list[list[dict]] → per-frame boxes per_frame_boxes = [_boxes_to_tensor(frame_boxes) if frame_boxes else None for frame_boxes in bboxes] # Pad to B if fewer frames provided while len(per_frame_boxes) < B: per_frame_boxes.append(per_frame_boxes[-1] if per_frame_boxes else None) elif isinstance(bboxes, list) and len(bboxes) > 0: # list[dict] → same boxes for all frames shared = _boxes_to_tensor(bboxes) per_frame_boxes = [shared] * B # Parse point prompts from JSON (KJNodes PointsEditor format: [{"x": int, "y": int}, ...]) pos_pts = json.loads(positive_coords) if positive_coords else [] neg_pts = json.loads(negative_coords) if negative_coords else [] has_points = len(pos_pts) > 0 or len(neg_pts) > 0 comfy.model_management.load_model_gpu(model) device = comfy.model_management.get_torch_device() dtype = model.model.get_dtype() sam3_model = model.model.diffusion_model # Build point inputs for tracker SAM decoder path point_inputs = None if has_points: all_coords = [[p["x"] / W * 1008, p["y"] / H * 1008] for p in pos_pts] + \ [[p["x"] / W * 1008, p["y"] / H * 1008] for p in neg_pts] all_labels = [1] * len(pos_pts) + [0] * len(neg_pts) point_inputs = { "point_coords": torch.tensor([all_coords], dtype=dtype, device=device), "point_labels": torch.tensor([all_labels], dtype=torch.int32, device=device), } cond_list = _extract_text_prompts(conditioning, device, dtype) if conditioning is not None and len(conditioning) > 0 else [] has_text = len(cond_list) > 0 # Run per-image through detector (text/boxes) and/or tracker (points) all_bbox_dicts = [] all_masks = [] pbar = comfy.utils.ProgressBar(B) for b in range(B): frame = image_in[b:b+1].to(device=device, dtype=dtype) b_boxes = None if per_frame_boxes is not None and per_frame_boxes[b] is not None: b_boxes = per_frame_boxes[b].to(device=device, dtype=dtype) frame_bbox_dicts = [] frame_masks = [] # Point prompts: tracker SAM decoder path with iterative refinement if point_inputs is not None: mask_logit = sam3_model.forward_segment(frame, point_inputs=point_inputs) for _ in range(max(0, refine_iterations - 1)): mask_logit = sam3_model.forward_segment(frame, mask_inputs=mask_logit) mask = F.interpolate(mask_logit, size=(H, W), mode="bilinear", align_corners=False) frame_masks.append((mask[0] > 0).float()) # Box prompts: SAM decoder path (segment inside each box) if b_boxes is not None and not has_text: for box_cxcywh in b_boxes[0]: cx, cy, bw, bh = box_cxcywh.tolist() # Convert cxcywh normalized → xyxy in 1008 space → [1, 2, 2] corners sam_box = torch.tensor([[[(cx - bw/2) * 1008, (cy - bh/2) * 1008], [(cx + bw/2) * 1008, (cy + bh/2) * 1008]]], device=device, dtype=dtype) mask_logit = sam3_model.forward_segment(frame, box_inputs=sam_box) for _ in range(max(0, refine_iterations - 1)): mask_logit = sam3_model.forward_segment(frame, mask_inputs=mask_logit) mask = F.interpolate(mask_logit, size=(H, W), mode="bilinear", align_corners=False) frame_masks.append((mask[0] > 0).float()) # Text prompts: run detector per text prompt (each detects one category) for text_embeddings, text_mask, max_det in cond_list: results = sam3_model( frame, text_embeddings=text_embeddings, text_mask=text_mask, boxes=b_boxes, threshold=threshold, orig_size=(H, W)) pred_boxes = results["boxes"][0] scores = results["scores"][0] masks = results["masks"][0] probs = scores.sigmoid() keep = probs > threshold kept_boxes = pred_boxes[keep].cpu() kept_scores = probs[keep].cpu() kept_masks = masks[keep] order = kept_scores.argsort(descending=True)[:max_det] kept_boxes = kept_boxes[order] kept_scores = kept_scores[order] kept_masks = kept_masks[order] for box, score in zip(kept_boxes, kept_scores): frame_bbox_dicts.append({ "x": float(box[0]), "y": float(box[1]), "width": float(box[2] - box[0]), "height": float(box[3] - box[1]), "score": float(score), }) for m, box in zip(kept_masks, kept_boxes): frame_masks.append(_refine_mask( sam3_model, image[b], m, box, H, W, device, dtype, refine_iterations)) all_bbox_dicts.append(frame_bbox_dicts) if len(frame_masks) > 0: combined = torch.cat(frame_masks, dim=0) # [N_obj, H, W] if individual_masks: all_masks.append(combined) else: all_masks.append((combined > 0).any(dim=0).float()) else: if individual_masks: all_masks.append(torch.zeros(0, H, W, device=comfy.model_management.intermediate_device())) else: all_masks.append(torch.zeros(H, W, device=comfy.model_management.intermediate_device())) pbar.update(1) idev = comfy.model_management.intermediate_device() all_masks = [m.to(idev) for m in all_masks] mask_out = torch.cat(all_masks, dim=0) if individual_masks else torch.stack(all_masks) return io.NodeOutput(mask_out, all_bbox_dicts) SAM3TrackData = io.Custom("SAM3_TRACK_DATA") class SAM3_VideoTrack(io.ComfyNode): """Track objects across video frames using SAM3's memory-based tracker.""" @classmethod def define_schema(cls): return io.Schema( node_id="SAM3_VideoTrack", display_name="SAM3 Video Track", category="image/detection", search_aliases=["sam3", "video", "track", "propagate"], inputs=[ io.Image.Input("images", display_name="images", tooltip="Video frames as batched images"), io.Model.Input("model", display_name="model"), io.Mask.Input("initial_mask", display_name="initial_mask", optional=True, tooltip="Mask(s) for the first frame to track (one per object)"), io.Conditioning.Input("conditioning", display_name="conditioning", optional=True, tooltip="Text conditioning for detecting new objects during tracking"), io.Float.Input("detection_threshold", display_name="detection_threshold", default=0.5, min=0.0, max=1.0, step=0.01, tooltip="Score threshold for text-prompted detection."), io.Int.Input("max_objects", display_name="max_objects", default=4, min=0, max=64, tooltip="Max tracked objects. Initial masks count toward this limit. 0 uses the internal cap of 64."), io.Int.Input("detect_interval", display_name="detect_interval", default=1, min=1, tooltip="Run detection every N frames (1=every frame). Higher values save compute."), ], outputs=[ SAM3TrackData.Output("track_data", display_name="track_data"), ], ) @classmethod def execute(cls, images, model, initial_mask=None, conditioning=None, detection_threshold=0.5, max_objects=0, detect_interval=1) -> io.NodeOutput: N, H, W, C = images.shape comfy.model_management.load_model_gpu(model) device = comfy.model_management.get_torch_device() dtype = model.model.get_dtype() sam3_model = model.model.diffusion_model frames_in = images[..., :3].movedim(-1, 1) init_masks = None if initial_mask is not None: init_masks = initial_mask.unsqueeze(1).to(device=device, dtype=dtype) pbar = comfy.utils.ProgressBar(N) text_prompts = None if conditioning is not None and len(conditioning) > 0: text_prompts = [(emb, mask) for emb, mask, _ in _extract_text_prompts(conditioning, device, dtype)] elif initial_mask is None: raise ValueError("Either initial_mask or conditioning must be provided") result = sam3_model.forward_video( images=frames_in, initial_masks=init_masks, pbar=pbar, text_prompts=text_prompts, new_det_thresh=detection_threshold, max_objects=max_objects, detect_interval=detect_interval, target_device=device, target_dtype=dtype) result["orig_size"] = (H, W) return io.NodeOutput(result) class SAM3_TrackPreview(io.ComfyNode): """Visualize tracked objects with distinct colors as a video preview. No tensor output — saves to temp video.""" @classmethod def define_schema(cls): return io.Schema( node_id="SAM3_TrackPreview", display_name="SAM3 Track Preview", category="image/detection", inputs=[ SAM3TrackData.Input("track_data", display_name="track_data"), io.Image.Input("images", display_name="images", optional=True), io.Float.Input("opacity", display_name="opacity", default=0.5, min=0.0, max=1.0, step=0.05), io.Float.Input("fps", display_name="fps", default=24.0, min=1.0, max=120.0, step=1.0), ], is_output_node=True, ) COLORS = [ (0.12, 0.47, 0.71), (1.0, 0.5, 0.05), (0.17, 0.63, 0.17), (0.84, 0.15, 0.16), (0.58, 0.4, 0.74), (0.55, 0.34, 0.29), (0.89, 0.47, 0.76), (0.5, 0.5, 0.5), (0.74, 0.74, 0.13), (0.09, 0.75, 0.81), (0.94, 0.76, 0.06), (0.42, 0.68, 0.84), ] # 5x3 bitmap font atlas for digits 0-9 [10, 5, 3] _glyph_cache = {} # (device, scale) -> (glyphs, outlines, gh, gw, oh, ow) @staticmethod def _get_glyphs(device, scale=3): key = (device, scale) if key in SAM3_TrackPreview._glyph_cache: return SAM3_TrackPreview._glyph_cache[key] atlas = torch.tensor([ [[1,1,1],[1,0,1],[1,0,1],[1,0,1],[1,1,1]], [[0,1,0],[1,1,0],[0,1,0],[0,1,0],[1,1,1]], [[1,1,1],[0,0,1],[1,1,1],[1,0,0],[1,1,1]], [[1,1,1],[0,0,1],[1,1,1],[0,0,1],[1,1,1]], [[1,0,1],[1,0,1],[1,1,1],[0,0,1],[0,0,1]], [[1,1,1],[1,0,0],[1,1,1],[0,0,1],[1,1,1]], [[1,1,1],[1,0,0],[1,1,1],[1,0,1],[1,1,1]], [[1,1,1],[0,0,1],[0,0,1],[0,0,1],[0,0,1]], [[1,1,1],[1,0,1],[1,1,1],[1,0,1],[1,1,1]], [[1,1,1],[1,0,1],[1,1,1],[0,0,1],[1,1,1]], ], dtype=torch.bool) glyphs, outlines = [], [] for d in range(10): g = atlas[d].repeat_interleave(scale, 0).repeat_interleave(scale, 1) padded = F.pad(g.float().unsqueeze(0).unsqueeze(0), (1,1,1,1)) o = (F.max_pool2d(padded, 3, stride=1, padding=1)[0, 0] > 0) glyphs.append(g.to(device)) outlines.append(o.to(device)) gh, gw = glyphs[0].shape oh, ow = outlines[0].shape SAM3_TrackPreview._glyph_cache[key] = (glyphs, outlines, gh, gw, oh, ow) return SAM3_TrackPreview._glyph_cache[key] @staticmethod def _draw_number_gpu(frame, number, cx, cy, color, scale=3): """Draw a number on a GPU tensor [H, W, 3] float 0-1 at (cx, cy) with outline.""" H, W = frame.shape[:2] device = frame.device glyphs, outlines, gh, gw, oh, ow = SAM3_TrackPreview._get_glyphs(device, scale) color_t = torch.tensor(color, device=device, dtype=frame.dtype) digs = [int(d) for d in str(number)] total_w = len(digs) * (gw + scale) - scale x0 = cx - total_w // 2 y0 = cy - gh // 2 for i, d in enumerate(digs): dx = x0 + i * (gw + scale) # Black outline oy0, ox0 = y0 - 1, dx - 1 osy1, osx1 = max(0, -oy0), max(0, -ox0) osy2, osx2 = min(oh, H - oy0), min(ow, W - ox0) if osy2 > osy1 and osx2 > osx1: fy1, fx1 = oy0 + osy1, ox0 + osx1 frame[fy1:fy1+(osy2-osy1), fx1:fx1+(osx2-osx1)][outlines[d][osy1:osy2, osx1:osx2]] = 0 # Colored fill sy1, sx1 = max(0, -y0), max(0, -dx) sy2, sx2 = min(gh, H - y0), min(gw, W - dx) if sy2 > sy1 and sx2 > sx1: fy1, fx1 = y0 + sy1, dx + sx1 frame[fy1:fy1+(sy2-sy1), fx1:fx1+(sx2-sx1)][glyphs[d][sy1:sy2, sx1:sx2]] = color_t @classmethod def execute(cls, track_data, images=None, opacity=0.5, fps=24.0) -> io.NodeOutput: from comfy.ldm.sam3.tracker import unpack_masks packed = track_data["packed_masks"] H, W = track_data["orig_size"] if images is not None: H, W = images.shape[1], images.shape[2] if packed is None: N, N_obj = track_data["n_frames"], 0 else: N, N_obj = packed.shape[0], packed.shape[1] import uuid gpu = comfy.model_management.get_torch_device() temp_dir = folder_paths.get_temp_directory() filename = f"sam3_track_preview_{uuid.uuid4().hex[:8]}.mp4" filepath = os.path.join(temp_dir, filename) with av.open(filepath, mode='w') as output: stream = output.add_stream('h264', rate=Fraction(round(fps * 1000), 1000)) stream.width = W stream.height = H stream.pix_fmt = 'yuv420p' frame_cpu = torch.empty(H, W, 3, dtype=torch.uint8) frame_np = frame_cpu.numpy() if N_obj > 0: colors_t = torch.tensor([cls.COLORS[i % len(cls.COLORS)] for i in range(N_obj)], device=gpu, dtype=torch.float32) grid_y = torch.arange(H, device=gpu).view(1, H, 1) grid_x = torch.arange(W, device=gpu).view(1, 1, W) for t in range(N): if images is not None and t < images.shape[0]: frame = images[t].clone() else: frame = torch.zeros(H, W, 3) if N_obj > 0: frame_binary = unpack_masks(packed[t:t+1].to(gpu)) # [1, N_obj, H, W] bool frame_masks = F.interpolate(frame_binary.float(), size=(H, W), mode="nearest")[0] frame_gpu = frame.to(gpu) bool_masks = frame_masks > 0.5 any_mask = bool_masks.any(dim=0) if any_mask.any(): obj_idx_map = bool_masks.to(torch.uint8).argmax(dim=0) color_overlay = colors_t[obj_idx_map] mask_3d = any_mask.unsqueeze(-1) frame_gpu = torch.where(mask_3d, frame_gpu * (1 - opacity) + color_overlay * opacity, frame_gpu) area = bool_masks.sum(dim=(-1, -2)).clamp_(min=1) cy = (bool_masks * grid_y).sum(dim=(-1, -2)) // area cx = (bool_masks * grid_x).sum(dim=(-1, -2)) // area has = area > 1 scores = track_data.get("scores", []) label_scale = max(3, H // 240) # Scale font with resolutio size_caps = (area.float().sqrt() / 15).clamp_(min=1).long().tolist() #cap per-object so the number doesn't dwarf small masks for obj_idx in range(N_obj): if has[obj_idx]: _cx, _cy = int(cx[obj_idx]), int(cy[obj_idx]) color = cls.COLORS[obj_idx % len(cls.COLORS)] obj_scale = min(label_scale, size_caps[obj_idx]) score_scale = max(1, obj_scale * 2 // 3) SAM3_TrackPreview._draw_number_gpu(frame_gpu, obj_idx, _cx, _cy, color, scale=obj_scale) if obj_idx < len(scores) and scores[obj_idx] < 1.0: SAM3_TrackPreview._draw_number_gpu(frame_gpu, int(scores[obj_idx] * 100), _cx, _cy + 5 * obj_scale + 3, color, scale=score_scale) frame_cpu.copy_(frame_gpu.clamp_(0, 1).mul_(255).byte()) else: frame_cpu.copy_(frame.clamp_(0, 1).mul_(255).byte()) vframe = av.VideoFrame.from_ndarray(frame_np, format='rgb24') output.mux(stream.encode(vframe.reformat(format='yuv420p'))) output.mux(stream.encode(None)) return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(filename, "", io.FolderType.temp)])) class SAM3_TrackToMask(io.ComfyNode): """Select tracked objects by index and output as mask.""" @classmethod def define_schema(cls): return io.Schema( node_id="SAM3_TrackToMask", display_name="SAM3 Track to Mask", category="image/detection", inputs=[ SAM3TrackData.Input("track_data", display_name="track_data"), io.String.Input("object_indices", display_name="object_indices", default="", tooltip="Comma-separated object indices to include (e.g. '0,2,3'). Empty = all objects."), ], outputs=[ io.Mask.Output("masks", display_name="masks"), ], ) @classmethod def execute(cls, track_data, object_indices="") -> io.NodeOutput: from comfy.ldm.sam3.tracker import unpack_masks packed = track_data["packed_masks"] H, W = track_data["orig_size"] if packed is None: N = track_data["n_frames"] return io.NodeOutput(torch.zeros(N, H, W, device=comfy.model_management.intermediate_device())) N, N_obj = packed.shape[0], packed.shape[1] if object_indices.strip(): indices = [int(i.strip()) for i in object_indices.split(",") if i.strip().isdigit()] indices = [i for i in indices if 0 <= i < N_obj] else: indices = list(range(N_obj)) if not indices: return io.NodeOutput(torch.zeros(N, H, W, device=comfy.model_management.intermediate_device())) union_packed = packed[:, indices[0]].clone() for i in indices[1:]: union_packed |= packed[:, i] union = unpack_masks(union_packed).unsqueeze(1).float() # [N, 1, Hm, Wm] mask_out = F.interpolate(union, size=(H, W), mode="bilinear", align_corners=False)[:, 0] return io.NodeOutput(mask_out) class SAM3Extension(ComfyExtension): @override async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ SAM3_Detect, SAM3_VideoTrack, SAM3_TrackPreview, SAM3_TrackToMask, ] async def comfy_entrypoint() -> SAM3Extension: return SAM3Extension()