diff --git a/src/bin/link_versions.rs b/src/bin/link_versions.rs new file mode 100644 index 0000000..580d424 --- /dev/null +++ b/src/bin/link_versions.rs @@ -0,0 +1,225 @@ +//! Link photo versions that were ingested before linking existed (#191). +//! +//! **Default is a dry run**: nothing is written. It prints a TSV of every +//! proposed link (`status method library_id version_path original_path +//! phash_distance kind ...`) and a summary on stderr. +//! +//! - `exact` rows (the PNG's XMP `DerivedFrom` hash, or the +//! `.[.].` convention next to an indexed original) are +//! the same links the watcher makes at ingest; `--apply` writes them +//! without review. +//! - `review` rows (ComfyUI prompt `LoadImage` input, Upscayl-style names, +//! nearest pHash) are heuristic. Save the dry-run output, delete the rows +//! you don't want, and pass the file as `--reviewed`; only those are +//! written, and each is checked against the database first so a stale file +//! can't link a path that now holds different bytes. +//! - `unresolved` rows are listed, never guessed. +//! +//! `--apply` first takes a WAL-safe snapshot of the database +//! (`VACUUM INTO`) and refuses to write if that fails. Prove it on a copy of +//! the database first. The library roots are only ever read. + +use std::collections::HashSet; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + +use clap::Parser; +use diesel::prelude::*; +use diesel::sql_types::{BigInt, Integer, Nullable, Text}; + +use image_api::database::version_dao::VersionDao; +use image_api::database::{SqliteVersionDao, connect}; +use image_api::libraries; +use image_api::versions::backfill::{self, FileRow, Status}; + +#[derive(Parser, Debug)] +#[command(name = "link_versions")] +#[command(about = "Propose (dry run) or write links from photo versions to their originals")] +struct Args { + /// Write the links. Without it nothing is changed. + #[arg(long)] + apply: bool, + + /// A reviewed copy of the dry-run TSV: its `review` rows are linked too. + /// Only meaningful with --apply. + #[arg(long, value_name = "FILE")] + reviewed: Option, + + /// Directory for the pre-apply snapshot (default: next to the database). + #[arg(long, value_name = "DIR")] + backup_dir: Option, + + /// Only consider this library id (repeatable). Default: every library. + #[arg(long = "library", value_name = "ID")] + libraries: Vec, +} + +#[derive(QueryableByName)] +struct DbRow { + #[diesel(sql_type = Integer)] + library_id: i32, + #[diesel(sql_type = Text)] + rel_path: String, + #[diesel(sql_type = Text)] + content_hash: String, + #[diesel(sql_type = Nullable)] + phash_64: Option, +} + +#[derive(QueryableByName)] +struct Count { + #[diesel(sql_type = BigInt)] + n: i64, +} + +#[derive(QueryableByName)] +struct Hash { + #[diesel(sql_type = Text)] + version_hash: String, +} + +fn main() -> anyhow::Result<()> { + env_logger::init(); + dotenv::dotenv().ok(); + let args = Args::parse(); + if args.reviewed.is_some() && !args.apply { + anyhow::bail!("--reviewed only applies together with --apply (the default is a dry run)"); + } + + let mut conn = connect(); + let has_table = diesel::sql_query( + "SELECT COUNT(*) AS n FROM sqlite_master WHERE type = 'table' AND name = 'photo_versions'", + ) + .get_result::(&mut conn)?; + if has_table.n == 0 { + anyhow::bail!( + "this database has no photo_versions table; start the server once (it runs the \ + migrations) or point DATABASE_URL at a migrated copy" + ); + } + + let mut libs = libraries::load_all(&mut conn); + if !args.libraries.is_empty() { + libs.retain(|l| args.libraries.contains(&l.id)); + } + if libs.is_empty() { + anyhow::bail!("no matching libraries"); + } + let lib_ids: HashSet = libs.iter().map(|l| l.id).collect(); + + let rows: Vec = diesel::sql_query( + "SELECT library_id, rel_path, content_hash, phash_64 FROM image_exif \ + WHERE content_hash IS NOT NULL ORDER BY library_id, rel_path", + ) + .load::(&mut conn)? + .into_iter() + .filter(|r| lib_ids.contains(&r.library_id)) + .map(|r| FileRow { + library_id: r.library_id, + rel_path: r.rel_path, + content_hash: r.content_hash, + phash: r.phash_64, + }) + .collect(); + let linked: HashSet = diesel::sql_query("SELECT version_hash FROM photo_versions") + .load::(&mut conn)? + .into_iter() + .map(|h| h.version_hash) + .collect(); + + eprintln!( + "{} hashed image/video rows, {} version links already recorded", + rows.len(), + linked.len() + ); + let proposals = backfill::plan(&libs, &rows, &linked); + + println!("{}", backfill::TSV_HEADER); + for p in &proposals { + println!("{}", p.to_tsv()); + } + let count = |s: Status| proposals.iter().filter(|p| p.status == s).count(); + eprintln!( + "\nproposals: {} exact (apply without review), {} review (need --reviewed), {} unresolved", + count(Status::Exact), + count(Status::Review), + count(Status::Unresolved) + ); + for method in ["xmp", "name", "comfy-loader", "name-prefix", "phash"] { + let n = proposals.iter().filter(|p| p.method == method).count(); + if n > 0 { + eprintln!(" {method}: {n}"); + } + } + + if !args.apply { + eprintln!("dry run: nothing written. Re-run with --apply to write the exact links."); + return Ok(()); + } + + let reviewed = match &args.reviewed { + Some(path) => { + let text = std::fs::read_to_string(path)?; + let (links, skipped) = backfill::parse_reviewed(&text); + eprintln!( + "reviewed file: {} review row(s) to link, {} other line(s) ignored", + links.len(), + skipped + ); + links + } + None => Vec::new(), + }; + + let backup = backup_database(&mut conn, args.backup_dir.as_deref())?; + eprintln!("snapshot written: {}", backup.display()); + drop(conn); + + let dao: Arc>> = + Arc::new(Mutex::new(Box::new(SqliteVersionDao::new()))); + let mut guard = dao.lock().unwrap(); + let report = backfill::apply(guard.as_mut(), &proposals, &reviewed, &rows); + let synced = guard + .sync_version_metadata() + .map_err(|e| anyhow::anyhow!("{e:?}"))?; + eprintln!( + "linked {} ({} already linked, {} refused, {} stale reviewed row(s)); \ + date/GPS/camera copied onto {} version row(s)", + report.linked, report.already_linked, report.refused, report.stale, synced + ); + Ok(()) +} + +/// Snapshot the live database with `VACUUM INTO`, which is consistent with +/// the WAL (a plain file copy is not). +fn backup_database( + conn: &mut diesel::SqliteConnection, + dir: Option<&std::path::Path>, +) -> anyhow::Result { + let url = dotenv::var("DATABASE_URL")?; + let db_path = PathBuf::from(url.trim_start_matches("sqlite://")); + let dir = match dir { + Some(d) => d.to_path_buf(), + None => db_path + .parent() + .filter(|p| !p.as_os_str().is_empty()) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(".")), + }; + let stem = db_path + .file_name() + .map(|n| n.to_string_lossy().into_owned()) + .unwrap_or_else(|| "database".into()); + let target = dir.join(format!( + "{stem}.pre-link_versions-{}.bak", + chrono::Utc::now().format("%Y%m%dT%H%M%SZ") + )); + if target.exists() { + anyhow::bail!("{} already exists", target.display()); + } + let escaped = target.to_string_lossy().replace('\'', "''"); + diesel::sql_query(format!("VACUUM INTO '{escaped}'")) + .execute(conn) + .map_err(|e| anyhow::anyhow!("snapshot failed, nothing was written: {e}"))?; + Ok(target) +} diff --git a/src/versions/backfill.rs b/src/versions/backfill.rs new file mode 100644 index 0000000..0d7af8b --- /dev/null +++ b/src/versions/backfill.rs @@ -0,0 +1,931 @@ +//! Backfill of photo versions that were ingested before linking existed +//! (the `link_versions` binary). +//! +//! Two kinds of proposal: +//! +//! - **Exact** (`xmp`, `name`): the PNG carries the producer's XMP +//! `DerivedFrom` hash, or its name follows `.[.].` next +//! to an indexed original. The ingest path would have made the same link, +//! so these apply without review. +//! - **Review** (`comfy-loader`, `name-prefix`, `phash`): files with no XMP, +//! matched by the ComfyUI prompt's `LoadImage` input, an Upscayl-style +//! name, or the nearest pHash. Heuristic, so they only go through a TSV a +//! human has read (`--reviewed`). Ambiguous or unresolvable cases are +//! listed, never guessed. +//! +//! [`plan`] only reads (database rows handed in, library files for XMP and +//! the prompt chunk); [`apply`] is the only writer and goes through +//! [`VersionDao`]. + +use std::collections::{HashMap, HashSet}; +use std::path::Path; +use std::sync::LazyLock; + +use regex::Regex; + +use super::xmp::{self, XMP_KEYWORD}; +use super::{VERSIONS_SUBFOLDER, original_candidates, parse_version_name}; +use crate::database::version_dao::{ + LinkOutcome, NewVersionLink, SOURCE_BACKFILL, SOURCE_NAME, SOURCE_XMP, VersionDao, + is_known_kind, +}; +use crate::file_types; +use crate::libraries::Library; + +/// A nearest-pHash match farther than this (of 64 bits) is not proposed. +pub const MAX_PHASH_DISTANCE: u32 = 16; + +/// One `image_exif` row, as far as planning needs it. +#[derive(Debug, Clone)] +pub struct FileRow { + pub library_id: i32, + pub rel_path: String, + pub content_hash: String, + pub phash: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Status { + /// Applies without review. + Exact, + /// Heuristic; applies only from a reviewed TSV. + Review, + /// Listed, not guessed. + Unresolved, +} + +impl Status { + fn as_str(self) -> &'static str { + match self { + Status::Exact => "exact", + Status::Review => "review", + Status::Unresolved => "unresolved", + } + } +} + +#[derive(Debug, Clone)] +pub struct Proposal { + pub status: Status, + /// `xmp` | `name` | `comfy-loader` | `name-prefix` | `phash` + pub method: &'static str, + pub library_id: i32, + pub version_path: String, + pub version_hash: String, + /// Empty for an unresolved case, or an XMP link whose original isn't indexed. + pub original_path: String, + pub original_hash: String, + pub kind: String, + pub distance: Option, + /// Hashes of reference photos (the other loader inputs). + pub references: Vec, + pub notes: String, +} + +pub const TSV_HEADER: &str = "status\tmethod\tlibrary_id\tversion_path\toriginal_path\tphash_distance\tkind\tversion_hash\toriginal_hash\treferences\tnotes"; + +impl Proposal { + pub fn to_tsv(&self) -> String { + let clean = |s: &str| s.replace(['\t', '\n', '\r'], " "); + [ + self.status.as_str().to_string(), + self.method.to_string(), + self.library_id.to_string(), + clean(&self.version_path), + clean(&self.original_path), + self.distance.map(|d| d.to_string()).unwrap_or_default(), + self.kind.clone(), + self.version_hash.clone(), + self.original_hash.clone(), + self.references.join(","), + clean(&self.notes), + ] + .join("\t") + } +} + +/// A row of a reviewed TSV: just enough to link. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ReviewedLink { + pub method: String, + pub library_id: i32, + pub version_path: String, + pub kind: String, + pub version_hash: String, + pub original_hash: String, + pub references: Vec, +} + +/// Parse a reviewed TSV (the dry-run output with rows deleted or kept). +/// Rows that aren't `review` status with both hashes are skipped; so are the +/// header, blanks and `#` comments. Returns `(links, skipped_lines)`. +pub fn parse_reviewed(tsv: &str) -> (Vec, usize) { + let mut out = Vec::new(); + let mut skipped = 0; + for line in tsv.lines() { + if line.trim().is_empty() || line.starts_with('#') || line.starts_with("status\t") { + continue; + } + let f: Vec<&str> = line.split('\t').collect(); + let usable = f.len() >= 10 + && f[0] == "review" + && f[2].parse::().is_ok() + && !f[7].is_empty() + && !f[8].is_empty(); + if !usable { + skipped += 1; + continue; + } + out.push(ReviewedLink { + method: f[1].to_string(), + library_id: f[2].parse().unwrap(), + version_path: f[3].to_string(), + kind: f[6].to_string(), + version_hash: f[7].to_string(), + original_hash: f[8].to_string(), + references: f[9] + .split(',') + .filter(|s| !s.is_empty()) + .map(str::to_string) + .collect(), + }); + } + (out, skipped) +} + +fn dir_of(rel: &str) -> &str { + rel.rsplit_once('/').map_or("", |(d, _)| d) +} + +fn name_of(rel: &str) -> &str { + rel.rsplit('/').next().unwrap_or(rel) +} + +fn parent_of_dir(dir: &str) -> &str { + dir.rsplit_once('/').map_or("", |(p, _)| p) +} + +fn is_aiedits_dir(dir: &str) -> bool { + name_of(dir).eq_ignore_ascii_case(VERSIONS_SUBFOLDER) +} + +/// Names that are themselves versions (convention or Upscayl output) are +/// never proposed as someone's original. +fn looks_like_a_version(file_name: &str) -> bool { + parse_version_name(file_name).is_some() || upscayl_stem(file_name).is_some() +} + +fn hamming(a: i64, b: i64) -> u32 { + ((a ^ b) as u64).count_ones() +} + +static UPSCAYL_SUFFIX: LazyLock = LazyLock::new(|| { + // `-4x.png`, `_upscayl_2x_ultramix_balanced.png` + Regex::new(r"(?i)^(?P.+?)(?:-\d+x|_upscayl_\d+(?:\.\d+)?x_[^.]+)\.png$").unwrap() +}); + +fn upscayl_stem(file_name: &str) -> Option { + UPSCAYL_SUFFIX + .captures(file_name) + .map(|c| c["stem"].to_string()) +} + +/// File names a ComfyUI prompt loads, in node order: the `image` input of +/// every `LoadImage*` node. +pub fn loader_names(prompt_json: &str) -> Vec { + let Ok(value) = serde_json::from_str::(prompt_json) else { + return Vec::new(); + }; + let Some(nodes) = value.as_object() else { + return Vec::new(); + }; + let mut ids: Vec<&String> = nodes.keys().collect(); + ids.sort_by_key(|k| k.parse::().unwrap_or(u64::MAX)); + let mut out = Vec::new(); + for id in ids { + let node = &nodes[id]; + let class = node["class_type"].as_str().unwrap_or(""); + if !class.starts_with("LoadImage") { + continue; + } + if let Some(raw) = node["inputs"]["image"].as_str() { + let name = raw.trim_end_matches(" [input]"); + let name = name.rsplit(['/', '\\']).next().unwrap_or(name); + if !name.is_empty() { + out.push(name.to_string()); + } + } + } + out +} + +struct Index<'a> { + by_path: HashMap<(i32, &'a str), &'a FileRow>, + by_name: HashMap<(i32, &'a str), Vec<&'a FileRow>>, + by_dir: HashMap<(i32, &'a str), Vec<&'a FileRow>>, +} + +impl<'a> Index<'a> { + fn new(rows: &'a [FileRow]) -> Self { + let mut idx = Index { + by_path: HashMap::new(), + by_name: HashMap::new(), + by_dir: HashMap::new(), + }; + for r in rows { + idx.by_path.insert((r.library_id, r.rel_path.as_str()), r); + idx.by_name + .entry((r.library_id, name_of(&r.rel_path))) + .or_default() + .push(r); + idx.by_dir + .entry((r.library_id, dir_of(&r.rel_path))) + .or_default() + .push(r); + } + idx + } +} + +/// Plan every link the backfill can propose. `rows` are all hashed image +/// rows; `linked` the version hashes that already have a link. +pub fn plan(libs: &[Library], rows: &[FileRow], linked: &HashSet) -> Vec { + let idx = Index::new(rows); + let root_of: HashMap = libs.iter().map(|l| (l.id, l.root_path.as_str())).collect(); + let mut out = Vec::new(); + + for row in rows { + if linked.contains(&row.content_hash) { + continue; + } + let name = name_of(&row.rel_path); + let dir = dir_of(&row.rel_path); + let is_png = name.to_ascii_lowercase().ends_with(".png"); + let name_conv = parse_version_name(name); + let upscayl = if is_png { upscayl_stem(name) } else { None }; + let in_aiedits = is_aiedits_dir(dir); + if !(is_png || name_conv.is_some()) { + continue; + } + if !file_types::is_image_file(Path::new(&row.rel_path)) { + continue; + } + + // Text chunks (XMP + ComfyUI prompt) of PNGs, read once. + let mut xmp_text = None; + let mut prompt = None; + if is_png && let Some(root) = root_of.get(&row.library_id) { + let abs = Path::new(root).join(&row.rel_path); + if let Ok(chunks) = xmp::read_png_text(&abs, &[XMP_KEYWORD, "prompt"]) { + for (k, v) in chunks { + if k == XMP_KEYWORD { + xmp_text = Some(v); + } else { + prompt = Some(v); + } + } + } + } + + // 1. XMP: exact, even when the original isn't indexed yet. + if let Some(x) = xmp_text.as_deref().and_then(xmp::parse_version_xmp) + && x.original_hash != row.content_hash + { + let original = rows.iter().find(|r| r.content_hash == x.original_hash); + let kind = x + .kind + .filter(|k| is_known_kind(k)) + .or_else(|| name_conv.as_ref().map(|n| n.kind.to_string())) + .unwrap_or_else(|| "edit".to_string()); + out.push(Proposal { + status: Status::Exact, + method: "xmp", + library_id: row.library_id, + version_path: row.rel_path.clone(), + version_hash: row.content_hash.clone(), + original_path: original.map(|o| o.rel_path.clone()).unwrap_or_default(), + original_hash: x.original_hash, + kind, + distance: distance(row, original), + references: x.references, + notes: if original.is_some() { + String::new() + } else { + "original not indexed (yet)".into() + }, + }); + continue; + } + + // 2. Name convention next to an indexed original: exact. + if let Some(n) = &name_conv { + let found = original_candidates(&row.rel_path) + .into_iter() + .find_map(|c| { + idx.by_path + .get(&(row.library_id, c.as_str())) + .copied() + .filter(|o| o.content_hash != row.content_hash) + }); + if let Some(o) = found { + out.push(Proposal { + status: Status::Exact, + method: "name", + library_id: row.library_id, + version_path: row.rel_path.clone(), + version_hash: row.content_hash.clone(), + original_path: o.rel_path.clone(), + original_hash: o.content_hash.clone(), + kind: n.kind.to_string(), + distance: distance(row, Some(o)), + references: Vec::new(), + notes: String::new(), + }); + continue; + } + // A convention name with no original found falls through to the + // heuristics only if it is also a ComfyUI output in aiedits/. + if !in_aiedits { + out.push(unresolved(row, "name", "no original with that name nearby")); + continue; + } + } + + // 3. Heuristics (PNG only; reviewed before applying). + if !is_png { + continue; + } + let loaders = prompt.as_deref().map(loader_names).unwrap_or_default(); + if in_aiedits && !loaders.is_empty() { + out.push(plan_comfy_loader(row, &loaders, &idx)); + } else if let Some(stem) = upscayl { + out.push(plan_prefix(row, &stem, &idx)); + } else if in_aiedits { + out.push(plan_phash(row, rows)); + } + } + out.sort_by(|a, b| { + (a.status as u8, a.library_id, &a.version_path).cmp(&( + b.status as u8, + b.library_id, + &b.version_path, + )) + }); + out +} + +fn distance(version: &FileRow, original: Option<&FileRow>) -> Option { + Some(hamming(version.phash?, original?.phash?)) +} + +fn unresolved(row: &FileRow, method: &'static str, why: &str) -> Proposal { + Proposal { + status: Status::Unresolved, + method, + library_id: row.library_id, + version_path: row.rel_path.clone(), + version_hash: row.content_hash.clone(), + original_path: String::new(), + original_hash: String::new(), + kind: "edit".into(), + distance: None, + references: Vec::new(), + notes: why.to_string(), + } +} + +/// Nearest row to `row` by pHash among `candidates` (needs both hashes). +fn nearest<'a>(row: &FileRow, candidates: &[&'a FileRow]) -> Option<(&'a FileRow, u32)> { + let p = row.phash?; + candidates + .iter() + .filter_map(|c| Some((*c, hamming(p, c.phash?)))) + .min_by_key(|(c, d)| (*d, c.rel_path.clone())) +} + +fn plan_comfy_loader(row: &FileRow, loaders: &[String], idx: &Index) -> Proposal { + let dir = dir_of(&row.rel_path); + let parent = parent_of_dir(dir); + // Per loader name: the matching rows from the first tier that has any, + // preferring the parent of aiedits/, then the same folder, then anywhere + // in the library. Version-looking files never count as originals. + let resolve = |loader: &str| -> Vec<&FileRow> { + let all: Vec<&FileRow> = idx + .by_name + .get(&(row.library_id, loader)) + .map(|v| { + v.iter() + .copied() + .filter(|r| { + r.content_hash != row.content_hash && !is_aiedits_dir(dir_of(&r.rel_path)) + }) + .collect() + }) + .unwrap_or_default(); + for tier in [parent, dir] { + let hit: Vec<&FileRow> = all + .iter() + .copied() + .filter(|r| dir_of(&r.rel_path) == tier) + .collect(); + if !hit.is_empty() { + return hit; + } + } + all + }; + + let mut resolved: Vec<(&String, Vec<&FileRow>)> = Vec::new(); + let mut missing: Vec<&String> = Vec::new(); + for l in loaders { + let hits = resolve(l); + if hits.is_empty() { + missing.push(l); + } else { + resolved.push((l, hits)); + } + } + if resolved.is_empty() { + let mut p = unresolved( + row, + "comfy-loader", + &format!( + "loader input(s) not found in the library: {}", + loaders.join(", ") + ), + ); + // A hint for the reviewer, not a proposal. + let siblings: Vec<&FileRow> = idx + .by_dir + .get(&(row.library_id, parent)) + .map(|v| { + v.iter() + .copied() + .filter(|r| !is_aiedits_dir(dir_of(&r.rel_path))) + .collect() + }) + .unwrap_or_default(); + if let Some((n, d)) = nearest(row, &siblings) { + p.notes + .push_str(&format!("; hint: nearest pHash {} (d={d})", n.rel_path)); + } + return p; + } + + // Primary: the nearest pHash across everything the loaders resolved to. + let flat: Vec<(&String, &FileRow)> = resolved + .iter() + .flat_map(|(l, hits)| hits.iter().map(move |h| (*l, *h))) + .collect(); + let rows_only: Vec<&FileRow> = flat.iter().map(|(_, r)| *r).collect(); + let (primary, dist) = match nearest(row, &rows_only) { + Some((p, d)) => (p, Some(d)), + None => (rows_only[0], None), + }; + let mut notes = Vec::new(); + if resolved.len() > 1 { + notes.push("two loader inputs: primary chosen by nearest pHash, the other recorded as a reference — check".to_string()); + } + if flat.len() > resolved.len() { + notes.push("a loader name matched several files — check".into()); + } + if !missing.is_empty() { + notes.push(format!( + "other loader input(s) not in the library: {}", + missing + .iter() + .map(|s| s.as_str()) + .collect::>() + .join(", ") + )); + } + if dist.is_none() { + notes.push("no pHash to compare".into()); + } + let mut references: Vec = Vec::new(); + for (_, hits) in &resolved { + for h in hits { + if h.content_hash != primary.content_hash && !references.contains(&h.content_hash) { + references.push(h.content_hash.clone()); + } + } + } + Proposal { + status: Status::Review, + method: "comfy-loader", + library_id: row.library_id, + version_path: row.rel_path.clone(), + version_hash: row.content_hash.clone(), + original_path: primary.rel_path.clone(), + original_hash: primary.content_hash.clone(), + kind: "edit".into(), + distance: dist, + references, + notes: notes.join("; "), + } +} + +fn plan_prefix(row: &FileRow, stem: &str, idx: &Index) -> Proposal { + let dir = dir_of(&row.rel_path); + let matches: Vec<&FileRow> = idx + .by_dir + .get(&(row.library_id, dir)) + .map(|v| { + v.iter() + .copied() + .filter(|r| { + r.content_hash != row.content_hash + && name_of(&r.rel_path) + .rsplit_once('.') + .is_some_and(|(s, _)| s == stem) + }) + .collect() + }) + .unwrap_or_default(); + match matches.as_slice() { + [] => unresolved( + row, + "name-prefix", + &format!("no file named {stem}.* beside it"), + ), + [only] => Proposal { + status: Status::Review, + method: "name-prefix", + library_id: row.library_id, + version_path: row.rel_path.clone(), + version_hash: row.content_hash.clone(), + original_path: only.rel_path.clone(), + original_hash: only.content_hash.clone(), + kind: "upscale".into(), + distance: distance(row, Some(only)), + references: Vec::new(), + notes: String::new(), + }, + several => { + let (best, d) = nearest(row, several).unwrap_or((several[0], u32::MAX)); + Proposal { + status: Status::Review, + method: "name-prefix", + library_id: row.library_id, + version_path: row.rel_path.clone(), + version_hash: row.content_hash.clone(), + original_path: best.rel_path.clone(), + original_hash: best.content_hash.clone(), + kind: "upscale".into(), + distance: (d != u32::MAX).then_some(d), + references: Vec::new(), + notes: format!("{} files share that name — check", several.len()), + } + } + } +} + +fn plan_phash(row: &FileRow, rows: &[FileRow]) -> Proposal { + let dir = dir_of(&row.rel_path); + let parent = parent_of_dir(dir); + let candidates: Vec<&FileRow> = rows + .iter() + .filter(|r| { + r.library_id == row.library_id + && r.content_hash != row.content_hash + && !is_aiedits_dir(dir_of(&r.rel_path)) + && !looks_like_a_version(name_of(&r.rel_path)) + && (dir_of(&r.rel_path) == parent || dir_of(&r.rel_path) == dir) + }) + .collect(); + match nearest(row, &candidates) { + Some((o, d)) if d <= MAX_PHASH_DISTANCE => Proposal { + status: Status::Review, + method: "phash", + library_id: row.library_id, + version_path: row.rel_path.clone(), + version_hash: row.content_hash.clone(), + original_path: o.rel_path.clone(), + original_hash: o.content_hash.clone(), + kind: "edit".into(), + distance: Some(d), + references: Vec::new(), + notes: "matched by nearest pHash only".into(), + }, + Some((o, d)) => unresolved( + row, + "phash", + &format!("nearest pHash is {} at distance {d}, too far", o.rel_path), + ), + None => unresolved(row, "phash", "no pHash candidates nearby"), + } +} + +/// What [`apply`] did. +#[derive(Debug, Default, PartialEq, Eq)] +pub struct ApplyReport { + pub linked: usize, + pub already_linked: usize, + pub refused: usize, + /// Reviewed rows whose version row no longer matches the file in the TSV. + pub stale: usize, +} + +/// Write the links. `exact` proposals apply as they are; `reviewed` come +/// from a TSV a human kept, and are checked against `rows` first so a stale +/// file can't link a path that now holds different bytes. +pub fn apply( + dao: &mut dyn VersionDao, + exact: &[Proposal], + reviewed: &[ReviewedLink], + rows: &[FileRow], +) -> ApplyReport { + let mut report = ApplyReport::default(); + let current: HashMap<(i32, &str), &str> = rows + .iter() + .map(|r| ((r.library_id, r.rel_path.as_str()), r.content_hash.as_str())) + .collect(); + + let mut links: Vec = Vec::new(); + for p in exact.iter().filter(|p| p.status == Status::Exact) { + links.push(NewVersionLink { + version_hash: p.version_hash.clone(), + original_hash: p.original_hash.clone(), + kind: p.kind.clone(), + source: if p.method == "xmp" { + SOURCE_XMP + } else { + SOURCE_NAME + } + .to_string(), + library_id: p.library_id, + rel_path: p.version_path.clone(), + references: p.references.clone(), + }); + } + for r in reviewed { + if current.get(&(r.library_id, r.version_path.as_str())) != Some(&r.version_hash.as_str()) { + report.stale += 1; + continue; + } + links.push(NewVersionLink { + version_hash: r.version_hash.clone(), + original_hash: r.original_hash.clone(), + kind: if is_known_kind(&r.kind) { + r.kind.clone() + } else { + "edit".into() + }, + source: SOURCE_BACKFILL.to_string(), + library_id: r.library_id, + rel_path: r.version_path.clone(), + references: r.references.clone(), + }); + } + for link in links { + match dao.link_version(&link) { + Ok(LinkOutcome::Linked) => report.linked += 1, + Ok(LinkOutcome::AlreadyLinked) => report.already_linked += 1, + Ok(LinkOutcome::Rejected(_)) | Err(_) => report.refused += 1, + } + } + report +} + +#[cfg(test)] +mod tests { + use super::super::xmp::test_png::{producer_xmp, write_png}; + use super::*; + use crate::database::SqliteVersionDao; + use crate::database::test::in_memory_db_connection; + use std::sync::{Arc, Mutex}; + + fn row(rel: &str, hash: &str, phash: Option) -> FileRow { + FileRow { + library_id: 1, + rel_path: rel.into(), + content_hash: hash.into(), + phash, + } + } + + fn lib(root: &Path) -> Library { + Library { + id: 1, + name: "main".into(), + root_path: root.to_string_lossy().into(), + enabled: true, + excluded_dirs: vec![], + } + } + + fn prompt(loaders: &[&str]) -> String { + let mut nodes = serde_json::Map::new(); + nodes.insert( + "1".into(), + serde_json::json!({"class_type": "UNETLoader", "inputs": {"unet_name": "m"}}), + ); + for (i, l) in loaders.iter().enumerate() { + nodes.insert( + (i + 2).to_string(), + serde_json::json!({"class_type": "LoadImage", "inputs": {"image": l}}), + ); + } + serde_json::Value::Object(nodes).to_string() + } + + fn write_comfy_png(root: &Path, rel: &str, loaders: &[&str], salt: u8) { + super::super::xmp::test_png::write_png(&root.join(rel), None, Some(&prompt(loaders)), salt); + } + + #[test] + fn loader_names_reads_only_loadimage_inputs_in_node_order() { + let p = prompt(&["a.jpg", "sub/b [input]"]); + assert_eq!(loader_names(&p), vec!["a.jpg", "b"]); + assert!(loader_names("not json").is_empty()); + } + + #[test] + fn already_ingested_xmp_files_are_proposed_as_exact_without_naming_help() { + let tmp = tempfile::tempdir().unwrap(); + let orig_hash = "d5".repeat(32); + write_png( + &tmp.path().join("2009/aiedits/019_19.JPG.restored.png"), + Some(&producer_xmp(&orig_hash, "restored", &[])), + None, + 1, + ); + // An XMP file whose name gives no hint (a renamed copy). + write_png( + &tmp.path().join("misc/renamed.png"), + Some(&producer_xmp( + &orig_hash, + "upscale", + &["ee".repeat(32).as_str()], + )), + None, + 2, + ); + let rows = vec![ + row("2009/019_19.JPG", &orig_hash, None), + row("2009/aiedits/019_19.JPG.restored.png", "v1", None), + row("misc/renamed.png", "v2", None), + ]; + let plan = plan(&[lib(tmp.path())], &rows, &HashSet::new()); + let by_path: HashMap<&str, &Proposal> = + plan.iter().map(|p| (p.version_path.as_str(), p)).collect(); + let a = by_path["2009/aiedits/019_19.JPG.restored.png"]; + assert_eq!((a.status, a.method), (Status::Exact, "xmp")); + assert_eq!(a.original_path, "2009/019_19.JPG"); + let b = by_path["misc/renamed.png"]; + assert_eq!( + (b.status, b.method, b.kind.as_str()), + (Status::Exact, "xmp", "upscale") + ); + assert_eq!(b.references.len(), 1); + } + + #[test] + fn comfy_outputs_are_resolved_by_loader_and_flagged_when_two_inputs() { + let tmp = tempfile::tempdir().unwrap(); + write_comfy_png(tmp.path(), "Amanda/aiedits/out_00001_.png", &["a.jpg"], 1); + write_comfy_png( + tmp.path(), + "Amanda/aiedits/out_00002_.png", + &["a.jpg", "b.jpg"], + 2, + ); + write_comfy_png( + tmp.path(), + "Amanda/aiedits/out_00003_.png", + &["typo.jpg"], + 3, + ); + let rows = vec![ + row("Amanda/a.jpg", "ha", Some(0b1111)), + row("Amanda/b.jpg", "hb", Some(-1)), + row("Amanda/aiedits/out_00001_.png", "v1", Some(0b1110)), + row("Amanda/aiedits/out_00002_.png", "v2", Some(0b1011)), + row("Amanda/aiedits/out_00003_.png", "v3", Some(0)), + ]; + let plan = plan(&[lib(tmp.path())], &rows, &HashSet::new()); + let get = |p: &str| plan.iter().find(|x| x.version_path == p).unwrap(); + + let one = get("Amanda/aiedits/out_00001_.png"); + assert_eq!((one.status, one.method), (Status::Review, "comfy-loader")); + assert_eq!(one.original_path, "Amanda/a.jpg"); + assert_eq!(one.distance, Some(1)); + assert!(one.references.is_empty()); + + let two = get("Amanda/aiedits/out_00002_.png"); + assert_eq!( + two.original_path, "Amanda/a.jpg", + "primary by nearest pHash" + ); + assert_eq!(two.references, vec!["hb"]); + assert!(two.notes.contains("two loader inputs")); + + let lost = get("Amanda/aiedits/out_00003_.png"); + assert_eq!(lost.status, Status::Unresolved); + assert!(lost.original_path.is_empty(), "listed, not guessed"); + } + + #[test] + fn upscayl_names_and_phash_fallback() { + let tmp = tempfile::tempdir().unwrap(); + write_png(&tmp.path().join("p/att_1-4x.png"), None, None, 1); + write_png( + &tmp.path().join("p/pic_upscayl_3x_ultramix_balanced.png"), + None, + None, + 2, + ); + write_png(&tmp.path().join("p/aiedits/mystery.png"), None, None, 3); + let rows = vec![ + row("p/att_1.jpg", "h1", Some(5)), + row("p/pic.jpeg", "h2", Some(5)), + row("p/other.jpg", "h3", Some(0b0111_1111_1111)), + row("p/att_1-4x.png", "u1", Some(5)), + row("p/pic_upscayl_3x_ultramix_balanced.png", "u2", Some(5)), + row("p/aiedits/mystery.png", "m1", Some(7)), + ]; + let plan = plan(&[lib(tmp.path())], &rows, &HashSet::new()); + let get = |p: &str| plan.iter().find(|x| x.version_path == p).unwrap(); + assert_eq!(get("p/att_1-4x.png").original_path, "p/att_1.jpg"); + assert_eq!(get("p/att_1-4x.png").kind, "upscale"); + assert_eq!( + get("p/pic_upscayl_3x_ultramix_balanced.png").original_path, + "p/pic.jpeg" + ); + let m = get("p/aiedits/mystery.png"); + assert_eq!(m.method, "phash"); + assert_eq!(m.status, Status::Review); + assert_eq!( + m.original_path, "p/att_1.jpg", + "distance 2 beats the others" + ); + } + + #[test] + fn linked_rows_are_not_proposed_again() { + let tmp = tempfile::tempdir().unwrap(); + write_png(&tmp.path().join("a.jpg.restored.png"), None, None, 1); + let rows = vec![ + row("a.jpg", "o", None), + row("a.jpg.restored.png", "v", None), + ]; + let linked = HashSet::from(["v".to_string()]); + assert!(plan(&[lib(tmp.path())], &rows, &linked).is_empty()); + let p = plan(&[lib(tmp.path())], &rows, &HashSet::new()); + assert_eq!((p[0].status, p[0].method), (Status::Exact, "name")); + } + + #[test] + fn tsv_round_trip_and_reviewed_application() { + let tmp = tempfile::tempdir().unwrap(); + write_comfy_png(tmp.path(), "d/aiedits/o_00001_.png", &["a.jpg", "b.jpg"], 1); + let rows = vec![ + row("d/a.jpg", "ha", Some(1)), + row("d/b.jpg", "hb", Some(-1)), + row("d/aiedits/o_00001_.png", "v", Some(1)), + row("d/x.jpg", "hx", None), + row("d/x.jpg.edit.png", "vx", None), + ]; + write_png(&tmp.path().join("d/x.jpg.edit.png"), None, None, 9); + let planned = plan(&[lib(tmp.path())], &rows, &HashSet::new()); + let tsv: String = std::iter::once(TSV_HEADER.to_string()) + .chain(planned.iter().map(Proposal::to_tsv)) + .collect::>() + .join("\n"); + let (reviewed, skipped) = parse_reviewed(&tsv); + assert_eq!( + reviewed.len(), + 1, + "only the review row; the exact one isn't re-applied from the file" + ); + assert_eq!(skipped, 1); + assert_eq!(reviewed[0].references, vec!["hb"]); + + let conn = Arc::new(Mutex::new(in_memory_db_connection())); + let mut dao = SqliteVersionDao::from_shared(conn); + // Without --reviewed: exact only. + let report = apply(&mut dao, &planned, &[], &rows); + assert_eq!(report.linked, 1); + assert!(dao.get_link("vx").unwrap().is_some()); + assert!( + dao.get_link("v").unwrap().is_none(), + "heuristic rows never auto-apply" + ); + // With the reviewed file. + let report = apply(&mut dao, &planned, &reviewed, &rows); + assert_eq!(report.linked, 1); + assert_eq!(report.already_linked, 1); + let l = dao.get_link("v").unwrap().unwrap(); + assert_eq!( + (l.original_hash.as_str(), l.source.as_str()), + ("ha", "backfill") + ); + assert_eq!(dao.list_references("v").unwrap(), vec!["hb"]); + + // A stale TSV (the path now holds other bytes) is refused. + let mut stale_rows = rows.clone(); + stale_rows[2].content_hash = "changed".into(); + let report = apply(&mut dao, &[], &reviewed, &stale_rows); + assert_eq!(report.stale, 1); + } +} diff --git a/src/versions/mod.rs b/src/versions/mod.rs index 1100890..2ef7803 100644 --- a/src/versions/mod.rs +++ b/src/versions/mod.rs @@ -14,6 +14,7 @@ //! original's `content_hash`, never a path: ComfyUI writes through a UNC //! share and ImageApi reads `/media/...`, and files get moved. +pub mod backfill; pub mod handlers; pub mod xmp;