updates UI

This commit is contained in:
mike
2026-06-29 14:10:58 +02:00
parent 139785d108
commit 61268de34b
4 changed files with 394 additions and 65 deletions

View File

@@ -75,6 +75,9 @@ Scenery:
- make the faceswap preview mode fixed size instead of 50% we dont know how big the original clip is. just 470p or something.
- include 2 filter(s) in the filmstrip for VIDEO and ORBIT only show them if the sub-set is not empty.
- add a 2nd visible filmstrip to easy move items from one to another strip in the UI
- on the landing page we now just play all images of an group. However if the orbit generaition is complete. I want to circulate those images , (this basically is a new filmstrip feature)) so we set the orbit in a filmstrip by default, but can modify it. Also we want the face to be visible in the landing page and the most prominent (most 3d alike full body with femine descriptors)) image
- still looking for that killer-feminine shot, one image that would show most characteristics of a nude female including eyes.
-
## refine
- when refresh page, we lose track of current jobs running.
@@ -106,3 +109,17 @@ videos
2) 3ref image MERGE
3) 3ref straight on- head-on - looking viewer
ok cool,
Let do some work in the "Info" tab.
Pose information
<usefull info is here>
GAZE: -
BODY: -
SUBJECTS: -
TRACKING
1) we basically always see GAZE: - BODY: - etc.. (perhaps we only mis it passing thru in the frontend)
2) We would like to see information about the face, completeness of the anatomy (perhaps an indication on how much to outpaint to make it anatomically correct), camera/view oriention. Objects found in the image. We use this information later on to refine the prompts
3) we can move some actions (or duplicate) the actions into the "studioAngleBar". It already got included the Expand canvas icon. we can ad the icons for Pose, Duplicate, Crop...(auto-crop) and Crop
4) add a new video filter, and hide video by default in the Group Active
-- feel free to furhter optimize the studio view "info" tab

View File

