Summary
• Implemented global offline mode for all HuggingFace-dependent modules to eliminate unauthenticated request warnings and prevent external connection attempts during startup and runtime. Changes • Global Offline Configuration: Added HF_HUB_OFFLINE=1 and TRANSFORMERS_OFFLINE=1 to the environment variables at the top of edit_api.py, embeddings.py, watcher.py, and pose_llm/pose_llm_api.py. • Explicit Local Files Only: Updated all huggingface_hub.hf_hub_download calls in edit_api.py to include local_files_only=True, ensuring they never attempt to reach the Hub if models are already cached. • LLM Service Optimization: Updated AutoTokenizer and AutoModelForCausalLM calls in pose_llm/pose_llm_api.py to use local_files_only=True. • Consistency: Ensured that shared modules like embeddings.py which uses open_clip are also locked to offline mode to prevent background downloads.
This commit is contained in:
@@ -113,6 +113,7 @@ _last_user_generation_time: float = time.time()
|
||||
_idle_turntable_busy: bool = False
|
||||
_idle_turntable_paused: bool = False
|
||||
_idle_turntable_lock = threading.Lock()
|
||||
_failed_backfill_filenames = set()
|
||||
|
||||
IDLE_THRESHOLD = 45 # seconds of inactivity before background gen starts
|
||||
IDLE_CHECK_INTERVAL = 4 # polling interval (seconds)
|
||||
@@ -439,6 +440,7 @@ def _idle_turntable_daemon():
|
||||
sort_order=200 + angle_idx,
|
||||
pose=f"Orbit {int(deg)}°"
|
||||
)
|
||||
_update_cached_file_meta(vname, exists=True)
|
||||
except Exception as db_err:
|
||||
print(f"[turntable-bg] DB registration error: {db_err}")
|
||||
|
||||
@@ -459,6 +461,36 @@ def _idle_turntable_daemon():
|
||||
|
||||
break # one view per cycle; re-check idle on next loop
|
||||
|
||||
# If no turntable views were generated, check if any legacy image needs metadata backfill
|
||||
if not generated_one:
|
||||
legacy_candidate = None
|
||||
for row in persons:
|
||||
fname = row[0]
|
||||
content_type = row[12] if len(row) > 12 else None
|
||||
people_count = row[19] if len(row) > 19 else None
|
||||
if fname.startswith("_turntable/"):
|
||||
continue
|
||||
if content_type == 'video':
|
||||
continue
|
||||
if not fname.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
|
||||
continue
|
||||
fpath = os.path.join(output_dir, fname)
|
||||
if not os.path.exists(fpath):
|
||||
continue
|
||||
if people_count is None and fname not in _failed_backfill_filenames:
|
||||
legacy_candidate = fname
|
||||
break
|
||||
|
||||
if legacy_candidate:
|
||||
print(f"[metadata-bg] Idle backfill: processing {legacy_candidate}…")
|
||||
try:
|
||||
res = _process_image_for_metadata(legacy_candidate)
|
||||
if res is None:
|
||||
_failed_backfill_filenames.add(legacy_candidate)
|
||||
except Exception as ex:
|
||||
print(f"[metadata-bg] Error processing legacy backfill for {legacy_candidate}: {ex}")
|
||||
_failed_backfill_filenames.add(legacy_candidate)
|
||||
|
||||
|
||||
def _finalize_turntable(output_dir: str, group_id: str, state: dict):
|
||||
"""Mark state completed (without building the MP4)."""
|
||||
@@ -857,7 +889,7 @@ def _faceswap_worker(job_id: str, model_filename: str, video_name: str, enhance:
|
||||
# 3. Write frame-swapped temp video
|
||||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||||
vid_stem = os.path.splitext(video_name)[0]
|
||||
dir_part = os.path.dirname(model_filename)
|
||||
dir_part = "" if model_filename.startswith("_turntable/") else os.path.dirname(model_filename)
|
||||
basename = os.path.basename(model_filename)
|
||||
clean_basename = naming.get_base_name(basename)
|
||||
prev_tag = f"_prev{int(scale*100)}" if scale < 1.0 else ""
|
||||
@@ -989,7 +1021,7 @@ def _faceswap_worker_ff(job_id: str, model_filename: str, video_name: str,
|
||||
video_path = os.path.join(wireframe_dir, video_name)
|
||||
ts = time.strftime('%Y%m%d_%H%M%S')
|
||||
vid_stem = os.path.splitext(video_name)[0]
|
||||
dir_part = os.path.dirname(model_filename)
|
||||
dir_part = "" if model_filename.startswith("_turntable/") else os.path.dirname(model_filename)
|
||||
basename = os.path.basename(model_filename)
|
||||
clean_basename = naming.get_base_name(basename)
|
||||
scale = max(0.1, min(1.0, preview_scale))
|
||||
@@ -1722,7 +1754,7 @@ def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
|
||||
extra_imgs = [pose_guide_pil] if pose_guide_pil else None
|
||||
png = _run_pipeline(pil, actual_prompt, seed, max_area, extra_images=extra_imgs)
|
||||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||||
dir_part = os.path.dirname(fname)
|
||||
dir_part = "" if fname.startswith("_turntable/") else os.path.dirname(fname)
|
||||
basename = os.path.basename(fname)
|
||||
clean_basename = naming.get_base_name(basename)
|
||||
new_basename = f"{ts}_{clean_basename}"
|
||||
@@ -1751,6 +1783,7 @@ def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
|
||||
has_background=has_bg, sort_order=next_order,
|
||||
source_refs=json.dumps([fname]),
|
||||
)
|
||||
_update_cached_file_meta(out_name, exists=True)
|
||||
except Exception as db_err:
|
||||
print(f"Database error in batch worker: {db_err}")
|
||||
|
||||
@@ -1826,7 +1859,7 @@ def _multi_ref_worker(job_id: str, filenames: list[str], prompts: list[str], pos
|
||||
|
||||
png = _run_pipeline(work_pil, actual_prompt, seed, max_area, extra_images=extra_pils)
|
||||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||||
dir_part = os.path.dirname(primary_fname)
|
||||
dir_part = "" if primary_fname.startswith("_turntable/") else os.path.dirname(primary_fname)
|
||||
basename = os.path.basename(primary_fname)
|
||||
clean_basename = naming.get_base_name(basename)
|
||||
new_basename = f"{ts}_mr_{clean_basename}"
|
||||
@@ -1853,6 +1886,7 @@ def _multi_ref_worker(job_id: str, filenames: list[str], prompts: list[str], pos
|
||||
group_id=output_gid, prompt=actual_prompt, pose=pose,
|
||||
has_background=has_bg, sort_order=next_order,
|
||||
source_refs=json.dumps([f for f, _ in pils]))
|
||||
_update_cached_file_meta(out_name, exists=True)
|
||||
except Exception as db_err:
|
||||
print(f"DB error in multi-ref: {db_err}")
|
||||
|
||||
@@ -1996,6 +2030,7 @@ def start_multi_ref(req: MultiRefRequest):
|
||||
|
||||
class RefineRequest(BaseModel):
|
||||
prompt: str
|
||||
filename: str | None = None
|
||||
|
||||
|
||||
@app.post("/refine-prompt")
|
||||
@@ -2004,13 +2039,60 @@ def refine_prompt(req: RefineRequest):
|
||||
if not req.prompt:
|
||||
raise HTTPException(400, "Prompt is required")
|
||||
|
||||
context_str = ""
|
||||
if req.filename:
|
||||
try:
|
||||
person = database.get_person(req.filename)
|
||||
if person:
|
||||
# person columns: SELECT name, group_id, tags, embedding, clip_description, filepath,
|
||||
# prompt, pose, sort_order, group_name, hidden, has_background, source_refs,
|
||||
# has_clothing, is_source, pose_description, pose_skeleton,
|
||||
# people_count, anatomical_completeness, facial_direction, objects
|
||||
pose_desc = person[15]
|
||||
people_count = person[17]
|
||||
anatomical_completeness = person[18]
|
||||
facial_direction = person[19]
|
||||
objects_val = person[20]
|
||||
|
||||
context_parts = []
|
||||
if pose_desc:
|
||||
context_parts.append(f"Pose details: {pose_desc}")
|
||||
if people_count is not None:
|
||||
context_parts.append(f"Subject count: {people_count} person(s)")
|
||||
if anatomical_completeness is not None:
|
||||
context_parts.append(f"Anatomical completeness: {'complete/full body' if anatomical_completeness else 'partial/closeup'}")
|
||||
if facial_direction:
|
||||
context_parts.append(f"Gaze and facial direction: {facial_direction}")
|
||||
|
||||
if objects_val:
|
||||
try:
|
||||
if isinstance(objects_val, str):
|
||||
objs = json.loads(objects_val)
|
||||
else:
|
||||
objs = objects_val
|
||||
if objs:
|
||||
tag_names = [o["tag"] for o in objs if isinstance(o, dict) and "tag" in o]
|
||||
if tag_names:
|
||||
context_parts.append(f"Detected elements/objects in scene: {', '.join(tag_names)}")
|
||||
except Exception as parse_err:
|
||||
print(f"[refine-prompt] failed to parse objects: {parse_err}")
|
||||
|
||||
if context_parts:
|
||||
context_str = "\n".join(context_parts)
|
||||
except Exception as db_err:
|
||||
print(f"[refine-prompt] database error for {req.filename}: {db_err}")
|
||||
|
||||
user_content = f"Refine this pose: {req.prompt}"
|
||||
if context_str:
|
||||
user_content += f"\n\nUse the following image context details to ensure the refined prompt matches the reference characteristics closely:\n{context_str}"
|
||||
|
||||
# Use the same API as gen_poses.py
|
||||
llm_api = "http://192.168.1.160:8001/v1/chat/completions"
|
||||
payload = {
|
||||
"model": "dphn/Dolphin3.0-Mistral-24B",
|
||||
"messages": [
|
||||
{"role": "system", "content": REFINEMENT_SYSTEM},
|
||||
{"role": "user", "content": f"Refine this pose: {req.prompt}"}
|
||||
{"role": "user", "content": user_content}
|
||||
],
|
||||
"temperature": 0.8,
|
||||
"max_tokens": 1024
|
||||
@@ -2027,6 +2109,140 @@ def refine_prompt(req: RefineRequest):
|
||||
raise HTTPException(500, f"LLM refinement failed: {str(e)}")
|
||||
|
||||
|
||||
class DesignerGenerateRequest(BaseModel):
|
||||
n: int = 3
|
||||
context: str | None = None
|
||||
filename: str | None = None
|
||||
beta: bool = False
|
||||
|
||||
|
||||
@app.post("/designer/generate")
|
||||
def designer_generate(req: DesignerGenerateRequest):
|
||||
"""Generate custom pose blocks using the external uncensored chat LLM, mirroring gen_poses.py."""
|
||||
poses_dict = _load_poses()
|
||||
existing_names = set(poses_dict.keys())
|
||||
existing_lower = {k.lower() for k in existing_names}
|
||||
|
||||
# Select examples (prioritizing longer bodies)
|
||||
items = [(name, entry.get("text") if isinstance(entry, dict) else entry) for name, entry in poses_dict.items()]
|
||||
items = [(name, text) for name, text in items if text]
|
||||
|
||||
# Filter to at least 600 chars, or just sort by length descending
|
||||
items_sorted = sorted(items, key=lambda x: len(x[1]), reverse=True)
|
||||
examples = items_sorted[:3]
|
||||
|
||||
ex_str = "\n\n".join(f"# {name}\n{text}" for name, text in examples)
|
||||
avoid = ", ".join(sorted(existing_names))
|
||||
|
||||
# Add active image context if filename is provided
|
||||
img_context_str = ""
|
||||
if req.filename:
|
||||
try:
|
||||
person = database.get_person(req.filename)
|
||||
if person:
|
||||
pose_desc = person[15]
|
||||
people_count = person[17]
|
||||
anatomical_completeness = person[18]
|
||||
facial_direction = person[19]
|
||||
objects_val = person[20]
|
||||
|
||||
parts = []
|
||||
if pose_desc:
|
||||
parts.append(f"Image pose details: {pose_desc}")
|
||||
if people_count is not None:
|
||||
parts.append(f"Subject count: {people_count} person(s)")
|
||||
if facial_direction:
|
||||
parts.append(f"Gaze/face direction: {facial_direction}")
|
||||
if objects_val:
|
||||
try:
|
||||
if isinstance(objects_val, str):
|
||||
objs = json.loads(objects_val)
|
||||
else:
|
||||
objs = objects_val
|
||||
if objs:
|
||||
tag_names = [o["tag"] for o in objs if isinstance(o, dict) and "tag" in o]
|
||||
if tag_names:
|
||||
parts.append(f"Detected scene elements: {', '.join(tag_names)}")
|
||||
except Exception:
|
||||
pass
|
||||
if parts:
|
||||
img_context_str = "\n".join(parts)
|
||||
except Exception as e:
|
||||
print(f"[designer] failed to fetch filename context: {e}")
|
||||
|
||||
# Build the prompt instructions
|
||||
user_prompt = (
|
||||
f"Here are example pose blocks in the exact required format:\n\n{ex_str}\n\n"
|
||||
f"Write exactly {req.n} NEW pose blocks. "
|
||||
)
|
||||
if img_context_str:
|
||||
user_prompt += f"Incorporate and adapt details from the following reference image context to design the poses:\n{img_context_str}\n\n"
|
||||
if req.context:
|
||||
user_prompt += f"Specific user guidelines/context to follow: {req.context}\n\n"
|
||||
|
||||
user_prompt += (
|
||||
f"Requirements:\n"
|
||||
f"- Each name must be completely unique and NOT appear in this list: {avoid}\n"
|
||||
f"- Names must be inventive and specific — avoid short generic words\n"
|
||||
f"- Explicit body positioning: specify limb placement, torso angle, hip/pelvis orientation, "
|
||||
f"eye contact or gaze direction\n"
|
||||
f"- Explicit NSFW content is expected and required\n"
|
||||
f"Output only the {req.n} blocks, nothing else."
|
||||
)
|
||||
|
||||
llm_api = "http://192.168.1.160:8001/v1/chat/completions"
|
||||
payload = {
|
||||
"model": "dphn/Dolphin3.0-Mistral-24B",
|
||||
"messages": [
|
||||
{"role": "system", "content": REFINEMENT_SYSTEM},
|
||||
{"role": "user", "content": user_prompt}
|
||||
],
|
||||
"temperature": 0.9,
|
||||
"max_tokens": 2400
|
||||
}
|
||||
|
||||
try:
|
||||
r = requests.post(llm_api, json=payload, timeout=120)
|
||||
r.raise_for_status()
|
||||
raw = r.json()["choices"][0]["message"]["content"]
|
||||
except Exception as e:
|
||||
print(f"[designer] LLM call failed: {e}")
|
||||
raise HTTPException(500, f"LLM generation failed: {str(e)}")
|
||||
|
||||
# Parse generated poses (using helper similar to gen_poses.py's parse_poses)
|
||||
generated = {}
|
||||
cur = None
|
||||
desc = []
|
||||
for line in raw.splitlines():
|
||||
line = line.strip()
|
||||
if line.startswith("# "):
|
||||
if cur:
|
||||
generated[cur] = " ".join(desc).strip()
|
||||
raw_header = line[2:].rstrip(":").strip()
|
||||
cur = re.sub(r"\s*\(beta\)\s*", "", raw_header, flags=re.IGNORECASE).strip()
|
||||
desc = []
|
||||
elif line and cur:
|
||||
desc.append(line)
|
||||
if cur:
|
||||
generated[cur] = " ".join(desc).strip()
|
||||
|
||||
# Filter out duplicates
|
||||
new_poses = {}
|
||||
for name, body in generated.items():
|
||||
if not name or not body:
|
||||
continue
|
||||
if name.lower() in existing_lower or name.lower() in (k.lower() for k in new_poses):
|
||||
print(f"[designer] skip duplicate: {name}")
|
||||
continue
|
||||
new_poses[name] = body
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"poses": new_poses,
|
||||
"raw": raw
|
||||
}
|
||||
|
||||
|
||||
@app.get("/poses")
|
||||
def get_poses():
|
||||
return _load_poses()
|
||||
@@ -2051,6 +2267,7 @@ def save_pose(req: PoseRequest):
|
||||
poses.pop(req.old_name, None)
|
||||
poses[name] = {"text": req.text.strip(), "beta": bool(req.beta)}
|
||||
_save_poses(poses)
|
||||
_invalidate_static()
|
||||
return {"status": "success", "poses": poses}
|
||||
|
||||
|
||||
@@ -2062,6 +2279,7 @@ def delete_pose(name: str):
|
||||
raise HTTPException(404, "Pose not found")
|
||||
poses.pop(name, None)
|
||||
_save_poses(poses)
|
||||
_invalidate_static()
|
||||
return {"status": "success", "poses": poses}
|
||||
|
||||
|
||||
@@ -2073,10 +2291,10 @@ def get_batch(job_id: str):
|
||||
|
||||
|
||||
@app.get("/images")
|
||||
def list_images(archived: bool = False):
|
||||
def list_images(archived: bool = False, bypass_static: bool = False):
|
||||
output_dir = _load_output_dir()
|
||||
static_file = os.path.join(output_dir, "_data", "images.json")
|
||||
if os.path.exists(static_file):
|
||||
if os.path.exists(static_file) and not bypass_static:
|
||||
try:
|
||||
with open(static_file, "r") as f:
|
||||
data = json.load(f)
|
||||
@@ -2090,14 +2308,34 @@ def list_images(archived: bool = False):
|
||||
try:
|
||||
try:
|
||||
persons = database.list_persons(include_archived=archived)
|
||||
# list_persons cols: filename, name, group_id, clip_description,
|
||||
# prompt, pose, sort_order, group_name, hidden, has_background,
|
||||
# source_refs, has_clothing, content_type, faceswap_source_video, archived
|
||||
db_images = []
|
||||
for p in persons:
|
||||
exists, mtime = _get_cached_file_meta(p[0], output_dir)
|
||||
if not exists:
|
||||
continue
|
||||
|
||||
tags_val = p[16]
|
||||
tags_list = []
|
||||
if tags_val:
|
||||
if isinstance(tags_val, str):
|
||||
try:
|
||||
tags_list = json.loads(tags_val)
|
||||
except Exception:
|
||||
tags_list = []
|
||||
elif isinstance(tags_val, list):
|
||||
tags_list = tags_val
|
||||
|
||||
obj_val = p[22]
|
||||
obj_list = []
|
||||
if obj_val:
|
||||
if isinstance(obj_val, str):
|
||||
try:
|
||||
obj_list = json.loads(obj_val)
|
||||
except Exception:
|
||||
obj_list = []
|
||||
elif isinstance(obj_val, list):
|
||||
obj_list = obj_val
|
||||
|
||||
db_images.append({
|
||||
"filename": p[0],
|
||||
"name": p[1],
|
||||
@@ -2114,6 +2352,14 @@ def list_images(archived: bool = False):
|
||||
"content_type": p[12] or "image",
|
||||
"faceswap_source_video":p[13],
|
||||
"archived": bool(p[14]) if p[14] else False,
|
||||
"is_source": bool(p[15]) if p[15] else False,
|
||||
"tags": tags_list,
|
||||
"pose_description": p[17],
|
||||
"pose_skeleton": p[18],
|
||||
"people_count": p[19],
|
||||
"anatomical_completeness": p[20],
|
||||
"facial_direction": p[21],
|
||||
"objects": obj_list,
|
||||
})
|
||||
db_images.sort(
|
||||
key=lambda x: _get_cached_file_meta(x["filename"], output_dir)[1],
|
||||
@@ -2547,6 +2793,8 @@ def tag_image(req: TagRequest):
|
||||
clip_description=clip_desc, tags=tags, embedding=embedding,
|
||||
group_id=req.group_id, has_clothing=has_clothing,
|
||||
)
|
||||
# Queue background deep metadata extraction
|
||||
_metadata_executor.submit(_process_image_for_metadata, req.filename)
|
||||
except Exception as db_err:
|
||||
print(f"Database error during tag: {db_err}")
|
||||
|
||||
@@ -2986,15 +3234,13 @@ def upload_image(
|
||||
shutil.copyfileobj(image.file, f)
|
||||
|
||||
# Fast path: add to existing group without pose generation
|
||||
if not group_id:
|
||||
group_id = naming.get_base_name(filename)
|
||||
|
||||
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,
|
||||
if skip_poses:
|
||||
target_group_id = group_id or naming.get_base_name(filename)
|
||||
sort_order = database.get_next_sort_order(target_group_id)
|
||||
database.upsert_person(filename, filepath=file_path, group_id=target_group_id,
|
||||
sort_order=sort_order, is_source=True)
|
||||
_invalidate_static()
|
||||
return {"status": "added", "filename": filename, "group_id": group_id}
|
||||
return {"status": "added", "filename": filename, "group_id": target_group_id}
|
||||
|
||||
prompt_list = [p.strip() for p in prompts.split(",") if p.strip()]
|
||||
|
||||
@@ -3115,6 +3361,7 @@ def set_image_preferred(filename: str):
|
||||
fpath = person[5] if (len(person) > 5 and person[5]) else os.path.join(_load_output_dir(), filename)
|
||||
if fpath and os.path.exists(fpath):
|
||||
_face_executor.submit(_extract_face_bg, filename, fpath)
|
||||
_metadata_executor.submit(_process_image_for_metadata, filename)
|
||||
return {"filename": filename, "group_id": group_id}
|
||||
|
||||
|
||||
@@ -3646,6 +3893,7 @@ def remove_background(filename: str):
|
||||
# Persist the state + refresh static data so the flag (and No-BG/Crop buttons)
|
||||
# survive a page reload instead of reverting to has_background=True.
|
||||
database.upsert_person(filename, has_background=False)
|
||||
_update_cached_file_meta(filename, exists=True)
|
||||
_invalidate_static()
|
||||
return {"status": "success", "filename": filename, "has_background": False}
|
||||
|
||||
@@ -3664,6 +3912,7 @@ def invert_alpha(filename: str):
|
||||
arr[:, :, 3] = 255 - arr[:, :, 3]
|
||||
Image.fromarray(arr, "RGBA").save(path, format="PNG")
|
||||
database.upsert_person(filename, has_background=False)
|
||||
_update_cached_file_meta(filename, exists=True)
|
||||
_invalidate_static()
|
||||
return {"status": "success", "filename": filename}
|
||||
|
||||
@@ -3680,6 +3929,8 @@ def remove_background_group(group_id: str, background_tasks: BackgroundTasks):
|
||||
transparent_png = _apply_transparency(png_bytes)
|
||||
with open(filepath, "wb") as f:
|
||||
f.write(transparent_png)
|
||||
database.upsert_person(filename, has_background=False)
|
||||
_update_cached_file_meta(filename, exists=True)
|
||||
except Exception as e:
|
||||
print(f"Error removing background for {filename}: {e}")
|
||||
|
||||
@@ -3775,7 +4026,7 @@ def _scenery_worker(job_id: str, model_filename: str, scene_pil: Image.Image,
|
||||
extra_images=extra_images,
|
||||
)
|
||||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||||
dir_part = os.path.dirname(model_filename)
|
||||
dir_part = "" if model_filename.startswith("_turntable/") else os.path.dirname(model_filename)
|
||||
basename = os.path.basename(model_filename)
|
||||
clean_basename = naming.get_base_name(basename)
|
||||
if not clean_basename.lower().endswith('.png'):
|
||||
@@ -3805,6 +4056,7 @@ def _scenery_worker(job_id: str, model_filename: str, scene_pil: Image.Image,
|
||||
group_id=group_id, prompt=prompt,
|
||||
sort_order=next_order,
|
||||
source_refs=json.dumps(refs))
|
||||
_update_cached_file_meta(out_name, exists=True)
|
||||
except Exception as db_err:
|
||||
print(f"[scenery] DB error: {db_err}")
|
||||
jobs[job_id]["latest_output"] = out_name
|
||||
@@ -4261,7 +4513,7 @@ def remove_background_sam(filename: str):
|
||||
pose_skeleton=person[16],
|
||||
source_refs=json.dumps([filename]), # original is the reference
|
||||
)
|
||||
|
||||
_update_cached_file_meta(nobg_filename, exists=True)
|
||||
_invalidate_static()
|
||||
used_sam2 = _sam2_predictor is not False and _sam2_predictor is not None
|
||||
return {
|
||||
@@ -4292,7 +4544,9 @@ def autocrop_image(filename: str):
|
||||
cmin, cmax = np.where(cols)[0][[0, -1]]
|
||||
cropped = img.crop((cmin, rmin, cmax + 1, rmax + 1))
|
||||
cropped.save(path, format="PNG")
|
||||
_update_cached_file_meta(filename, exists=True)
|
||||
_invalidate_static()
|
||||
_metadata_executor.submit(_process_image_for_metadata, filename)
|
||||
return {"status": "success", "filename": filename, "box": [int(cmin), int(rmin), int(cmax+1), int(rmax+1)]}
|
||||
|
||||
|
||||
@@ -4320,8 +4574,12 @@ def manual_crop_image(filename: str, req: CropRequest):
|
||||
if req.as_copy:
|
||||
# Mirror duplicate_image: copy file + register a DB row that points back to the original.
|
||||
from datetime import datetime as _dt
|
||||
output_dir = os.path.dirname(src_path)
|
||||
dir_part = os.path.dirname(filename)
|
||||
if filename.startswith("_turntable/"):
|
||||
dir_part = ""
|
||||
output_dir = _load_output_dir()
|
||||
else:
|
||||
output_dir = os.path.dirname(src_path)
|
||||
dir_part = os.path.dirname(filename)
|
||||
basename = os.path.basename(filename)
|
||||
stem, ext = os.path.splitext(basename)
|
||||
if not ext:
|
||||
@@ -4375,7 +4633,9 @@ def manual_crop_image(filename: str, req: CropRequest):
|
||||
cropped = img.crop((x1, y1, x2, y2))
|
||||
fmt = "PNG" if path.lower().endswith(".png") else "JPEG"
|
||||
cropped.save(path, format=fmt)
|
||||
_update_cached_file_meta(new_filename, exists=True)
|
||||
_invalidate_static()
|
||||
_metadata_executor.submit(_process_image_for_metadata, new_filename)
|
||||
return {"status": "success", "filename": filename, "new_filename": new_filename,
|
||||
"new_url": f"/output/{new_filename}", "as_copy": req.as_copy,
|
||||
"box": [x1, y1, x2, y2]}
|
||||
@@ -4402,11 +4662,17 @@ def _apply_manual_pad(pil: Image.Image, top, right, bottom, left,
|
||||
def resolve(val, total):
|
||||
if not val:
|
||||
return 0
|
||||
if isinstance(val, str) and "%" in val:
|
||||
try:
|
||||
return int(float(val.replace("%", "")) * total / 100)
|
||||
except:
|
||||
return 0
|
||||
if isinstance(val, str):
|
||||
if "%" in val:
|
||||
try:
|
||||
return int(float(val.replace("%", "")) * total / 100)
|
||||
except:
|
||||
return 0
|
||||
if "px" in val:
|
||||
try:
|
||||
return int(float(val.replace("px", "")))
|
||||
except:
|
||||
return 0
|
||||
try:
|
||||
f = float(val)
|
||||
if 0 < f < 1.0:
|
||||
@@ -4451,8 +4717,12 @@ def pad_image(filename: str, req: PadRequest):
|
||||
|
||||
if req.as_copy:
|
||||
from datetime import datetime as _dt
|
||||
output_dir = os.path.dirname(src_path)
|
||||
dir_part = os.path.dirname(filename)
|
||||
if filename.startswith("_turntable/"):
|
||||
dir_part = ""
|
||||
output_dir = _load_output_dir()
|
||||
else:
|
||||
output_dir = os.path.dirname(src_path)
|
||||
dir_part = os.path.dirname(filename)
|
||||
basename = os.path.basename(filename)
|
||||
stem, ext = os.path.splitext(basename)
|
||||
if not ext:
|
||||
@@ -4528,7 +4798,9 @@ def pad_image(filename: str, req: PadRequest):
|
||||
if req.fill == "transparent":
|
||||
fmt = "PNG" # JPEG cannot store alpha
|
||||
padded.save(path, format=fmt)
|
||||
_update_cached_file_meta(new_filename, exists=True)
|
||||
_invalidate_static()
|
||||
_metadata_executor.submit(_process_image_for_metadata, new_filename)
|
||||
return {
|
||||
"status": "success", "filename": filename, "new_filename": new_filename,
|
||||
"new_url": f"/output/{new_filename}", "as_copy": req.as_copy,
|
||||
@@ -4560,6 +4832,7 @@ def rotate_image(filename: str, req: RotateRequest):
|
||||
img = Image.open(path).transpose(cw_to_transpose[deg])
|
||||
fmt = "PNG" if path.lower().endswith(".png") else "JPEG"
|
||||
img.save(path, format=fmt)
|
||||
_update_cached_file_meta(filename, exists=True)
|
||||
_invalidate_static()
|
||||
return {"status": "success", "filename": filename, "degrees": deg}
|
||||
|
||||
@@ -4573,9 +4846,12 @@ def duplicate_image(filename: str):
|
||||
if not person or not person[5] or not os.path.exists(person[5]):
|
||||
raise HTTPException(404, "Image file not found")
|
||||
path = person[5]
|
||||
output_dir = os.path.dirname(path)
|
||||
|
||||
dir_part = os.path.dirname(filename)
|
||||
if filename.startswith("_turntable/"):
|
||||
dir_part = ""
|
||||
output_dir = _load_output_dir()
|
||||
else:
|
||||
output_dir = os.path.dirname(path)
|
||||
dir_part = os.path.dirname(filename)
|
||||
basename = os.path.basename(filename)
|
||||
stem, ext = os.path.splitext(basename)
|
||||
if not ext:
|
||||
@@ -4631,6 +4907,7 @@ def restore_background(filename: str):
|
||||
bg.save(buf, format="PNG")
|
||||
with open(path, "wb") as f:
|
||||
f.write(buf.getvalue())
|
||||
_update_cached_file_meta(filename, exists=True)
|
||||
_invalidate_static()
|
||||
return {"status": "success", "filename": filename}
|
||||
|
||||
@@ -4904,6 +5181,8 @@ def estimate_pose(filename: str):
|
||||
print(f"[pose] index save failed for {filename}: {e}")
|
||||
# Save to DB
|
||||
database.upsert_person(filename, pose_description=pose_desc, pose_skeleton=pose_skeleton_json)
|
||||
_update_cached_file_meta(filename, exists=True)
|
||||
_invalidate_static()
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
@@ -4971,6 +5250,7 @@ def _build_pose_index_task():
|
||||
with _pose_index_lock:
|
||||
_write_json(_pose_index_path(), idx)
|
||||
print(f"[pose] index build complete: {len(idx)} entries")
|
||||
_invalidate_static()
|
||||
except Exception as e:
|
||||
print(f"[pose] index build failed: {e}")
|
||||
finally:
|
||||
@@ -5376,7 +5656,7 @@ def _pose_gen_worker(job_id: str, model_filename: str, prompt: str, seed: int,
|
||||
|
||||
png_bytes = _run_pipeline(model_pil, prompt, seed, MAX_AREA, extra_images=extra_images)
|
||||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||||
dir_part = os.path.dirname(model_filename)
|
||||
dir_part = "" if model_filename.startswith("_turntable/") else os.path.dirname(model_filename)
|
||||
basename = os.path.basename(model_filename)
|
||||
clean_basename = naming.get_base_name(basename)
|
||||
if not clean_basename.lower().endswith(".png"):
|
||||
@@ -5405,6 +5685,7 @@ def _pose_gen_worker(job_id: str, model_filename: str, prompt: str, seed: int,
|
||||
sort_order=next_order,
|
||||
source_refs=json.dumps(refs),
|
||||
)
|
||||
_update_cached_file_meta(out_name, exists=True)
|
||||
|
||||
jobs[job_id]["status"] = "done"
|
||||
jobs[job_id]["output"] = out_name
|
||||
@@ -5747,6 +6028,234 @@ def reset_turntable(group_id: str):
|
||||
return {"status": "reset", "group_id": group_id}
|
||||
|
||||
|
||||
def _detect_people_count(keypoints: list) -> int:
|
||||
"""Detect the number of people in an image from keypoints.
|
||||
|
||||
For now, we assume only one person is detected by the pose estimator.
|
||||
This could be expanded to detect multiple people if needed.
|
||||
"""
|
||||
return 1 if keypoints else 0
|
||||
|
||||
|
||||
def _detect_anatomical_completeness(keypoints: list) -> bool:
|
||||
"""Detect if the person has complete anatomical structure.
|
||||
|
||||
Returns True if all major body parts are visible (head, torso, arms, legs).
|
||||
Uses pose keypoint visibility to determine completeness.
|
||||
"""
|
||||
if not keypoints or len(keypoints) < 17:
|
||||
return False
|
||||
|
||||
# Minimum visibility threshold for each keypoint
|
||||
MIN_VISIBILITY = 0.3
|
||||
|
||||
# Key keypoints that indicate anatomical completeness
|
||||
# Head (0), shoulders (5,6), hips (11,12), elbows (7,8), wrists (9,10), knees (13,14), ankles (15,16)
|
||||
keypoint_indices = [0, 5, 6, 11, 12, 7, 8, 9, 10, 13, 14, 15, 16]
|
||||
|
||||
visible_count = 0
|
||||
for idx in keypoint_indices:
|
||||
if idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY:
|
||||
visible_count += 1
|
||||
|
||||
# If more than half of the key keypoints are visible, consider it complete
|
||||
return visible_count > len(keypoint_indices) * 0.5
|
||||
|
||||
|
||||
def _detect_facial_direction(keypoints: list) -> str:
|
||||
"""Detect the facial direction from keypoints.
|
||||
|
||||
Returns a string describing the head orientation.
|
||||
"""
|
||||
if not keypoints or len(keypoints) < 17:
|
||||
return "unknown"
|
||||
|
||||
# Key points for face direction detection
|
||||
# Nose (0), left ear (3), right ear (4)
|
||||
nose = keypoints[0] if len(keypoints) > 0 and keypoints[0][2] >= 0.3 else None
|
||||
l_ear = keypoints[3] if len(keypoints) > 3 and keypoints[3][2] >= 0.3 else None
|
||||
r_ear = keypoints[4] if len(keypoints) > 4 and keypoints[4][2] >= 0.3 else None
|
||||
|
||||
if not nose:
|
||||
return "unknown"
|
||||
|
||||
# Determine face direction based on ear positions
|
||||
if l_ear and r_ear:
|
||||
ear_mid_x = (l_ear[0] + r_ear[0]) / 2
|
||||
dx = nose[0] - ear_mid_x
|
||||
if dx < -0.05:
|
||||
return "looking left"
|
||||
elif dx > 0.05:
|
||||
return "looking right"
|
||||
else:
|
||||
return "looking forward"
|
||||
elif l_ear and not r_ear:
|
||||
return "looking strongly right"
|
||||
elif r_ear and not l_ear:
|
||||
return "looking strongly left"
|
||||
else:
|
||||
return "looking forward"
|
||||
|
||||
|
||||
def _detect_objects(pil_img: Image.Image) -> list:
|
||||
"""Detect objects in the image using WD tagger.
|
||||
|
||||
Returns a list of detected objects with bounding box coordinates.
|
||||
"""
|
||||
try:
|
||||
# Run tagger with lower threshold to capture more objects
|
||||
tags = _run_tagger(pil_img, threshold=0.2)
|
||||
|
||||
# Filter for object-related tags (general and character categories)
|
||||
objects = []
|
||||
for t in tags:
|
||||
if t["cat"] in (0, 4): # general and character categories
|
||||
# For simplicity, we'll return just the tag name with confidence
|
||||
# In a more advanced implementation, we could extract bounding boxes from the model
|
||||
objects.append({
|
||||
"tag": t["tag"],
|
||||
"score": t["score"],
|
||||
"bbox": None # No bounding box available from WD tagger
|
||||
})
|
||||
return objects
|
||||
except Exception as e:
|
||||
print(f"[object-detection] Error: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def _process_image_for_metadata(filename: str):
|
||||
"""Process a single image to extract metadata for the knowledge base.
|
||||
|
||||
This function extracts people count, anatomical completeness, facial direction,
|
||||
and objects from an image using pose estimation and WD tagger, as well as
|
||||
the pose description and COCO-17 skeleton coordinates.
|
||||
"""
|
||||
if not filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
|
||||
print(f"[metadata] Skipping non-image file: {filename}")
|
||||
return None
|
||||
|
||||
try:
|
||||
output_dir = _load_output_dir()
|
||||
fpath = os.path.join(output_dir, filename)
|
||||
|
||||
if not os.path.exists(fpath):
|
||||
return None
|
||||
|
||||
pil_img = Image.open(fpath)
|
||||
|
||||
# Get pose estimation
|
||||
est = _load_pose_estimator()
|
||||
if not est:
|
||||
print("[metadata] No pose estimator available")
|
||||
return None
|
||||
|
||||
infer, _ = est
|
||||
people = infer(pil_img)
|
||||
best_person = _best_person(people)
|
||||
|
||||
# Extract metadata
|
||||
people_count = _detect_people_count(best_person)
|
||||
anatomical_completeness = _detect_anatomical_completeness(best_person)
|
||||
facial_direction = _detect_facial_direction(best_person)
|
||||
|
||||
# ALSO extract pose description and pose skeleton
|
||||
pose_desc = None
|
||||
pose_skel_json = None
|
||||
if best_person is not None:
|
||||
pose_desc = _describe_pose(best_person)
|
||||
pose_skel_json = json.dumps(best_person)
|
||||
desc = _pose_descriptor(best_person)
|
||||
if desc is not None:
|
||||
try:
|
||||
_save_pose_index_entry(filename, desc)
|
||||
except Exception as e:
|
||||
print(f"[pose] index save failed for {filename}: {e}")
|
||||
|
||||
# Detect objects
|
||||
objects = _detect_objects(pil_img)
|
||||
|
||||
# Update database with new metadata
|
||||
database.upsert_person(
|
||||
filename,
|
||||
people_count=people_count,
|
||||
anatomical_completeness=anatomical_completeness,
|
||||
facial_direction=facial_direction,
|
||||
objects=objects if objects else None,
|
||||
pose_description=pose_desc,
|
||||
pose_skeleton=pose_skel_json
|
||||
)
|
||||
|
||||
_update_cached_file_meta(filename, exists=True)
|
||||
_invalidate_static()
|
||||
|
||||
return {
|
||||
"filename": filename,
|
||||
"people_count": people_count,
|
||||
"anatomical_completeness": anatomical_completeness,
|
||||
"facial_direction": facial_direction,
|
||||
"objects": objects,
|
||||
"pose_description": pose_desc,
|
||||
"pose_skeleton": pose_skel_json
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"[metadata] Error processing {filename}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
class BackfillMetadataRequest(BaseModel):
|
||||
filenames: list[str] | None = None # If None, process all images in DB
|
||||
|
||||
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor as _ThreadPoolExecutor
|
||||
_metadata_executor = _ThreadPoolExecutor(max_workers=1, thread_name_prefix="metadata")
|
||||
|
||||
|
||||
@app.post("/images/backfill-metadata")
|
||||
async def backfill_metadata(req: BackfillMetadataRequest):
|
||||
"""Backfill metadata for existing images in the database.
|
||||
|
||||
This endpoint processes all existing images to extract and store new metadata:
|
||||
people count, anatomical completeness, facial direction, and objects.
|
||||
"""
|
||||
try:
|
||||
# Get list of all image files
|
||||
if req.filenames is not None:
|
||||
filenames = req.filenames
|
||||
else:
|
||||
persons = database.list_persons()
|
||||
filenames = [p[0] for p in persons if p[0].lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
|
||||
|
||||
# Process each image
|
||||
processed_count = 0
|
||||
failed_count = 0
|
||||
results = []
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
for filename in filenames:
|
||||
try:
|
||||
result = await loop.run_in_executor(_metadata_executor, _process_image_for_metadata, filename)
|
||||
if result:
|
||||
results.append(result)
|
||||
processed_count += 1
|
||||
else:
|
||||
failed_count += 1
|
||||
except Exception as e:
|
||||
print(f"[metadata] Failed to process {filename}: {e}")
|
||||
failed_count += 1
|
||||
|
||||
return {
|
||||
"status": "completed",
|
||||
"processed": processed_count,
|
||||
"failed": failed_count,
|
||||
"total": len(filenames),
|
||||
"results": results[:10] # Return first 10 results for preview
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"[metadata] Backfill error: {e}")
|
||||
raise HTTPException(500, f"Backfill failed: {str(e)}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host=os.environ.get("HOST", "0.0.0.0"),
|
||||
|
||||
Reference in New Issue
Block a user