updates UI
This commit is contained in:
mike
0
.trash/watcher.lock
Normal file
0
.trash/watcher.lock
Normal file
334
.trash/watcher.py
Normal file
334
.trash/watcher.py
Normal file
@@ -0,0 +1,334 @@
|
||||
import os
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||||
import time
|
||||
import json
|
||||
import shutil
|
||||
import requests
|
||||
from PIL import Image
|
||||
import logging
|
||||
import hashlib
|
||||
import sys
|
||||
import fcntl
|
||||
import re
|
||||
|
||||
try:
|
||||
from . import database
|
||||
from . import embeddings
|
||||
except ImportError:
|
||||
import database
|
||||
import embeddings
|
||||
|
||||
# Load configuration
|
||||
CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json")
|
||||
|
||||
def load_config():
|
||||
with open(CONFIG_PATH, 'r') as f:
|
||||
conf = json.load(f)
|
||||
# Resolve relative paths relative to this script's directory
|
||||
base_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
for key in ["stage_dir", "output_dir", "failed_dir", "processed_file", "log_file"]:
|
||||
if not os.path.isabs(conf[key]):
|
||||
conf[key] = os.path.normpath(os.path.join(base_dir, "..", conf[key]))
|
||||
return conf
|
||||
|
||||
CONF = load_config()
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(levelname)s - %(message)s',
|
||||
handlers=[
|
||||
logging.FileHandler(CONF["log_file"]),
|
||||
logging.StreamHandler()
|
||||
]
|
||||
)
|
||||
|
||||
def get_processed_files():
|
||||
if os.path.exists(CONF["processed_file"]):
|
||||
try:
|
||||
with open(CONF["processed_file"], 'r') as f:
|
||||
data = json.load(f)
|
||||
if isinstance(data, list):
|
||||
# Migration: convert old list format to dict
|
||||
return {name: None for name in data}
|
||||
return data
|
||||
except Exception as e:
|
||||
logging.error(f"Error reading processed file: {e}")
|
||||
return {}
|
||||
return {}
|
||||
|
||||
def save_processed_files(processed):
|
||||
try:
|
||||
p = CONF["processed_file"]
|
||||
tmp = p + ".tmp"
|
||||
with open(tmp, 'w') as f:
|
||||
json.dump(processed, f, indent=2)
|
||||
os.replace(tmp, p)
|
||||
except Exception as e:
|
||||
logging.error(f"Error saving processed file: {e}")
|
||||
|
||||
def calculate_hash(filepath):
|
||||
"""Calculate MD5 hash of a file."""
|
||||
hasher = hashlib.md5()
|
||||
try:
|
||||
with open(filepath, 'rb') as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hasher.update(chunk)
|
||||
return hasher.hexdigest()
|
||||
except Exception as e:
|
||||
logging.error(f"Error calculating hash for {filepath}: {e}")
|
||||
return None
|
||||
|
||||
def crop_to_bbox(image_path, margin, top_margin=None, headroom=0.0):
|
||||
try:
|
||||
img = Image.open(image_path)
|
||||
if img.mode != 'RGBA':
|
||||
logging.info(f"Image {image_path} is mode {img.mode}, not RGBA. Skipping crop.")
|
||||
return img
|
||||
|
||||
alpha = img.split()[-1]
|
||||
bbox = alpha.getbbox()
|
||||
if not bbox:
|
||||
logging.info(f"No non-transparent bbox found for {image_path}. Returning original.")
|
||||
return img
|
||||
|
||||
if top_margin is None:
|
||||
top_margin = margin
|
||||
|
||||
# Add margin
|
||||
left, upper, right, lower = bbox
|
||||
left = max(0, left - margin)
|
||||
upper = max(0, upper - top_margin)
|
||||
right = min(img.width, right + margin)
|
||||
lower = min(img.height, lower + margin)
|
||||
|
||||
logging.info(f"Cropping {image_path} to {left, upper, right, lower} (margin={margin}, top_margin={top_margin})")
|
||||
cropped = img.crop((left, upper, right, lower))
|
||||
|
||||
if headroom > 0:
|
||||
h_px = int(cropped.height * headroom)
|
||||
if h_px > 0:
|
||||
logging.info(f"Adding {h_px}px headroom to {image_path}")
|
||||
new_img = Image.new("RGBA", (cropped.width, cropped.height + h_px), (0, 0, 0, 0))
|
||||
new_img.paste(cropped, (0, h_px))
|
||||
return new_img
|
||||
|
||||
return cropped
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to crop {image_path}: {e}")
|
||||
raise
|
||||
|
||||
def is_file_stable(filepath):
|
||||
"""Check if file size is stable for at least 1 second."""
|
||||
try:
|
||||
size1 = os.path.getsize(filepath)
|
||||
time.sleep(1)
|
||||
size2 = os.path.getsize(filepath)
|
||||
return size1 == size2 and size1 > 0
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
def flag_image(filename):
|
||||
input_path = os.path.join(CONF["stage_dir"], filename)
|
||||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
failed_filename = f"{timestamp}_{filename}"
|
||||
failed_path = os.path.join(CONF["failed_dir"], failed_filename)
|
||||
try:
|
||||
os.makedirs(CONF["failed_dir"], exist_ok=True)
|
||||
logging.info(f"Flagging image {filename} (moving to failed directory as {failed_filename})")
|
||||
shutil.move(input_path, failed_path)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to move {filename} to failed directory: {e}")
|
||||
|
||||
def process_image(filename):
|
||||
# Reload config in case it changed
|
||||
global CONF
|
||||
try:
|
||||
CONF = load_config()
|
||||
except:
|
||||
pass
|
||||
|
||||
input_path = os.path.join(CONF["stage_dir"], filename)
|
||||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
output_filename = f"{timestamp}_{filename}"
|
||||
output_path = os.path.join(CONF["output_dir"], output_filename)
|
||||
|
||||
temp_path = input_path + ".tmp.png"
|
||||
try:
|
||||
logging.info(f"Starting processing for {filename}...")
|
||||
cropped_img = crop_to_bbox(
|
||||
input_path,
|
||||
CONF["margin"],
|
||||
top_margin=CONF.get("top_margin"),
|
||||
headroom=CONF.get("headroom", 0.0)
|
||||
)
|
||||
|
||||
# Save temporary cropped image for upload
|
||||
cropped_img.save(temp_path, format="PNG")
|
||||
|
||||
prompt = CONF.get("prompt")
|
||||
if not prompt:
|
||||
bp = CONF.get("base_prompts", [])
|
||||
if bp and isinstance(bp, list) and len(bp) > 0:
|
||||
prompt = bp[0]
|
||||
else:
|
||||
prompt = "high quality, masterpiece"
|
||||
|
||||
with open(temp_path, 'rb') as f:
|
||||
files = {'image': (filename, f, 'image/png')}
|
||||
data = {
|
||||
'prompt': prompt,
|
||||
'seed': CONF.get("seed", -1),
|
||||
'max_area': CONF.get("max_area", 0)
|
||||
}
|
||||
logging.info(f"Calling API for {filename} -> {output_filename} with prompt: {prompt}")
|
||||
response = requests.post(CONF["api_url"], files=files, data=data, timeout=600)
|
||||
|
||||
if response.status_code == 200:
|
||||
with open(output_path, 'wb') as f:
|
||||
f.write(response.content)
|
||||
logging.info(f"Successfully processed {filename} -> {output_path}")
|
||||
|
||||
# Register in DB
|
||||
try:
|
||||
embedding = embeddings.generate_embedding(output_path)
|
||||
gid = filename
|
||||
database.upsert_person(output_filename, filepath=output_path, embedding=embedding, group_id=gid, is_source=True)
|
||||
|
||||
# Also trigger tagging to get auto-name and clip description
|
||||
tag_url = CONF["api_url"].replace("/edit", "/tag")
|
||||
try:
|
||||
requests.post(tag_url, json={"filename": output_filename, "group_id": gid}, timeout=30)
|
||||
except Exception as tag_err:
|
||||
logging.error(f"Error triggering tagging for {output_filename}: {tag_err}")
|
||||
except Exception as db_err:
|
||||
logging.error(f"Database error registering {output_filename}: {db_err}")
|
||||
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
return True
|
||||
else:
|
||||
logging.error(f"API failed for {filename}: {response.status_code} - {response.text}")
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
return False
|
||||
except requests.exceptions.ConnectionError as e:
|
||||
logging.error(f"Connection error while processing {filename}: {e}")
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
return None
|
||||
except Exception as e:
|
||||
logging.error(f"Error processing {filename}: {str(e)}", exc_info=True)
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
return False
|
||||
|
||||
def update_car_html():
|
||||
output_dir = CONF["output_dir"]
|
||||
car_html_path = os.path.join(output_dir, "car.html")
|
||||
if not os.path.exists(car_html_path):
|
||||
logging.warning(f"car.html not found at {car_html_path}")
|
||||
return
|
||||
|
||||
try:
|
||||
# Use database to list only non-archived images
|
||||
persons = database.list_persons(include_archived=False)
|
||||
db_images = {p[0] for p in persons}
|
||||
|
||||
# List images in output_dir
|
||||
extensions = ('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.svg')
|
||||
images = [f for f in os.listdir(output_dir)
|
||||
if f.lower().endswith(extensions) and f != "car.html" and f in db_images]
|
||||
|
||||
# Sort by mtime, newest first
|
||||
images.sort(key=lambda x: os.path.getmtime(os.path.join(output_dir, x)), reverse=True)
|
||||
|
||||
with open(car_html_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
images_json = json.dumps(images, indent=12).strip()
|
||||
# Ensure it looks nice in the JS
|
||||
images_json = images_json.replace('\n', '\n ')
|
||||
|
||||
pattern = r'// --- HYDRATION_START ---.*?// --- HYDRATION_END ---'
|
||||
replacement = f'// --- HYDRATION_START ---\n const PRELOADED_IMAGES = {images_json};\n // --- HYDRATION_END ---'
|
||||
|
||||
new_content = re.sub(pattern, replacement, content, flags=re.DOTALL)
|
||||
|
||||
tmp_path = car_html_path + ".tmp"
|
||||
with open(tmp_path, 'w') as f:
|
||||
f.write(new_content)
|
||||
os.replace(tmp_path, car_html_path)
|
||||
logging.info(f"Updated {car_html_path} atomically with {len(images)} images")
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to update car.html: {e}")
|
||||
|
||||
def main():
|
||||
# Prevent multiple instances
|
||||
lock_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "watcher.lock")
|
||||
fp = open(lock_file, 'w')
|
||||
try:
|
||||
fcntl.lockf(fp, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
except IOError:
|
||||
print("Another instance of watcher.py is already running. Exiting.")
|
||||
sys.exit(1)
|
||||
|
||||
processed = get_processed_files()
|
||||
|
||||
# Ensure directories exist
|
||||
os.makedirs(CONF["stage_dir"], exist_ok=True)
|
||||
os.makedirs(CONF["output_dir"], exist_ok=True)
|
||||
os.makedirs(CONF["failed_dir"], exist_ok=True)
|
||||
|
||||
logging.info(f"Watcher started. Monitoring {CONF['stage_dir']}...")
|
||||
logging.info(f"Output directory: {CONF['output_dir']}")
|
||||
logging.info(f"API URL: {CONF['api_url']}")
|
||||
|
||||
while True:
|
||||
try:
|
||||
files = [f for f in os.listdir(CONF["stage_dir"])
|
||||
if f.lower().endswith(('.png', '.jpg', '.jpeg'))
|
||||
and not f.endswith('.tmp.png')]
|
||||
|
||||
for f in files:
|
||||
input_path = os.path.join(CONF["stage_dir"], f)
|
||||
|
||||
# Check if file is stable (not still being copied)
|
||||
if not is_file_stable(input_path):
|
||||
continue
|
||||
|
||||
# Calculate current file hash
|
||||
current_hash = calculate_hash(input_path)
|
||||
if not current_hash:
|
||||
continue
|
||||
|
||||
# Check if already processed
|
||||
if f in processed:
|
||||
stored_hash = processed[f]
|
||||
if stored_hash == current_hash:
|
||||
continue
|
||||
if stored_hash is None:
|
||||
# Migration case: filename exists but no hash.
|
||||
# Skip to avoid mass re-processing, but update the hash.
|
||||
logging.info(f"Updating hash for previously processed {f}")
|
||||
processed[f] = current_hash
|
||||
save_processed_files(processed)
|
||||
continue
|
||||
|
||||
res = process_image(f)
|
||||
if res is True:
|
||||
processed[f] = current_hash
|
||||
save_processed_files(processed)
|
||||
update_car_html()
|
||||
elif res is False:
|
||||
flag_image(f)
|
||||
# We don't add to processed here so that if the user
|
||||
# moves the file back to stage, it will be retried.
|
||||
except Exception as e:
|
||||
logging.error(f"Main loop error: {e}")
|
||||
|
||||
time.sleep(CONF["poll_interval"])
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user