1075 lines
38 KiB
Rust
1075 lines
38 KiB
Rust
//! 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 `<name>.<kind>[.<n>].<ext>` 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<i64>,
|
|
}
|
|
|
|
#[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<u32>,
|
|
/// Hashes of reference photos (the other loader inputs).
|
|
pub references: Vec<String>,
|
|
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,
|
|
/// The original's path as the dry run proposed it (may be empty).
|
|
pub original_path: String,
|
|
pub kind: String,
|
|
pub version_hash: String,
|
|
pub original_hash: String,
|
|
pub references: Vec<String>,
|
|
}
|
|
|
|
/// 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<ReviewedLink>, 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::<i32>().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(),
|
|
original_path: f[4].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<Regex> = LazyLock::new(|| {
|
|
// `<name>-4x.png`, `<name>_upscayl_2x_ultramix_balanced.png`
|
|
Regex::new(r"(?i)^(?P<stem>.+?)(?:-\d+x|_upscayl_\d+(?:\.\d+)?x_[^.]+)\.png$").unwrap()
|
|
});
|
|
|
|
fn upscayl_stem(file_name: &str) -> Option<String> {
|
|
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<String> {
|
|
let Ok(value) = serde_json::from_str::<serde_json::Value>(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::<u64>().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<String>) -> Vec<Proposal> {
|
|
let idx = Index::new(rows);
|
|
let root_of: HashMap<i32, &str> = 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();
|
|
// A ComfyUI output names its inputs wherever it was saved (early test
|
|
// runs went to folders like `edited/` or sat beside the original), so
|
|
// the loader match isn't limited to aiedits/. It's review-only anyway.
|
|
if !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<u32> {
|
|
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:
|
|
// for an output in aiedits/ its parent then the same folder; for one
|
|
// elsewhere the same folder (saved beside the original) then the parent
|
|
// (e.g. `edited/`); then anywhere in the library. Version-looking files
|
|
// never count as originals.
|
|
let tiers = if is_aiedits_dir(dir) {
|
|
[parent, dir]
|
|
} else {
|
|
[dir, parent]
|
|
};
|
|
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 tiers {
|
|
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::<Vec<_>>()
|
|
.join(", ")
|
|
));
|
|
}
|
|
if dist.is_none() {
|
|
notes.push("no pHash to compare".into());
|
|
}
|
|
let mut references: Vec<String> = 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,
|
|
/// One line per stale reviewed row, saying what no longer matches.
|
|
pub stale_details: Vec<String>,
|
|
}
|
|
|
|
/// 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 live_hashes: HashSet<&str> = rows.iter().map(|r| r.content_hash.as_str()).collect();
|
|
|
|
let mut links: Vec<NewVersionLink> = 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 {
|
|
let stale = if current.get(&(r.library_id, r.version_path.as_str()))
|
|
!= Some(&r.version_hash.as_str())
|
|
{
|
|
Some("version file now holds different bytes (or is gone)")
|
|
} else if !live_hashes.contains(r.original_hash.as_str()) {
|
|
Some("original_hash is not in the index any more")
|
|
} else if !r.original_path.is_empty()
|
|
&& current
|
|
.get(&(r.library_id, r.original_path.as_str()))
|
|
.is_some_and(|h| *h != r.original_hash.as_str())
|
|
{
|
|
Some("original path now holds different bytes than original_hash")
|
|
} else {
|
|
None
|
|
};
|
|
if let Some(why) = stale {
|
|
report.stale += 1;
|
|
report.stale_details.push(format!(
|
|
"{} -> {}: {why}",
|
|
r.version_path,
|
|
if r.original_path.is_empty() {
|
|
&r.original_hash
|
|
} else {
|
|
&r.original_path
|
|
}
|
|
));
|
|
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<i64>) -> 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(), "album_a/aiedits/out_00001_.png", &["a.jpg"], 1);
|
|
write_comfy_png(
|
|
tmp.path(),
|
|
"album_a/aiedits/out_00002_.png",
|
|
&["a.jpg", "b.jpg"],
|
|
2,
|
|
);
|
|
write_comfy_png(
|
|
tmp.path(),
|
|
"album_a/aiedits/out_00003_.png",
|
|
&["typo.jpg"],
|
|
3,
|
|
);
|
|
let rows = vec![
|
|
row("album_a/a.jpg", "ha", Some(0b1111)),
|
|
row("album_a/b.jpg", "hb", Some(-1)),
|
|
row("album_a/aiedits/out_00001_.png", "v1", Some(0b1110)),
|
|
row("album_a/aiedits/out_00002_.png", "v2", Some(0b1011)),
|
|
row("album_a/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("album_a/aiedits/out_00001_.png");
|
|
assert_eq!((one.status, one.method), (Status::Review, "comfy-loader"));
|
|
assert_eq!(one.original_path, "album_a/a.jpg");
|
|
assert_eq!(one.distance, Some(1));
|
|
assert!(one.references.is_empty());
|
|
|
|
let two = get("album_a/aiedits/out_00002_.png");
|
|
assert_eq!(
|
|
two.original_path, "album_a/a.jpg",
|
|
"primary by nearest pHash"
|
|
);
|
|
assert_eq!(two.references, vec!["hb"]);
|
|
assert!(two.notes.contains("two loader inputs"));
|
|
|
|
let lost = get("album_a/aiedits/out_00003_.png");
|
|
assert_eq!(lost.status, Status::Unresolved);
|
|
assert!(lost.original_path.is_empty(), "listed, not guessed");
|
|
}
|
|
|
|
#[test]
|
|
fn comfy_outputs_outside_aiedits_resolve_from_their_own_folder_then_the_parent() {
|
|
let tmp = tempfile::tempdir().unwrap();
|
|
// Saved into a sibling folder: the original is one level up.
|
|
write_comfy_png(tmp.path(), "album_b/edited/out_00016_.png", &["a.jpg"], 1);
|
|
// Saved beside the original; a same-named file at the root must lose.
|
|
write_comfy_png(tmp.path(), "album_c/x-aiedit_00001_.png", &["x.jpg"], 2);
|
|
// A ComfyUI output with no loader (txt2img) is left alone.
|
|
write_comfy_png(tmp.path(), "album_c/generated_00001_.png", &[], 3);
|
|
let rows = vec![
|
|
row("album_b/a.jpg", "ha", Some(0)),
|
|
row("album_b/edited/out_00016_.png", "v1", Some(1)),
|
|
row("x.jpg", "hroot", Some(0)),
|
|
row("album_c/x.jpg", "hx", Some(0)),
|
|
row("album_c/x-aiedit_00001_.png", "v2", Some(1)),
|
|
row("album_c/generated_00001_.png", "g", Some(0)),
|
|
];
|
|
let plan = plan(&[lib(tmp.path())], &rows, &HashSet::new());
|
|
let get = |p: &str| plan.iter().find(|x| x.version_path == p);
|
|
|
|
let edited = get("album_b/edited/out_00016_.png").unwrap();
|
|
assert_eq!(
|
|
(edited.status, edited.method),
|
|
(Status::Review, "comfy-loader")
|
|
);
|
|
assert_eq!(edited.original_path, "album_b/a.jpg");
|
|
|
|
let beside = get("album_c/x-aiedit_00001_.png").unwrap();
|
|
assert_eq!(beside.status, Status::Review);
|
|
assert_eq!(
|
|
beside.original_path, "album_c/x.jpg",
|
|
"same folder beats the parent"
|
|
);
|
|
|
|
assert!(
|
|
get("album_c/generated_00001_.png").is_none(),
|
|
"no loader, not a version"
|
|
);
|
|
}
|
|
|
|
#[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::<Vec<_>>()
|
|
.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);
|
|
|
|
// So is one whose original changed since the dry run: its hash is
|
|
// gone from the index, or its path holds other bytes.
|
|
let mut gone = rows.clone();
|
|
gone[0].content_hash = "ha-edited".into();
|
|
let report = apply(&mut dao, &[], &reviewed, &gone);
|
|
assert_eq!(report.stale, 1);
|
|
assert_eq!(report.linked + report.already_linked, 0);
|
|
assert!(report.stale_details[0].contains("d/aiedits/o_00001_.png"));
|
|
let mut swapped = rows.clone();
|
|
swapped[0].content_hash = "hb".into(); // d/a.jpg now holds hb's bytes
|
|
let report = apply(&mut dao, &[], &reviewed, &swapped);
|
|
assert_eq!(report.stale, 1);
|
|
}
|
|
#[test]
|
|
fn apply_drops_derived_data_of_an_already_ingested_version() {
|
|
let conn = Arc::new(Mutex::new(in_memory_db_connection()));
|
|
let exec = |sql: &str| {
|
|
diesel::RunQueryDsl::execute(diesel::sql_query(sql), &mut *conn.lock().unwrap())
|
|
.unwrap_or_else(|e| panic!("{sql}: {e}"));
|
|
};
|
|
exec(
|
|
"INSERT INTO image_exif (library_id, rel_path, created_time, last_modified, content_hash, clip_embedding, clip_model_version) VALUES (1, 'v.png', 0, 0, 'v', x'00', 'm')",
|
|
);
|
|
exec(
|
|
"INSERT INTO face_detections (library_id, content_hash, rel_path, bbox_x, embedding, source, status, model_version, created_at) VALUES (1, 'v', 'v.png', 0.1, x'00', 'auto', 'detected', 'm', 0)",
|
|
);
|
|
exec(
|
|
"INSERT INTO clip_failures (library_id, rel_path, content_hash, error, failed_at) VALUES (1, 'v.png', 'v', 'boom', 0)",
|
|
);
|
|
let mut dao = SqliteVersionDao::from_shared(conn.clone());
|
|
let reviewed = vec![ReviewedLink {
|
|
method: "phash".into(),
|
|
library_id: 1,
|
|
version_path: "v.png".into(),
|
|
original_path: "o.jpg".into(),
|
|
kind: "edit".into(),
|
|
version_hash: "v".into(),
|
|
original_hash: "o".into(),
|
|
references: vec![],
|
|
}];
|
|
let rows = vec![row("v.png", "v", None), row("o.jpg", "o", None)];
|
|
assert_eq!(apply(&mut dao, &[], &reviewed, &rows).linked, 1);
|
|
|
|
#[derive(diesel::QueryableByName)]
|
|
struct N {
|
|
#[diesel(sql_type = diesel::sql_types::BigInt)]
|
|
n: i64,
|
|
}
|
|
for sql in [
|
|
"SELECT COUNT(*) AS n FROM face_detections WHERE content_hash = 'v'",
|
|
"SELECT COUNT(*) AS n FROM clip_failures WHERE content_hash = 'v'",
|
|
"SELECT COUNT(*) AS n FROM image_exif WHERE content_hash = 'v' AND (clip_embedding IS NOT NULL OR clip_model_version IS NOT NULL)",
|
|
] {
|
|
let n = diesel::RunQueryDsl::get_result::<N>(
|
|
diesel::sql_query(sql),
|
|
&mut *conn.lock().unwrap(),
|
|
)
|
|
.unwrap()
|
|
.n;
|
|
assert_eq!(n, 0, "{sql}");
|
|
}
|
|
}
|
|
}
|