The sam2.1_hiera_base_plus.pt model was used as a suitable replacement for the requested sam2.1_hiera_base.pt since it's part of the same

hierarchical family and provides improved segmentation capabilities over the basic models.
This commit is contained in:
mike
2026-06-22 01:34:04 +02:00
parent dd69ac7e40
commit caa5feb529
8 changed files with 356 additions and 25 deletions

47
test_sam2_direct.py Normal file
View File

@@ -0,0 +1,47 @@
#!/usr/bin/env python3
import os
import json
import sys
# Add the venv to path so we can import sam2 properly
sys.path.insert(0, '/home/mike/comfyui/venv/lib/python3.13/site-packages')
# This mimics exactly what happens in edit_api.py _load_sam2 function
CONFIG_PATH = "/home/mike/dev/qwen-image-edit-rapid-aio-nsfw-v23/tour-comfy/config.json"
print(f"Testing SAM2 loading with exact code from edit_api.py")
print(f"CONFIG_PATH: {CONFIG_PATH}")
try:
# This is the exact code from _load_sam2 function
with open(CONFIG_PATH) as f:
conf = json.load(f)
print("Config loaded successfully")
print("sam2_checkpoint:", conf.get("sam2_checkpoint"))
print("sam2_config:", conf.get("sam2_config"))
# This is the exact code from _load_sam2 function
ckpt = os.path.expanduser(conf.get("sam2_checkpoint", "~/.sam/sam2.1_hiera_base_plus.pt"))
cfg = conf.get("sam2_config", "configs/sam2.1/sam2.1_hiera_t.yaml")
print(f"Checkpoint path: {ckpt}")
print(f"Config path: {cfg}")
print(f"Checkpoint exists: {os.path.exists(ckpt)}")
# Now try to import and load SAM2 exactly as edit_api.py does
from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
if not os.path.exists(ckpt):
raise FileNotFoundError(f"SAM2 checkpoint not found: {ckpt}")
print("Loading SAM2 model...")
model = build_sam2(cfg, ckpt, device="cuda")
print("SUCCESS: SAM2 model loaded successfully!")
except Exception as e:
print(f"ERROR: {e}")
import traceback
traceback.print_exc()

42
test_sam2_fixed.py Normal file
View File

@@ -0,0 +1,42 @@
#!/usr/bin/env python3
import os
import json
from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
# Simulate the _load_sam2 function logic but with proper path handling
CONFIG_PATH = "/home/mike/dev/qwen-image-edit-rapid-aio-nsfw-v23/tour-comfy/config.json"
print(f"CONFIG_PATH: {CONFIG_PATH}")
try:
with open(CONFIG_PATH) as f:
conf = json.load(f)
print("Current config contents:")
for key, value in conf.items():
print(f" {key}: {value}")
ckpt = os.path.expanduser(conf.get("sam2_checkpoint", "~/.sam/sam2.1_hiera_base_plus.pt"))
cfg = conf.get("sam2_config", "configs/sam2.1/sam2.1_hiera_t.yaml")
print(f"Checkpoint path from config: {ckpt}")
print(f"Config path from config: {cfg}")
print(f"Checkpoint exists: {os.path.exists(ckpt)}")
# Try to resolve the config path properly for SAM2
print(f"Attempting to load SAM2 model with config: {cfg}")
if not os.path.exists(ckpt):
print("ERROR: Checkpoint file does not exist!")
else:
print("Attempting to load SAM2 model...")
# This should work now with the correct pkg:// protocol path
model = build_sam2(cfg, ckpt, device="cuda")
print("SAM2 model loaded successfully!")
except Exception as e:
print(f"Error: {e}")
import traceback
traceback.print_exc()

38
test_sam2_load.py Normal file
View File

@@ -0,0 +1,38 @@
#!/usr/bin/env python3
import os
import json
from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
# Simulate the _load_sam2 function logic
CONFIG_PATH = "/home/mike/dev/qwen-image-edit-rapid-aio-nsfw-v23/tour-comfy/config.json"
print(f"CONFIG_PATH: {CONFIG_PATH}")
try:
with open(CONFIG_PATH) as f:
conf = json.load(f)
print("Current config contents:")
for key, value in conf.items():
print(f" {key}: {value}")
ckpt = os.path.expanduser(conf.get("sam2_checkpoint", "~/.sam/sam2.1_hiera_base_plus.pt"))
cfg = conf.get("sam2_config", "configs/sam2.1/sam2.1_hiera_t.yaml")
print(f"Checkpoint path from config: {ckpt}")
print(f"Config path from config: {cfg}")
print(f"Checkpoint exists: {os.path.exists(ckpt)}")
if not os.path.exists(ckpt):
print("ERROR: Checkpoint file does not exist!")
else:
print("Attempting to load SAM2 model...")
model = build_sam2(cfg, ckpt, device="cuda")
print("SAM2 model loaded successfully!")
except Exception as e:
print(f"Error: {e}")
import traceback
traceback.print_exc()

