Resume reclass jobs and expand metadata detection

backend/app.py

@@ -25,7 +25,7 @@ from fastapi import BackgroundTasks, FastAPI, HTTPException, Request
 from fastapi.responses import FileResponse, Response, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
-from PIL import Image
+from PIL import ExifTags, Image
 from pydantic import BaseModel
 
 TURSO_DB_URL = os.getenv("TURSO_DB_URL", "").strip()
@@ -49,7 +49,18 @@ AI_TAGS = {
     "ai",
     "comfyui",
 }
-EXIF_TYPE_ORDER = ("novelai", "comfy", "sd", "mj", "celsys", "photoshop", "stealth")
+EXIF_TYPE_ORDER = (
+    "novelai",
+    "comfy",
+    "sd",
+    "mj",
+    "celsys",
+    "photoshop",
+    "stealth",
+    "invoke",
+    "fooocus",
+    "swarm",
+)
 EXIF_TYPE_TO_CODE = {name: idx + 1 for idx, name in enumerate(EXIF_TYPE_ORDER)}
 EXIF_CODE_TO_TYPE = {idx + 1: name for idx, name in enumerate(EXIF_TYPE_ORDER)}
 POST_SCAN_LIMIT = 64
@@ -58,6 +69,7 @@ EXIF_RANGE_BYTES = 16384
 FULL_IMAGE_LIMIT = 32
 PAGE_SIZE = 60
 SEARCH_PAGE_SIZE = 10
+RECLASS_BATCH_SIZE = 64
 SEARCH_LIST_WHERE = "post_ids != '[]' AND found_exif > 0"
 POST_SEARCH_LIMIT = 1000
 RECLASS_LOG_KEY = "TF#g9FA@UgKdN'usVtm3dTz*yoqbria`"
@@ -179,6 +191,8 @@ async def init_db():
 
 async def migrate_exif_order():
     order_value = ",".join(EXIF_TYPE_ORDER)
+    swapped_order = "novelai,comfy,sd,mj,celsys,photoshop,stealth"
+    legacy_order = "novelai,sd,comfy,mj,celsys,photoshop,stealth"
     resp = await turso_execute(
         [
             {
@@ -187,22 +201,31 @@ async def migrate_exif_order():
         ]
     )
     rows = first_rows(resp)
-    if rows and cell_value(rows[0], 0, "") == order_value:
+    prev_order = cell_value(rows[0], 0, "") if rows else ""
+    if prev_order == order_value:
         return
-    await turso_execute(
-        [
-            {
-                "sql": "UPDATE pi_scans SET exif_type = CASE exif_type "
-                "WHEN 2 THEN -2 WHEN 3 THEN -3 ELSE exif_type END"
-            },
-            {"sql": "UPDATE pi_scans SET exif_type = 3 WHERE exif_type = -2"},
-            {"sql": "UPDATE pi_scans SET exif_type = 2 WHERE exif_type = -3"},
+    stmts = []
+    if prev_order == legacy_order:
+        stmts.extend(
+            [
+                {
+                    "sql": "UPDATE pi_scans SET exif_type = CASE exif_type "
+                    "WHEN 2 THEN -2 WHEN 3 THEN -3 ELSE exif_type END"
+                },
+                {"sql": "UPDATE pi_scans SET exif_type = 3 WHERE exif_type = -2"},
+                {"sql": "UPDATE pi_scans SET exif_type = 2 WHERE exif_type = -3"},
+            ]
+        )
+    if prev_order in ("", swapped_order, legacy_order):
+        stmts.append(
             {
-                "sql": "INSERT OR REPLACE INTO pi_meta (key, value) VALUES ('exif_order', ?)",
+                "sql": "INSERT OR REPLACE INTO pi_meta "
+                "(key, value) VALUES ('exif_order', ?)",
                 "args": [{"type": "text", "value": order_value}],
-            },
-        ]
-    )
+            }
+        )
+    if stmts:
+        await turso_execute(stmts)
 
 
 async def discord_notify(msg):