@@ -2334,12 +2334,24 @@
// --- HYDRATION_START ---
const PRELOADED_IMAGES = [
"20260619_040244_5_20260619_040135_image.png",
"cg_f85e7ae1_face.png",
"20260618_181804_8_20260618_181600_image.png",
"b1.png_face.png",
"_turntable/b1.png/views/view_023_345deg.nobg.png",
"cg_e6ba3260_face.png",
"cg_4e2b6449_face.png",
"cg_ad73335e_face.png",
"cg_e964fb19_face.png",
"20260629_133926_6_20260621_202521_image.png",
"out4.png_face.png",
"20260629_132244_sc_0_20260619_124038_image.png",
"20260629_132101_sc_0_20260619_124038_image.png",
"20260629_131834_sc_1_20260619_123958_image.png",
"20260629_124948_sc_view_023_345deg.png",
"20260629_124523_sc_view_023_345deg.png",
"20260629_124433_sc_view_023_345deg.png",
"20260629_122907_fs_Step Sis will do anything to make me Delete this Videos_1633s-2269s_image.png_prev50.mp4",
"20260629_123127_test_upload_image.png",
"20260629_123101_test_upload_image.png",
"20260629_120459_fs_bitch insemination_dup_20260626_003125_dup_Pasted image (9).png.mp4",
"20260629_114451_sc_sc_image.nobg.png",
"20260629_114346_sc_sc_image.nobg.png",
@@ -5364,7 +5376,6 @@
"20260626_005504_Pasted image (9).png",
"20260626_005322_dup_20260626_003125_dup_Pasted image (9).png",
"20260626_003125_dup_Pasted image (9).png",
"cg_ad73335e_face.png",
"20260625_232346_sc_image.png",
"20260625_232242_sc_image.png",
"20260625_232204_sc_image.png",
@@ -6162,8 +6173,8 @@
"20260622_193411_1_20260622_193118_image.png",
"20260622_193354_0_20260622_193118_image.png",
"20260622_193118_image.png",
"20260622_212012_dup_20260622_181910_0_20260619_124038_image.png",
"20260622_181910_0_20260619_124038_image.png",
"20260622_212012_dup_20260622_181910_0_20260619_124038_image.png",
"20260622_181734_0_20260619_124038_image.png",
"20260622_181720_0_20260619_124038_image.png",
"20260622_181634_0_20260619_124038_image.png",
@@ -6955,7 +6966,6 @@
"20260619_040319_8_20260619_040135_image.png",
"20260619_040307_7_20260619_040135_image.png",
"20260619_040255_6_20260619_040135_image.png",
"20260619_040244_5_20260619_040135_image.png",
"20260619_040233_4_20260619_040135_image.png",
"20260619_040221_3_20260619_040135_image.png",
"20260619_040135_image.png",
@@ -7048,7 +7058,6 @@
"20260618_181839_4_20260618_181715_image.png",
"20260618_181827_3_20260618_181715_image.png",
"20260618_181815_2_20260618_181715_image.png",
"20260618_181804_8_20260618_181600_image.png",
"20260618_181752_1_20260618_181715_image.png",
"20260618_181742_7_20260618_181600_image.png",
"20260618_181730_0_20260618_181715_image.png",
@@ -8676,10 +8685,31 @@
if (refsDiv) {
if (refs && refs.length > 0) {
refsDiv.style.display = 'flex';
refsDiv.innerHTML = refs.map(r =>
`<img src="${IMAGE_FOLDER + r}" loading="lazy" title="${r}"
onclick="jumpToImage('${r.replace(/'/g,"\\'")}');">`
).join('');
refsDiv.innerHTML = refs.map(r => {
let srcUrl = '';
let isGallery = true;
let titleText = r;
if (r.startsWith('video:')) {
const vidName = r.substring(6);
srcUrl = `${API}/wireframe/frame/${encodeURIComponent(vidName)}?t=0`;
isGallery = false;
titleText = `Scene Video: ${vidName}`;
} else if (r.startsWith('wireframe:')) {
const wfName = r.substring(10);
srcUrl = `${API}/wireframe/${encodeURIComponent(wfName)}`;
isGallery = false;
titleText = `Scene Image: ${wfName}`;
} else {
srcUrl = IMAGE_FOLDER + r;
}
const clickAction = isGallery
? `jumpToImage('${r.replace(/'/g,"\\'")}');`
: `showToast('${titleText.replace(/'/g,"\\'")}','info');`;
return `<img src="${srcUrl}" loading="lazy" decoding="async" title="${escHtml(titleText)}"
onerror="this.style.opacity=.3"
onclick="${clickAction}">`;
}).join('');
} else {
refsDiv.style.display = 'none';
}
@@ -11712,12 +11742,21 @@
}
}
if (e.altKey && e.key === 'Enter') {
if (document.getElementById('studio').classList.contains('open') && _activeSidebarTab === 'generate') {
const btn = document.getElementById('sbGenBtn');
if (btn && !btn.disabled) {
btn.click();
e.preventDefault();
return;
if (document.getElementById('studio').classList.contains('open')) {
if (_activeSidebarTab === 'generate') {
const btn = document.getElementById('sbGenBtn');
if (btn && !btn.disabled) {
btn.click();
e.preventDefault();
return;
}
} else if (_activeSidebarTab === 'scenery') {
const btn = document.getElementById('sceneGenBtn');
if (btn && !btn.disabled) {
btn.click();
e.preventDefault();
return;
}
}
}
}
@@ -13353,24 +13392,27 @@
// --- Selected video: player + slider + capture (always at top) ---
if (_sceneVideo) {
if (!gridOpen) {
html += `
<div class="scene-frame-preview" id="sceneFramePreview" style="display:block">
<video id="sceneVideoEl" muted preload="auto" playsinline
style="${_sceneFrameBytes?'display:none':'display:block'};width:100%;height:100%;object-fit:contain"></video>
<img id="sceneFrameImg" style="${_sceneFrameBytes?'display:block':'display:none'};width:100%;height:100%;object-fit:contain"
${_sceneFrameBytes?`src="data:image/png;base64,${_sceneFrameBytes}"`:''}/>
<div class="scene-captured-badge" id="sceneCapturedBadge" style="${_sceneFrameBytes?'display:block':'display:none'}">✓ frame captured</div>
</div>
<input type="range" class="scene-scrubber" id="sceneScrubber" min="0" max="1000" value="0"
oninput="sceneScrub(this.value)" ${_sceneFrameBytes?'disabled':''}>
<div style="display:flex;justify-content:space-between;align-items:center;margin-bottom:8px">
<span id="sceneTimeLabel" style="font-size:11px;color:#888">0:00.0</span>
<span style="font-size:10px;color:#555">${escHtml(_sceneVideo.replace(/\.[^.]+$/,''))}</span>
</div>
<button class="sb-btn primary" id="sceneCaptureBtn" onclick="sceneExtractFrame()"
style="${_sceneFrameBytes?'display:none':'display:block'};width:100%;margin-bottom:8px">📷 Capture this frame</button>
<button class="sb-btn" id="sceneReleaseBtn" onclick="sceneReleaseFrame()"
style="${_sceneFrameBytes?'display:block':'display:none'};width:100%;margin-bottom:8px">✖ Pick a different frame</button>`;
}
html += `
<div class="scene-frame-preview" id="sceneFramePreview" style="display:block">
<video id="sceneVideoEl" muted preload="auto" playsinline
style="${_sceneFrameBytes?'display:none':'display:block'};width:100%;height:100%;object-fit:contain"></video>
<img id="sceneFrameImg" style="${_sceneFrameBytes?'display:block':'display:none'};width:100%;height:100%;object-fit:contain"
${_sceneFrameBytes?`src="data:image/png;base64,${_sceneFrameBytes}"`:''}/>
<div class="scene-captured-badge" id="sceneCapturedBadge" style="${_sceneFrameBytes?'display:block':'display:none'}">✓ frame captured</div>
</div>
<input type="range" class="scene-scrubber" id="sceneScrubber" min="0" max="1000" value="0"
oninput="sceneScrub(this.value)" ${_sceneFrameBytes?'disabled':''}>
<div style="display:flex;justify-content:space-between;align-items:center;margin-bottom:8px">
<span id="sceneTimeLabel" style="font-size:11px;color:#888">0:00.0</span>
<span style="font-size:10px;color:#555">${escHtml(_sceneVideo.replace(/\.[^.]+$/,''))}</span>
</div>
<button class="sb-btn primary" id="sceneCaptureBtn" onclick="sceneExtractFrame()"
style="${_sceneFrameBytes?'display:none':'display:block'};width:100%;margin-bottom:8px">📷 Capture this frame</button>
<button class="sb-btn" id="sceneReleaseBtn" onclick="sceneReleaseFrame()"
style="${_sceneFrameBytes?'display:block':'display:none'};width:100%;margin-bottom:8px">✖ Pick a different frame</button>
<button class="sb-btn" onclick="sceneToggleGrid()" style="width:100%;margin-bottom:10px">
${gridOpen?'▲ Hide video list':'▼ Change background video'}</button>`;
}
@@ -13475,7 +13517,7 @@
const borderStyle = isSelected ? 'border:2px solid #8b5cf6' : 'border:1px solid #27272a';
html += `
<div style="cursor:pointer;position:relative" title="${escHtml(frame)}" onclick="sceneSelectWireframeImage('${fSafe}')">
<img src="${API}/wireframe/${encodeURIComponent(frame)}" style="width:40px;height:40px;object-fit:cover;border-radius:4px;${borderStyle}">
<img src="${API}/wireframe/${encodeURIComponent(frame)}" loading="lazy" decoding="async" style="width:40px;height:40px;object-fit:cover;border-radius:4px;${borderStyle}" onerror="this.style.opacity=.3">
</div>`;
});
@@ -13512,7 +13554,7 @@
const borderStyle = isSelected ? 'border:2px solid #2563eb' : 'border:1px solid #333';
html += `
<div style="cursor:pointer;position:relative" title="${escHtml(img)}" onclick="sceneSelectWireframeImage('${imgSafe}')">
<img src="${API}/wireframe/${encodeURIComponent(img)}" style="width:48px;height:48px;object-fit:cover;border-radius:4px;${borderStyle}">
<img src="${API}/wireframe/${encodeURIComponent(img)}" loading="lazy" decoding="async" style="width:48px;height:48px;object-fit:cover;border-radius:4px;${borderStyle}" onerror="this.style.opacity=.3">
</div>`;
});
html += '</div>';
@@ -13608,26 +13650,51 @@
} else if (!_sceneLib.groups.length && !_sceneLib.ungrouped.length) {
html += `<div style="font-size:11px;color:#555;padding:4px 0">No previous scenery found.</div>`;
} else {
const renderLibItems = (items) => items.map(it => {
const url = IMAGE_FOLDER + encodeURIComponent(it.filename);
const gLabel = it.group_id ? ` <span style="color:#555">${escHtml(it.group_id)}</span>` : '';
return `<div style="display:inline-block;margin:2px;cursor:pointer;position:relative;vertical-align:top"
title="Use as background · ${escHtml(it.filename)}${it.group_id?' ('+escHtml(it.group_id)+')':''}"
onclick="sceneLibUseAsBackground(${JSON.stringify(it.filename)})">
<img src="${url}" style="width:52px;height:52px;object-fit:cover;border-radius:4px;border:1px solid #333">
</div>`;
}).join('');
if (!window._expandedSceneryGroups) {
window._expandedSceneryGroups = {};
}
const renderLibGroup = (groupKey, items) => {
const isExpanded = !!window._expandedSceneryGroups[groupKey];
const showCount = isExpanded ? items.length : 12;
const visibleItems = items.slice(0, showCount);
const hasMore = items.length > 12;
let groupHtml = `<div style="display:flex;flex-wrap:wrap;gap:4px">`;
groupHtml += visibleItems.map(it => {
const url = `${API}/output/` + encodeURIComponent(it.filename);
return `<div style="display:inline-block;cursor:pointer;position:relative;vertical-align:top"
title="Use as background · ${escHtml(it.filename)}${it.group_id?' ('+escHtml(it.group_id)+')':''}"
onclick="sceneLibUseAsBackground(${JSON.stringify(it.filename)})">
<img src="${url}" style="width:48px;height:48px;object-fit:cover;border-radius:4px;border:1px solid #27272a;background:#09090b"
onerror="this.style.opacity=.3">
</div>`;
}).join('');
groupHtml += `</div>`;
if (hasMore) {
groupHtml += `
<div style="margin-top:4px">
<button class="sb-btn" onclick="toggleSceneryGroup('${groupKey.replace(/'/g, "\\'")}')"
style="font-size:10px;padding:2px 6px;line-height:1;background:#18181b;color:#a1a1aa;border-color:#27272a">
${isExpanded ? 'Show less' : `Show all (+${items.length - 12} more)`}
</button>
</div>`;
}
return groupHtml;
};
_sceneLib.groups.forEach(g => {
const vName = g.video.replace(/\.[^.]+$/, '');
html += `<div style="margin-bottom:8px">
<div style="font-size:10px;color:#888;margin-bottom:3px">📹 ${escHtml(vName)}</div>
<div>${renderLibItems(g.items)}</div>
const groupKey = 'group_' + g.video;
html += `<div style="margin-bottom:12px;background:#141414;border:1px solid #222;border-radius:6px;padding:8px">
<div style="font-size:11px;font-weight:bold;color:#ccc;margin-bottom:6px">📹 ${escHtml(vName)}</div>
${renderLibGroup(groupKey, g.items)}
</div>`;
});
if (_sceneLib.ungrouped.length) {
html += `<div style="margin-bottom:8px">
<div style="font-size:10px;color:#888;margin-bottom:3px">📷 uploaded / other</div>
<div>${renderLibItems(_sceneLib.ungrouped)}</div>
html += `<div style="margin-bottom:12px;background:#141414;border:1px solid #222;border-radius:6px;padding:8px">
<div style="font-size:11px;font-weight:bold;color:#ccc;margin-bottom:6px">📷 uploaded / other</div>
${renderLibGroup('ungrouped', _sceneLib.ungrouped)}
</div>`;
}
}
@@ -13649,18 +13716,20 @@
if (_sceneVideo) {
const vid = document.getElementById('sceneVideoEl');
vid.src = `${API}/wireframe/${encodeURIComponent(_sceneVideo)}`;
// Once metadata is known, size the slider so 1000 steps span the whole clip.
vid.addEventListener('loadedmetadata', () => {
_sceneDuration = vid.duration || 0;
const lbl = document.getElementById('sceneTimeLabel');
if (lbl) lbl.textContent = formatSecs(vid.currentTime || 0);
}, { once: true });
// Keep the time label in sync if the frame is seeked any other way.
vid.addEventListener('seeked', () => {
const lbl = document.getElementById('sceneTimeLabel');
if (lbl) lbl.textContent = formatSecs(vid.currentTime || 0);
});
if (vid) {
vid.src = `${API}/wireframe/${encodeURIComponent(_sceneVideo)}`;
// Once metadata is known, size the slider so 1000 steps span the whole clip.
vid.addEventListener('loadedmetadata', () => {
_sceneDuration = vid.duration || 0;
const lbl = document.getElementById('sceneTimeLabel');
if (lbl) lbl.textContent = formatSecs(vid.currentTime || 0);
}, { once: true });
// Keep the time label in sync if the frame is seeked any other way.
vid.addEventListener('seeked', () => {
const lbl = document.getElementById('sceneTimeLabel');
if (lbl) lbl.textContent = formatSecs(vid.currentTime || 0);
});
}
}
panel.scrollTop = savedScrollTop;
@@ -13691,16 +13760,22 @@
renderSidebarScenery();
}
function toggleSceneryGroup(groupKey) {
if (!window._expandedSceneryGroups) window._expandedSceneryGroups = {};
window._expandedSceneryGroups[groupKey] = !window._expandedSceneryGroups[groupKey];
renderSidebarScenery();
}
// Use a library scenery image as slot 0 background (fetched to bytes).
async function sceneLibUseAsBackground(filename) {
try {
const url = IMAGE_FOLDER + encodeURIComponent(filename);
const url = `${API}/output/` + encodeURIComponent(filename);
const b64 = await _fetchAsBase64(url);
_sceneRefs[0] = { kind: 'bytes', data: b64, filename: filename };
_sceneVideo = null;
_sceneFrameBytes = null;
_sceneFrameBytes = b64; // Set this as well to align with other captured/set backgrounds!
_sceneGridOpen = false;
renderSidebarScenery();
_sceneRefresh();
showToast('Background set from previous scenery', 'success');
} catch (e) {
showToast('Could not load image: ' + e, 'error');
@@ -13724,6 +13799,9 @@
}
async function _fetchAsBase64(url) {
const r = await fetch(url);
if (!r.ok) {
throw new Error(`HTTP error! status: ${r.status}`);
}
const blob = await r.blob();
return await new Promise((res, rej) => {
const fr = new FileReader();

View File

@@ -2444,6 +2444,57 @@ def _get_grouped_wireframes(wireframe_dir: str) -> dict:
"frames": data["frames"]
})
# Find last used wireframe filenames from DB
last_used = []
try:
conn = database.get_db_connection()
cur = conn.cursor()
try:
cur.execute("""
SELECT source_refs
FROM person
WHERE archived IS NOT TRUE
AND filename LIKE '%_sc_%'
AND source_refs IS NOT NULL
ORDER BY filename DESC
""")
rows = cur.fetchall()
seen_refs = set()
for (source_refs_raw,) in rows:
try:
refs = json.loads(source_refs_raw) if source_refs_raw else []
except Exception:
refs = []
for r in refs:
if r.startswith("wireframe:"):
wf_name = r[len("wireframe:"):]
if wf_name and wf_name not in seen_refs:
seen_refs.add(wf_name)
last_used.append(wf_name)
finally:
cur.close()
database._put_db_connection(conn)
except Exception as db_err:
print(f"[wireframe] Failed to query last used scenes: {db_err}")
last_used_map = {name: idx for idx, name in enumerate(last_used)}
def standalone_sort_key(img):
fpath = os.path.join(wireframe_dir, img)
mtime = 0.0
try:
if os.path.exists(fpath):
mtime = os.path.getmtime(fpath)
except Exception:
pass
if img in last_used_map:
return (0, last_used_map[img], -mtime)
else:
return (1, 0, -mtime)
standalone_images.sort(key=standalone_sort_key)
standalone_images = standalone_images[:60]
return {
"videos": videos,
"groups": groups,
@@ -4071,6 +4122,7 @@ def _scenery_worker(job_id: str, model_filename: str, scene_pil: Image.Image,
sort_order=next_order,
source_refs=json.dumps(refs))
_update_cached_file_meta(out_name, exists=True)
_metadata_executor.submit(_process_image_for_metadata, out_name)
except Exception as db_err:
print(f"[scenery] DB error: {db_err}")
jobs[job_id]["latest_output"] = out_name
@@ -4149,6 +4201,7 @@ def generate_scenery(req: SceneryRequest):
@app.get("/scenery/library")
def scenery_library():
"""Return all scenery images grouped by source video reference."""
output_dir = _load_output_dir()
conn = database.get_db_connection()
cur = conn.cursor()
try:
@@ -4168,6 +4221,11 @@ def scenery_library():
by_video: dict[str, list] = {}
ungrouped: list = []
for filename, group_id, source_refs_raw in rows:
# Filter out files that don't exist on disk (avoid broken links)
exists, _ = _get_cached_file_meta(filename, output_dir)
if not exists:
continue
try:
refs = json.loads(source_refs_raw) if source_refs_raw else []
except Exception:

View File

@@ -623,5 +623,181 @@ class TestAPIRegression(unittest.TestCase):
self.assertEqual(dance_grp["video"], "dance.mp4")
self.assertIn("dance_10s-20s.mp4", dance_grp["clips"])
def test_12_standalone_scenery_sorting_and_existence_check(self):
from unittest.mock import patch
import edit_api
import database
# Insert some mock scenery generation records so we can test last_used sorting
conn = database.get_db_connection()
cur = conn.cursor()
try:
cur.execute("""
INSERT INTO person (filename, group_id, source_refs)
VALUES ('20260629_130000_sc_test.png', 'g1', '["test_person.png", "wireframe:used_scenery_2.png"]')
""")
cur.execute("""
INSERT INTO person (filename, group_id, source_refs)
VALUES ('20260629_130500_sc_test.png', 'g1', '["test_person.png", "wireframe:used_scenery_1.png"]')
""")
conn.commit()
finally:
cur.close()
database._put_db_connection(conn)
try:
# Mock wireframe folder and output directories
with patch("edit_api._load_wireframe_dir", return_value="/tmp/test_wireframe"), \
patch("edit_api._load_output_dir", return_value=self.output_dir), \
patch("os.path.isdir", return_value=True), \
patch("os.path.exists", return_value=True), \
patch("os.path.getmtime", return_value=12345.0), \
patch("os.listdir", return_value=["used_scenery_1.png", "used_scenery_2.png", "unused_scenery_3.png"]):
response = self.client.get("/videos?refresh=true")
self.assertEqual(response.status_code, 200)
data = response.json()
# Check standalone images exist in the response
self.assertIn("standalone_images", data)
standalone = data["standalone_images"]
# Should contain our wireframes sorted with last used on top:
# 'used_scenery_1.png' was used in the most recent scenery generation (at 13:05:00)
# 'used_scenery_2.png' was used in the older scenery generation (at 13:00:00)
# 'unused_scenery_3.png' was not used
self.assertEqual(standalone[0], "used_scenery_1.png")
self.assertEqual(standalone[1], "used_scenery_2.png")
self.assertEqual(standalone[2], "unused_scenery_3.png")
# Also verify that /scenery/library filters out missing files
# Let's mock _get_cached_file_meta: '20260629_130500_sc_test.png' exists, but '20260629_130000_sc_test.png' is missing
def mock_cached_meta(fname, out_dir):
if fname == "20260629_130500_sc_test.png":
return True, 12345.0
return False, 0.0
with patch("edit_api._get_cached_file_meta", side_effect=mock_cached_meta):
response = self.client.get("/scenery/library")
self.assertEqual(response.status_code, 200)
lib_data = response.json()
# Verify that only the existing scenery item '20260629_130500_sc_test.png' is returned
all_filenames = []
for g in lib_data["groups"]:
for item in g["items"]:
all_filenames.append(item["filename"])
for item in lib_data["ungrouped"]:
all_filenames.append(item["filename"])
self.assertIn("20260629_130500_sc_test.png", all_filenames)
self.assertNotIn("20260629_130000_sc_test.png", all_filenames)
finally:
# Clean up test DB rows
conn = database.get_db_connection()
cur = conn.cursor()
try:
cur.execute("DELETE FROM person WHERE filename IN ('20260629_130000_sc_test.png', '20260629_130500_sc_test.png')")
conn.commit()
finally:
cur.close()
database._put_db_connection(conn)
def test_13_scenery_metadata_extraction_and_refs(self):
from unittest.mock import patch, MagicMock
import edit_api
import database
import json
import os
from PIL import Image
# Write dummy model image on disk
model_filename = "model_image_123.png"
model_path = os.path.join(self.output_dir, model_filename)
Image.new("RGB", (10, 10), color="red").save(model_path)
# Set up a test person record in DB for the model
conn = database.get_db_connection()
cur = conn.cursor()
try:
cur.execute("""
INSERT INTO person (filename, group_id)
VALUES (%s, 'model_group_123')
""", (model_filename,))
conn.commit()
finally:
cur.close()
database._put_db_connection(conn)
try:
# Let's mock _run_pipeline, generation of embeddings, etc.
# We want to verify that _scenery_worker correctly structures 'source_refs'
# and triggers metadata extraction.
mock_png = b"dummy_png_bytes"
edit_api.jobs["test_scenery_job"] = {"status": "running"}
with patch("edit_api._run_pipeline", return_value=mock_png), \
patch("edit_api._load_output_dir", return_value=self.output_dir), \
patch("edit_api.embeddings.generate_embedding", return_value=None), \
patch("edit_api._metadata_executor.submit") as mock_submit:
# Run the scenery worker synchronously for the test
dummy_scene = Image.new("RGB", (100, 100), color="blue")
edit_api._scenery_worker(
job_id="test_scenery_job",
model_filename=model_filename,
scene_pil=dummy_scene,
prompt="test scenery prompt",
seed=42,
extra_pils=[],
scene_video="test_video.mp4",
scene_image="test_image.png",
extra_filename="test_extra.png"
)
# Find the newly created scenery filename
scenery_filename = edit_api.jobs["test_scenery_job"]["output"]
self.assertIsNotNone(scenery_filename)
# Check database record to verify refs are correct
person = database.get_person(scenery_filename)
self.assertIsNotNone(person)
self.assertEqual(person[1], "model_group_123") # Correctly pulled group_id
refs = json.loads(person[12]) # source_refs is at index 12 in tuple or get_person response
self.assertIn(model_filename, refs)
self.assertIn("video:test_video.mp4", refs)
self.assertIn("wireframe:test_image.png", refs)
self.assertIn("test_extra.png", refs)
# Verify that _metadata_executor.submit was called to trigger background metadata extraction
mock_submit.assert_called_once_with(edit_api._process_image_for_metadata, scenery_filename)
finally:
# Clean up all created files and database entries
if os.path.exists(model_path):
os.remove(model_path)
conn = database.get_db_connection()
cur = conn.cursor()
try:
cur.execute("DELETE FROM person WHERE filename = 'model_image_123.png'")
cur.execute("DELETE FROM person WHERE filename LIKE '%_sc_model_image_123.png'")
conn.commit()
finally:
cur.close()
database._put_db_connection(conn)
# Clean up generated scenery file on disk
try:
scenery_filename = edit_api.jobs.get("test_scenery_job", {}).get("output")
if scenery_filename:
out_path = os.path.join(self.output_dir, scenery_filename)
if os.path.exists(out_path):
os.remove(out_path)
except Exception:
pass
if __name__ == "__main__":
unittest.main()