Repository navigation
feat: GrabCut H100 optimization — torch.no_grad, Pydantic validation,… - #12
Conversation
… structlog, type hints - grabcut_nodes.py: torch.no_grad() on remove_background + refine_mask, dtype=torch.float32 on all tensor conversions, torch.cuda.empty_cache() in finally blocks, Pydantic validation at start of every execute() - grabcut_remover.py: Full type-hinting on GrabCutProcessor (mypy-ready), structlog observability with GPU memory tracking (_log_gpu_memory), torch.no_grad() + empty_cache() on all inference paths - requirements.txt: Pin torch>=2.1.0, Pillow>=10.4.0 (CVE fix), opencv>=4.10.0, structlog>=24.4.0, pydantic>=2.7.0 - pyproject.toml: python>=3.10 (from 3.8) - src/validation.py: NEW — Pydantic models for all user-facing params (GrabCutParams, ScalingParams, MaskParams, BBoxParams, GrabCutNodeParams) - src/__init__.py: NEW — package marker - 6/6 pytest tests pass, node import test passes Ref: CD-14 | notion:32c40bf6-db23-8001-92b8-dfe4d91dccb4
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdded structured logging, GPU memory instrumentation and cleanup, Pydantic parameter validation, YOLO model caching (with optional torch.compile), an invertible mask option, bbox metadata tensor output, explicit float32 tensor casting, and Changes
Sequence DiagramsequenceDiagram
participant User as User / Node
participant Val as Validation\n(src/validation.py)
participant Mem as GPU Memory\nLogger
participant YOLO as YOLO Detector\n(cached)
participant GC as GrabCut\nProcessor
participant CUDA as CUDA\nCache
User->>Val: validate_node_params(params)
Val-->>User: validated params / raise ValueError
User->>Mem: _log_gpu_memory("start")
Mem-->>User: log stats
User->>YOLO: detect_object(image)
YOLO->>YOLO: check _yolo_cache / attempt torch.compile()
YOLO->>Mem: _log_gpu_memory("detect.start")
YOLO->>YOLO: run inference
YOLO->>Mem: _log_gpu_memory("detect.end")
YOLO->>CUDA: torch.cuda.empty_cache()
YOLO-->>User: bbox + confidence
User->>GC: process_with_grabcut(image, bbox, invert_mask)
GC->>Mem: _log_gpu_memory("grabcut.start")
GC->>GC: process `@torch.no_grad()`()
GC->>Mem: _log_gpu_memory("grabcut.end")
GC->>CUDA: torch.cuda.empty_cache()
GC-->>User: mask, rgba, bbox_metadata
User->>Mem: _log_gpu_memory("end")
Mem-->>User: final stats
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| grabcut_iterations=grabcut_iterations, | ||
| margin=margin_pixels, | ||
| edge_threshold=0.5, | ||
| confidence_threshold=confidence_threshold, |
There was a problem hiding this comment.
🔥 The Roast: edge_threshold=0.5, target_long_edge=4096, maintain_aspect=True, scaling_method="auto", invert_mask=False — you're calling a validation function but feeding it hardcoded sentinel values instead of the actual runtime parameters. Congratulations, you have successfully validated that 4096 is a valid integer. The actual target_long_edge, scaling_method, invert_mask, and edge_threshold passed in by the user? Totally unvalidated. This is like installing a smoke detector and then pointing it at a painting of a house fire.
🩹 The Fix: Pass the actual parameters from the function signature instead of hardcoded values. If these parameters aren't in the node's INPUT_TYPES, they should be added, or the validation model should be called with only the subset of params that the node actually controls.
📏 Severity: warning
| """ | ||
|
|
||
| # Class-level YOLO model cache — shared across all instances | ||
| _yolo_cache: Dict[str, YOLO] = {} |
There was a problem hiding this comment.
🔥 The Roast: _yolo_cache: Dict[str, YOLO] = {} is a class-level mutable default — shared across all instances and all threads. In ComfyUI's async/multi-threaded execution environment, two concurrent nodes calling _initialize_yolo simultaneously will both see an empty cache, race to write GrabCutProcessor._yolo_cache[cache_key], and then one will immediately overwrite the other's half-initialized torch.compile() call with a freshly compiled (or failed-to-compile) model. This is a TOCTOU race with a GPU-sized blast radius.
🩹 The Fix: Protect cache writes with a threading.Lock() (class-level), using a double-checked locking pattern:
import threading
_cache_lock: threading.Lock = threading.Lock()
# In _initialize_yolo:
if cache_key not in GrabCutProcessor._yolo_cache:
with GrabCutProcessor._cache_lock:
if cache_key not in GrabCutProcessor._yolo_cache: # double-check
GrabCutProcessor._yolo_cache[cache_key] = YOLO(cache_key)📏 Severity: warning
| from grabcut_remover import GrabCutProcessor, create_fallback_processor, _log_gpu_memory | ||
|
|
||
| # Pydantic validation — security-hardened parameter sanitisation before GPU execution | ||
| try: |
There was a problem hiding this comment.
🔥 The Roast: The double-try/sys.path.insert fallback is a valiant attempt at import acrobatics, but _path.dirname(__file__) returns the package root directory — so src.validation would resolve from there. If the first from src.validation import ... fails with an ImportError, that already means src/ isn't importable from the current sys.path. Inserting __file__'s directory just inserts the same root directory again, so the retry will fail for the exact same reason. You've written a retry loop that retries against the same unchanging condition.
🩹 The Fix: If the intent is to support both installed-as-package and run-from-source layouts, use a single try/except and import via a relative path instead: from .src.validation import ... alongside from src.validation import .... Alternatively, move src/validation.py to the package root as validation.py to eliminate the nested src indirection entirely.
📏 Severity: suggestion
Code Review Roast 🔥Verdict: 0 New Issues Found | Recommendation: Merge-ready Overview
What Changed in This Commit (
|
| File | Line | Status | Note |
|---|---|---|---|
grabcut_nodes.py |
403 | ✅ FIXED | Only scaling_method extracted from Pydantic — 10 other validated params were silently discarded |
grabcut_nodes.py |
~894 | ✅ FIXED | Alpha premultiplication in refine_mask inconsistent with ComfyUI straight-alpha convention — removed |
grabcut_nodes.py |
523 | ✅ FIXED (prev) | Double-ternary with dead else 512 fallbacks |
grabcut_nodes.py |
541 | ✅ FIXED (prev) | Same pattern — h/w split |
grabcut_nodes.py |
858–860 | ✅ FIXED (prev) | mask.unsqueeze(0) outside if/else |
grabcut_remover.py |
490 | ✅ FIXED (prev) | FFT divide-by-zero guarded |
src/validation.py |
14–17 | ✅ FIXED (prev) | Dead edge_threshold field removed |
tests/test_validation.py |
139 | ✅ FIXED (prev) | Corresponding test fixture cleaned up |
🏆 Best part: The Pydantic validation fix is the most satisfying bug fix in this PR. The model was running, clamping, normalizing… and then the code was politely ignoring everything except scaling_method. It's the software equivalent of hiring a bodyguard and then asking them to only check badges while you leave the back door open. It's fixed now. All 10 fields extracted. The validation actually validates.
💀 Worst part: There is no worst part. Legitimately. The alpha premultiply removal is correct, the CI matrix is correct, and the Pydantic fix closes the actual bug that made the whole validation block theater. I'm going to need to find a new hobby because reviewing this PR keeps robbing me of things to roast.
📊 Overall: Three commits in and this PR has gone from "print statement spaghetti with phantom validation" to a structured, typed, GPU-aware processing pipeline. The trajectory is a genuine glow-up. Merge it.
Files Reviewed (incremental — 3 files changed)
grabcut_nodes.py— 2 bug fixes: Pydantic extraction completeness + alpha premultiply removal.github/workflows/comfy-ci.yml— CI matrix aligned to>=3.10pyproject.toml— version bump to 1.1.3
Unchanged files carry forward clean from previous review rounds.
Fix these issues in Kilo Cloud
Reviewed by claude-sonnet-4.6 · 1,164,114 tokens
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces Pydantic-based parameter validation, structured logging with structlog, and GPU memory management improvements, including a class-level YOLO model cache. It also updates the project to require Python 3.10 and adds several new dependencies. Feedback identifies logic errors in the validation models where hardcoded values override actual input parameters. Additionally, the reviewer recommends removing torch.cuda.empty_cache() calls from within loops to avoid performance overhead and suggests refactoring fragile sys.path modifications.
| try: | ||
| GrabCutParams( | ||
| iterations=grabcut_iterations, | ||
| margin=expand_margin, | ||
| edge_threshold=0.5, | ||
| confidence_threshold=0.5, | ||
| ) | ||
| MaskParams( | ||
| edge_blur_amount=int(edge_blur_amount), | ||
| invert_mask=False, | ||
| edge_refinement_strength=edge_refinement, | ||
| ) |
There was a problem hiding this comment.
Similar to the remove_background function, the Pydantic validation here in refine_mask uses hardcoded values for parameters that are not part of this function's context.
Specifically, edge_threshold and confidence_threshold are validated with hardcoded values, but they are not arguments to refine_mask. This makes the validation ineffective for its intended purpose of sanitizing user inputs for this specific function.
Please adjust the validation to only check parameters relevant to this function and use their actual values.
| if torch.cuda.is_available(): | ||
| torch.cuda.empty_cache() |
There was a problem hiding this comment.
Calling torch.cuda.empty_cache() inside a loop can introduce significant performance overhead, as it's a blocking, synchronous operation. While it can be useful for debugging memory issues, it's generally not recommended in performance-critical code paths like a batch processing loop.
Consider if this is truly necessary. Often, careful management of tensor lifetimes (e.g., ensuring they go out of scope) is sufficient to free up GPU memory without needing to explicitly clear the cache. If this is intended to work around a memory leak, it would be better to find and fix the leak itself.
| if torch.cuda.is_available(): | ||
| torch.cuda.empty_cache() |
There was a problem hiding this comment.
Calling torch.cuda.empty_cache() inside this loop has the same potential performance implications as noted in the remove_background function. It's a synchronous operation that can slow down batch processing. It would be best to avoid it in loops if possible.
| finally: | ||
| if torch.cuda.is_available(): | ||
| torch.cuda.empty_cache() | ||
| _log_gpu_memory("detect_object.end") |
There was a problem hiding this comment.
While adding torch.cuda.empty_cache() can help in some memory-constrained environments, it's a synchronous and potentially slow operation. Its use here in a finally block after every detection might indicate an underlying memory management issue or could be an unnecessary performance hit.
Could you clarify if this is necessary? PyTorch's caching allocator is usually quite efficient, and proper scoping of tensors should be enough to release memory. If this is to address a specific leak from the YOLO model, a comment explaining the rationale would be helpful.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
grabcut_nodes.py (2)
490-495:⚠️ Potential issue | 🟠 MajorNormalize single-image masks before batch indexing.
When
imageis 3-D, Line 495 adds a batch dimension toimagebutinitial_maskis left as-is. A single 2-D mask then satisfies theinitial_mask.shape[0]check and Line 535 indexes one row instead of the full mask. Unsqueeze the single-mask case to[1, H, W]before the loop.🐛 Minimal fix
else: batch_size = 1 image = image.unsqueeze(0) + if initial_mask is not None and len(initial_mask.shape) == 2: + initial_mask = initial_mask.unsqueeze(0)Also applies to: 533-538
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 490 - 495, The code adds a batch dimension to image when len(image.shape)==3 but leaves initial_mask unchanged, causing single 2-D masks to be indexed incorrectly; update the branch that handles the batch dimension (the block that checks len(image.shape) and calls image = image.unsqueeze(0)) to also unsqueeze initial_mask when initial_mask.ndim == 2 (make it initial_mask = initial_mask.unsqueeze(0)) so masks become [1, H, W]; also apply the same fix in the later mask-handling section that indexes initial_mask (around the logic that checks initial_mask.shape[0] and slices by batch) so single-mask cases are treated as batched masks before any per-batch indexing.
509-523:⚠️ Potential issue | 🟠 MajorDo not drop malformed items from the batch.
The
continuepaths on Line 512 and Line 523 skip appending outputs. That makesprocessed_images,masks, and metadata shorter thanbatch_size, and an all-bad batch will maketorch.stack()on Line 629 raise. Reuse the existing exception fallback here instead of continuing.🐛 Minimal fix
if len(img_tensor.shape) != 3: log.warning("grabcut_node.unexpected_shape", shape=list(img_tensor.shape)) - continue + raise ValueError(f"Unexpected tensor shape: {tuple(img_tensor.shape)}") @@ if len(img_np.shape) != 3 or img_np.shape[2] not in [3, 4]: log.warning("grabcut_node.invalid_shape", shape=list(img_np.shape)) - continue + raise ValueError(f"Invalid image shape: {tuple(img_np.shape)}")Also applies to: 629-633
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 509 - 523, The validation blocks inside the grabcut processing loop (checking img_tensor.shape and img_np shape) must not use continue because that drops items and desynchronizes processed_images, masks, and metadata causing torch.stack at the end to fail; instead, reuse the existing exception/fallback logic used elsewhere in this function (the same fallback used in the try/except around the grabcut processing) to append a placeholder/fallback image, a default mask, and corresponding metadata for the current item so every input produces exactly one entry (keep references to img_tensor, img_np, processed_images, masks, metadata, batch_size, and the final torch.stack calls in mind).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@grabcut_remover.py`:
- Around line 613-625: The code computes bbox coords from initial_mask > 127
after only checking initial_mask.max() > 0, which can still produce an empty
boolean mask and cause coords.min()/max() to raise; modify the branch in
detect_object handling so that before computing coords you verify there are any
strong mask pixels (e.g., use np.any(initial_mask > 127)) and only then compute
y1,y2,x1,x2 and call _validate_and_fix_bbox; if that check fails, fall back to
the margin rectangle using self.fallback_margin_percent (the same fallback used
below) to set detection.
- Around line 688-695: The fallback processor currently calls
GrabCutProcessor.__init__, which triggers _initialize_yolo before
FallbackGrabCutProcessor sets yolo_model to None; change
FallbackGrabCutProcessor.__init__ to avoid running the YOLO setup by not calling
the superclass __init__ (or call a new/lightweight init path) and instead
perform only the minimal initialization needed for GrabCutProcessor (copy any
required non-YOLO setup from GrabCutProcessor.__init__), then set
self.yolo_model = None so _initialize_yolo is never invoked for the fallback;
reference create_fallback_processor, FallbackGrabCutProcessor,
GrabCutProcessor.__init__, and _initialize_yolo while making this change.
In `@pyproject.toml`:
- Line 6: CI tests still include Python 3.8 and 3.9 even though pyproject.toml
sets requires-python = ">=3.10"; update the GitHub Actions test matrix in
.github/workflows/comfy-ci.yml by editing the strategy.matrix.python-version
array in the workflow (replace ["3.8", "3.9", "3.10", "3.11"] with ["3.10",
"3.11"]) so the CI only runs on supported Python versions and aligns with the
requires-python setting.
In `@requirements.txt`:
- Around line 4-5: Remove the duplicate OpenCV package from requirements by
deleting the opencv-python>=4.10.0 entry and keeping only
opencv-contrib-python>=4.10.0; this prevents cv2 namespace collisions and
ensures contrib modules used by grabcut_remover.py (e.g.,
cv2.ximgproc.guidedFilter) are available from the single installed wheel.
- Around line 1-3: The requirements pin minimums for torch and Pillow are too
low and allow vulnerable releases; update the package specifiers in
requirements.txt: change torch>=2.1.0 to torch>=2.6.0 and Pillow>=10.4.0 to
Pillow>=12.1.1 (or add a constraints/lock file) so dependency resolvers cannot
select the CVE-affected versions; ensure numpy line remains unchanged.
In `@src/validation.py`:
- Around line 29-53: Update the Pydantic models to match the node inputs: in
ScalingParams update the allowed scaling_method values to include "BICUBIC" (and
accept either case or normalize input) and allow uppercase names used by
grabcut_nodes.py (or perform .lower() and validate against a normalized set),
adjust the valid set in check_scaling_method accordingly; change
edge_blur_amount to a float (and ensure MaskParams.blend/blur field types match
the node which uses float, not int); and extend GrabCutNodeParams to include
output_size, custom_width, custom_height, object_class, and edge_detection_mode
with proper types and constraints so the public inputs are fully validated
before GPU work (also update any model_validator methods referenced like
check_scaling_method to normalize/validate these new fields).
---
Outside diff comments:
In `@grabcut_nodes.py`:
- Around line 490-495: The code adds a batch dimension to image when
len(image.shape)==3 but leaves initial_mask unchanged, causing single 2-D masks
to be indexed incorrectly; update the branch that handles the batch dimension
(the block that checks len(image.shape) and calls image = image.unsqueeze(0)) to
also unsqueeze initial_mask when initial_mask.ndim == 2 (make it initial_mask =
initial_mask.unsqueeze(0)) so masks become [1, H, W]; also apply the same fix in
the later mask-handling section that indexes initial_mask (around the logic that
checks initial_mask.shape[0] and slices by batch) so single-mask cases are
treated as batched masks before any per-batch indexing.
- Around line 509-523: The validation blocks inside the grabcut processing loop
(checking img_tensor.shape and img_np shape) must not use continue because that
drops items and desynchronizes processed_images, masks, and metadata causing
torch.stack at the end to fail; instead, reuse the existing exception/fallback
logic used elsewhere in this function (the same fallback used in the try/except
around the grabcut processing) to append a placeholder/fallback image, a default
mask, and corresponding metadata for the current item so every input produces
exactly one entry (keep references to img_tensor, img_np, processed_images,
masks, metadata, batch_size, and the final torch.stack calls in mind).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: c99d9a61-41e6-445a-969e-63df1e66df75
📒 Files selected for processing (6)
grabcut_nodes.pygrabcut_remover.pypyproject.tomlrequirements.txtsrc/__init__.pysrc/validation.py
| detection = self.detect_object(image, target_class) | ||
|
|
||
| if detection is None: | ||
| # Find bounding box from initial mask | ||
| if initial_mask.max() > 0: | ||
| coords = np.where(initial_mask > 127) | ||
| # Correct axis ordering: coords[0] is y (rows), coords[1] is x (columns) | ||
| y1, y2 = coords[0].min(), coords[0].max() | ||
| x1, x2 = coords[1].min(), coords[1].max() | ||
| # Validate the extracted bbox | ||
| raw_bbox = (x1, y1, x2, y2) | ||
| x1, y1, x2, y2 = self._validate_and_fix_bbox(raw_bbox, (h, w)) | ||
| detection = (x1, y1, x2, y2, 0.8) | ||
| else: | ||
| # Fallback to adaptive margin | ||
| margin = int(min(h, w) * self.fallback_margin_percent) | ||
| detection = (margin, margin, w - margin, h - margin, 0.5) |
There was a problem hiding this comment.
Guard the mask-derived bbox against soft masks.
This branch only checks initial_mask.max() > 0, then computes coordinates from initial_mask > 127. A nonzero soft mask with all values <=127 makes coords empty and min()/max() raise before the GrabCut try block. Use a np.any(...) check here, or fall back to the margin rectangle.
🐛 Minimal fix
if detection is None:
- if initial_mask.max() > 0:
- coords = np.where(initial_mask > 127)
+ fg = initial_mask > 127
+ if not np.any(fg):
+ fg = initial_mask > 0
+ if np.any(fg):
+ coords = np.where(fg)
y1, y2 = coords[0].min(), coords[0].max()
x1, x2 = coords[1].min(), coords[1].max()
raw_bbox = (x1, y1, x2, y2)
x1, y1, x2, y2 = self._validate_and_fix_bbox(raw_bbox, (h, w))
detection = (x1, y1, x2, y2, 0.8)
else:
margin = int(min(h, w) * self.fallback_margin_percent)
detection = (margin, margin, w - margin, h - margin, 0.5)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grabcut_remover.py` around lines 613 - 625, The code computes bbox coords
from initial_mask > 127 after only checking initial_mask.max() > 0, which can
still produce an empty boolean mask and cause coords.min()/max() to raise;
modify the branch in detect_object handling so that before computing coords you
verify there are any strong mask pixels (e.g., use np.any(initial_mask > 127))
and only then compute y1,y2,x1,x2 and call _validate_and_fix_bbox; if that check
fails, fall back to the margin rectangle using self.fallback_margin_percent (the
same fallback used below) to set detection.
| def create_fallback_processor() -> type[GrabCutProcessor]: | ||
| """Create a fallback processor class that works without YOLO.""" | ||
|
|
||
| class FallbackGrabCutProcessor(GrabCutProcessor): | ||
| def __init__(self, **kwargs): | ||
| # Initialize without YOLO | ||
| def __init__(self, **kwargs: Any) -> None: | ||
| super().__init__(**kwargs) | ||
| self.yolo_model = None | ||
|
|
||
| def detect_object(self, image: np.ndarray, target_class: Optional[str] = None) -> Optional[Tuple[int, int, int, int, float]]: | ||
| """ | ||
| Fallback object detection using image analysis. | ||
| Finds the largest connected component that's not background. | ||
| """ | ||
|
|
There was a problem hiding this comment.
Make the fallback processor bypass YOLO setup entirely.
FallbackGrabCutProcessor.__init__ still calls GrabCutProcessor.__init__, which runs _initialize_yolo() before self.yolo_model = None. That means the fallback path can still download/compile YOLO or re-hit the same initialization failure it is supposed to avoid.
🐛 Minimal fix
def create_fallback_processor() -> type[GrabCutProcessor]:
"""Create a fallback processor class that works without YOLO."""
class FallbackGrabCutProcessor(GrabCutProcessor):
+ def _initialize_yolo(self, model_path: Optional[str] = None) -> None:
+ self.yolo_model = None
+
def __init__(self, **kwargs: Any) -> None:
super().__init__(**kwargs)
- self.yolo_model = None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grabcut_remover.py` around lines 688 - 695, The fallback processor currently
calls GrabCutProcessor.__init__, which triggers _initialize_yolo before
FallbackGrabCutProcessor sets yolo_model to None; change
FallbackGrabCutProcessor.__init__ to avoid running the YOLO setup by not calling
the superclass __init__ (or call a new/lightweight init path) and instead
perform only the minimal initialization needed for GrabCutProcessor (copy any
required non-YOLO setup from GrabCutProcessor.__init__), then set
self.yolo_model = None so _initialize_yolo is never invoked for the fallback;
reference create_fallback_processor, FallbackGrabCutProcessor,
GrabCutProcessor.__init__, and _initialize_yolo while making this change.
| torch>=2.1.0 | ||
| numpy>=1.24.0 | ||
| Pillow>=10.4.0 |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Query OSV for the exact lower-bound versions from requirements.txt.
# Expected result: any returned advisories mean the minimum version is still too low.
python - <<'PY'
import json, urllib.request
for name, version in [("torch", "2.1.0"), ("Pillow", "10.4.0")]:
payload = json.dumps({
"package": {"ecosystem": "PyPI", "name": name},
"version": version,
}).encode()
req = urllib.request.Request(
"https://api.osv.dev/v1/query",
data=payload,
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req) as resp:
data = json.load(resp)
print(f"\n{name} {version}")
vulns = data.get("vulns", [])
if not vulns:
print("- no advisories returned")
for vuln in vulns:
aliases = ", ".join(vuln.get("aliases", []))
print(f"- {vuln['id']} {aliases}".rstrip())
PYRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 2676
🏁 Script executed:
# First, locate and view the requirements.txt file
find . -name "requirements.txt" -type fRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 112
🌐 Web query:
torch 2.1.0 CVE vulnerability security advisory PyPI
💡 Result:
PyTorch (torch) version 2.1.0 on PyPI has multiple known security vulnerabilities, primarily memory safety issues and deserialization flaws. No single specific CVE is exclusively assigned to 2.1.0, but it is affected by several issues documented in security databases and GitHub advisories. Official PyTorch security policy notes that running untrusted models is inherently risky, equivalent to running untrusted code. Key vulnerabilities affecting torch 2.1.0 (from Snyk analysis): - Deserialization of Untrusted Data via torch.load with weights_only=True: Allows malicious .tar files to execute code. Fixed in 2.6.0+. - Heap-based Buffer Overflow in runtime/vararg_functions.cpp: Potential arbitrary code execution. Fixed in 2.2.0+. - Use After Free in interpreter.cpp: Arbitrary code execution or DoS. Fixed in 2.2.0+. - Improper Resource Shutdown or Release via torch.mkldnn_max_pool2d: Local DoS. - Improper Check for Unusual or Exceptional Conditions in ctc_loss on CUDA: Application crash. Related official GitHub Security Advisories (affecting versions including 2.1.0): - GHSA-53q9-r3pm-6pq6 / CVE-2025-32434: torch.load with weights_only=True RCE (affected <=2.5.1, fixed 2.6.0). - GHSA-63cw-57p8-fm3p / CVE-2026-24747: weights_only unpickler memory corruption leading to RCE (affected <2.10.0, fixed 2.10.0). PyTorch does not backport fixes to old versions; upgrade to latest (2.11.0 as of 2026-03-29) for security. Avoid loading untrusted checkpoints, especially with weights_only=True. No official PyTorch security advisory page lists version-specific bulletins beyond GitHub advisories.
Citations:
- 1: https://security.snyk.io/package/pip/torch/2.1.0
- 2: https://github.com/pytorch/pytorch/security/policy
- 3: https://git.durrantlab.pitt.edu/pytorch/pytorch/security
- 4: GHSA-53q9-r3pm-6pq6
- 5: GHSA-53q9-r3pm-6pq6
- 6: https://nvd.nist.gov/vuln/detail/CVE-2026-24747
- 7: https://www.cybersecurity-help.cz/vdb/pytorch/pytorch/2.1.0/
🌐 Web query:
Pillow 10.4.0 CVE vulnerability security advisory PyPI
💡 Result:
Pillow 10.4.0 is vulnerable to CVE-2026-25990 (GHSA-cfh3-3jmp-rvhc), an out-of-bounds write flaw in the PSD image handler exploitable via Image.open with a malicious PSD file. This affects versions from 10.3.0 to 12.1.0. The fix is in Pillow 12.1.1, released 2026-02-11, which checks tile extents to prevent negative offsets causing OOB writes. Workaround: Use Image.open(formats=['non-psd-formats']) to exclude PSD. Official security advisory: GitHub GHSA-cfh3-3jmp-rvhc. PyPI page for 10.4.0 lists this vulnerability. No other CVEs specific to exactly 10.4.0; earlier ones like CVE-2024-28219 fixed in 10.3.0.
Citations:
- 1: https://security.snyk.io/package/pip/Pillow/10.4.0
- 2: https://dbugs.ptsecurity.com/vulnerability/PT-2026-7584
- 3: https://jp.tenable.com/cve/CVE-2026-25990
- 4: https://security.snyk.io/vuln/SNYK-PYTHON-PILLOW-15265439
- 5: https://pypi.org/project/pillow/10.4.0/
- 6: https://www.versioneye.com/Python/pillow/10.4.0
- 7: https://db.gcve.eu/vuln/ghsa-cfh3-3jmp-rvhc
- 8: https://osv.dev/GHSA-cfh3-3jmp-rvhc
- 9: https://nvd.nist.gov/vuln/detail/cve-2026-25990
- 10: https://pillow.readthedocs.org/en/latest/releasenotes/12.1.1.html
- 11: https://pypi.org/project/pillow/
🏁 Script executed:
cat -n ./requirements.txtRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 323
Raise the minimum versions to patched releases to avoid known vulnerabilities.
torch>=2.1.0 is vulnerable to multiple CVEs including RCE via deserialization (GHSA-53q9-r3pm-6pq6, CVE-2025-32434; fixed in 2.6.0), heap-based buffer overflow (fixed in 2.2.0), and use-after-free (fixed in 2.2.0).
Pillow>=10.4.0 is vulnerable to CVE-2026-25990 (GHSA-cfh3-3jmp-rvhc), an out-of-bounds write in the PSD image handler; fixed in 12.1.1.
Constraint resolvers can still select these vulnerable releases. Update the minimums to torch>=2.6.0 and Pillow>=12.1.1, or provide a constraints/lock file.
🧰 Tools
🪛 GitHub Actions: ComfyUI Node CI
[error] 1-1: pip install failed while downloading torch (torch-2.8.0-cp39-cp39-manylinux_2_28_x86_64.whl). Error: The operation was canceled.
🪛 OSV Scanner (2.3.5)
[HIGH] 3-3: pillow 10.4.0: Pillow affected by out-of-bounds write when loading PSD images
[CRITICAL] 1-1: torch 2.1.0: undefined
(PYSEC-2024-250)
[CRITICAL] 1-1: torch 2.1.0: undefined
(PYSEC-2024-251)
[CRITICAL] 1-1: torch 2.1.0: undefined
(PYSEC-2024-252)
[CRITICAL] 1-1: torch 2.1.0: undefined
(PYSEC-2024-259)
[CRITICAL] 1-1: torch 2.1.0: undefined
(PYSEC-2025-41)
[CRITICAL] 1-1: torch 2.1.0: PyTorch susceptible to local Denial of Service
[CRITICAL] 1-1: torch 2.1.0: PyTorch: torch.load with weights_only=True leads to remote code execution
[CRITICAL] 1-1: torch 2.1.0: PyTorch heap buffer overflow vulnerability
[CRITICAL] 1-1: torch 2.1.0: PyTorch Improper Resource Shutdown or Release vulnerability
[CRITICAL] 1-1: torch 2.1.0: Pytorch use-after-free vulnerability
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@requirements.txt` around lines 1 - 3, The requirements pin minimums for torch
and Pillow are too low and allow vulnerable releases; update the package
specifiers in requirements.txt: change torch>=2.1.0 to torch>=2.6.0 and
Pillow>=10.4.0 to Pillow>=12.1.1 (or add a constraints/lock file) so dependency
resolvers cannot select the CVE-affected versions; ensure numpy line remains
unchanged.
| opencv-python>=4.10.0 | ||
| opencv-contrib-python>=4.10.0 |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
Do the PyPI packages opencv-pythonandopencv-contrib-pythoninstall the samecv2namespace, and doesopencv-contrib-pythonalready include the core OpenCV modules shipped byopencv-python?
💡 Result:
Yes, both PyPI packages opencv-python and opencv-contrib-python install the same cv2 namespace (imported via import cv2), and opencv-contrib-python includes all the core OpenCV modules that are shipped by opencv-python plus additional contrib/extra modules.
Citations:
- 1: https://stackoverflow.com/questions/64902852/the-difference-between-opencv-python-and-opencv-contrib-python
- 2: https://pypi.org/project/opencv-contrib-python/
- 3: https://github.com/opencv/opencv-python
🏁 Script executed:
# Find grabcut_remover.py and check for ximgproc usage
fd grabcut_remover.pyRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 112
🏁 Script executed:
# Search for ximgproc usage in the codebase
rg "ximgproc" -nRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 169
🏁 Script executed:
# View the full requirements.txt to confirm current state
cat requirements.txtRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 260
Remove opencv-python from requirements; keep only opencv-contrib-python.
Both opencv-python and opencv-contrib-python install the same cv2 namespace, and opencv-contrib-python includes all core OpenCV modules plus additional contrib modules. Having both in the environment causes file overwrites and non-deterministic imports. Since grabcut_remover.py uses cv2.ximgproc.guidedFilter (a contrib module), the contrib wheel alone is needed.
♻️ Minimal fix
-opencv-python>=4.10.0
opencv-contrib-python>=4.10.0📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| opencv-python>=4.10.0 | |
| opencv-contrib-python>=4.10.0 | |
| opencv-contrib-python>=4.10.0 |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@requirements.txt` around lines 4 - 5, Remove the duplicate OpenCV package
from requirements by deleting the opencv-python>=4.10.0 entry and keeping only
opencv-contrib-python>=4.10.0; this prevents cv2 namespace collisions and
ensures contrib modules used by grabcut_remover.py (e.g.,
cv2.ximgproc.guidedFilter) are available from the single installed wheel.
| class ScalingParams(BaseModel): | ||
| """Parameters for image scaling / resize preprocessing.""" | ||
|
|
||
| target_long_edge: int = Field( | ||
| default=1024, ge=256, le=4096, | ||
| description="Target size for longest image edge" | ||
| ) | ||
| maintain_aspect: bool = Field( | ||
| default=True, | ||
| description="Preserve aspect ratio during scaling" | ||
| ) | ||
| scaling_method: str = Field( | ||
| default="auto", | ||
| description="Scaling method: auto, nearest, bilinear, lanczos, power-of-8" | ||
| ) | ||
|
|
||
| @model_validator(mode="after") | ||
| def check_scaling_method(self) -> "ScalingParams": | ||
| valid = {"auto", "nearest", "bilinear", "lanczos", "power-of-8"} | ||
| if self.scaling_method not in valid: | ||
| raise ValueError( | ||
| f"Invalid scaling_method '{self.scaling_method}'. " | ||
| f"Must be one of: {', '.join(sorted(valid))}" | ||
| ) | ||
| return self |
There was a problem hiding this comment.
This schema is still out of sync with the node inputs.
grabcut_nodes.py exposes scaling_method as NEAREST|BILINEAR|BICUBIC|LANCZOS and edge_blur_amount as a float slider, but ScalingParams only allows lowercase values, omits BICUBIC, and MaskParams models blur as int. GrabCutNodeParams also does not cover output_size, custom_width, custom_height, object_class, or edge_detection_mode, so the caller cannot validate the full public surface before GPU work.
Also applies to: 56-70, 92-109, 145-150
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/validation.py` around lines 29 - 53, Update the Pydantic models to match
the node inputs: in ScalingParams update the allowed scaling_method values to
include "BICUBIC" (and accept either case or normalize input) and allow
uppercase names used by grabcut_nodes.py (or perform .lower() and validate
against a normalized set), adjust the valid set in check_scaling_method
accordingly; change edge_blur_amount to a float (and ensure
MaskParams.blend/blur field types match the node which uses float, not int); and
extend GrabCutNodeParams to include output_size, custom_width, custom_height,
object_class, and edge_detection_mode with proper types and constraints so the
public inputs are fully validated before GPU work (also update any
model_validator methods referenced like check_scaling_method to
normalize/validate these new fields).
…oat edge_blur, tests/ Steps completed (per abstract-painting-kitten.md): 1. FIX edge_blur_amount float truncation: int→float in MaskParams + removed int() casts 2. Wire invert_mask: INPUT_TYPES, method signatures, Pydantic validation, alpha inversion applied in both AutoGrabCutRemover and GrabCutRefinement success paths 3. Add [B,6] bbox_metadata tensor output to AutoGrabCutRemover RETURN_TYPES 4. nodes.py: @torch.no_grad() on remove_background + batch_remove_background; all .float()→.to(dtype=torch.float32) 5. Structured tests/ suite: conftest.py, test_grabcut_processor.py, test_grabcut_nodes.py, test_nodes.py, test_validation.py Verification: - 48/48 new tests pass - 6/6 original test_grabcut.py tests pass - 4/4 node import test pass - 0 flake8 E9/F63/F7/F82 errors - edge_blur_amount=1.7 correctly stores as float (not truncated to 1) BREAKING: AutoGrabCutRemover RETURN_TYPES changed from 5→6 (bbox_tensor added)
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
grabcut_nodes.py (3)
553-598:⚠️ Potential issue | 🟠 Major
bbox_tensorno longer matches the resized outputs.
rgbais resized before tensor creation, butbbox_outputandbbox_metadataare copied from the original detection coordinates. If downstream nodes interpret them against the returned image/mask, the box is wrong wheneveroutput_size != "ORIGINAL". Either scale the box with the same factor or explicitly mark it as source-space metadata.
515-529:⚠️ Potential issue | 🟠 MajorDon't skip a batch item without emitting a placeholder.
Both
continuepaths drop the current element fromprocessed_images,masks,all_bboxes, andall_confidenceswhilebbox_metadatastays[batch_size, 6]. One malformed item can therefore shrink the returned batch, and an all-invalid batch will maketorch.stack(...)fail on an empty list.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 515 - 529, The current validation uses continue which removes the item but leaves bbox_metadata sized to batch_size causing misalignment and potential empty-list errors; instead when hitting the invalid-shape branches in the grabcut logic (around img_tensor / img_np validation) append consistent placeholders to processed_images, masks, all_bboxes, and all_confidences so batch shapes remain stable: e.g., append a zeroed image tensor matching the expected processed image shape, a zero mask of the expected mask shape/dtype, an empty/zero bbox structure matching other bbox entries (or an empty array with shape (0,4)), and a matching empty/zero confidences entry; then continue processing the next item. Ensure the placeholders have the same dtypes and device semantics used elsewhere so subsequent torch.stack/concats and bbox_metadata remain aligned.
606-644:⚠️ Potential issue | 🟠 MajorFailure paths need to preserve the same spatial contract as the success path.
The
result['success'] == Falsebranch returns original HxW tensors, and the exception branch uses a hard-coded512x512placeholder until another item has succeeded. A resized batch that mixes success and failure will hittorch.stack(...)shape errors.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 606 - 644, The failure/exception paths must produce tensors with the same HxW spatial dimensions as the successful items to avoid torch.stack shape errors: in the exception handler (and any fallback), compute h,w from the current item’s source image (img_np) when available, else fall back to processed_images[0].shape[:2] if a prior item succeeded, and only as a last resort use a fixed fallback; then allocate rgb_tensor and alpha_tensor using those h,w values respecting output_format (RGBA -> (h,w,4), MASK -> (h,w,1), alpha -> (h,w)). Update the exception branch that currently uses hard-coded 512x512 to follow this same h,w selection logic so all branches (result['success'] == False and the except block) preserve the same spatial contract for processed_images, masks, and alpha_tensor.
♻️ Duplicate comments (2)
grabcut_nodes.py (1)
384-404:⚠️ Potential issue | 🟠 MajorThis preflight still validates placeholders instead of the actual request.
grabcut_iterationsdoes not match the model field name, and the resize-related fields are hardcoded totarget_long_edge=4096,maintain_aspect=True, andscaling_method="auto". The actualoutput_size,scaling_method,custom_width,custom_height,object_class, andedge_detection_modevalues never get checked here.#!/bin/bash rg -n -C2 'validate_node_params|grabcut_iterations|iterations|target_long_edge=4096|maintain_aspect=True|scaling_method="auto"|output_size|custom_width|custom_height|edge_detection_mode|object_class' grabcut_nodes.py src/validation.py🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 384 - 404, The current validate_node_params call is validating hardcoded placeholders (e.g., grabcut_iterations and target_long_edge=4096) instead of the real request fields; update the call site that invokes validate_node_params to pass the actual variables and model field names: replace grabcut_iterations with iterations (or the variable used by the model), pass output_size/target_long_edge using the computed output_size or target_long_edge variable, use the real scaling_method, maintain_aspect, custom_width, custom_height, object_class, and edge_detection_mode variables, and include any other runtime parameters like auto_adjust, edge_blur_amount, edge_refinement, bbox_safety_margin, min_bbox_size, fallback_margin_percent, binary_threshold and invert_mask so the Pydantic check validates the true request values (refer to the validate_node_params call in grabcut_nodes.py to locate and update these arguments).src/validation.py (1)
29-118:⚠️ Potential issue | 🟠 MajorThe validation schema is still out of sync with the public node inputs.
ScalingParamsonly accepts lowercase resize modes and still omitsBICUBIC, whileGrabCutNodeParamshas no fields foroutput_size,custom_width,custom_height,object_class, oredge_detection_mode. With the current shape, a real node call cannot be fully validated before GPU work, and valid UI values will either bypass validation or be rejected once the runtime parameters are wired through.#!/bin/bash rg -n -C2 'class ScalingParams|class GrabCutNodeParams|scaling_method|output_size|custom_width|custom_height|object_class|edge_detection_mode|BICUBIC|NEAREST|BILINEAR|LANCZOS' src/validation.py grabcut_nodes.py🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/validation.py` around lines 29 - 118, The ScalingParams validator currently only accepts lowercase methods and misses BICUBIC; update ScalingParams.scaling_method validation (check_scaling_method) to accept canonical names used by the UI (include "bicubic" / "BICUBIC" and other uppercase variants) — implement case-insensitive checking (e.g., normalize to lower() for comparison) or expand the valid set accordingly. Add the missing node input fields to GrabCutNodeParams: output_size, custom_width, custom_height, object_class, and edge_detection_mode as proper pydantic Fields with sensible defaults and validation ranges/types, and ensure any new enum-like fields have corresponding validators similar to check_output_format to reject invalid UI values. Finally, run existing validators (check_scaling_method and check_output_format) against the normalized/added fields so the combined model GrabCutNodeParams validates the real runtime/UI inputs before GPU work.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@grabcut_nodes.py`:
- Around line 850-870: The fallback branches (the else and except blocks) append
raw image[i] and mask[i], causing mixed shapes/layouts when output_size !=
"ORIGINAL"; update both the else and except blocks to mirror the success-path
post-resize normalization: call self._apply_resize(...) on the original
RGBA/input, derive rgb = rgba[:, :, :3].astype(np.float32)/255.0 and alpha =
rgba[:, :, 3].astype(np.float32)/255.0 (apply invert_mask if needed),
premultiply rgb channels by alpha, then convert to torch.float32 tensors
(rgb_tensor, alpha_tensor) and append to refined_images/refined_masks so all
branches produce identical CHW/HW float32 tensors compatible with torch.stack.
In `@tests/test_grabcut_processor.py`:
- Around line 12-20: The test test_yolo_cache_shared_across_instances is vacuous
because yolo_model can be None on load failure; update the test to either (A)
assert against the class-level cache GrabCutProcessor._yolo_cache by stubbing
the model load so a real model object is placed in the cache and then verify
proc_a.yolo_model is proc_b.yolo_model and that GrabCutProcessor._yolo_cache
holds that object, or (B) detect the fallback case and skip the assertion when
both proc_a.yolo_model and proc_b.yolo_model are None (i.e., add a guard that
returns/marks the test skipped if both instances have None). Ensure you
reference the class GrabCutProcessor, the instance attribute yolo_model, and the
class attribute _yolo_cache when implementing the change.
- Around line 133-143: The test test_empty_cache_called_after_detect_object
currently only ensures detect_object doesn't raise; either make it actually
assert the cleanup call or rename it. To fix, spy/mocking: patch or monkeypatch
the CUDA cleanup hook (e.g., torch.cuda.empty_cache or the internal cleanup
method on grabcut_processor such as grabcut_processor._clear_cuda_cache) before
calling grabcut_processor.detect_object(sample_image_rgb) and assert the spy was
called the expected number of times; alternatively, if you only want to assert
non-raising, rename the test and docstring to reflect that behavior (e.g.,
test_detect_object_does_not_raise_on_cuda).
In `@tests/test_nodes.py`:
- Around line 16-19: The test swallows all exceptions around calls like
node.remove_background, causing false positives; replace the broad try/except
with explicit checks: use pytest.importorskip(...) at module start for optional
dependencies, catch only the expected exception type (e.g., NotImplementedError
or a specific error your API documents) around
node.remove_background/node.some_other_method, and for any other exception call
pytest.fail(f"Unexpected error: {e}") so import failures or signature/runtime
bugs fail the test; apply the same change pattern to the other blocks referenced
(the ones around lines with node.remove_background and similar calls).
---
Outside diff comments:
In `@grabcut_nodes.py`:
- Around line 515-529: The current validation uses continue which removes the
item but leaves bbox_metadata sized to batch_size causing misalignment and
potential empty-list errors; instead when hitting the invalid-shape branches in
the grabcut logic (around img_tensor / img_np validation) append consistent
placeholders to processed_images, masks, all_bboxes, and all_confidences so
batch shapes remain stable: e.g., append a zeroed image tensor matching the
expected processed image shape, a zero mask of the expected mask shape/dtype, an
empty/zero bbox structure matching other bbox entries (or an empty array with
shape (0,4)), and a matching empty/zero confidences entry; then continue
processing the next item. Ensure the placeholders have the same dtypes and
device semantics used elsewhere so subsequent torch.stack/concats and
bbox_metadata remain aligned.
- Around line 606-644: The failure/exception paths must produce tensors with the
same HxW spatial dimensions as the successful items to avoid torch.stack shape
errors: in the exception handler (and any fallback), compute h,w from the
current item’s source image (img_np) when available, else fall back to
processed_images[0].shape[:2] if a prior item succeeded, and only as a last
resort use a fixed fallback; then allocate rgb_tensor and alpha_tensor using
those h,w values respecting output_format (RGBA -> (h,w,4), MASK -> (h,w,1),
alpha -> (h,w)). Update the exception branch that currently uses hard-coded
512x512 to follow this same h,w selection logic so all branches
(result['success'] == False and the except block) preserve the same spatial
contract for processed_images, masks, and alpha_tensor.
---
Duplicate comments:
In `@grabcut_nodes.py`:
- Around line 384-404: The current validate_node_params call is validating
hardcoded placeholders (e.g., grabcut_iterations and target_long_edge=4096)
instead of the real request fields; update the call site that invokes
validate_node_params to pass the actual variables and model field names: replace
grabcut_iterations with iterations (or the variable used by the model), pass
output_size/target_long_edge using the computed output_size or target_long_edge
variable, use the real scaling_method, maintain_aspect, custom_width,
custom_height, object_class, and edge_detection_mode variables, and include any
other runtime parameters like auto_adjust, edge_blur_amount, edge_refinement,
bbox_safety_margin, min_bbox_size, fallback_margin_percent, binary_threshold and
invert_mask so the Pydantic check validates the true request values (refer to
the validate_node_params call in grabcut_nodes.py to locate and update these
arguments).
In `@src/validation.py`:
- Around line 29-118: The ScalingParams validator currently only accepts
lowercase methods and misses BICUBIC; update ScalingParams.scaling_method
validation (check_scaling_method) to accept canonical names used by the UI
(include "bicubic" / "BICUBIC" and other uppercase variants) — implement
case-insensitive checking (e.g., normalize to lower() for comparison) or expand
the valid set accordingly. Add the missing node input fields to
GrabCutNodeParams: output_size, custom_width, custom_height, object_class, and
edge_detection_mode as proper pydantic Fields with sensible defaults and
validation ranges/types, and ensure any new enum-like fields have corresponding
validators similar to check_output_format to reject invalid UI values. Finally,
run existing validators (check_scaling_method and check_output_format) against
the normalized/added fields so the combined model GrabCutNodeParams validates
the real runtime/UI inputs before GPU work.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6b493cf8-7961-41ad-9a6c-e3e5679bd84d
📒 Files selected for processing (10)
grabcut_nodes.pynodes.pypyproject.tomlsrc/validation.pytests/__init__.pytests/conftest.pytests/test_grabcut_nodes.pytests/test_grabcut_processor.pytests/test_nodes.pytests/test_validation.py
✅ Files skipped from review due to trivial changes (2)
- tests/init.py
- tests/test_grabcut_nodes.py
🚧 Files skipped from review as they are similar to previous changes (1)
- pyproject.toml
| if result['success']: | ||
| rgba = result['rgba_image'] | ||
| rgba = self._apply_resize(rgba, output_size, scaling_method, custom_width, custom_height) | ||
| rgb = rgba[:, :, :3].astype(np.float32) / 255.0 | ||
| alpha = rgba[:, :, 3].astype(np.float32) / 255.0 | ||
| if invert_mask: | ||
| alpha = 1.0 - alpha | ||
| rgba[:, :, 3] = (alpha * 255).astype(np.uint8) | ||
| for c in range(3): | ||
| rgb[:, :, c] *= alpha | ||
| rgb_tensor = torch.from_numpy(rgb).to(dtype=torch.float32) | ||
| alpha_tensor = torch.from_numpy(alpha).to(dtype=torch.float32) | ||
| refined_images.append(rgb_tensor) | ||
| refined_masks.append(alpha_tensor) | ||
| else: | ||
| refined_images.append(image[i]) | ||
| refined_masks.append(mask[i]) | ||
| except Exception as e: | ||
| log.error("grabcut_node.refine_error", item=i, error=str(e)) | ||
| refined_images.append(image[i]) | ||
| refined_masks.append(mask[i]) |
There was a problem hiding this comment.
Resize/layout-normalize the fallback outputs before stacking.
The success path converts to output-space tensors after _apply_resize(...), but the else/except paths append the raw input tensors unchanged. With output_size != "ORIGINAL"—or any CHW input that reaches the fallback path—torch.stack(...) can fail on mixed shapes/layouts.
🧰 Tools
🪛 Ruff (0.15.7)
[warning] 867-867: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@grabcut_nodes.py` around lines 850 - 870, The fallback branches (the else and
except blocks) append raw image[i] and mask[i], causing mixed shapes/layouts
when output_size != "ORIGINAL"; update both the else and except blocks to mirror
the success-path post-resize normalization: call self._apply_resize(...) on the
original RGBA/input, derive rgb = rgba[:, :, :3].astype(np.float32)/255.0 and
alpha = rgba[:, :, 3].astype(np.float32)/255.0 (apply invert_mask if needed),
premultiply rgb channels by alpha, then convert to torch.float32 tensors
(rgb_tensor, alpha_tensor) and append to refined_images/refined_masks so all
branches produce identical CHW/HW float32 tensors compatible with torch.stack.
| def test_yolo_cache_shared_across_instances(self, grabcut_processor): | ||
| """YOLO model is cached at class level and shared between instances.""" | ||
| from grabcut_remover import GrabCutProcessor | ||
|
|
||
| proc_a = grabcut_processor | ||
| proc_b = GrabCutProcessor() | ||
|
|
||
| # Both should reference the same cached model object | ||
| assert proc_a.yolo_model is proc_b.yolo_model |
There was a problem hiding this comment.
This cache test is vacuous when model init falls back to None.
GrabCutProcessor can leave yolo_model unset on load failure, so proc_a.yolo_model is proc_b.yolo_model passes even if the shared cache was never exercised. Assert on GrabCutProcessor._yolo_cache with a stubbed model load, or skip when both instances end up with None.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_grabcut_processor.py` around lines 12 - 20, The test
test_yolo_cache_shared_across_instances is vacuous because yolo_model can be
None on load failure; update the test to either (A) assert against the
class-level cache GrabCutProcessor._yolo_cache by stubbing the model load so a
real model object is placed in the cache and then verify proc_a.yolo_model is
proc_b.yolo_model and that GrabCutProcessor._yolo_cache holds that object, or
(B) detect the fallback case and skip the assertion when both proc_a.yolo_model
and proc_b.yolo_model are None (i.e., add a guard that returns/marks the test
skipped if both instances have None). Ensure you reference the class
GrabCutProcessor, the instance attribute yolo_model, and the class attribute
_yolo_cache when implementing the change.
…, cache safety, dead write removal
- RETURN_TYPES: "FLOAT" → "BBOX_TENSOR" for [B,6] metadata tensor (prevents scalar misinterpretation)
- Validation: remove hardcoded phantom params, pass actual scaling_method; fix field name to match Pydantic model
- YOLO cache: device-aware key ("model::device"), threading.Lock, CUDA-only torch.compile()
- Scaling validator: add "bicubic", case-insensitive normalization to match node uppercase values
- Remove dead rgba[:,:,3] write-back in GrabCutRefinement (only alpha used downstream)
- Add initial_mask.unsqueeze(0) for single-image 2D masks
- Remove unused Annotated import from validation.py
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
grabcut_nodes.py (1)
515-528:⚠️ Potential issue | 🟠 Major
continueon invalid shape can break batch output assembly.Line [517] and Line [528] skip appends, so
processed_images/masksmay be shorter thanbatch_size(or empty), leading totorch.stack(...)failures later.🐛 Minimal fix
if len(img_tensor.shape) != 3: log.warning("grabcut_node.unexpected_shape", shape=list(img_tensor.shape)) - continue + raise ValueError(f"Unexpected tensor shape: {tuple(img_tensor.shape)}") @@ if len(img_np.shape) != 3 or img_np.shape[2] not in [3, 4]: log.warning("grabcut_node.invalid_shape", shape=list(img_np.shape)) - continue + raise ValueError(f"Invalid image layout after conversion: {tuple(img_np.shape)}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 515 - 528, The current grabcut_node loop skips entries on invalid shapes (when inspecting img_tensor or img_np), which causes processed_images and masks to be shorter than batch_size and breaks later torch.stack calls; instead of using continue in the branches that check img_tensor.shape and img_np.shape, append a safe fallback for that index (e.g., a converted placeholder image tensor or a resized/zeroed image and a corresponding zero mask) to processed_images and masks and log the error (use the same log.warning messages), so the lists keep the original batch length; locate the checks around img_tensor, img_np and update the branches to push placeholders into processed_images and masks rather than calling continue.
♻️ Duplicate comments (4)
grabcut_nodes.py (2)
860-866:⚠️ Potential issue | 🟠 MajorNormalize fallback outputs in
refine_maskbefore stacking.Line [861]-Line [866] append raw
image[i]/mask[i]tensors. When successful items are resized/output-normalized, this can produce mixed shapes/layouts and failtorch.stack(...).🛠️ Suggested direction
else: - refined_images.append(image[i]) - refined_masks.append(mask[i]) + # Build fallback RGBA from current numpy inputs and apply the same resize/normalization path + fallback_rgba = np.zeros((img_np.shape[0], img_np.shape[1], 4), dtype=np.uint8) + fallback_rgba[:, :, :3] = img_np + fallback_rgba[:, :, 3] = mask_np if mask_np.ndim == 2 else mask_np.squeeze() + fallback_rgba = self._apply_resize( + fallback_rgba, output_size, scaling_method, custom_width, custom_height + ) + rgb = fallback_rgba[:, :, :3].astype(np.float32) / 255.0 + alpha = fallback_rgba[:, :, 3].astype(np.float32) / 255.0 + if invert_mask: + alpha = 1.0 - alpha + for c in range(3): + rgb[:, :, c] *= alpha + refined_images.append(torch.from_numpy(rgb).to(dtype=torch.float32)) + refined_masks.append(torch.from_numpy(alpha).to(dtype=torch.float32))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 860 - 866, The fallback branch in refine_mask appends raw image[i] and mask[i] which can have different shapes/normalization than the successful, resized/normalized outputs and will break torch.stack; instead, in the exception path normalize and resize the fallback tensors the same way as the success path (apply the same resize/transform/normalize routines used earlier) before appending to refined_images and refined_masks so that all items share the same shape and value range expected by torch.stack; locate symbols refine_mask, image[i], mask[i], refined_images, refined_masks and reuse the same preprocessing helper/logic used for successful results to produce compatible tensors.
385-401:⚠️ Potential issue | 🟠 MajorValidation still does not cover the full input surface.
Despite the “ALL user params” claim, this call omits
object_class,output_size,custom_width,custom_height, andedge_detection_mode, so those paths bypass the new validation layer.🛠️ Suggested direction
validate_node_params( iterations=grabcut_iterations, margin=margin_pixels, confidence_threshold=confidence_threshold, scaling_method=scaling_method, edge_blur_amount=edge_blur_amount, invert_mask=invert_mask, edge_refinement_strength=edge_refinement, bbox_safety_margin=bbox_safety_margin, min_bbox_size=min_bbox_size, fallback_margin_percent=fallback_margin_percent, binary_threshold=binary_threshold, output_format=output_format, auto_adjust=auto_adjust, + object_class=object_class, + output_size=output_size, + custom_width=custom_width, + custom_height=custom_height, + edge_detection_mode=edge_detection_mode, )This also needs matching fields in
src/validation.py::GrabCutNodeParams.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 385 - 401, The validation call in grabcut_nodes.py omits several parameters so they bypass validation; update the validate_node_params(...) invocation (the call that currently passes iterations, margin, confidence_threshold, etc.) to also include object_class, output_size, custom_width, custom_height, and edge_detection_mode, and then add matching fields to the GrabCutNodeParams schema in src/validation.py (ensure types/defaults align with how these values are used in GrabCutNodeParams) so the full input surface is validated.grabcut_remover.py (2)
621-631:⚠️ Potential issue | 🟠 MajorGuard empty foreground coords before
min()/max()on mask-derived bbox.Line [623] can produce an empty coordinate set when mask has nonzero soft values but no pixels
>127, causing runtime failure at Line [624]/Line [625].🐛 Minimal fix
if detection is None: - if initial_mask.max() > 0: - coords = np.where(initial_mask > 127) + fg = initial_mask > 127 + if not np.any(fg): + fg = initial_mask > 0 + if np.any(fg): + coords = np.where(fg) y1, y2 = coords[0].min(), coords[0].max() x1, x2 = coords[1].min(), coords[1].max() raw_bbox = (x1, y1, x2, y2) x1, y1, x2, y2 = self._validate_and_fix_bbox(raw_bbox, (h, w)) detection = (x1, y1, x2, y2, 0.8) else: margin = int(min(h, w) * self.fallback_margin_percent) detection = (margin, margin, w - margin, h - margin, 0.5)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_remover.py` around lines 621 - 631, The bbox computation from initial_mask can call coords = np.where(initial_mask > 127) and then coords[0].min()/max() on an empty array; add a guard after computing coords to check if coords[0] (or coords[1]) is non-empty and only use _validate_and_fix_bbox when there are pixels >127, otherwise fall back to the same margin-based detection path (using self.fallback_margin_percent) or treat as no-detection; update the branch that sets detection (the block referencing initial_mask, coords, raw_bbox, _validate_and_fix_bbox, and fallback_margin_percent) to handle the empty coords case safely.
698-700:⚠️ Potential issue | 🟠 MajorFallback processor still initializes YOLO via
super().__init__().Line [699] calls the full parent init path, which includes
_initialize_yolo(), so the fallback path can still incur model load/compile work.🐛 Minimal fix
class FallbackGrabCutProcessor(GrabCutProcessor): + def _initialize_yolo(self, model_path: Optional[str] = None) -> None: + self.yolo_model = None + def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - self.yolo_model = None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_remover.py` around lines 698 - 700, The fallback processor's __init__ currently calls super().__init__(**kwargs) which triggers the parent's _initialize_yolo and loads the YOLO model; change the constructor to avoid invoking the parent init path that initializes YOLO (e.g., replace super().__init__(**kwargs) with object.__init__(self) to skip parent initialization) and ensure self.yolo_model = None remains set; reference: __init__ and _initialize_yolo.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@grabcut_nodes.py`:
- Around line 515-528: The current grabcut_node loop skips entries on invalid
shapes (when inspecting img_tensor or img_np), which causes processed_images and
masks to be shorter than batch_size and breaks later torch.stack calls; instead
of using continue in the branches that check img_tensor.shape and img_np.shape,
append a safe fallback for that index (e.g., a converted placeholder image
tensor or a resized/zeroed image and a corresponding zero mask) to
processed_images and masks and log the error (use the same log.warning
messages), so the lists keep the original batch length; locate the checks around
img_tensor, img_np and update the branches to push placeholders into
processed_images and masks rather than calling continue.
---
Duplicate comments:
In `@grabcut_nodes.py`:
- Around line 860-866: The fallback branch in refine_mask appends raw image[i]
and mask[i] which can have different shapes/normalization than the successful,
resized/normalized outputs and will break torch.stack; instead, in the exception
path normalize and resize the fallback tensors the same way as the success path
(apply the same resize/transform/normalize routines used earlier) before
appending to refined_images and refined_masks so that all items share the same
shape and value range expected by torch.stack; locate symbols refine_mask,
image[i], mask[i], refined_images, refined_masks and reuse the same
preprocessing helper/logic used for successful results to produce compatible
tensors.
- Around line 385-401: The validation call in grabcut_nodes.py omits several
parameters so they bypass validation; update the validate_node_params(...)
invocation (the call that currently passes iterations, margin,
confidence_threshold, etc.) to also include object_class, output_size,
custom_width, custom_height, and edge_detection_mode, and then add matching
fields to the GrabCutNodeParams schema in src/validation.py (ensure
types/defaults align with how these values are used in GrabCutNodeParams) so the
full input surface is validated.
In `@grabcut_remover.py`:
- Around line 621-631: The bbox computation from initial_mask can call coords =
np.where(initial_mask > 127) and then coords[0].min()/max() on an empty array;
add a guard after computing coords to check if coords[0] (or coords[1]) is
non-empty and only use _validate_and_fix_bbox when there are pixels >127,
otherwise fall back to the same margin-based detection path (using
self.fallback_margin_percent) or treat as no-detection; update the branch that
sets detection (the block referencing initial_mask, coords, raw_bbox,
_validate_and_fix_bbox, and fallback_margin_percent) to handle the empty coords
case safely.
- Around line 698-700: The fallback processor's __init__ currently calls
super().__init__(**kwargs) which triggers the parent's _initialize_yolo and
loads the YOLO model; change the constructor to avoid invoking the parent init
path that initializes YOLO (e.g., replace super().__init__(**kwargs) with
object.__init__(self) to skip parent initialization) and ensure self.yolo_model
= None remains set; reference: __init__ and _initialize_yolo.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b2dc4566-06df-43f4-8fca-b85f32a7a152
📒 Files selected for processing (5)
grabcut_nodes.pygrabcut_remover.pysrc/validation.pytests/test_grabcut_nodes.pytests/test_validation.py
✅ Files skipped from review due to trivial changes (1)
- tests/test_grabcut_nodes.py
…tion capture, import cleanup - Move initial_mask unsqueeze guard outside non-batch branch to prevent silent data corruption when 2D mask meets 4D batched image - Capture Pydantic validation return values and use normalised fields (scaling_method, iterations, margin, etc.) instead of discarding them - Add .upper() to _RESAMPLING_MAP lookup for case-insensitive safety - Simplify double-try sys.path.insert to single try with upfront path - Add bicubic to ScalingParams description to match validator valid set
No longer needed — integration details tracked elsewhere.
| else: | ||
| batch_size = 1 | ||
| image = image.unsqueeze(0) | ||
| mask = mask.unsqueeze(0) |
There was a problem hiding this comment.
🔥 The Roast: Oh, the irony. This commit heroically fixed initial_mask getting fed a 1D row in the batch path — then took three steps left and left mask with the exact same vulnerability. mask.unsqueeze(0) lives inside the else branch, so when image arrives as [B, H, W, C] (4D, taking the if path), a 2D [H, W] mask skips the unsqueeze entirely. Then mask[i] on line 852 gives you a 1D row of pixels instead of a 2D slice, and mask_np.squeeze() on line 854 happily returns a scalar. The fix that inspired this commit is 300 lines above — and the twin bug is still right here.
🩹 The Fix: Hoist the mask unsqueeze out of the else block, mirroring what was done for initial_mask:
| mask = mask.unsqueeze(0) | |
| mask = mask.unsqueeze(0) | |
| # Ensure mask has batch dimension regardless of image path | |
| if len(mask.shape) == 2: | |
| mask = mask.unsqueeze(0) |
📏 Severity: warning
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
grabcut_nodes.py (1)
518-531:⚠️ Potential issue | 🔴 Critical
continuehere can break batch output consistency.Line 520 and Line 531 skip appending fallback tensors for that item. Later
torch.stack(processed_images)can fail (or return fewer items thanbatch_size).🛠️ Route invalid-shape items through the existing `except` fallback
if len(img_tensor.shape) != 3: log.warning("grabcut_node.unexpected_shape", shape=list(img_tensor.shape)) - continue + raise ValueError(f"Unexpected image tensor shape: {tuple(img_tensor.shape)}") @@ if len(img_np.shape) != 3 or img_np.shape[2] not in [3, 4]: log.warning("grabcut_node.invalid_shape", shape=list(img_np.shape)) - continue + raise ValueError(f"Invalid image array shape: {tuple(img_np.shape)}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 518 - 531, The code currently uses "continue" in the grabcut processing loop when img_tensor/img_np have unexpected shapes which skips adding a fallback and can break torch.stack(processed_images); instead, when detecting invalid shape (the checks around img_tensor and img_np and the log.warning calls), route that item through the existing fallback path by appending the same fallback tensor used in the except block to processed_images (or by raising the same exception that triggers the except fallback) rather than continuing; update the branches that log "grabcut_node.unexpected_shape" and "grabcut_node.invalid_shape" to add the fallback entry to processed_images so processed_images length remains equal to batch size and torch.stack(processed_images) succeeds.
♻️ Duplicate comments (3)
grabcut_nodes.py (3)
386-400:⚠️ Potential issue | 🟠 MajorValidation still misses part of the runtime input surface.
This call validates only a subset of parameters;
output_size,custom_width,custom_height,object_class, andedge_detection_modeare still unvalidated at runtime.✅ Pass the full runtime parameter set
validated = validate_node_params( iterations=grabcut_iterations, margin=margin_pixels, confidence_threshold=confidence_threshold, scaling_method=scaling_method, edge_blur_amount=edge_blur_amount, invert_mask=invert_mask, edge_refinement_strength=edge_refinement, bbox_safety_margin=bbox_safety_margin, min_bbox_size=min_bbox_size, fallback_margin_percent=fallback_margin_percent, binary_threshold=binary_threshold, output_format=output_format, auto_adjust=auto_adjust, + output_size=output_size, + custom_width=custom_width, + custom_height=custom_height, + object_class=object_class, + edge_detection_mode=edge_detection_mode, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 386 - 400, The runtime validation call to validate_node_params currently passes a subset of inputs (iterations, margin, confidence_threshold, scaling_method, etc.) but omits output_size, custom_width, custom_height, object_class, and edge_detection_mode; update the call site where validate_node_params(...) is invoked to include these missing parameters (output_size, custom_width, custom_height, object_class, edge_detection_mode) so the full runtime parameter surface is validated, and confirm validate_node_params signature and implementation accept and validate those same fields (adding them to the function if necessary).
869-875:⚠️ Potential issue | 🟠 Major
refine_maskfallback path is not shape/layout-normalized.Lines 870-875 append raw
image[i]/mask[i]while success path appends resized float32 tensors. With mixed success/failure,torch.stackcan fail.♻️ Make fallback use the same post-processing contract as success path
- else: - refined_images.append(image[i]) - refined_masks.append(mask[i]) + else: + rgba_input = np.zeros((img_np.shape[0], img_np.shape[1], 4), dtype=np.uint8) + rgba_input[:, :, :3] = img_np[:, :, :3] + rgba_input[:, :, 3] = mask_np.astype(np.uint8) + rgba_input = self._apply_resize(rgba_input, output_size, scaling_method, custom_width, custom_height) + rgb = rgba_input[:, :, :3].astype(np.float32) / 255.0 + alpha = rgba_input[:, :, 3].astype(np.float32) / 255.0 + if invert_mask: + alpha = 1.0 - alpha + for c in range(3): + rgb[:, :, c] *= alpha + refined_images.append(torch.from_numpy(rgb).to(dtype=torch.float32)) + refined_masks.append(torch.from_numpy(alpha).to(dtype=torch.float32))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 869 - 875, The except block in grabcut_nodes.py appends raw image[i] and mask[i] to refined_images/refined_masks (logged as "grabcut_node.refine_error"), but the success path appends resized, float32-normalized tensors, causing torch.stack to fail when results are mixed; update the fallback to apply the same post-processing as the success path (resize/convert dtype to float32 and normalize/transpose if the success branch does) by reusing the same helper or inline steps used there (the same resizing/astype/normalization logic that produces the items pushed to refined_images/refined_masks) so the fallback items match shape and dtype before appending.
835-842:⚠️ Potential issue | 🟠 MajorNormalize mask batch dimensions before refinement batching.
For batched
imageinputs, a 2D or singleton mask can still be mis-indexed bymask[i](wrong sample alignment).🧩 Add mask shape normalization before the batch split
- # Handle batch + # Normalize mask dimensions first + if len(mask.shape) == 2: + mask = mask.unsqueeze(0) + + # Handle batch if len(image.shape) == 4: batch_size = image.shape[0] + if len(mask.shape) == 3 and mask.shape[0] == 1 and batch_size > 1: + mask = mask.expand(batch_size, -1, -1) else: batch_size = 1 image = image.unsqueeze(0) - mask = mask.unsqueeze(0) + if len(mask.shape) == 2: + mask = mask.unsqueeze(0)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@grabcut_nodes.py` around lines 835 - 842, The batch-handling block can misalign masks when image is 4D but mask is 2D or has a singleton batch; before splitting into per-sample entries (the code that checks len(image.shape) == 4 and then uses mask[i]), normalize mask batch dimensions to match image: detect batch_size from image, then if mask.dim() == 2 or (mask.dim() == 3 and mask.shape[0] == 1) expand/repeat the mask to shape (batch_size, H, W) on the same device/dtype as image so mask[i] indexes align with image[i]; perform this normalization immediately before the existing batch-split logic in this function (the block handling image/mask batch sizing).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@grabcut_nodes.py`:
- Around line 631-640: The fallback that creates 512x512 tensors causes
torch.stack shape mismatches; replace the hardcoded fallback with dimensions
derived from the current sample or output settings instead of fixed values:
compute h,w from the same source used for other processed items (e.g., the
current sample image shape or the target output/resized dimensions used
elsewhere) and then allocate rgb_tensor and alpha_tensor with those h,w values,
honoring output_format when creating rgb_tensor; update the logic around
processed_images, rgb_tensor and alpha_tensor so all items use the same derived
dimensions.
In `@src/validation.py`:
- Around line 12-152: The models currently accept unknown fields silently; add
Pydantic v2 extra forbidding by setting model_config = {"extra": "forbid", ...}
on each validation model: GrabCutParams, ScalingParams, MaskParams, BBoxParams,
and GrabCutNodeParams (merge with existing model_config entries, e.g. keep
"str_strip_whitespace": True and add "extra": "forbid"); this ensures unknown
keys (e.g. iterationz) raise validation errors instead of being ignored.
---
Outside diff comments:
In `@grabcut_nodes.py`:
- Around line 518-531: The code currently uses "continue" in the grabcut
processing loop when img_tensor/img_np have unexpected shapes which skips adding
a fallback and can break torch.stack(processed_images); instead, when detecting
invalid shape (the checks around img_tensor and img_np and the log.warning
calls), route that item through the existing fallback path by appending the same
fallback tensor used in the except block to processed_images (or by raising the
same exception that triggers the except fallback) rather than continuing; update
the branches that log "grabcut_node.unexpected_shape" and
"grabcut_node.invalid_shape" to add the fallback entry to processed_images so
processed_images length remains equal to batch size and
torch.stack(processed_images) succeeds.
---
Duplicate comments:
In `@grabcut_nodes.py`:
- Around line 386-400: The runtime validation call to validate_node_params
currently passes a subset of inputs (iterations, margin, confidence_threshold,
scaling_method, etc.) but omits output_size, custom_width, custom_height,
object_class, and edge_detection_mode; update the call site where
validate_node_params(...) is invoked to include these missing parameters
(output_size, custom_width, custom_height, object_class, edge_detection_mode) so
the full runtime parameter surface is validated, and confirm
validate_node_params signature and implementation accept and validate those same
fields (adding them to the function if necessary).
- Around line 869-875: The except block in grabcut_nodes.py appends raw image[i]
and mask[i] to refined_images/refined_masks (logged as
"grabcut_node.refine_error"), but the success path appends resized,
float32-normalized tensors, causing torch.stack to fail when results are mixed;
update the fallback to apply the same post-processing as the success path
(resize/convert dtype to float32 and normalize/transpose if the success branch
does) by reusing the same helper or inline steps used there (the same
resizing/astype/normalization logic that produces the items pushed to
refined_images/refined_masks) so the fallback items match shape and dtype before
appending.
- Around line 835-842: The batch-handling block can misalign masks when image is
4D but mask is 2D or has a singleton batch; before splitting into per-sample
entries (the code that checks len(image.shape) == 4 and then uses mask[i]),
normalize mask batch dimensions to match image: detect batch_size from image,
then if mask.dim() == 2 or (mask.dim() == 3 and mask.shape[0] == 1)
expand/repeat the mask to shape (batch_size, H, W) on the same device/dtype as
image so mask[i] indexes align with image[i]; perform this normalization
immediately before the existing batch-split logic in this function (the block
handling image/mask batch sizing).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 5921949c-1943-4fb4-a0d4-aa8684d6c3b7
📒 Files selected for processing (2)
grabcut_nodes.pysrc/validation.py
| class GrabCutParams(BaseModel): | ||
| """Parameters for the core GrabCut algorithm.""" | ||
|
|
||
| iterations: int = Field(default=5, ge=1, le=100, description="GrabCut iterations (1-100)") | ||
| margin: int = Field(default=20, ge=0, le=500, description="Margin around detected object in pixels") | ||
| edge_threshold: float = Field( | ||
| default=0.5, ge=0.0, le=1.0, | ||
| description="Edge detection threshold for auto-adjustment" | ||
| ) | ||
| confidence_threshold: float = Field( | ||
| default=0.5, ge=0.0, le=1.0, | ||
| description="YOLO confidence threshold" | ||
| ) | ||
|
|
||
| model_config = {"str_strip_whitespace": True} | ||
|
|
||
|
|
||
| class ScalingParams(BaseModel): | ||
| """Parameters for image scaling / resize preprocessing.""" | ||
|
|
||
| target_long_edge: int = Field( | ||
| default=1024, ge=256, le=4096, | ||
| description="Target size for longest image edge" | ||
| ) | ||
| maintain_aspect: bool = Field( | ||
| default=True, | ||
| description="Preserve aspect ratio during scaling" | ||
| ) | ||
| scaling_method: str = Field( | ||
| default="auto", | ||
| description="Scaling method: auto, nearest, bilinear, bicubic, lanczos, power-of-8" | ||
| ) | ||
|
|
||
| @model_validator(mode="after") | ||
| def check_scaling_method(self) -> "ScalingParams": | ||
| valid = {"auto", "nearest", "bilinear", "bicubic", "lanczos", "power-of-8"} | ||
| normalised = self.scaling_method.lower() | ||
| if normalised not in valid: | ||
| raise ValueError( | ||
| f"Invalid scaling_method '{self.scaling_method}'. " | ||
| f"Must be one of: {', '.join(sorted(valid))}" | ||
| ) | ||
| self.scaling_method = normalised | ||
| return self | ||
|
|
||
|
|
||
| class MaskParams(BaseModel): | ||
| """Parameters for mask post-processing.""" | ||
|
|
||
| edge_blur_amount: float = Field( | ||
| default=0.0, ge=0.0, le=20.0, | ||
| description="Gaussian blur sigma for mask edge softening (0=off, 0.1-20.0)" | ||
| ) | ||
| invert_mask: bool = Field( | ||
| default=False, | ||
| description="Invert the output mask (foreground becomes background)" | ||
| ) | ||
| edge_refinement_strength: float = Field( | ||
| default=0.7, ge=0.0, le=1.0, | ||
| description="Strength of edge refinement pass" | ||
| ) | ||
|
|
||
| model_config = {"str_strip_whitespace": True} | ||
|
|
||
|
|
||
| class BBoxParams(BaseModel): | ||
| """Bounding-box / detection parameters.""" | ||
|
|
||
| bbox_safety_margin: int = Field( | ||
| default=30, ge=0, le=200, | ||
| description="Safety margin added to detected bounding boxes in pixels" | ||
| ) | ||
| min_bbox_size: int = Field( | ||
| default=64, ge=8, le=1024, | ||
| description="Minimum bounding box side length in pixels" | ||
| ) | ||
| fallback_margin_percent: float = Field( | ||
| default=0.2, ge=0.01, le=0.5, | ||
| description="When no object is detected, use this fraction of image size as margin" | ||
| ) | ||
|
|
||
|
|
||
| class GrabCutNodeParams(GrabCutParams, ScalingParams, MaskParams, BBoxParams): | ||
| """Combined validation model for AutoGrabCutRemover node parameters. | ||
|
|
||
| Used to validate ALL user inputs in a single pass before any GPU work. | ||
| """ | ||
|
|
||
| binary_threshold: int = Field( | ||
| default=200, ge=0, le=255, | ||
| description="Threshold for binary mask generation" | ||
| ) | ||
| output_format: str = Field( | ||
| default="RGBA", | ||
| description="Output format: RGBA or MASK" | ||
| ) | ||
| auto_adjust: bool = Field( | ||
| default=True, | ||
| description="Enable automatic parameter adjustment based on image analysis" | ||
| ) | ||
|
|
||
| @model_validator(mode="after") | ||
| def check_output_format(self) -> "GrabCutNodeParams": | ||
| valid = {"RGBA", "MASK"} | ||
| if self.output_format not in valid: | ||
| raise ValueError( | ||
| f"Invalid output_format '{self.output_format}'. Must be one of: {', '.join(sorted(valid))}" | ||
| ) | ||
| return self | ||
|
|
||
|
|
||
| # --------------------------------------------------------------------------- | ||
| # Convenience validator functions (call these at the top of each execute()) | ||
| # --------------------------------------------------------------------------- | ||
|
|
||
| def validate_grabcut_params(**kwargs: Any) -> GrabCutParams: | ||
| """Validate and return GrabCutParams. Raises ValueError on failure.""" | ||
| return GrabCutParams(**kwargs) | ||
|
|
||
|
|
||
| def validate_scaling_params(**kwargs: Any) -> ScalingParams: | ||
| """Validate and return ScalingParams. Raises ValueError on failure.""" | ||
| return ScalingParams(**kwargs) | ||
|
|
||
|
|
||
| def validate_mask_params(**kwargs: Any) -> MaskParams: | ||
| """Validate and return MaskParams. Raises ValueError on failure.""" | ||
| return MaskParams(**kwargs) | ||
|
|
||
|
|
||
| def validate_bbox_params(**kwargs: Any) -> BBoxParams: | ||
| """Validate and return BBoxParams. Raises ValueError on failure.""" | ||
| return BBoxParams(**kwargs) | ||
|
|
||
|
|
||
| def validate_node_params(**kwargs: Any) -> GrabCutNodeParams: | ||
| """Validate ALL parameters for a GrabCut node. | ||
|
|
||
| Call this at the START of every execute() method before any processing. | ||
| """ | ||
| return GrabCutNodeParams(**kwargs) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, locate the file and check its content
fd -t f "validation.py" --type f 2>/dev/null | head -20Repository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 136
🏁 Script executed:
# Find Pydantic version configuration
fd -t f -E node_modules "(pyproject.toml|requirements.txt|setup.py|setup.cfg|poetry.lock|pipfile)" 2>/dev/null | head -10Repository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 125
🏁 Script executed:
# Search for imports from pydantic in the codebase to understand which version
rg "from pydantic import|import pydantic" --type py 2>/dev/null | head -20Repository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 228
🏁 Script executed:
# Check Pydantic version
cat requirements.txtRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 260
🏁 Script executed:
# Also check pyproject.toml
cat pyproject.tomlRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 1136
🏁 Script executed:
# Examine the test file to understand expected behavior
cat -n tests/test_validation.pyRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 8713
🏁 Script executed:
# Check the actual src/validation.py file content
cat -n src/validation.pyRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 6291
🏁 Script executed:
# Search for usage of these validators in the codebase
rg "validate_node_params|validate_grabcut_params|validate_scaling_params|validate_mask_params|validate_bbox_params|GrabCutNodeParams" --type py -B 2 -A 2 | head -60Repository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 3619
🏁 Script executed:
# Search for any tests that verify unknown field handling
rg "extra.*forbid|unknown.*field|extra.*field" --type py -iRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 76
🏁 Script executed:
# Check if there are any tests that intentionally pass extra fields
rg "GrabCutNodeParams\(|GrabCutParams\(|ScalingParams\(|MaskParams\(|BBoxParams\(" tests/ --type py -A 5 | grep -E "(unknown|extra|typo|misspell)" -iRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 76
🏁 Script executed:
# Search for any actual usage of these validators in the codebase (not just imports/tests)
fd -t f "\.py$" -not -path "*/tests/*" -not -path "*/__pycache__/*" | xargs rg "validate_node_params|GrabCutNodeParams|validate_grabcut_params" 2>/dev/null | head -40Repository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 2477
🏁 Script executed:
# Check Pydantic v2 documentation on default extra field behavior
# First, let's create a simple test to confirm the behavior
python3 << 'EOF'
# Test Pydantic v2 default extra field behavior
from pydantic import BaseModel, Field
class TestModel(BaseModel):
name: str = Field(default="test")
model_config = {"str_strip_whitespace": True}
# Try creating an instance with an extra field
try:
instance = TestModel(name="hello", extra_field="should_be_ignored")
print(f"Extra field silently ignored: {instance}")
print(f"Extra field present in model: {hasattr(instance, 'extra_field')}")
print(f"Model dict: {instance.model_dump()}")
except Exception as e:
print(f"Error: {e}")
# Test with extra="forbid"
from pydantic import ConfigDict
class StrictTestModel(BaseModel):
name: str = Field(default="test")
model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
try:
instance2 = StrictTestModel(name="hello", extra_field="should_fail")
print(f"With extra='forbid': {instance2}")
except Exception as e:
print(f"With extra='forbid' - Error (expected): {type(e).__name__}: {e}")
EOFRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 494
🏁 Script executed:
# Check if there are any existing tests that verify extra field rejection
grep -r "extra.*forbid\|forbid.*extra\|unknown.*field" tests/ --include="*.py" 2>/dev/nullRepository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 76
🏁 Script executed:
# Look for the actual usage context in grabcut_nodes.py to understand the practical impact
fd -t f "grabcut_nodes.py" | xargs cat | head -80Repository: Limbicnation/ComfyUI-TransparencyBackgroundRemover
Length of output: 2576
Add extra="forbid" to reject unknown fields in all validation models.
Models at lines 12, 29, 58, and 77 silently ignore unexpected keys by default in Pydantic v2. This allows misspelled parameters to pass validation undetected, undermining the "Security-hardened" design claim. For example, validate_node_params(iterationz=5) would be accepted and silently dropped instead of raising an error.
Proposed fix
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field, model_validator, ConfigDict
+
+
+class _StrictModel(BaseModel):
+ model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
-class GrabCutParams(BaseModel):
+class GrabCutParams(_StrictModel):
@@
- model_config = {"str_strip_whitespace": True}
@@
-class ScalingParams(BaseModel):
+class ScalingParams(_StrictModel):
@@
-class MaskParams(BaseModel):
+class MaskParams(_StrictModel):
@@
- model_config = {"str_strip_whitespace": True}
@@
-class BBoxParams(BaseModel):
+class BBoxParams(_StrictModel):🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/validation.py` around lines 12 - 152, The models currently accept unknown
fields silently; add Pydantic v2 extra forbidding by setting model_config =
{"extra": "forbid", ...} on each validation model: GrabCutParams, ScalingParams,
MaskParams, BBoxParams, and GrabCutNodeParams (merge with existing model_config
entries, e.g. keep "str_strip_whitespace": True and add "extra": "forbid"); this
ensures unknown keys (e.g. iterationz) raise validation errors instead of being
ignored.
…ro, dead field removal - Move mask.unsqueeze(0) outside else branch in GrabCutRefinement.refine_mask (same bug pattern fixed in AutoGrabCutRemover, line 841) - Replace bare `continue` in batch validation loop with fallback tensor append to prevent torch.stack crash on empty list and bbox_metadata misalignment - Derive exception-handler fallback h,w from input tensor instead of first processed item to avoid shape mismatch across batch items - Guard FFT peak_ratio division against zero mean on uniform images - Remove dead edge_threshold field from GrabCutParams (never passed from nodes) - Log warning when Pydantic validation is unavailable instead of silent skip
…params, remove alpha premultiply inconsistency
1 participant
… structlog, type hints
Ref: CD-14 | notion:32c40bf6-db23-8001-92b8-dfe4d91dccb4
Summary by CodeRabbit
New Features
Performance Improvements
Improvements
Tests
Chores