@@ -688,6 +711,8 @@ def decode_meta_text(data):
 
 
 def parse_png_metadata(data):
+    if not data.startswith(b"\x89PNG\r\n\x1a\n"):
+        return None
     fields = {}
     index = 8
     while index < len(data):
@@ -723,6 +748,99 @@ def parse_png_metadata(data):
     return fields or None
 
 
+def decode_exif_text(value):
+    if value is None:
+        return ""
+    if isinstance(value, bytes):
+        prefixes = (
+            (b"ASCII\x00\x00\x00", "utf-8"),
+            (b"UNICODE\x00", "utf-16-be"),
+            (b"JIS\x00\x00\x00\x00\x00", "shift_jis"),
+        )
+        for prefix, enc in prefixes:
+            if value.startswith(prefix):
+                try:
+                    return value[len(prefix) :].rstrip(b"\0").decode(enc).strip()
+                except UnicodeDecodeError:
+                    break
+        return decode_meta_text(value.rstrip(b"\0"))
+    return str(value).strip()
+
+
+def parse_pil_metadata(data):
+    fields = {}
+    try:
+        image = Image.open(io.BytesIO(data))
+        for key, value in (image.info or {}).items():
+            if isinstance(value, (str, bytes)):
+                fields[str(key)] = decode_exif_text(value)
+        for tag, value in image.getexif().items():
+            name = ExifTags.TAGS.get(tag, str(tag))
+            fields[name] = decode_exif_text(value)
+    except Exception:
+        pass
+    return {k: v for k, v in fields.items() if v} or None
+
+
+def parse_jpeg_comments(data):
+    if not data.startswith(b"\xff\xd8"):
+        return None
+    fields = {}
+    index = 2
+    n = 1
+    while index + 4 <= len(data):
+        if data[index] != 0xFF:
+            index += 1
+            continue
+        marker = data[index + 1]
+        index += 2
+        if marker == 0xDA:
+            break
+        if marker in (0x01, *range(0xD0, 0xD9)):
+            continue
+        size = int.from_bytes(data[index : index + 2], "big")
+        if size < 2 or index + size > len(data):
+            break
+        content = data[index + 2 : index + size]
+        if marker == 0xFE:
+            fields[f"Comment{n}"] = decode_meta_text(content)
+            n += 1
+        index += size
+    return fields or None
+
+
+def parse_xmp_metadata(data):
+    start = data.find(b"<x:xmpmeta")
+    if start < 0:
+        start = data.find(b"<?xpacket")
+    if start < 0:
+        return None
+    end = data.find(b"</x:xmpmeta>", start)
+    if end >= 0:
+        end += len(b"</x:xmpmeta>")
+    else:
+        end = data.find(b"?>", start)
+        if end >= 0:
+            end += 2
+    if end < 0:
+        return None
+    return {"XMP": decode_meta_text(data[start:end])}
+
+
+def parse_image_metadata(data):
+    fields = {}
+    for parser in (
+        parse_png_metadata,
+        parse_pil_metadata,
+        parse_jpeg_comments,
+        parse_xmp_metadata,
+    ):
+        parsed = parser(data)
+        if parsed:
+            fields.update(parsed)
+    return fields or None
+
+
 def parse_json_maybe(value):
     try:
         return json.loads(value)
@@ -771,6 +889,76 @@ def has_comfy_metadata(fields):
     return False
 
 
