335 lines
12 KiB
Python
335 lines
12 KiB
Python
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()
|