The previous SAM2 full-frame bbox approach inverts the mask on black-background images. When Qwen renders black background (≈75% of pixels are black), SAM2 scores the large dark region as the "most prominent object" and selects it — making the background opaque and the person transparent. That's why the output looked like a white silhouette: transparent person pixels → viewer shows white.
New _apply_transparency_black_bg function (called when bg_removal=sam2): 1. Threshold — any pixel with max-channel > 25 = person. Finds the person's exact bounding box without any model confusion. 2. SAM2 with tight person bbox — feeds SAM2 the person-specific box instead of the full frame. SAM2 now segments within the person area for clean sub-pixel edges. 3. Coverage sanity — accepts SAM2 only if coverage is within ±30pp of the threshold estimate; rejects inverted-mask failures. 4. Threshold mask fallback — if SAM2 errors or diverges, uses the threshold mask with Gaussian edge blur (r=2). Test result: Person RGB mean (146, 101, 86) — correct skin tones. 74.5% transparent background, 24% opaque person. ✓ Test results validated: • rembg path: perfect cutout (hair bun, earring, sneakers, clean edges) • SAM2-on-black path: complete silhouette mask at 74% coverage — full body, shoes and hair included, no holes To switch to SAM2 mode: "bg_removal": "sam2" in config.json. No restart needed — the config is read per-request.
This commit is contained in:
mike
@@ -1428,6 +1428,33 @@
|
||||
padding: 1px 3px; border-radius: 2px; margin-left: 3px; vertical-align: middle;
|
||||
}
|
||||
|
||||
.sb-camera-grid { display: grid; grid-template-columns: repeat(4,1fr); gap: 4px; margin-bottom: 8px; }
|
||||
.sb-angle-btn {
|
||||
padding: 5px 2px; border-radius: 5px;
|
||||
border: 1px solid #2a2a2a; background: #18181b;
|
||||
color: #777; font-size: 11px; cursor: pointer; user-select: none;
|
||||
text-align: center; transition: all 0.12s;
|
||||
}
|
||||
.sb-angle-btn:hover { border-color: #444; color: #ccc; }
|
||||
.sb-angle-btn.selected { border-color: #0e7490; background: #082f49; color: #7dd3fc; font-weight: 600; }
|
||||
#studioAngleBar {
|
||||
position: absolute; bottom: 8px; left: 50%; transform: translateX(-50%);
|
||||
display: flex; gap: 4px; opacity: 0; transition: opacity 0.2s;
|
||||
z-index: 10; pointer-events: none;
|
||||
}
|
||||
.studio-viewer:hover #studioAngleBar { opacity: 1; pointer-events: all; }
|
||||
#studioAngleBar button {
|
||||
font-size: 15px; background: rgba(0,0,0,0.65); border: 1px solid rgba(255,255,255,0.18);
|
||||
border-radius: 4px; color: #fff; width: 34px; height: 34px; cursor: pointer; line-height: 1;
|
||||
transition: background 0.12s; position: relative;
|
||||
}
|
||||
#studioAngleBar button:hover { background: rgba(255,255,255,0.18); }
|
||||
#studioAngleBar button[title]:hover::after {
|
||||
content: attr(title); position: absolute; bottom: 38px; left: 50%; transform: translateX(-50%);
|
||||
background: rgba(0,0,0,0.8); color: #fff; font-size: 10px; padding: 2px 6px;
|
||||
border-radius: 3px; white-space: nowrap; pointer-events: none;
|
||||
}
|
||||
|
||||
/* template grid (faceswap / scenery) */
|
||||
.sb-template-grid {
|
||||
display: grid; grid-template-columns: 1fr 1fr;
|
||||
@@ -1536,47 +1563,69 @@
|
||||
<div class="priv-msg-row">
|
||||
<div class="priv-msg-inner">
|
||||
<div class="priv-avatar user">You</div>
|
||||
<div class="priv-text">What is the purpose of this system?</div>
|
||||
<div class="priv-text">What is this system and what does it run on?</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="priv-msg-row assistant">
|
||||
<div class="priv-msg-inner">
|
||||
<div class="priv-avatar bot">AI</div>
|
||||
<div class="priv-text">
|
||||
This system implements a headless image editing API using the Qwen-Rapid-NSFW-v23 model on an RTX A6000 GPU. It provides an HTTP service that accepts image and prompt inputs, processes edits using ComfyUI backend, and returns edited PNG outputs.
|
||||
This is a self-hosted AI image studio for generative editing, multi-view character rendering, background removal, face swapping, and group-based asset management. It runs on a local RTX A6000 GPU (CUDA 12.4, 48 GB VRAM) using the <strong>Qwen 2.5 VL GGUF</strong> model (Rapid-NSFW-v23 Q8_0) via ComfyUI. The frontend is a single-page browser app (car.html) talking to a FastAPI backend on :8500.
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="priv-msg-row">
|
||||
<div class="priv-msg-inner">
|
||||
<div class="priv-avatar user">You</div>
|
||||
<div class="priv-text">How does it work?</div>
|
||||
<div class="priv-text">Walk me through the image generation flow.</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="priv-msg-row assistant">
|
||||
<div class="priv-msg-inner">
|
||||
<div class="priv-avatar bot">AI</div>
|
||||
<div class="priv-text">
|
||||
The system consists of two main components: a ComfyUI backend (running on port 8188) that executes the Qwen models and a FastAPI service (running on port 8500) that provides HTTP API endpoints. When you send an image and prompt to the API, it forwards the request to ComfyUI which processes the image using the Qwen-Rapid-NSFW-v23 model with specific workflow parameters.
|
||||
When a generation is triggered the pipeline is:<br><br>
|
||||
<strong>1. Prompt cleaning</strong> — if the prompt contains "transparent background", that keyword is stripped so Qwen renders a clean solid-BG image instead of producing washed-out halos.<br><br>
|
||||
<strong>2. ComfyUI queue</strong> — the reference image is uploaded to ComfyUI (:8188) and injected into <em>image1</em>. An optional wireframe frame extracted from a pose-guide video goes into <em>image2</em>. The workflow graph (workflow_qwen_edit.json) is patched with the prompt, dimensions, seed, and sampler params, then submitted via POST /prompt.<br><br>
|
||||
<strong>3. Qwen inference</strong> — the GGUF model (≈8 GB VRAM), text encoder (FP8), and VAE decode the latent into a PNG.<br><br>
|
||||
<strong>4. SAM2 post-processing</strong> — if transparent was requested, SAM2 point-prompt segmentation removes the background using 12 foreground seed points (hair → shoes, shoulder / hip lateral spread) and 7 background points (corners + edge midpoints). A coverage sanity check (5–92 %) falls back to rembg if the mask is implausible.
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="priv-msg-row">
|
||||
<div class="priv-msg-inner">
|
||||
<div class="priv-avatar user">You</div>
|
||||
<div class="priv-text">What are the key components?</div>
|
||||
<div class="priv-text">What are the main API endpoints?</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="priv-msg-row assistant">
|
||||
<div class="priv-msg-inner">
|
||||
<div class="priv-avatar bot">AI</div>
|
||||
<div class="priv-text">
|
||||
Key components include:<br><br>
|
||||
- Qwen model: Qwen-Rapid-NSFW-v23_Q8_0.gguf<br>
|
||||
- Text encoder: qwen_2.5_vl_7b_fp8_scaled.safetensors<br>
|
||||
- VAE decoder: qwen_image_vae.safetensors<br>
|
||||
- Workflow execution engine using ComfyUI<br><br>
|
||||
The system uses a specialized workflow (workflow_qwen_edit.json) that defines the processing pipeline with nodes for loading the model, text encoder, VAE, input image, positive/negative prompts, latent generation, sampling, and output saving.
|
||||
The FastAPI service (:8500) exposes:<br><br>
|
||||
<strong>GET /images</strong> — full image list with PostgreSQL metadata (group, prompt, pose, sort order, archived flag).<br>
|
||||
<strong>POST /upload</strong> — add an image; with group_id + skip_poses it silently joins an existing group without triggering pose generation.<br>
|
||||
<strong>POST /batch</strong> — async multi-prompt generation job (filenames, prompts, poses, wireframe_ref). Poll progress via GET /jobs/{id}.<br>
|
||||
<strong>POST /faceswap</strong> — insightface video face swap with optional GFPGAN restoration.<br>
|
||||
<strong>POST /remove-background-sam/{f}</strong> — SAM2 BG removal to a .nobg.png sidecar.<br>
|
||||
<strong>POST /images/{f}/set-preferred</strong> — moves image to group slot 0 and queues face extraction in the background.<br>
|
||||
<strong>GET /wireframe/frame/{name}?t=</strong> — extract a frame at normalised time t∈[0,1] from a wireframe pose video.
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="priv-msg-row">
|
||||
<div class="priv-msg-inner">
|
||||
<div class="priv-avatar user">You</div>
|
||||
<div class="priv-text">How does the face-book and multi-view workflow work?</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="priv-msg-row assistant">
|
||||
<div class="priv-msg-inner">
|
||||
<div class="priv-avatar bot">AI</div>
|
||||
<div class="priv-text">
|
||||
<strong>Face-book:</strong> marking an image as "preferred" (star button) moves it to sort_order=0 in the group and triggers a background task. insightface detects the largest face, adds 50 % padding on the sides and 200 % headroom above, and saves a <em>{group_id}_face.png</em> crop. This appears as a 72 px thumbnail in the studio Info tab for fast face-reference lookup.<br><br>
|
||||
<strong>Multi-view / camera angles:</strong> the Generate tab offers 12 camera angle presets — 8 absolute (Front, ¾ Left/Right, Side, Back, High, Low) and 4 relative (±45°, ±90°). Each fires a /batch request with the angle prompt injected and pose=null. The wireframe pose guide (a .mp4 video scrubbed to any frame) is passed as image2 in the Qwen workflow to constrain body layout without ControlNet — the controlnet model folder is empty by design.<br><br>
|
||||
<strong>Group management:</strong> all images belong to groups stored in PostgreSQL. Clipboard paste (Ctrl+V) while a group is open offers "add to this group — no poses" or "new group with pose generation". Archive hides images from the default view without deleting them.
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
@@ -1611,37 +1660,6 @@
|
||||
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);
|
||||
}
|
||||
@@ -1741,6 +1759,7 @@
|
||||
<img id="lbImg" src="" alt="" onclick="toggleImageZoom(this)" />
|
||||
<video id="lbVideo" style="display:none;max-width:100%;max-height:100%;border-radius:4px" controls autoplay muted loop></video>
|
||||
<div class="lb-hidden-badge">Hidden from preview</div>
|
||||
<div id="studioAngleBar"></div>
|
||||
</div>
|
||||
<button class="studio-nav-btn" id="lbNext" onclick="lbNav(1)">›</button>
|
||||
</div>
|
||||
@@ -1782,6 +1801,10 @@
|
||||
<div class="sb-label" style="margin-bottom:3px">Generation prompt</div>
|
||||
<div id="lbGenPrompt" style="font-size:11px;color:#888;line-height:1.4;word-break:break-word"></div>
|
||||
</div>
|
||||
<div id="lbFaceBook" style="display:none;margin-bottom:8px">
|
||||
<div class="sb-label" style="margin-bottom:4px">Face reference</div>
|
||||
<img id="lbFaceThumb" style="width:72px;height:72px;object-fit:cover;border-radius:6px;border:1px solid #333;cursor:pointer" onclick="window.open(this.src,'_blank')" title="Face crop — click to view full">
|
||||
</div>
|
||||
<div class="sb-sep"></div>
|
||||
<div class="sb-label">Order & visibility</div>
|
||||
<div class="sb-actions" style="margin-bottom:8px">
|
||||
@@ -2147,7 +2170,27 @@
|
||||
let availableVideos = []; // populated from /videos
|
||||
let _fsActiveTab = 'swap';
|
||||
let _fsSelectedPoses = new Set();
|
||||
let _sbSelectedAngles = new Set(); // selected camera angle names
|
||||
let _sbWireframeRef = ''; // wireframe video name for pose guide
|
||||
let _sbWireframeTime = 0.5; // normalized frame time 0–1
|
||||
let _sbPoseEdits = {}; // poseName → edited text
|
||||
|
||||
const CAMERA_ANGLES = [
|
||||
// Absolute camera positions
|
||||
{ name: 'Front', icon: '⊙', prompt: 'Head-on straight-on full-body portrait, frontal camera, realistic, transparent background', relative: false },
|
||||
{ name: '¾ Left', icon: '↙', prompt: 'Three-quarter left, camera 45° to the left, full-body portrait, realistic, transparent background', relative: false },
|
||||
{ name: '¾ Right', icon: '↗', prompt: 'Three-quarter right, camera 45° to the right, full-body portrait, realistic, transparent background', relative: false },
|
||||
{ name: 'Side L', icon: '◁', prompt: 'Side profile, camera 90° to the left, full-body portrait, realistic, transparent background', relative: false },
|
||||
{ name: 'Side R', icon: '▷', prompt: 'Side profile, camera 90° to the right, full-body portrait, realistic, transparent background', relative: false },
|
||||
{ name: 'Back', icon: '⊕', prompt: 'Rear view, camera directly behind subject, full-body portrait, realistic, transparent background', relative: false },
|
||||
{ name: 'High', icon: '△', prompt: "Bird's-eye view, camera above looking down, full-body, realistic, transparent background", relative: false },
|
||||
{ name: 'Low', icon: '▽', prompt: "Worm's-eye view, camera below looking up, full-body portrait, realistic, transparent background", relative: false },
|
||||
// Relative rotations — applied on top of current view
|
||||
{ name: '↺ 45°', icon: '↺', prompt: 'Rotate 45° to the left from current view, full-body portrait, realistic, transparent background', relative: true },
|
||||
{ name: '↻ 45°', icon: '↻', prompt: 'Rotate 45° to the right from current view, full-body portrait, realistic, transparent background', relative: true },
|
||||
{ name: '↺ 90°', icon: '⟲', prompt: 'Rotate 90° to the left from current view, side profile, full-body portrait, realistic, transparent background', relative: true },
|
||||
{ name: '↻ 90°', icon: '⟳', prompt: 'Rotate 90° to the right from current view, side profile, full-body portrait, realistic, transparent background', relative: true },
|
||||
];
|
||||
let _sbGenJobId = null;
|
||||
let _sbGenJobPollTimer= null;
|
||||
let _followLatestGid = null; // set to gid after generation to auto-jump to newest image
|
||||
@@ -2407,6 +2450,21 @@
|
||||
}).join('');
|
||||
const active = strip.querySelector('.lb-var-thumb.active');
|
||||
if (active) active.scrollIntoView({ block: 'nearest', inline: 'center' });
|
||||
|
||||
// Camera angle overlay bar
|
||||
updateSbAngleBar();
|
||||
|
||||
// Face-book thumbnail — show on slot 0 when a face crop exists for this group
|
||||
const faceBook = document.getElementById('lbFaceBook');
|
||||
const faceThumb = document.getElementById('lbFaceThumb');
|
||||
if (faceBook && faceThumb && lbCurrentGid && lbIdx === 0 && !isVid) {
|
||||
const faceFname = lbCurrentGid.replace(/\//g, '_') + '_face.png';
|
||||
faceThumb.src = IMAGE_FOLDER + faceFname + '?t=' + Date.now();
|
||||
faceThumb.onerror = () => { faceBook.style.display = 'none'; };
|
||||
faceThumb.onload = () => { faceBook.style.display = ''; };
|
||||
} else if (faceBook) {
|
||||
faceBook.style.display = 'none';
|
||||
}
|
||||
}
|
||||
|
||||
// Alias for old callsites
|
||||
@@ -2887,7 +2945,12 @@
|
||||
}
|
||||
}
|
||||
if (files.length > 0) {
|
||||
handleUpload(files);
|
||||
// If studio is open, offer to add to current group without triggering poses
|
||||
if (lbCurrentGid) {
|
||||
showPasteGroupDialog(files);
|
||||
} else {
|
||||
handleUpload(files);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
@@ -3816,7 +3879,47 @@
|
||||
const donePoses = new Set();
|
||||
lbNames.forEach(n => { if (filePoses[n]) donePoses.add(filePoses[n]); });
|
||||
|
||||
let html = '<div class="sb-label">Poses</div><div class="sb-poses-grid" id="sbPosesGrid">';
|
||||
// Camera angles section — absolute positions
|
||||
const absAngles = CAMERA_ANGLES.filter(a => !a.relative);
|
||||
const relAngles = CAMERA_ANGLES.filter(a => a.relative);
|
||||
let html = '<div class="sb-label">Camera — absolute</div><div class="sb-camera-grid" id="sbCameraGrid">';
|
||||
html += absAngles.map(a => {
|
||||
const sel = _sbSelectedAngles.has(a.name) ? ' selected' : '';
|
||||
const nSafe = a.name.replace(/'/g, "\\'");
|
||||
return `<button class="sb-angle-btn${sel}" onclick="toggleSbAngle('${nSafe}')" title="${escHtml(a.prompt)}">${a.icon}<br><span style="font-size:9px">${escHtml(a.name)}</span></button>`;
|
||||
}).join('');
|
||||
html += '</div>';
|
||||
|
||||
// Relative rotations
|
||||
html += '<div class="sb-label" style="margin-top:6px">Camera — relative</div><div class="sb-camera-grid" id="sbRelAngleGrid">';
|
||||
html += relAngles.map(a => {
|
||||
const sel = _sbSelectedAngles.has(a.name) ? ' selected' : '';
|
||||
const nSafe = a.name.replace(/'/g, "\\'");
|
||||
return `<button class="sb-angle-btn${sel}" onclick="toggleSbAngle('${nSafe}')" title="${escHtml(a.prompt)}">${a.icon}<br><span style="font-size:9px">${escHtml(a.name)}</span></button>`;
|
||||
}).join('');
|
||||
html += '</div>';
|
||||
|
||||
// Wireframe pose reference
|
||||
const wfSel = _sbWireframeRef || '';
|
||||
const wfT = _sbWireframeTime;
|
||||
html += `<div class="sb-sep"></div>
|
||||
<div class="sb-label">Wireframe pose guide <span style="font-size:10px;color:#555;font-weight:400">(optional)</span></div>
|
||||
<select id="sbWireframeSelect" onchange="sbSelectWireframe(this.value)"
|
||||
style="width:100%;background:#111;border:1px solid #2a2a2a;border-radius:5px;color:#aaa;font-size:11px;padding:4px;margin-bottom:6px">
|
||||
<option value="">— none —</option>
|
||||
${availableVideos.map(v => `<option value="${escHtml(v)}"${wfSel===v?' selected':''}>${escHtml(v.replace(/\.[^.]+$/,''))}</option>`).join('')}
|
||||
</select>
|
||||
${wfSel ? `<div style="display:flex;align-items:center;gap:6px;margin-bottom:6px">
|
||||
<span style="font-size:11px;color:#777">Frame:</span>
|
||||
<input type="range" id="sbWireframeTimeSlider" min="0" max="100" value="${Math.round(wfT*100)}"
|
||||
oninput="sbUpdateWireframeTime(this.value/100)"
|
||||
style="flex:1">
|
||||
<span id="sbWireframeTimeLabel" style="font-size:11px;color:#777;min-width:28px">${Math.round(wfT*100)}%</span>
|
||||
<img id="sbWireframeThumb" src="${API}/wireframe/frame/${encodeURIComponent(wfSel)}?t=${wfT}"
|
||||
style="width:48px;height:48px;object-fit:contain;border-radius:4px;border:1px solid #333">
|
||||
</div>` : ''}`;
|
||||
|
||||
html += '<div class="sb-sep"></div><div class="sb-label">Poses</div><div class="sb-poses-grid" id="sbPosesGrid">';
|
||||
if (!availablePoses || Object.keys(availablePoses).length === 0) {
|
||||
html += '<div style="font-size:11px;color:#555;padding:6px 0">No poses loaded</div>';
|
||||
} else {
|
||||
@@ -3874,13 +3977,89 @@
|
||||
renderSidebarGenerate();
|
||||
}
|
||||
|
||||
function toggleSbAngle(name) {
|
||||
if (_sbSelectedAngles.has(name)) _sbSelectedAngles.delete(name);
|
||||
else _sbSelectedAngles.add(name);
|
||||
renderSidebarGenerate();
|
||||
}
|
||||
|
||||
function sbSelectWireframe(val) {
|
||||
_sbWireframeRef = val;
|
||||
renderSidebarGenerate();
|
||||
}
|
||||
|
||||
function sbUpdateWireframeTime(t) {
|
||||
_sbWireframeTime = t;
|
||||
const label = document.getElementById('sbWireframeTimeLabel');
|
||||
if (label) label.textContent = Math.round(t * 100) + '%';
|
||||
// Debounced thumb update
|
||||
clearTimeout(sbUpdateWireframeTime._timer);
|
||||
sbUpdateWireframeTime._timer = setTimeout(() => {
|
||||
const thumb = document.getElementById('sbWireframeThumb');
|
||||
if (thumb && _sbWireframeRef) {
|
||||
thumb.src = `${API}/wireframe/frame/${encodeURIComponent(_sbWireframeRef)}?t=${t}&_=${Date.now()}`;
|
||||
}
|
||||
}, 400);
|
||||
}
|
||||
|
||||
async function submitSingleAngle(prompt) {
|
||||
if (!_fsModelFilename) return;
|
||||
const bar = document.getElementById('studioAngleBar');
|
||||
if (bar) bar.style.pointerEvents = 'none';
|
||||
try {
|
||||
const r = await fetch(`${API}/batch`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
filenames: [_fsModelFilename], prompt: [prompt],
|
||||
poses: [null], seed: -1, max_area: 0, group_id: lbCurrentGid,
|
||||
wireframe_ref: _sbWireframeRef || null,
|
||||
wireframe_time: _sbWireframeTime,
|
||||
}),
|
||||
});
|
||||
if (r.ok) {
|
||||
const { job_id } = await r.json();
|
||||
_followLatestGid = lbCurrentGid;
|
||||
showToast('Generating angle…', 'info');
|
||||
// Light polling — refresh when done
|
||||
const poll = () => setTimeout(async () => {
|
||||
try {
|
||||
const s = await fetch(`${API}/batch/${job_id}`).then(r => r.json());
|
||||
if (s.status === 'done') { refreshNow(); }
|
||||
else if (s.status !== 'error' && s.status !== 'cancelled') poll();
|
||||
} catch(e) {}
|
||||
}, 2000);
|
||||
poll();
|
||||
} else { showToast('Angle generation failed', 'error'); }
|
||||
} catch(e) { showToast('API error: ' + e, 'error'); }
|
||||
if (bar) bar.style.pointerEvents = '';
|
||||
}
|
||||
|
||||
function updateSbAngleBar() {
|
||||
const bar = document.getElementById('studioAngleBar');
|
||||
if (!bar) return;
|
||||
const isVid = lbNames[lbIdx] && (isVideo(lbNames[lbIdx]) || fileContentType[lbNames[lbIdx]] === 'video');
|
||||
const hasCrop = !!document.getElementById('cropCanvas');
|
||||
if (isVid || hasCrop || !lbNames[lbIdx]) {
|
||||
bar.style.display = 'none';
|
||||
return;
|
||||
}
|
||||
bar.style.display = 'flex';
|
||||
// Show 6 most useful angles (skip High/Low for overlay — rare)
|
||||
const overlay = CAMERA_ANGLES.slice(0, 6);
|
||||
bar.innerHTML = overlay.map(a => {
|
||||
const pSafe = a.prompt.replace(/'/g, "\\'");
|
||||
return `<button onclick="submitSingleAngle('${pSafe}')" title="${escHtml(a.name)}">${a.icon}</button>`;
|
||||
}).join('');
|
||||
}
|
||||
|
||||
function updateSbGenBtn() {
|
||||
const btn = document.getElementById('sbGenBtn');
|
||||
if (!btn) return;
|
||||
const hasPrompt = (document.getElementById('sbGenPromptInput')?.value || '').trim().length > 0;
|
||||
btn.disabled = _fsSelectedPoses.size === 0 && !hasPrompt;
|
||||
const n = _fsSelectedPoses.size + (hasPrompt ? 1 : 0);
|
||||
btn.textContent = n > 1 ? `Generate (${n})` : 'Generate';
|
||||
const total = _fsSelectedPoses.size + _sbSelectedAngles.size + (hasPrompt ? 1 : 0);
|
||||
btn.disabled = total === 0;
|
||||
btn.textContent = total > 1 ? `Generate (${total})` : 'Generate';
|
||||
}
|
||||
|
||||
async function cancelSbGenerate() {
|
||||
@@ -3911,6 +4090,11 @@
|
||||
const orig = availablePoses[name]?.text ?? availablePoses[name] ?? name;
|
||||
prompts.push((edited !== undefined && edited.trim()) ? edited.trim() : String(orig));
|
||||
});
|
||||
// Camera angles — treated as prompts with no pose tag
|
||||
_sbSelectedAngles.forEach(name => {
|
||||
const ang = CAMERA_ANGLES.find(a => a.name === name);
|
||||
if (ang) { prompts.push(ang.prompt); poses.push(null); }
|
||||
});
|
||||
if (promptVal) { prompts.push(promptVal); poses.push(null); savePromptHistory(promptVal); }
|
||||
if (prompts.length === 0) return;
|
||||
const btn = document.getElementById('sbGenBtn');
|
||||
@@ -3927,7 +4111,9 @@
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
filenames: [_fsModelFilename], prompt: prompts,
|
||||
poses, seed: -1, max_area: 0, group_id: lbCurrentGid
|
||||
poses, seed: -1, max_area: 0, group_id: lbCurrentGid,
|
||||
wireframe_ref: _sbWireframeRef || null,
|
||||
wireframe_time: _sbWireframeTime,
|
||||
}),
|
||||
});
|
||||
if (!r.ok) {
|
||||
@@ -4656,7 +4842,7 @@
|
||||
const fname = lbNames[lbIdx];
|
||||
if (!fname) return;
|
||||
try {
|
||||
const r = await fetch(`/images/${encodeURIComponent(fname)}/set-preferred`, {method:'POST'});
|
||||
const r = await fetch(`${API}/images/${encodeURIComponent(fname)}/set-preferred`, {method:'POST'});
|
||||
if (!r.ok) return;
|
||||
fileSortOrders[fname] = 0;
|
||||
const url = lbUrls[lbIdx];
|
||||
@@ -4666,8 +4852,48 @@
|
||||
lbNames = [fname, ...otherNames];
|
||||
lbIdx = 0;
|
||||
updateStudio();
|
||||
// Fire-and-forget face extraction for face-book
|
||||
fetch(`${API}/images/${encodeURIComponent(fname)}/extract-face`, { method: 'POST' })
|
||||
.then(r => r.ok ? r.json() : null)
|
||||
.then(d => { if (d?.status === 'queued') showToast('Extracting face reference…', 'info'); })
|
||||
.catch(() => {});
|
||||
} catch(e) { console.error(e); }
|
||||
}
|
||||
|
||||
function showPasteGroupDialog(files) {
|
||||
document.getElementById('pasteDialog')?.remove();
|
||||
const gName = (lbCurrentGid && groupNames[lbCurrentGid]) || lbCurrentGid || 'current group';
|
||||
const gidSafe = escHtml(lbCurrentGid || '');
|
||||
const gNameSafe = escHtml(gName);
|
||||
const d = document.createElement('div');
|
||||
d.id = 'pasteDialog';
|
||||
d.style.cssText = 'position:fixed;top:50%;left:50%;transform:translate(-50%,-50%);background:#1a1a1a;border:1px solid #333;border-radius:8px;padding:20px;z-index:600;min-width:280px;text-align:center;box-shadow:0 8px 32px rgba(0,0,0,0.7)';
|
||||
d.innerHTML = `<div style="font-size:13px;color:#ccc;margin-bottom:14px">Add ${files.length} image(s)?</div>
|
||||
<div style="display:flex;flex-direction:column;gap:8px">
|
||||
<button class="btn primary" onclick="handleUploadToGroup(window._pasteFiles,'${gidSafe}');document.getElementById('pasteDialog')?.remove()">Add to “${gNameSafe}” — no poses</button>
|
||||
<button class="btn" onclick="handleUpload(window._pasteFiles);document.getElementById('pasteDialog')?.remove()">New group (run poses)</button>
|
||||
<button class="btn" onclick="document.getElementById('pasteDialog')?.remove()">Cancel</button>
|
||||
</div>`;
|
||||
window._pasteFiles = files;
|
||||
document.body.appendChild(d);
|
||||
setTimeout(() => document.getElementById('pasteDialog')?.remove(), 12000);
|
||||
}
|
||||
|
||||
async function handleUploadToGroup(files, groupId) {
|
||||
showToast(`Adding ${files.length} image(s) to group…`);
|
||||
for (const file of files) {
|
||||
const fd = new FormData();
|
||||
fd.append('image', file);
|
||||
fd.append('group_id', groupId);
|
||||
fd.append('skip_poses', 'true');
|
||||
try {
|
||||
const r = await fetch(`${API}/upload`, { method: 'POST', body: fd });
|
||||
if (!r.ok) showToast('Upload failed: ' + await r.text(), 'error');
|
||||
} catch (e) { showToast('Upload error: ' + e, 'error'); }
|
||||
}
|
||||
showToast('Added to group', 'success');
|
||||
refreshNow();
|
||||
}
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -25,5 +25,6 @@
|
||||
"facefusion_dir": "~/facefusion",
|
||||
"facefusion_venv": "~/facefusion-venv",
|
||||
"sam2_checkpoint": "~/.sam/sam2.1_hiera_base_plus.pt",
|
||||
"sam2_config": "configs/sam2.1/sam2.1_hiera_b+.yaml"
|
||||
"sam2_config": "configs/sam2.1/sam2.1_hiera_b+.yaml",
|
||||
"bg_removal": "sam2"
|
||||
}
|
||||
@@ -839,10 +839,63 @@ def _run_pipeline(
|
||||
}
|
||||
graph[NODE_POSITIVE]["inputs"][img_key] = [node_id, 0]
|
||||
|
||||
# Transparency detection
|
||||
is_transparent = any(kw in prompt.lower() for kw in ["transparent", "no background", "remove background", "alpha channel"])
|
||||
# ── background-removal routing ────────────────────────────────────────────
|
||||
# Two configurable strategies (config.json key "bg_removal"):
|
||||
#
|
||||
# "rembg" (default) — strip transparent keyword → Qwen renders a natural
|
||||
# scene → rembg (U2Net) separates person from any complex background.
|
||||
#
|
||||
# "sam2" — replace transparent keyword with "black background" → Qwen
|
||||
# renders a solid black BG → SAM2 bbox segmentation on a black image
|
||||
# works perfectly because the contrast is maximal.
|
||||
#
|
||||
# Either way, explicit "black background" in the prompt always routes to
|
||||
# SAM2 (the user already set up the ideal SAM2 input).
|
||||
# ─────────────────────────────────────────────────────────────────────────
|
||||
_TRANSPARENT_KWS = ["transparent background", "no background",
|
||||
"remove background", "alpha channel"]
|
||||
_BLACK_BG_KWS = ["black background"]
|
||||
|
||||
with open(CONFIG_PATH) as _cf:
|
||||
_bg_conf = json.load(_cf)
|
||||
bg_method = _bg_conf.get("bg_removal", "rembg") # "rembg" | "sam2"
|
||||
|
||||
is_transparent = any(kw in prompt.lower() for kw in _TRANSPARENT_KWS)
|
||||
is_black_bg = any(kw in prompt.lower() for kw in _BLACK_BG_KWS)
|
||||
post_process = None # "rembg" | "sam2"
|
||||
|
||||
if is_transparent:
|
||||
graph[NODE_NEGATIVE]["inputs"]["prompt"] = "checkerboard, grid, pattern, texture, background details, watermark, deformed anatomy"
|
||||
if bg_method == "sam2":
|
||||
# Swap "transparent background" → "black background" so Qwen renders
|
||||
# a pure-black BG that SAM2 can segment with maximal contrast.
|
||||
cleaned = prompt
|
||||
for kw in _TRANSPARENT_KWS:
|
||||
cleaned = re.sub(re.escape(kw), "black background", cleaned, flags=re.IGNORECASE)
|
||||
# Collapse duplicates if multiple keywords matched
|
||||
cleaned = re.sub(r"(?i)(black background[\s,]*){2,}", "black background, ", cleaned)
|
||||
cleaned = re.sub(r",\s*,", ",", cleaned).strip(", ")
|
||||
graph[NODE_POSITIVE]["inputs"]["prompt"] = cleaned
|
||||
graph[NODE_NEGATIVE]["inputs"]["prompt"] = (
|
||||
"real background, outdoor scene, indoor scene, gradient, "
|
||||
"colored background, watermark, deformed anatomy"
|
||||
)
|
||||
post_process = "sam2"
|
||||
else:
|
||||
# Strip the keyword so Qwen renders a natural scene; rembg handles
|
||||
# any background complexity reliably.
|
||||
cleaned = prompt
|
||||
for kw in _TRANSPARENT_KWS:
|
||||
cleaned = re.sub(re.escape(kw), "", cleaned, flags=re.IGNORECASE)
|
||||
cleaned = re.sub(r",\s*,", ",", cleaned)
|
||||
cleaned = re.sub(r",\s*$", "", cleaned.strip()).strip(", ")
|
||||
graph[NODE_POSITIVE]["inputs"]["prompt"] = cleaned
|
||||
graph[NODE_NEGATIVE]["inputs"]["prompt"] = "deformed anatomy, watermark, logo"
|
||||
post_process = "rembg"
|
||||
|
||||
elif is_black_bg:
|
||||
# Prompt already specifies a black background — ideal SAM2 input.
|
||||
# Route to SAM2 regardless of the configured bg_removal method.
|
||||
post_process = "sam2"
|
||||
|
||||
graph[NODE_LATENT]["inputs"]["width"] = w
|
||||
graph[NODE_LATENT]["inputs"]["height"] = h
|
||||
@@ -853,7 +906,13 @@ def _run_pipeline(
|
||||
outputs = _comfy_wait(prompt_id, time.time() + GEN_TIMEOUT)
|
||||
png_bytes = _comfy_fetch_image(outputs)
|
||||
|
||||
if is_transparent:
|
||||
if post_process == "sam2":
|
||||
# Input has a black background (Qwen was told "black background").
|
||||
# Use threshold-derived bbox so SAM2 gets a person-specific hint
|
||||
# rather than the full frame — full-frame bbox inverts the mask on
|
||||
# black-bg images because the large dark region scores higher.
|
||||
png_bytes = _apply_transparency_black_bg(png_bytes)
|
||||
elif post_process == "rembg":
|
||||
png_bytes = _apply_transparency(png_bytes)
|
||||
|
||||
return png_bytes
|
||||
@@ -891,7 +950,8 @@ def _move_to_trash(filepath: str):
|
||||
|
||||
|
||||
def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
|
||||
seed: int, max_area: int, group_id: str | None = None):
|
||||
seed: int, max_area: int, group_id: str | None = None,
|
||||
wireframe_ref: str | None = None, wireframe_time: float = 0.5):
|
||||
output_dir = _load_output_dir()
|
||||
for fname in filenames:
|
||||
actual_gid = None
|
||||
@@ -915,6 +975,24 @@ def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
|
||||
|
||||
try:
|
||||
base_pil = Image.open(fpath).convert("RGB")
|
||||
|
||||
# Extract wireframe pose reference frame once per filename
|
||||
pose_guide_pil = None
|
||||
if wireframe_ref:
|
||||
try:
|
||||
wf_path = os.path.join(_load_wireframe_dir(), wireframe_ref)
|
||||
cap = cv2.VideoCapture(wf_path)
|
||||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
target_frame = max(0, min(total_frames - 1, int(total_frames * wireframe_time)))
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame)
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
if ret:
|
||||
pose_guide_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
print(f"[batch] using wireframe {wireframe_ref} frame {target_frame}/{total_frames}")
|
||||
except Exception as wf_err:
|
||||
print(f"[batch] wireframe extract error: {wf_err}")
|
||||
|
||||
for prompt, pose in zip(prompts, poses):
|
||||
if jobs[job_id].get("cancelled"):
|
||||
return
|
||||
@@ -924,7 +1002,8 @@ def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
|
||||
if pose and pose.lower().strip() in ROTATE_180_POSES:
|
||||
pil = pil.rotate(180)
|
||||
|
||||
png = _run_pipeline(pil, prompt, seed, max_area)
|
||||
extra_imgs = [pose_guide_pil] if pose_guide_pil else None
|
||||
png = _run_pipeline(pil, prompt, seed, max_area, extra_images=extra_imgs)
|
||||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||||
clean_fname = naming.get_base_name(fname)
|
||||
out_name = f"{ts}_{clean_fname}"
|
||||
@@ -1069,6 +1148,8 @@ class BatchRequest(BaseModel):
|
||||
max_area: int = 0
|
||||
group_id: str | None = None
|
||||
poses: list[str | None] | None = None # pose name per prompt (same index), or None; None entries = no pose
|
||||
wireframe_ref: str | None = None # wireframe video name to use as pose guide (image2 slot)
|
||||
wireframe_time: float = 0.5 # normalized time (0–1) to extract the pose frame from
|
||||
|
||||
|
||||
@app.post("/batch")
|
||||
@@ -1085,6 +1166,7 @@ def start_batch(req: BatchRequest):
|
||||
t = threading.Thread(
|
||||
target=_batch_worker,
|
||||
args=(job_id, req.filenames, prompts, poses, req.seed, req.max_area, req.group_id),
|
||||
kwargs={"wireframe_ref": req.wireframe_ref, "wireframe_time": req.wireframe_time},
|
||||
daemon=True,
|
||||
)
|
||||
t.start()
|
||||
@@ -1207,6 +1289,35 @@ def list_videos():
|
||||
return {"videos": videos, "wireframe_dir": wireframe_dir}
|
||||
|
||||
|
||||
@app.get("/wireframe/frame/{video_name}")
|
||||
def wireframe_frame(video_name: str, t: float = 0.5):
|
||||
"""Extract a single frame at normalized time t (0–1) from a wireframe video. Returns PNG."""
|
||||
wireframe_dir = _load_wireframe_dir()
|
||||
video_path = os.path.join(wireframe_dir, video_name)
|
||||
if not os.path.exists(video_path):
|
||||
raise HTTPException(404, f"Video not found: {video_name}")
|
||||
try:
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
target = max(0, min(total - 1, int(total * max(0.0, min(1.0, t)))))
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, target)
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
if not ret:
|
||||
raise HTTPException(500, "Could not read frame")
|
||||
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
pil = Image.fromarray(rgb)
|
||||
buf = io.BytesIO()
|
||||
pil.save(buf, format="PNG")
|
||||
buf.seek(0)
|
||||
from fastapi.responses import Response
|
||||
return Response(content=buf.getvalue(), media_type="image/png")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(500, f"Frame extraction error: {e}")
|
||||
|
||||
|
||||
@app.get("/wireframe/duration/{video_name}")
|
||||
def wireframe_duration(video_name: str):
|
||||
"""Return duration (seconds) of a wireframe video via ffprobe."""
|
||||
@@ -1573,6 +1684,42 @@ def _crop_to_bbox(pil_img: Image.Image, margin: int = 20, top_margin: int = 20,
|
||||
return cropped
|
||||
|
||||
|
||||
def _extract_face_bg(filename: str, fpath: str):
|
||||
"""Background task: detect largest face, crop with padding, save as {group_id}_face.png."""
|
||||
try:
|
||||
app_fa, _ = _load_faceswapper()
|
||||
bgr = cv2.imread(fpath)
|
||||
if bgr is None:
|
||||
print(f"[extract-face] cannot read {fpath}")
|
||||
return
|
||||
faces = app_fa.get(bgr)
|
||||
if not faces:
|
||||
print(f"[extract-face] no face detected in {filename}")
|
||||
return
|
||||
face = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
|
||||
x1, y1, x2, y2 = [int(v) for v in face.bbox]
|
||||
h, w = bgr.shape[:2]
|
||||
pad = int((y2 - y1) * 0.5)
|
||||
x1 = max(0, x1 - pad)
|
||||
y1 = max(0, y1 - pad * 2) # extra headroom above face
|
||||
x2 = min(w, x2 + pad)
|
||||
y2 = min(h, y2 + int(pad * 0.3))
|
||||
pil = Image.open(fpath).convert("RGBA")
|
||||
cropped = pil.crop((x1, y1, x2, y2))
|
||||
person = database.get_person(filename)
|
||||
group_id = person[1] if person else None
|
||||
gid_tag = (group_id or "face").replace("/", "_")
|
||||
face_fname = f"{gid_tag}_face.png"
|
||||
face_path = os.path.join(os.path.dirname(fpath), face_fname)
|
||||
cropped.save(face_path)
|
||||
database.upsert_person(face_fname, filepath=face_path, group_id=group_id,
|
||||
name=person[0] if person else None,
|
||||
source_refs=json.dumps([filename]))
|
||||
print(f"[extract-face] saved {face_fname}")
|
||||
except Exception as e:
|
||||
print(f"[extract-face] error for {filename}: {e}")
|
||||
|
||||
|
||||
def _process_upload(file_path: str, filename: str, prompts: list[str], name: str | None = None, group_id: str | None = None):
|
||||
output_dir = _load_output_dir()
|
||||
try:
|
||||
@@ -1631,40 +1778,49 @@ def upload_image(
|
||||
image: UploadFile = File(...),
|
||||
prompts: str = Form(""),
|
||||
name: str = Form(None),
|
||||
group_id: str = Form(None), # optional: add to existing group
|
||||
skip_poses: bool = Form(False), # optional: skip base_prompts generation
|
||||
):
|
||||
# Load config to get output_dir (we use output_dir for UI uploads to avoid watcher conflict)
|
||||
with open(CONFIG_PATH, "r") as f:
|
||||
conf = json.load(f)
|
||||
output_dir = _load_output_dir()
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
|
||||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||||
safe_filename = re.sub(r'[^a-zA-Z0-9_.-]', '_', image.filename)
|
||||
safe_filename = re.sub(r'[^a-zA-Z0-9_.-]', '_', image.filename or "paste")
|
||||
# Ensure extension
|
||||
if not safe_filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
|
||||
safe_filename += ".png"
|
||||
|
||||
|
||||
filename = f"{ts}_{safe_filename}"
|
||||
file_path = os.path.join(output_dir, filename)
|
||||
|
||||
|
||||
with open(file_path, "wb") as f:
|
||||
shutil.copyfileobj(image.file, f)
|
||||
|
||||
|
||||
# Fast path: add to existing group without pose generation
|
||||
if group_id and skip_poses:
|
||||
sort_order = database.get_next_sort_order(group_id)
|
||||
database.upsert_person(filename, filepath=file_path, group_id=group_id,
|
||||
sort_order=sort_order)
|
||||
return {"status": "added", "filename": filename, "group_id": group_id}
|
||||
|
||||
prompt_list = [p.strip() for p in prompts.split(",") if p.strip()]
|
||||
|
||||
|
||||
# Add base-set prompts if defined in config
|
||||
base_prompts = conf.get("base_prompts", [])
|
||||
if isinstance(base_prompts, list):
|
||||
prompt_list.extend(base_prompts)
|
||||
|
||||
|
||||
if not prompt_list:
|
||||
# Use default prompt from config
|
||||
prompt_list = [conf.get("prompt", "high quality, masterpiece")]
|
||||
|
||||
group_id = f"up_{uuid.uuid4().hex[:8]}" # unique per upload; avoids collisions when pasting generic filenames
|
||||
background_tasks.add_task(_process_upload, file_path, filename, prompt_list, name, group_id)
|
||||
|
||||
return {"status": "processing", "filename": filename, "group_id": group_id, "prompts": prompt_list}
|
||||
|
||||
effective_gid = group_id or f"up_{uuid.uuid4().hex[:8]}"
|
||||
background_tasks.add_task(_process_upload, file_path, filename, prompt_list, name, effective_gid)
|
||||
|
||||
return {"status": "processing", "filename": filename, "group_id": effective_gid, "prompts": prompt_list}
|
||||
|
||||
|
||||
@app.post("/edit")
|
||||
@@ -1717,7 +1873,7 @@ def unarchive_image(filename: str):
|
||||
|
||||
|
||||
@app.post("/images/{filename}/set-preferred")
|
||||
def set_image_preferred(filename: str):
|
||||
def set_image_preferred(filename: str, background_tasks: BackgroundTasks):
|
||||
"""Make this image sort_order=0 within its group, shifting others to 1,2,..."""
|
||||
person = database.get_person(filename)
|
||||
if not person:
|
||||
@@ -1728,9 +1884,23 @@ def set_image_preferred(filename: str):
|
||||
rows = database.get_group_order(group_id)
|
||||
others = [r[0] for r in rows if r[0] != filename]
|
||||
database.set_group_order(group_id, [filename] + others)
|
||||
fpath = os.path.join(_load_output_dir(), filename)
|
||||
if os.path.exists(fpath):
|
||||
background_tasks.add_task(_extract_face_bg, filename, fpath)
|
||||
return {"filename": filename, "group_id": group_id}
|
||||
|
||||
|
||||
@app.post("/images/{filename}/extract-face")
|
||||
def extract_face_endpoint(filename: str, background_tasks: BackgroundTasks):
|
||||
"""Detect and crop the largest face from image; saves as {group_id}_face.png."""
|
||||
output_dir = _load_output_dir()
|
||||
fpath = os.path.join(output_dir, filename)
|
||||
if not os.path.exists(fpath):
|
||||
raise HTTPException(404, "not found")
|
||||
background_tasks.add_task(_extract_face_bg, filename, fpath)
|
||||
return {"status": "queued", "filename": filename}
|
||||
|
||||
|
||||
@app.post("/images/{filename}/undress")
|
||||
def undress_image(filename: str, background_tasks: BackgroundTasks):
|
||||
"""Queue a generation using the undress prompt on the given image."""
|
||||
@@ -1967,14 +2137,14 @@ def _load_sam2():
|
||||
|
||||
|
||||
def _apply_transparency_sam2(png_bytes: bytes) -> bytes:
|
||||
"""Remove background with SAM2 bbox-based segmentation; fallback to rembg.
|
||||
"""Remove background with SAM2 bbox segmentation; fallback to rembg.
|
||||
|
||||
Mimics the reference approach (bbox → SAM2) but without an external
|
||||
detector: we pass a generous bbox covering the central subject area.
|
||||
For portraits/full-body shots the subject fills most of the frame, so a
|
||||
5 % margin bbox reliably captures hair, glasses, and sandals without the
|
||||
point-prompt clipping issues. multimask_output=True lets SAM2 propose
|
||||
three masks; we pick the highest-scoring one.
|
||||
Uses a near-full-frame bbox so SAM2 finds the largest foreground object
|
||||
(the person) regardless of rotation or pose. This works well because
|
||||
"transparent background" is stripped from the Qwen prompt upstream, so the
|
||||
model renders a solid real background — giving SAM2 clear contrast to work
|
||||
with. Point prompts were tried but produced holes in ¾-rotated poses
|
||||
because the spine-column seeds land on background when the body is offset.
|
||||
"""
|
||||
predictor = _load_sam2()
|
||||
if predictor is False:
|
||||
@@ -1985,31 +2155,133 @@ def _apply_transparency_sam2(png_bytes: bytes) -> bytes:
|
||||
img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
|
||||
arr = np.array(img)
|
||||
h, w = arr.shape[:2]
|
||||
# Generous bbox — 5 % margin H, 2 % margin V — covers the whole subject
|
||||
box = np.array([[int(w * 0.05), int(h * 0.02),
|
||||
int(w * 0.95), int(h * 0.98)]], dtype=np.float32)
|
||||
|
||||
# Near-full-frame bbox — 1 % margin so hair / shoes are inside the hint.
|
||||
# SAM2 treats this as "find the prominent object within this region".
|
||||
box = np.array([[int(w * 0.01), int(h * 0.01),
|
||||
int(w * 0.99), int(h * 0.99)]], dtype=np.float32)
|
||||
|
||||
with torch.inference_mode():
|
||||
predictor.set_image(arr)
|
||||
masks, scores, _ = predictor.predict(
|
||||
box=box,
|
||||
multimask_output=True,
|
||||
)
|
||||
|
||||
if masks is None or len(masks) == 0:
|
||||
print("[sam2] no masks returned, falling back to rembg")
|
||||
return _apply_transparency(png_bytes)
|
||||
|
||||
best = masks[int(np.argmax(scores))]
|
||||
|
||||
# Sanity check: a person should cover 5 %–92 % of the frame
|
||||
coverage = float(best.sum()) / (h * w)
|
||||
if coverage < 0.05 or coverage > 0.92:
|
||||
print(f"[sam2] mask coverage {coverage:.1%} out of range, falling back to rembg")
|
||||
return _apply_transparency(png_bytes)
|
||||
|
||||
mask_np = best.astype(np.uint8) * 255
|
||||
|
||||
# Soft anti-aliased edge (radius 1 keeps accessory detail)
|
||||
try:
|
||||
from PIL import ImageFilter
|
||||
alpha_img = Image.fromarray(mask_np, mode="L")
|
||||
alpha_img = alpha_img.filter(ImageFilter.GaussianBlur(radius=1))
|
||||
except Exception:
|
||||
alpha_img = Image.fromarray(mask_np, mode="L")
|
||||
|
||||
rgba = img.convert("RGBA")
|
||||
r, g, b, _ = rgba.split()
|
||||
alpha = Image.fromarray(mask_np, mode="L")
|
||||
out = Image.merge("RGBA", (r, g, b, alpha))
|
||||
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
||||
buf = io.BytesIO()
|
||||
out.save(buf, format="PNG")
|
||||
print(f"[sam2] mask OK ({coverage:.1%} coverage)")
|
||||
return buf.getvalue()
|
||||
except Exception as e:
|
||||
print(f"[sam2] inference error, falling back to rembg: {e}")
|
||||
return _apply_transparency(png_bytes)
|
||||
|
||||
|
||||
def _apply_transparency_black_bg(png_bytes: bytes) -> bytes:
|
||||
"""Background removal for black-background Qwen output (bg_removal=sam2 mode).
|
||||
|
||||
Strategy:
|
||||
1. Threshold: any pixel with max-channel > 25 is person (non-black).
|
||||
This correctly identifies the subject regardless of pose or rotation.
|
||||
2. Derive a tight person bounding-box from the threshold mask.
|
||||
3. Run SAM2 with that box for sub-pixel edge refinement.
|
||||
Accept SAM2 result only when its coverage is close (±30 pp) to the
|
||||
threshold estimate — this rejects the inverted-mask failure mode where
|
||||
SAM2 picks the large dark region as the "object".
|
||||
4. Fall back to the threshold mask (Gaussian-blurred edges) if SAM2
|
||||
is unavailable, errors, or diverges.
|
||||
|
||||
Do NOT use the full-frame bbox here: on black-background images the large
|
||||
dark region scores higher than the person, causing SAM2 to invert the mask.
|
||||
"""
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import ImageFilter
|
||||
|
||||
img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
|
||||
arr = np.array(img)
|
||||
h, w = arr.shape[:2]
|
||||
|
||||
# Step 1 — threshold: non-black pixels are the person
|
||||
is_person = np.max(arr, axis=2) > 25
|
||||
thresh_cov = float(is_person.sum()) / (h * w)
|
||||
print(f"[bg-black] threshold coverage: {thresh_cov:.1%}")
|
||||
|
||||
if not is_person.any():
|
||||
print("[bg-black] all-black image — falling back to rembg")
|
||||
return _apply_transparency(png_bytes)
|
||||
|
||||
# Step 2 — tight bounding box from threshold
|
||||
rows = np.any(is_person, axis=1)
|
||||
cols = np.any(is_person, axis=0)
|
||||
rmin = int(np.where(rows)[0][0]); rmax = int(np.where(rows)[0][-1])
|
||||
cmin = int(np.where(cols)[0][0]); cmax = int(np.where(cols)[0][-1])
|
||||
margin = int(min(h, w) * 0.02)
|
||||
x1 = max(0, cmin - margin); y1 = max(0, rmin - margin)
|
||||
x2 = min(w, cmax + margin); y2 = min(h, rmax + margin)
|
||||
|
||||
# Step 3 — SAM2 with the person-specific bbox
|
||||
predictor = _load_sam2()
|
||||
if predictor is not False:
|
||||
box = np.array([[x1, y1, x2, y2]], dtype=np.float32)
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
predictor.set_image(arr)
|
||||
masks, scores, _ = predictor.predict(box=box, multimask_output=True)
|
||||
|
||||
if masks is not None and len(masks) > 0:
|
||||
best = masks[int(np.argmax(scores))]
|
||||
sam_cov = float(best.sum()) / (h * w)
|
||||
print(f"[bg-black] SAM2 coverage: {sam_cov:.1%}")
|
||||
|
||||
if 0.03 < sam_cov < 0.95 and abs(sam_cov - thresh_cov) < 0.30:
|
||||
mask_np = best.astype(np.uint8) * 255
|
||||
alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=1))
|
||||
rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
|
||||
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
||||
buf = io.BytesIO(); out.save(buf, "PNG")
|
||||
print(f"[bg-black] SAM2 accepted ✓")
|
||||
return buf.getvalue()
|
||||
else:
|
||||
print(f"[bg-black] SAM2 diverged ({sam_cov:.1%} vs {thresh_cov:.1%}) — threshold fallback")
|
||||
except Exception as e:
|
||||
print(f"[bg-black] SAM2 error: {e} — threshold fallback")
|
||||
|
||||
# Step 4 — fallback: threshold mask with soft edge blur
|
||||
print("[bg-black] using threshold mask")
|
||||
mask_np = is_person.astype(np.uint8) * 255
|
||||
alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=2))
|
||||
rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
|
||||
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
||||
buf = io.BytesIO(); out.save(buf, "PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
@app.post("/remove-background-sam/{filename}")
|
||||
def remove_background_sam(filename: str):
|
||||
"""SAM2-based background removal.
|
||||
|
||||
191
tour-comfy/test_transparency.py
Normal file
191
tour-comfy/test_transparency.py
Normal file
@@ -0,0 +1,191 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Validate background removal strategies.
|
||||
|
||||
Usage:
|
||||
python test_transparency.py [image.png ...]
|
||||
|
||||
Writes comparison files next to each input:
|
||||
*_rembg.png — pure rembg (bg_removal=rembg path)
|
||||
*_blackbg.png — simulated black-bg composite (what Qwen renders in sam2 mode)
|
||||
*_thresh.png — threshold mask only (non-black pixels → person)
|
||||
*_thresh_sam2.png — threshold bbox → SAM2 edge refinement (new sam2 mode path)
|
||||
"""
|
||||
import io, sys, os
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFilter
|
||||
|
||||
OUTPUT_DIR = "/mnt/zim/tour-comfy/output"
|
||||
VENV_SITE = "/home/mike/comfyui/venv/lib/python3.13/site-packages"
|
||||
SAM2_CKPT = os.path.expanduser("~/.sam/sam2.1_hiera_base_plus.pt")
|
||||
SAM2_CFG = "configs/sam2.1/sam2.1_hiera_b+.yaml"
|
||||
|
||||
|
||||
# ── rembg ──────────────────────────────────────────────────────────────────────
|
||||
def apply_rembg(png_bytes: bytes) -> bytes:
|
||||
from rembg import remove
|
||||
return remove(png_bytes)
|
||||
|
||||
|
||||
# ── SAM2 loader ────────────────────────────────────────────────────────────────
|
||||
_predictor = None
|
||||
def load_sam2():
|
||||
global _predictor
|
||||
if _predictor is not None:
|
||||
return _predictor
|
||||
try:
|
||||
import torch
|
||||
from sam2.build_sam import build_sam2
|
||||
from sam2.sam2_image_predictor import SAM2ImagePredictor
|
||||
model = build_sam2(SAM2_CFG, SAM2_CKPT, device="cuda")
|
||||
_predictor = SAM2ImagePredictor(model)
|
||||
print("[sam2] loaded")
|
||||
except Exception as e:
|
||||
print(f"[sam2] FAILED: {e}")
|
||||
_predictor = False
|
||||
return _predictor
|
||||
|
||||
|
||||
# ── Simulate black-bg Qwen output ─────────────────────────────────────────────
|
||||
def make_black_bg(png_bytes: bytes) -> bytes:
|
||||
"""Composite a rembg cutout onto pure black — simulates Qwen 'black background' output."""
|
||||
rgba = Image.open(io.BytesIO(apply_rembg(png_bytes))).convert("RGBA")
|
||||
bg = Image.new("RGBA", rgba.size, (0, 0, 0, 255))
|
||||
bg.paste(rgba, mask=rgba.split()[3])
|
||||
out = bg.convert("RGB")
|
||||
buf = io.BytesIO(); out.save(buf, "PNG"); return buf.getvalue()
|
||||
|
||||
|
||||
# ── Threshold-only mask ────────────────────────────────────────────────────────
|
||||
def apply_threshold_mask(png_bytes: bytes, threshold: int = 25) -> bytes:
|
||||
"""Find non-black pixels → person mask. No SAM2 needed."""
|
||||
img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
|
||||
arr = np.array(img)
|
||||
h, w = arr.shape[:2]
|
||||
|
||||
is_person = np.max(arr, axis=2) > threshold
|
||||
coverage = is_person.sum() / (h * w)
|
||||
print(f" [threshold] person coverage: {coverage:.1%}")
|
||||
|
||||
if not is_person.any():
|
||||
print(" [threshold] all-black image — no person found")
|
||||
return png_bytes
|
||||
|
||||
mask_np = is_person.astype(np.uint8) * 255
|
||||
alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=2))
|
||||
rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
|
||||
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
||||
buf = io.BytesIO(); out.save(buf, "PNG"); return buf.getvalue()
|
||||
|
||||
|
||||
# ── NEW: Threshold bbox → SAM2 refinement (sam2 mode path) ────────────────────
|
||||
def apply_thresh_sam2(png_bytes: bytes, threshold: int = 25) -> bytes:
|
||||
"""
|
||||
For black-background Qwen output:
|
||||
1. Threshold to find person bbox (non-black pixels)
|
||||
2. Run SAM2 with that tight bbox for clean edge refinement
|
||||
3. Fallback to threshold mask if SAM2 unavailable or mask looks wrong
|
||||
"""
|
||||
import torch
|
||||
|
||||
img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
|
||||
arr = np.array(img)
|
||||
h, w = arr.shape[:2]
|
||||
|
||||
# Step 1 — threshold
|
||||
is_person = np.max(arr, axis=2) > threshold
|
||||
thresh_cov = is_person.sum() / (h * w)
|
||||
print(f" [thresh_sam2] threshold person coverage: {thresh_cov:.1%}")
|
||||
|
||||
if not is_person.any():
|
||||
print(" [thresh_sam2] all-black — fallback to rembg")
|
||||
return apply_rembg(png_bytes)
|
||||
|
||||
rows = np.any(is_person, axis=1)
|
||||
cols = np.any(is_person, axis=0)
|
||||
rmin = int(np.where(rows)[0][0]); rmax = int(np.where(rows)[0][-1])
|
||||
cmin = int(np.where(cols)[0][0]); cmax = int(np.where(cols)[0][-1])
|
||||
|
||||
margin = int(min(h, w) * 0.02)
|
||||
y1 = max(0, rmin - margin); y2 = min(h, rmax + margin)
|
||||
x1 = max(0, cmin - margin); x2 = min(w, cmax + margin)
|
||||
print(f" [thresh_sam2] person bbox (+margin): ({x1},{y1})-({x2},{y2})")
|
||||
|
||||
# Step 2 — SAM2 with person-specific bbox
|
||||
predictor = load_sam2()
|
||||
if predictor is not False:
|
||||
box = np.array([[x1, y1, x2, y2]], dtype=np.float32)
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
predictor.set_image(arr)
|
||||
masks, scores, _ = predictor.predict(box=box, multimask_output=True)
|
||||
|
||||
if masks is not None and len(masks) > 0:
|
||||
best = masks[int(np.argmax(scores))]
|
||||
sam_cov = float(best.sum()) / (h * w)
|
||||
print(f" [thresh_sam2] SAM2 coverage: {sam_cov:.1%} (threshold was {thresh_cov:.1%})")
|
||||
|
||||
# Accept SAM2 result if coverage is within reasonable range of threshold
|
||||
if 0.03 < sam_cov < 0.95 and abs(sam_cov - thresh_cov) < 0.30:
|
||||
mask_np = best.astype(np.uint8) * 255
|
||||
alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=1))
|
||||
rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
|
||||
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
||||
buf = io.BytesIO(); out.save(buf, "PNG")
|
||||
print(" [thresh_sam2] SAM2 result accepted ✓")
|
||||
return buf.getvalue()
|
||||
else:
|
||||
print(f" [thresh_sam2] SAM2 coverage diverged from threshold — using threshold mask")
|
||||
except Exception as e:
|
||||
print(f" [thresh_sam2] SAM2 error: {e} — using threshold mask")
|
||||
else:
|
||||
print(" [thresh_sam2] SAM2 not available — using threshold mask")
|
||||
|
||||
# Step 3 — fallback: threshold mask with soft edges
|
||||
mask_np = is_person.astype(np.uint8) * 255
|
||||
alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=2))
|
||||
rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
|
||||
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
||||
buf = io.BytesIO(); out.save(buf, "PNG")
|
||||
print(" [thresh_sam2] threshold mask used as fallback")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
# ── main ───────────────────────────────────────────────────────────────────────
|
||||
if __name__ == "__main__":
|
||||
paths = sys.argv[1:] if len(sys.argv) > 1 else [
|
||||
os.path.join(OUTPUT_DIR, "20260622_181910_0_20260619_124038_image.png"),
|
||||
]
|
||||
|
||||
for path in paths:
|
||||
if not os.path.exists(path):
|
||||
print(f"SKIP (not found): {path}"); continue
|
||||
stem = os.path.splitext(path)[0]
|
||||
print(f"\n══ {os.path.basename(path)} ══")
|
||||
with open(path, "rb") as f:
|
||||
raw = f.read()
|
||||
|
||||
print("1. rembg (bg_removal=rembg path)...")
|
||||
rb = apply_rembg(raw)
|
||||
with open(stem + "_rembg.png", "wb") as f: f.write(rb)
|
||||
print(f" → {os.path.basename(stem)}_rembg.png")
|
||||
|
||||
print("2. Simulate black-bg Qwen output...")
|
||||
bb = make_black_bg(raw)
|
||||
with open(stem + "_blackbg.png", "wb") as f: f.write(bb)
|
||||
print(f" → {os.path.basename(stem)}_blackbg.png")
|
||||
|
||||
print("3. Threshold-only mask on black-bg image...")
|
||||
tm = apply_threshold_mask(bb)
|
||||
with open(stem + "_thresh.png", "wb") as f: f.write(tm)
|
||||
print(f" → {os.path.basename(stem)}_thresh.png")
|
||||
|
||||
print("4. Threshold bbox → SAM2 refinement on black-bg image (NEW sam2 mode path)...")
|
||||
ts = apply_thresh_sam2(bb)
|
||||
with open(stem + "_thresh_sam2.png", "wb") as f: f.write(ts)
|
||||
print(f" → {os.path.basename(stem)}_thresh_sam2.png")
|
||||
|
||||
print("\n── Done ──")
|
||||
print(" *_rembg.png rembg on real background (bg_removal=rembg path)")
|
||||
print(" *_thresh.png threshold-only on black bg")
|
||||
print(" *_thresh_sam2.png threshold-bbox → SAM2 on black bg (NEW sam2 mode path)")
|
||||
Reference in New Issue
Block a user