+def field_json_values(fields):
+    for value in fields.values():
+        data = parse_json_maybe(value)
+        if isinstance(data, dict):
+            yield data
+
+
+def has_fooocus_metadata(fields):
+    scheme = fields.get("fooocus_scheme") or fields.get("MakerNote")
+    if str(scheme).casefold() in ("fooocus", "a1111"):
+        return True
+    software = fields.get("Software", "").casefold()
+    if "fooocus" in software:
+        return True
+    for data in field_json_values(fields):
+        keys = set(data)
+        if {
+            "adaptive_cfg",
+            "base_model_hash",
+            "full_negative_prompt",
+            "full_prompt",
+            "loras",
+            "refiner_model_hash",
+        } & keys:
+            return True
+    return False
+
+
+def has_invoke_metadata(fields):
+    if any(k.casefold() in ("dream", "sd-metadata") for k in fields):
+        return True
+    joined = "\n".join(str(v) for v in fields.values()).casefold()
+    if "invokeai" in joined or "invoke ai" in joined:
+        return True
+    for data in field_json_values(fields):
+        app = str(data.get("app_id", "")).casefold()
+        if "invoke" in app:
+            return True
+    return False
+
+
+def has_swarm_metadata(fields):
+    for data in field_json_values(fields):
+        if {"sui_image_params", "sui_extra_data", "sui_models"} & set(data):
+            return True
+    return "sui_image_params" in "\n".join(str(v) for v in fields.values()).casefold()
+
+
+def has_sd_metadata(fields):
+    if "parameters" in fields:
+        return True
+    for value in fields.values():
+        text = str(value)
+        folded = text.casefold()
+        if (
+            "steps:" in folded
+            and "seed:" in folded
+            and ("sampler:" in folded or "cfg scale:" in folded or "model:" in folded)
+        ):
+            return True
+        data = parse_json_maybe(value)
+        if isinstance(data, dict):
+            keys = set(data)
+            if "prompt" in keys and "seed" in keys and (
+                "steps" in keys or "sampler" in keys or "cfg_scale" in keys
+            ):
+                return True
+    return False
+
+
 def determine_exif_type(fields):
     if not fields:
         return None
@@ -789,7 +977,13 @@ def determine_exif_type(fields):
         return "novelai"
     if has_comfy_metadata(fields):
         return "comfy"
-    if "parameters" in fields or joined.startswith("parameters"):
+    if has_fooocus_metadata(fields):
+        return "fooocus"
+    if has_invoke_metadata(fields):
+        return "invoke"
+    if has_swarm_metadata(fields):
+        return "swarm"
+    if has_sd_metadata(fields) or joined.startswith("parameters"):
         return "sd"
     if "midjourney" in joined:
         return "mj"
@@ -829,7 +1023,12 @@ async def scan_post(post_id, session, post_sem, exif_sem, img_sem):
     async with post_sem:
         try:
             pages = await get_post_pages(post_id, session)
-            image_urls = [p["original"] for p in pages if "png" in p["original"]]
+            image_urls = [
+                p["original"]
+                for p in pages
+                if urlsplit(p["original"]).path.rsplit(".", 1)[-1].casefold()
+                in ("png", "jpg", "jpeg", "webp")
+            ]
             for url in image_urls:
                 metadata = await get_exif_range(url, session, exif_sem)
                 exif_type = determine_exif_type(metadata)
@@ -839,7 +1038,7 @@ async def scan_post(post_id, session, post_sem, exif_sem, img_sem):
             for url in image_urls:
                 img_data = await fetch_image(session, url, img_sem)
                 if img_data and has_stealth_png_bytes(img_data):
-                    return post_id, url, EXIF_TYPE_TO_CODE.get("comfy")
+                    return post_id, url, EXIF_TYPE_TO_CODE.get("stealth")
             return post_id, None, None
         except Exception:
             return post_id, None, None
@@ -857,7 +1056,7 @@ async def get_exif_range(url, session, sem):
     else:
         async with session.get(url, headers=hdrs) as r:
             data = await r.read()
-    return parse_png_metadata(data)
+    return parse_image_metadata(data)
 
 
 async def fetch_image(session, url, sem):
@@ -869,7 +1068,7 @@ async def fetch_image(session, url, sem):
         return await r.read()
 
 