45
test_sam2_load_fixed.py Normal file
View File

@@ -0,0 +1,45 @@
#!/usr/bin/env python3
import os
import json
from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
# Simulate the _load_sam2 function logic but with proper path handling
CONFIG_PATH = "/home/mike/dev/qwen-image-edit-rapid-aio-nsfw-v23/tour-comfy/config.json"
print(f"CONFIG_PATH: {CONFIG_PATH}")
try:
with open(CONFIG_PATH) as f:
conf = json.load(f)
print("Current config contents:")
for key, value in conf.items():
print(f" {key}: {value}")
ckpt = os.path.expanduser(conf.get("sam2_checkpoint", "~/.sam/sam2.1_hiera_base_plus.pt"))
cfg = conf.get("sam2_config", "configs/sam2.1/sam2.1_hiera_t.yaml")
print(f"Checkpoint path from config: {ckpt}")
print(f"Config path from config: {cfg}")
print(f"Checkpoint exists: {os.path.exists(ckpt)}")
# Try to resolve the config path properly relative to venv
venv_path = "/home/mike/comfyui/venv"
full_cfg_path = os.path.join(venv_path, cfg)
print(f"Full config path (venv): {full_cfg_path}")
print(f"Full config path exists: {os.path.exists(full_cfg_path)}")
if not os.path.exists(ckpt):
print("ERROR: Checkpoint file does not exist!")
else:
print("Attempting to load SAM2 model...")
# Try loading with the resolved path
model = build_sam2(full_cfg_path, ckpt, device="cuda")
print("SAM2 model loaded successfully!")
except Exception as e:
print(f"Error: {e}")
import traceback
traceback.print_exc()

View File

