diff --git a/src/versions/backfill.rs b/src/versions/backfill.rs index 32d965b..fb9cb26 100644 --- a/src/versions/backfill.rs +++ b/src/versions/backfill.rs @@ -356,7 +356,10 @@ pub fn plan(libs: &[Library], rows: &[FileRow], linked: &HashSet) -> Vec continue; } let loaders = prompt.as_deref().map(loader_names).unwrap_or_default(); - if in_aiedits && !loaders.is_empty() { + // 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)); @@ -406,9 +409,16 @@ fn nearest<'a>(row: &FileRow, candidates: &[&'a FileRow]) -> Option<(&'a FileRow 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. + // 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 @@ -422,7 +432,7 @@ fn plan_comfy_loader(row: &FileRow, loaders: &[String], idx: &Index) -> Proposal .collect() }) .unwrap_or_default(); - for tier in [parent, dir] { + for tier in tiers { let hit: Vec<&FileRow> = all .iter() .copied() @@ -783,48 +793,88 @@ mod tests { #[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(), "album_a/aiedits/out_00001_.png", &["a.jpg"], 1); write_comfy_png( tmp.path(), - "Amanda/aiedits/out_00002_.png", + "album_a/aiedits/out_00002_.png", &["a.jpg", "b.jpg"], 2, ); write_comfy_png( tmp.path(), - "Amanda/aiedits/out_00003_.png", + "album_a/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)), + 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("Amanda/aiedits/out_00001_.png"); + let one = get("album_a/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.original_path, "album_a/a.jpg"); assert_eq!(one.distance, Some(1)); assert!(one.references.is_empty()); - let two = get("Amanda/aiedits/out_00002_.png"); + let two = get("album_a/aiedits/out_00002_.png"); assert_eq!( - two.original_path, "Amanda/a.jpg", + 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("Amanda/aiedits/out_00003_.png"); + 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();