-async def run_scan(post_ids, phpsessid, task_id=None, save_live=False):
+async def run_scan(post_ids, phpsessid, task_id=None, save_live=False, done_offset=0):
     post_sem = asyncio.Semaphore(POST_SCAN_LIMIT)
     exif_sem = asyncio.Semaphore(EXIF_RANGE_LIMIT)
     img_sem = asyncio.Semaphore(FULL_IMAGE_LIMIT)
@@ -889,7 +1088,7 @@ async def run_scan(post_ids, phpsessid, task_id=None, save_live=False):
                 results.append(result)
                 pending.append(result)
                 if task_id and task_id in ACTIVE_TASKS:
-                    ACTIVE_TASKS[task_id]["done"] = len(results)
+                    ACTIVE_TASKS[task_id]["done"] = done_offset + len(results)
                 if save_live and len(pending) >= 20:
                     if task_canceled(task_id):
                         raise TaskCanceled()
@@ -952,6 +1151,12 @@ async def get_reclass_targets(allowed):
 
 
 async def save_reclass_job(job_id, post_ids, filters):
+    summary = {
+        "scanned": 0,
+        "changed": 0,
+        "updated": 0,
+        "filters": [exif_label(code) for code in filters],
+    }
     await turso_execute(
         [
             {
@@ -963,7 +1168,7 @@ async def save_reclass_job(job_id, post_ids, filters):
                     {"type": "text", "value": json.dumps(post_ids)},
                     {"type": "text", "value": json.dumps(filters)},
                     {"type": "integer", "value": str(len(post_ids))},
-                    {"type": "text", "value": json.dumps({})},
+                    {"type": "text", "value": json.dumps(summary)},
                     {"type": "text", "value": encrypt_log([])},
                     {"type": "integer", "value": str(int(time.time()))},
                 ],
@@ -1252,51 +1457,75 @@ async def bg_search_and_scan_task(search_id, url, pages, mode, phpsessid):
         await finish_task(search_id)
 
 
-async def bg_reclass_task(job_id, targets, filters, phpsessid):
+async def bg_reclass_task(
+    job_id,
+    targets,
+    filters,
+    phpsessid,
+    total=None,
+    scanned=0,
+    changes=None,
+    updated=0,
+):
+    total = total or len(targets)
+    changes = list(changes or [])
     ACTIVE_TASKS[job_id] = {
         "type": "reclass",
         "phase": "scanning",
-        "total": len(targets),
-        "done": 0,
+        "total": total,
+        "done": scanned,
+        "changed": len(changes),
     }
-    await discord_notify(f"`{job_id}` reclass started ({len(targets)} posts)")
-    changes = []
-    changed_results = []
+    action = "resumed" if scanned else "started"
+    await discord_notify(f"`{job_id}` reclass {action} ({total - scanned} posts)")
     try:
-        old = {item["post_id"]: item for item in targets}
-        post_ids = list(old)
-        results = await run_scan(post_ids, phpsessid, task_id=job_id, save_live=False)
-        for post_id, url, exif_type in results:
-            item = old.get(str(post_id))
-            if not item:
-                continue
-            short_url = url.replace(IMG_BASE, "", 1) if url else ""
-            old_type = item["exif_type"]
-            old_url = item["url"]
-            if old_type != exif_type:
-                changes.append(
-                    {
-                        "post_id": str(post_id),
-                        "from": exif_label(old_type),
-                        "to": exif_label(exif_type),
-                    }
+        for i in range(0, len(targets), RECLASS_BATCH_SIZE):
+            if task_canceled(job_id):
+                raise TaskCanceled()
+            batch = targets[i : i + RECLASS_BATCH_SIZE]
+            old = {item["post_id"]: item for item in batch}
+            results = await run_scan(
+                list(old),
+                phpsessid,
+                task_id=job_id,
+                save_live=False,
+                done_offset=scanned,
+            )
+            changed_results = []
+            for post_id, url, exif_type in results:
+                item = old.get(str(post_id))
+                if not item:
+                    continue
+                short_url = url.replace(IMG_BASE, "", 1) if url else ""
+                old_type = item["exif_type"]
+                old_url = item["url"]
+                if old_type != exif_type:
+                    changes.append(
+                        {
+                            "post_id": str(post_id),
+                            "from": exif_label(old_type),
+                            "to": exif_label(exif_type),
+                        }
+                    )
+                if old_type != exif_type or old_url != short_url:
+                    changed_results.append((post_id, url, exif_type))
+            if changed_results:
+                await save_scan_results(changed_results)
+                await update_searches_for_post_ids(
+                    [str(pid) for pid, _, _ in changed_results]
                 )
-            if old_type != exif_type or old_url != short_url:
-                changed_results.append((post_id, url, exif_type))
-        if changed_results:
-            await save_scan_results(changed_results)
-            await update_searches_for_post_ids([str(pid) for pid, _, _ in changed_results])
-        await finish_reclass_job(
-            job_id, len(results), changes, len(changed_results), filters
-        )
+                updated += len(changed_results)
+            scanned += len(results)
+            ACTIVE_TASKS[job_id].update({"done": scanned, "changed": len(changes)})
+            await finish_reclass_job(job_id, scanned, changes, updated, filters)
         await discord_notify(f"`{job_id}` Done")
     except TaskCanceled:
         done = ACTIVE_TASKS.get(job_id, {}).get("done", 0)
-        await finish_reclass_job(job_id, done, changes, len(changed_results), filters)
+        await finish_reclass_job(job_id, done, changes, updated, filters)
         await discord_notify(f"`{job_id}` reclass canceled")
     except Exception as e:
         done = ACTIVE_TASKS.get(job_id, {}).get("done", 0)
-        await finish_reclass_job(job_id, done, changes, len(changed_results), filters)
+        await finish_reclass_job(job_id, done, changes, updated, filters)
         await discord_notify(f"`{job_id}` reclass failed: {e}")
     finally:
         await finish_task(job_id)
@@ -1339,6 +1568,7 @@ async def startup():
         await init_db()
         asyncio.create_task(backfill_search_stats())
         asyncio.create_task(resume_incomplete_scans())
+        asyncio.create_task(resume_incomplete_reclasses())
     except Exception as e:
         print(f"WARN: DB init failed ({e}), will retry on first request")
 
@@ -1384,6 +1614,85 @@ async def resume_incomplete_scans():
         print(f"WARN: scan resume failed: {e}")
 
 
+async def get_incomplete_reclass_jobs():
+    resp = await turso_execute(
+        [
+            {
+                "sql": "SELECT id, post_ids, filters, total, scanned, summary, log_enc "
+                "FROM pi_reclass_jobs WHERE scanned < total ORDER BY id",
+            }
+        ]
+    )
+    jobs = []
+    for row in first_rows(resp):
+        job_id = cell_value(row, 0, "")
+        post_ids = json.loads(cell_value(row, 1, "[]"))
+        filters = json.loads(cell_value(row, 2, "[]"))
+        total = int(cell_value(row, 3, "0"))
+        scanned = min(int(cell_value(row, 4, "0")), total)
+        summary = parse_json_maybe(cell_value(row, 5, "{}")) or {}
+        changes = decrypt_log(cell_value(row, 6, ""))
+        remaining = post_ids[scanned:]
+        current = await get_scanned_post_ids(remaining)
+        targets = []
+        for post_id in remaining:
+            item = current.get(str(post_id), {})
+            targets.append(
+                {
+                    "post_id": str(post_id),
+                    "url": item.get("url", ""),
+                    "exif_type": item.get("exif_type"),
+                }
+            )
+        jobs.append(
+            {
+                "id": job_id,
+                "targets": targets,
+                "filters": filters,
+                "total": total or len(post_ids),
+                "scanned": scanned,
+                "changes": changes,
+                "updated": int(summary.get("updated") or 0),
+            }
+        )
+    return jobs
+
+
+async def resume_incomplete_reclasses():
+    try:
+        started = 0
+        for job in await get_incomplete_reclass_jobs():
+            job_id = job["id"]
+            if job_id in ACTIVE_TASKS:
+                continue
+            if job["targets"]:
+                asyncio.create_task(
+                    bg_reclass_task(
+                        job_id,
+                        job["targets"],
+                        job["filters"],
+                        PHPSESSID,
+                        total=job["total"],
+                        scanned=job["scanned"],
+                        changes=job["changes"],
+                        updated=job["updated"],
+                    )
+                )
+                started += 1
+            else:
+                await finish_reclass_job(
+                    job_id,
+                    job["scanned"],
+                    job["changes"],
+                    job["updated"],
+                    job["filters"],
+                )
+        if started:
+            print(f"Resumed {started} incomplete reclasses")
+    except Exception as e:
+        print(f"WARN: reclass resume failed: {e}")
+
+
 @app.post("/api/submit")
 async def submit_search(req: SearchRequest, bg: BackgroundTasks):
     search_id = base26_time()
@@ -1488,9 +1797,11 @@ async def get_reclass(job_id: str):
     active = ACTIVE_TASKS.get(job_id)
     total = int(cell_value(row, 3, "0"))
     scanned = int(cell_value(row, 4, "0"))
+    changed = int(cell_value(row, 5, "0"))
     if active:
         total = int(active.get("total") or total)
         scanned = int(active.get("done") or scanned)
+        changed = int(active.get("changed") or changed)
     return {
         "id": cell_value(row, 0, ""),
         "kind": "reclass",
@@ -1498,7 +1809,7 @@ async def get_reclass(job_id: str):
         "filters": json.loads(cell_value(row, 2, "[]")),
         "total": total,
         "scanned": scanned,
-        "changed": int(cell_value(row, 5, "0")),
+        "changed": changed,
         "summary": summary,
         "log": decrypt_log(cell_value(row, 7, "")),
         "created_at": int(cell_value(row, 8, "0")),

frontend/app.js

@@ -1,7 +1,7 @@
 const $ = s => document.querySelector(s)
 const $$ = s => document.querySelectorAll(s)
-const EXIF_NAMES = { 1: "novelai", 2: "comfy", 3: "sd", 4: "mj", 5: "celsys", 6: "photoshop", 7: "stealth" }
-const EXIF_CODES = [1, 2, 3, 4, 5, 6, 7]
+const EXIF_NAMES = { 1: "novelai", 2: "comfy", 3: "sd", 4: "mj", 5: "celsys", 6: "photoshop", 7: "stealth", 8: "invoke", 9: "fooocus", 10: "swarm" }
+const EXIF_CODES = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
 const CATEGORY_NAMES = { ...EXIF_NAMES, 0: "None" }
 const FILTER_CODES = [...EXIF_CODES, 0]
 const EXIF_FILTER_KEY = "pixif2-exif-types"
@@ -225,7 +225,7 @@ async function loadSearches(page = 1, silent = false) {
     const savedHtml = searches.map(s => {
       const d = new Date(parseInt(s.created_at) * 1000)
       const ts = d.toLocaleString()
-      const count = s.kind === "reclass" ? `${s.changed || 0}/${s.total || 0}` : `${s.found_exif || 0}/${s.total_searched || 0}`
+      const count = s.kind === "reclass" ? `${s.total_searched || 0}/${s.total || 0}` : `${s.found_exif || 0}/${s.total_searched || 0}`
       const kind = s.kind === "reclass" ? "reclass " : ""
       return `<div class="search-item" data-id="${esc(s.id)}">
         <span class="id">${esc(s.id)}</span>