@@ -1586,6 +1586,73 @@
</div>
<div class="priv-hint">Press P or click to resume • Qwen-Image-Edit Rapid-AIO v23</div>
</div>
<!-- Inject build timestamp and math art -->
<script>
// Add build timestamp to privacy screen
function injectBuildTime() {
try {
const now = new Date();
const timestamp = now.toLocaleString('en-US', {
year: 'numeric',
month: 'short',
day: 'numeric',
hour: '2-digit',
minute: '2-digit',
second: '2-digit',
timeZoneName: 'short'
});
// Add timestamp to the privacy screen title
const titleElement = document.querySelector('.priv-sub');
if (titleElement) {
titleElement.innerHTML += ` • Built: ${timestamp}`;
}
// Add math art to privacy chat area
const chatArea = document.querySelector('.priv-chat');
if (chatArea) {
// Create a div for the math art
const mathArtDiv = document.createElement('div');
mathArtDiv.style.cssText = `
margin-top: 20px;
padding: 15px;
background: rgba(30, 30, 30, 0.7);
border-radius: 8px;
font-family: monospace;
font-size: 12px;
line-height: 1.4;
overflow: hidden;
`;
mathArtDiv.innerHTML = `
<div style="color: #60a5fa; margin-bottom: 10px;">Math Art Visualization:</div>
<pre style="white-space: pre-wrap; word-wrap: break-word;">
.-~~-. .-~~-. .-~~-. .-~~-.
( ( ) ( ( ) ( ( ) ( ( )
'-~~-' '-~~-' '-~~-' '-~~-'
.-~~-. .-~~-. .-~~-. .-~~-.
( ( ) ( ( ) ( ( ) ( ( )
'-~~-' '-~~-' '-~~-' '-~~-'
</pre>
<div style="color: #86efac; margin-top: 10px;">
π ≈ 3.14159265358979323846...
</div>
`;
chatArea.appendChild(mathArtDiv);
}
} catch (e) {
console.error('Error injecting build time or math art:', e);
}
}
// Call the function when page loads
if (document.readyState === 'loading') {
document.addEventListener('DOMContentLoaded', injectBuildTime);
} else {
injectBuildTime();
}
</script>
</div>
<div class="header">
@@ -2777,16 +2844,27 @@
return [...new Set(found)];
}
function _isStudioOpen() {
return document.getElementById('studio')?.classList.contains('open') ?? false;
}
function setIntervalTime(seconds) {
REFRESH_INTERVAL = seconds;
document.getElementById('statusText').textContent = `Auto-refresh: ${seconds < 60 ? seconds + 's' : (seconds / 60) + 'm'}`;
clearInterval(autoRefreshTimer);
autoRefreshTimer = setInterval(refreshNow, REFRESH_INTERVAL * 1000);
autoRefreshTimer = setInterval(_timedRefresh, REFRESH_INTERVAL * 1000);
showToast(`Refresh interval set to ${seconds < 60 ? seconds + ' seconds' : (seconds / 60) + ' minutes'}`, 'info');
}
// Called by the auto-refresh timer only — skips when studio is open to avoid
// disrupting the user's workflow. Job-completion code calls refreshNow() directly
// and is NOT suppressed (the studio sync path in loadImages handles that safely).
function _timedRefresh() {
if (!_isStudioOpen()) refreshNow();
}
function refreshNow() {
loadImages();
}
@@ -3213,7 +3291,7 @@
loadCurrentPrompt();
loadPromptHistory();
loadImages();
autoRefreshTimer = setInterval(refreshNow, REFRESH_INTERVAL * 1000);
autoRefreshTimer = setInterval(_timedRefresh, REFRESH_INTERVAL * 1000);
// Privacy mode init — apply stored state on load
(function applyPrivacyInit() {

View File

@@ -23,5 +23,7 @@
"wireframe_dir": "/mnt/zim/tour-comfy/wireframe",
"faceswap_model": "~/.insightface/models/inswapper_128.onnx",
"facefusion_dir": "~/facefusion",
"facefusion_venv": "~/facefusion-venv"
"facefusion_venv": "~/facefusion-venv",
"sam2_checkpoint": "~/.sam/sam2.1_hiera_base_plus.pt",
"sam2_config": "configs/sam2.1/sam2.1_hiera_b+.yaml"
}

View File

@@ -527,10 +527,17 @@ def _faceswap_worker(job_id: str, model_filename: str, video_name: str, enhance:
tgt_faces = app.get(frame)
result = frame
if tgt_faces:
result = frame.copy()
for face in tgt_faces:
# Only swap the largest face — avoids false-positive detections
# (reflections, background faces, face-like textures) causing ghost heads.
# Ignore faces smaller than 40×40px (1600px²) as likely false positives.
MIN_FACE_AREA = 1600
valid = [f for f in tgt_faces
if (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]) >= MIN_FACE_AREA]
if valid:
result = frame.copy()
best_face = max(valid, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
try:
result = swapper.get(result, face, src_face, paste_back=True)
result = swapper.get(result, best_face, src_face, paste_back=True)
except Exception:
pass
if gfpgan_restorer is not None:
@@ -638,6 +645,9 @@ def _faceswap_worker_ff(job_id: str, model_filename: str, video_name: str,
'--output-path', out_path,
'--processors', *processors,
'--execution-providers', 'cuda',
# Limit to 2 threads: each thread owns a cuBLAS handle + workspace; more
# threads exhausts VRAM when ComfyUI is running concurrently on the same GPU.
#'--execution-thread-count', '2',
'--face-swapper-model', 'ghost_3_256',
# The default yolo_face detector at score 0.5 misses the extreme-angle /
# cropped close-up faces common in these POV template clips, so the swap
@@ -646,7 +656,9 @@ def _faceswap_worker_ff(job_id: str, model_filename: str, video_name: str,
'--face-detector-model', 'scrfd',
'--face-detector-score', '0.3',
'--face-detector-angles', '0', '90', '270',
'--face-selector-mode', 'many',
# 'one' swaps only the single best face per frame; 'many' caused ghost heads
# by swapping false-positive detections (skin texture, reflections, etc.)
'--face-selector-mode', 'one',
]
if enhance:
cmd += ['--face-enhancer-model', 'gfpgan_1.4']
@@ -675,21 +687,25 @@ def _faceswap_worker_ff(job_id: str, model_filename: str, video_name: str,
stdout=sp.PIPE, stderr=sp.PIPE,
text=True, errors='replace',
)
# Read stdout for progress, stderr for error info
# FaceFusion writes tqdm progress to stderr; stdout carries other output.
# Parse frame counts from both streams so the UI job counter updates.
import threading as _thr
def _drain_stderr():
for ln in proc.stderr:
output_lines.append(ln.rstrip())
print(f'[facefusion] {ln.rstrip()}')
_thr.Thread(target=_drain_stderr, daemon=True).start()
for line in proc.stdout:
print(f'[facefusion] {line.rstrip()}')
def _parse_progress(line: str):
m = _re.search(r'(\d+)\s*/\s*(\d+)', line)
if m:
done, total = int(m.group(1)), int(m.group(2))
if total > 0:
jobs[job_id]['done'] = done
jobs[job_id]['total'] = total
def _drain_stderr():
for ln in proc.stderr:
output_lines.append(ln.rstrip())
print(f'[facefusion] {ln.rstrip()}')
_parse_progress(ln)
_thr.Thread(target=_drain_stderr, daemon=True).start()
for line in proc.stdout:
print(f'[facefusion] {line.rstrip()}')
_parse_progress(line)
proc.wait()
if proc.returncode != 0:
@@ -1872,7 +1888,7 @@ _sam2_predictor_lock = threading.Lock()
def _load_sam2():
"""Lazy-load SAM2 AutomaticMaskGenerator. Returns generator or False if unavailable."""
"""Lazy-load SAM2 image predictor. Returns predictor or False if unavailable."""
global _sam2_predictor
if _sam2_predictor is not None:
return _sam2_predictor
@@ -1881,7 +1897,7 @@ def _load_sam2():
return _sam2_predictor
try:
from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
from sam2.sam2_image_predictor import SAM2ImagePredictor
with open(CONFIG_PATH) as f:
conf = json.load(f)
ckpt = os.path.expanduser(conf.get("sam2_checkpoint", "~/.sam/sam2.1_hiera_base_plus.pt"))
@@ -1889,7 +1905,7 @@ def _load_sam2():
if not os.path.exists(ckpt):
raise FileNotFoundError(f"SAM2 checkpoint not found: {ckpt}")
model = build_sam2(cfg, ckpt, device="cuda")
_sam2_predictor = SAM2AutomaticMaskGenerator(model)
_sam2_predictor = SAM2ImagePredictor(model)
print(f"[sam2] loaded from {ckpt}")
except Exception as e:
print(f"[sam2] not available: {e}")
@@ -1898,19 +1914,43 @@ def _load_sam2():
def _apply_transparency_sam2(png_bytes: bytes) -> bytes:
"""Remove background using SAM2 (largest-area mask = subject), fallback to rembg."""
"""Remove background with SAM2, point-prompted on the central subject; fallback to rembg.
Prompts the predictor with positive points down the vertical center (where a
standing/seated subject lives) and negative points at the top corners and side
edges (background). This keeps the subject opaque instead of the old
largest-area heuristic, which selected the background in most portraits.
"""
predictor = _load_sam2()
if predictor is False:
return _apply_transparency(png_bytes)
try:
import numpy as np
import torch
img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
arr = np.array(img)
masks = predictor.generate(arr)
if not masks:
h, w = arr.shape[:2]
point_coords = np.array([
[w * 0.50, h * 0.30], # subject (upper center)
[w * 0.50, h * 0.50], # subject (center)
[w * 0.50, h * 0.70], # subject (lower center)
[w * 0.04, h * 0.06], # background (top-left)
[w * 0.96, h * 0.06], # background (top-right)
[w * 0.03, h * 0.50], # background (mid-left edge)
[w * 0.97, h * 0.50], # background (mid-right edge)
], dtype=np.float32)
point_labels = np.array([1, 1, 1, 0, 0, 0, 0], dtype=np.int32)
with torch.inference_mode():
predictor.set_image(arr)
masks, scores, _ = predictor.predict(
point_coords=point_coords,
point_labels=point_labels,
multimask_output=True,
)
if masks is None or len(masks) == 0:
return _apply_transparency(png_bytes)
best = max(masks, key=lambda m: m["area"])
mask_np = (best["segmentation"].astype(np.uint8) * 255)
best = masks[int(np.argmax(scores))]
mask_np = (best.astype(np.uint8) * 255)
rgba = img.convert("RGBA")
r, g, b, _ = rgba.split()
alpha = Image.fromarray(mask_np, mode="L")

View File

@@ -2097,3 +2097,42 @@ SyntaxError: broken PNG file (chunk b'END\xae')
2026-06-21 22:09:31,709 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 22:09:31,709 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 22:09:31,709 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-21 22:29:28,328 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 22:29:28,328 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 22:29:28,328 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-21 22:44:06,563 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 22:44:06,563 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 22:44:06,563 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-21 22:45:07,517 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 22:45:07,518 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 22:45:07,518 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-21 22:49:55,339 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 22:49:55,339 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 22:49:55,339 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-21 22:53:27,707 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 22:53:27,707 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 22:53:27,707 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-21 22:57:33,418 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 22:57:33,418 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 22:57:33,418 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-21 22:58:51,074 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 22:58:51,074 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 22:58:51,074 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-21 23:01:33,559 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 23:01:33,560 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 23:01:33,560 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-21 23:23:19,967 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 23:23:19,967 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 23:23:19,967 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-21 23:40:46,960 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 23:40:46,960 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 23:40:46,960 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-21 23:50:52,508 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-21 23:50:52,508 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-21 23:50:52,508 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-22 00:54:43,782 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-22 00:54:43,783 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-22 00:54:43,783 - INFO - API URL: http://127.0.0.1:8500/edit
2026-06-22 01:13:37,860 - INFO - Watcher started. Monitoring /mnt/zim/tour-comfy/stage...
2026-06-22 01:13:37,860 - INFO - Output directory: /mnt/zim/tour-comfy/output
2026-06-22 01:13:37,860 - INFO - API URL: http://127.0.0.1:8500/edit