Enhance on demand: ComfyUI templates, jobs, approve as version (#192) #196

Merged
cameron merged 3 commits from feat/192-enhance into master 2026-10-30 19:41:43 +00:00
22 changed files with 5159 additions and 113 deletions

No files matched your search

+1
View File
@@ -7,6 +7,7 @@ database/target
.env
# Server-local LLM config (llm.example.toml is the template)
/llm.toml
/enhance.toml
# Server-local TTS pronunciation overrides (tts_pronunciations.example.json is the template)
/tts_pronunciations.json
/tmp
+89
View File
@@ -424,6 +424,94 @@ for the `link_versions` binary).
against the current hash; `unresolved` rows are listed, never guessed. `--apply`
takes a `VACUUM INTO` snapshot first.
### Enhance on demand (#192)
An **Enhance** action runs a ComfyUI workflow (a *template*) on one photo; the
user previews the result in the app and **approves** it, and it lands in the
library as a photo version (#191). Nothing reaches the library without an
approval. Code: `src/enhance/` (`config.rs`, `path_map.rs`, `template.rs`,
`comfy.rs`, `service.rs`, `handlers.rs`), table `enhance_jobs`
(`src/database/enhance_dao.rs`). The apps only talk to ImageApi; ComfyUI has no
auth and stays LAN-only, so ImageApi relays previews too.
- **Config: `enhance.toml`** (path from `ENHANCE_CONFIG_PATH`, else `./enhance.toml`;
gitignored; `enhance.example.toml` is the annotated schema and a test parses it).
No secrets live in it. Without `url` (or without the file) Enhance is off:
`GET /enhance/templates` answers `available: false`. A malformed file fails startup;
a *template* that doesn't validate is logged and skipped, the rest load.
Keys: `url`, `timeout_secs`, `probe_ttl_secs`, `poll_secs`, `idle_free_secs`,
`unreachable_grace_secs`, `default_template`, `warn_above_megapixels`,
`[[path_map]]` (`from`, `to`), `[[template]]` (`name`, `label`, `description`,
`file`, `prompt_mode`, `prompt_join`, `output_megapixels`).
- **Templates** are API-format workflow exports, read-only (never edit them from
here). Nodes are found by `_meta.title`: `imageapi:source` (required, gets the
photo path), `imageapi:save` (required, forced to `require_approval = true`),
`imageapi:prompt` (optional; fills the first of `text` / `prompt` / `value`),
`imageapi:reference` (optional). `prompt_mode = "append"`: the user's text is joined
after the template's own prompt with `prompt_join` (default one space); empty text
leaves the prompt as is. `"replace"`: the text *is* the prompt and is required (empty
= 400); it needs a prompt node. With no `reference_path` the reference node **and
every `["<id>", n]` input linking to it** are removed (its leftover export path must
never be used); if ComfyUI then rejects the prompt with `node_errors` the error says
"this template needs a reference photo". Every `seed`, `noise_seed` and `*.seed`
input gets a fresh random value per submit, or a retry would return ComfyUI's cached
result. A stray `SaveImage` is only a startup warning.
- **Path map**: library prefix <-> ComfyUI prefix (`/media/nas/` <-> `Z:\`), longest
prefix wins, separators converted, case-insensitive on the ComfyUI side; used on
submit and to recognise the path `/imageapi/approve` returns. No entries = identity.
A mapped drive only exists in the logged-in session; use UNC if ComfyUI becomes a
service.
- **Jobs** (`enhance_jobs`, keyed by the submitted file's `content_hash`): status
`queued | running | ready | approved | failed | expired | unavailable`. The poller
(`EnhanceService::spawn`, every `poll_secs`, active jobs only) reads `/queue` then
`/history/{id}`: in queue -> queued/running with `queue_position`; finished ->
`ready` (staged results in `pending`) or `failed` (ComfyUI's message); known to
neither -> **`expired`** (ComfyUI restarted and cleared its temp folder). An
unreachable ComfyUI for `unreachable_grace_secs` -> `unavailable`. A preview or
approve that finds the staged result gone also expires the job (HTTP 410).
- **Idle `/free`**: every 30 s, if the queue is empty, none of our jobs is active and
ComfyUI's newest `/history` entry has not changed for `idle_free_secs` (it includes
hand-run workflows), `POST /free {unload_models, free_memory}` is sent once per quiet
spell. "Not changed for" is measured on ImageApi's clock (no cross-host skew); the
first observation after ImageApi starts counts as activity. `idle_free_secs = 0` = off.
- **Approve** (`POST /enhance/jobs/{id}/approve`): ComfyUI's `/imageapi/approve` writes
the file (ImageApi writes nothing under library roots). The returned path is mapped
back, then a **targeted ingest** (`watcher::register_new_files` + `versions::link_new_files`)
registers and links just that file, so it is in `/image/versions` before the response.
Idempotent: an already-saved result returns its version without calling ComfyUI.
- **Output size**: Turbo output is about 1 MP regardless of the original, so
`output_megapixels` / `warn_above_megapixels` let clients warn above ~1.5 MP.
Endpoints (all need a bearer token; errors are `{"error": msg, "code": ...}`; `library`
is accepted as a query or body field like the versions endpoints):
```
GET /enhance/templates
-> {"available": true, "reason"?: "ComfyUI isn't running" | "Enhance isn't configured",
"default": "restore", "warn_above_megapixels": 1.5,
"templates": [{"name","label","description","accepts_prompt",
"prompt_mode": "append"|"replace", "prompt_required",
"accepts_reference","default","output_megapixels",
"warn_above_megapixels"}]}
POST /enhance {"path","template","prompt"?,"reference_path"?,"library"?} -> job
GET /enhance/jobs/{id} -> job
GET /enhance/jobs?path=&library= -> {"jobs": [job, ...]} newest first, discarded hidden
GET /enhance/jobs/{id}/preview/{n} -> the staged PNG (relayed from ComfyUI /view)
POST /enhance/jobs/{id}/approve {"n"} -> a /image/versions `versions[]` entry
POST /enhance/jobs/{id}/discard -> job (hidden from the list; ComfyUI untouched)
job = {"id","status","queue_position"?,"error"?,"path","library_id","template",
"prompt"?,"reference_path"?,"created_at","updated_at",
"pending": [{"n","preview": "/enhance/jobs/{id}/preview/{n}","approved"}],
"saved": [{"n","path","library_id"}]}
```
Error statuses: 400 bad request / empty required prompt / job not ready; 404 unknown
template, photo or job; 410 `expired`; 422 ComfyUI rejected the workflow; 502 upstream
or an approved file outside every library/path_map; 503 `unavailable` (ComfyUI down) or
`disabled`. `pending` is empty unless the job is `ready`/`approved`. Approve answers 202
`{"linked": false, ...}` if the file was saved but couldn't be ingested synchronously.
### File Processing Pipeline
**Thumbnail Generation:**
@@ -632,6 +720,7 @@ GET /image/versions?path=...&library=... (path of the original or any version
{ original:{path,library_id,hash}, versions:[{path,
library_id,hash,kind,source,created_at,references,
width,height}], display_hash })
POST /enhance ... (Enhance on demand; see "Enhance on demand")
PUT /image/versions/display (body { path, display_hash|null, library? }; global choice)
// /photos hides versions behind their original (include_versions=true lists them) and adds
// parallel arrays photo_display_paths / photo_version_counts; memories items, search hits and
+65
View File
@@ -0,0 +1,65 @@
# ImageApi Enhance configuration (#192). Copy to `enhance.toml` next to
# `.env` (or point ENHANCE_CONFIG_PATH at it). Every key is optional: delete
# a line to get its default (trailing `default:` comment). Without `url`,
# Enhance is off and GET /enhance/templates reports `available: false`.
#
# ComfyUI has no auth and stays LAN-only; the apps only ever talk to
# ImageApi, which relays previews. Nothing in this file is a secret.
#
# Templates are API-format workflows (ComfyUI -> Workflow -> Export (API)).
# ImageApi finds the nodes it fills in by `_meta.title`, so re-exporting
# after editing the graph keeps working:
# imageapi:source required Load Image From Path; gets the photo's path
# imageapi:save required Save Image Version; forced to require_approval
# imageapi:prompt optional text node; its `text`, `prompt` or `value`
# input (first one present) gets the user's text
# imageapi:reference optional second loader for a context photo. With no
# reference in the request the node and every
# link to it are removed from the prompt.
# A template missing a required title (or with a prompt node that has none of
# the three inputs) is rejected at startup with a log line; the rest still load.
# Never edit the template files from here: ImageApi only reads them.
url = "http://192.168.10.218:8188" # default: unset (Enhance disabled)
timeout_secs = 15 # default: 15; per ComfyUI request
probe_ttl_secs = 10 # default: 10; how long `available` is cached
poll_secs = 2 # default: 2; job poll interval
idle_free_secs = 300 # default: 300; ComfyUI idle (empty queue, newest
# history entry this old) -> POST /free once. 0 = off
unreachable_grace_secs = 120 # default: 120; active job -> `unavailable` after this
default_template = "restore" # default: the first template
warn_above_megapixels = 1.5 # default: 1.5; clients warn above this original size
# Library path prefix <-> ComfyUI path prefix; longest prefix wins and
# separators are converted, in both directions. With no entries, paths pass
# through unchanged. A mapped drive (Z:\) only exists in the logged-in
# session; use a UNC prefix (\\nas\) if ComfyUI ever runs as a service.
[[path_map]]
from = "/media/nas/"
to = 'Z:\'
[[template]]
name = "restore" # stable id clients send
label = "Restore"
description = "Sharpen and clean up the photo (about 3 s)."
file = "/media/nas/image_fibo2_restore_turbo_imageapi_API.json" # relative = beside this file
# append: the template's prompt is a full base instruction; the user's text
# is appended after prompt_join. Empty text leaves the prompt as is.
prompt_mode = "append" # default: "append" | "replace"
prompt_join = " " # default: " " (single space); e.g. "\n\nAdditional instructions: "
output_megapixels = 1.0 # default: 1.0; ~1 MP output, for the client's size warning
[[template]]
name = "edit"
label = "Edit"
description = "Change the photo as you describe."
file = "/media/nas/image_fibo2_edit_turbo_imageapi_API.json"
# replace: the user's text IS the prompt, so it is required (empty = 400).
prompt_mode = "replace"
# The Dev variant (50 steps, ~50 s) is no better on the archive test set, so
# it is not the default; configure it if you want it:
# [[template]]
# name = "restore-dev"
# label = "Restore (Dev, slow)"
# file = "/media/nas/image_fibo2_restore_imageapi_API.json"
@@ -0,0 +1 @@
DROP TABLE IF EXISTS enhance_jobs;
@@ -0,0 +1,31 @@
-- Enhance on demand (#192): one row per ComfyUI run submitted for a photo.
--
-- Keyed to the photo by `original_hash` (the blake3 content_hash of the file
-- that was submitted), like every other derived table, so a job survives a
-- rename or move; `library_id` + `rel_path` say which file it was at submit
-- time (and are what the API echoes back as `path`).
CREATE TABLE enhance_jobs (
id TEXT PRIMARY KEY NOT NULL, -- uuid, ImageApi's own
original_hash TEXT NOT NULL,
library_id INTEGER NOT NULL,
rel_path TEXT NOT NULL,
template TEXT NOT NULL,
prompt TEXT,
reference_hash TEXT,
reference_path TEXT, -- library-relative, same library
comfy_prompt_id TEXT NOT NULL,
-- queued | running | ready | approved | failed | expired | unavailable
status TEXT NOT NULL,
error TEXT,
queue_position INTEGER,
-- json: [{id, view: {filename, subfolder, type}}], ComfyUI's staged results
pending TEXT NOT NULL DEFAULT '[]',
-- json: [{n, comfy_path, library_id, rel_path}], approved results
saved TEXT NOT NULL DEFAULT '[]',
-- 1 = the user threw the result away; hidden from the per-photo list
discarded INTEGER NOT NULL DEFAULT 0,
created_at BIGINT NOT NULL,
updated_at BIGINT NOT NULL
);
CREATE INDEX idx_enhance_jobs_hash ON enhance_jobs(original_hash, created_at);
CREATE INDEX idx_enhance_jobs_status ON enhance_jobs(status);
+343
View File
@@ -0,0 +1,343 @@
//! Enhance jobs (#192): one row per ComfyUI run submitted for a photo.
//!
//! The row is the only state: the poller advances it from ComfyUI's
//! `/queue` and `/history`, the endpoints read it, approve appends to
//! `saved`. `pending` / `saved` hold small JSON documents owned by
//! `crate::enhance`; this layer treats them as opaque text.
use std::sync::{Arc, Mutex};
use diesel::prelude::*;
use diesel::sql_query;
use diesel::sql_types::{BigInt, Integer, Nullable, Text};
use diesel::sqlite::SqliteConnection;
use crate::database::{DbError, DbErrorKind, connect};
use crate::otel::trace_db_call;
pub const STATUS_QUEUED: &str = "queued";
pub const STATUS_RUNNING: &str = "running";
pub const STATUS_READY: &str = "ready";
pub const STATUS_APPROVED: &str = "approved";
pub const STATUS_FAILED: &str = "failed";
pub const STATUS_EXPIRED: &str = "expired";
pub const STATUS_UNAVAILABLE: &str = "unavailable";
#[derive(Debug, Clone, PartialEq, Eq, QueryableByName)]
pub struct EnhanceJobRow {
#[diesel(sql_type = Text)]
pub id: String,
#[diesel(sql_type = Text)]
pub original_hash: String,
#[diesel(sql_type = Integer)]
pub library_id: i32,
#[diesel(sql_type = Text)]
pub rel_path: String,
#[diesel(sql_type = Text)]
pub template: String,
#[diesel(sql_type = Nullable<Text>)]
pub prompt: Option<String>,
#[diesel(sql_type = Nullable<Text>)]
pub reference_hash: Option<String>,
#[diesel(sql_type = Nullable<Text>)]
pub reference_path: Option<String>,
#[diesel(sql_type = Text)]
pub comfy_prompt_id: String,
#[diesel(sql_type = Text)]
pub status: String,
#[diesel(sql_type = Nullable<Text>)]
pub error: Option<String>,
#[diesel(sql_type = Nullable<Integer>)]
pub queue_position: Option<i32>,
#[diesel(sql_type = Text)]
pub pending: String,
#[diesel(sql_type = Text)]
pub saved: String,
#[diesel(sql_type = Integer)]
pub discarded: i32,
#[diesel(sql_type = BigInt)]
pub created_at: i64,
#[diesel(sql_type = BigInt)]
pub updated_at: i64,
}
const COLUMNS: &str = "id, original_hash, library_id, rel_path, template, prompt, reference_hash, \
reference_path, comfy_prompt_id, status, error, queue_position, pending, saved, \
discarded, created_at, updated_at";
/// What the poller / approve changes on an existing job.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JobUpdate {
pub status: String,
pub error: Option<String>,
pub queue_position: Option<i32>,
/// `Some` replaces the `pending` document.
pub pending: Option<String>,
/// `Some` replaces the `saved` document.
pub saved: Option<String>,
}
impl JobUpdate {
pub fn status(status: &str) -> Self {
JobUpdate {
status: status.to_string(),
error: None,
queue_position: None,
pending: None,
saved: None,
}
}
}
pub trait EnhanceJobDao: Sync + Send {
fn insert(&mut self, job: &EnhanceJobRow) -> Result<(), DbError>;
fn get(&mut self, id: &str) -> Result<Option<EnhanceJobRow>, DbError>;
/// Newest first.
fn list_for_hash(
&mut self,
original_hash: &str,
limit: i64,
) -> Result<Vec<EnhanceJobRow>, DbError>;
/// Jobs the poller still has to advance (`queued` / `running`).
fn list_active(&mut self) -> Result<Vec<EnhanceJobRow>, DbError>;
fn update(&mut self, id: &str, update: &JobUpdate) -> Result<(), DbError>;
/// Hide a job from `list_for_hash`.
fn discard(&mut self, id: &str) -> Result<(), DbError>;
/// Delete jobs last touched before `cutoff` (unix seconds) that are not
/// active. Returns the number removed.
fn prune(&mut self, cutoff: i64) -> Result<usize, DbError>;
}
pub struct SqliteEnhanceJobDao {
connection: Arc<Mutex<SqliteConnection>>,
}
impl Default for SqliteEnhanceJobDao {
fn default() -> Self {
Self::new()
}
}
impl SqliteEnhanceJobDao {
pub fn new() -> Self {
SqliteEnhanceJobDao {
connection: Arc::new(Mutex::new(connect())),
}
}
/// Test-only: share a connection with other DAOs.
#[cfg(test)]
pub fn from_shared(connection: Arc<Mutex<SqliteConnection>>) -> Self {
SqliteEnhanceJobDao { connection }
}
}
impl EnhanceJobDao for SqliteEnhanceJobDao {
fn insert(&mut self, job: &EnhanceJobRow) -> Result<(), DbError> {
trace_db_call("insert", "enhance_job_insert", |_span| {
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
sql_query(format!(
"INSERT INTO enhance_jobs ({COLUMNS}) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)"
))
.bind::<Text, _>(&job.id)
.bind::<Text, _>(&job.original_hash)
.bind::<Integer, _>(job.library_id)
.bind::<Text, _>(&job.rel_path)
.bind::<Text, _>(&job.template)
.bind::<Nullable<Text>, _>(&job.prompt)
.bind::<Nullable<Text>, _>(&job.reference_hash)
.bind::<Nullable<Text>, _>(&job.reference_path)
.bind::<Text, _>(&job.comfy_prompt_id)
.bind::<Text, _>(&job.status)
.bind::<Nullable<Text>, _>(&job.error)
.bind::<Nullable<Integer>, _>(job.queue_position)
.bind::<Text, _>(&job.pending)
.bind::<Text, _>(&job.saved)
.bind::<Integer, _>(job.discarded)
.bind::<BigInt, _>(job.created_at)
.bind::<BigInt, _>(job.updated_at)
.execute(&mut *conn)
.map(|_| ())
.map_err(|e| anyhow::anyhow!("Insert error: {}", e))
})
.map_err(|e| DbError::log(DbErrorKind::InsertError, e))
}
fn get(&mut self, id: &str) -> Result<Option<EnhanceJobRow>, DbError> {
trace_db_call("query", "enhance_job_get", |_span| {
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
sql_query(format!("SELECT {COLUMNS} FROM enhance_jobs WHERE id = ?"))
.bind::<Text, _>(id)
.load::<EnhanceJobRow>(&mut *conn)
.map(|rows| rows.into_iter().next())
.map_err(|e| anyhow::anyhow!("Query error: {}", e))
})
.map_err(|e| DbError::log(DbErrorKind::QueryError, e))
}
fn list_for_hash(
&mut self,
original_hash: &str,
limit: i64,
) -> Result<Vec<EnhanceJobRow>, DbError> {
trace_db_call("query", "enhance_job_list_for_hash", |_span| {
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
sql_query(format!(
"SELECT {COLUMNS} FROM enhance_jobs WHERE original_hash = ? AND discarded = 0 \
ORDER BY created_at DESC, rowid DESC LIMIT ?"
))
.bind::<Text, _>(original_hash)
.bind::<BigInt, _>(limit)
.load::<EnhanceJobRow>(&mut *conn)
.map_err(|e| anyhow::anyhow!("Query error: {}", e))
})
.map_err(|e| DbError::log(DbErrorKind::QueryError, e))
}
fn list_active(&mut self) -> Result<Vec<EnhanceJobRow>, DbError> {
trace_db_call("query", "enhance_job_list_active", |_span| {
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
sql_query(format!(
"SELECT {COLUMNS} FROM enhance_jobs WHERE status IN ('queued', 'running') \
ORDER BY created_at, rowid"
))
.load::<EnhanceJobRow>(&mut *conn)
.map_err(|e| anyhow::anyhow!("Query error: {}", e))
})
.map_err(|e| DbError::log(DbErrorKind::QueryError, e))
}
fn update(&mut self, id: &str, u: &JobUpdate) -> Result<(), DbError> {
trace_db_call("update", "enhance_job_update", |_span| {
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
sql_query(
"UPDATE enhance_jobs SET status = ?, error = ?, queue_position = ?, \
pending = COALESCE(?, pending), saved = COALESCE(?, saved), updated_at = ? \
WHERE id = ?",
)
.bind::<Text, _>(&u.status)
.bind::<Nullable<Text>, _>(&u.error)
.bind::<Nullable<Integer>, _>(u.queue_position)
.bind::<Nullable<Text>, _>(&u.pending)
.bind::<Nullable<Text>, _>(&u.saved)
.bind::<BigInt, _>(chrono::Utc::now().timestamp())
.bind::<Text, _>(id)
.execute(&mut *conn)
.map(|_| ())
.map_err(|e| anyhow::anyhow!("Update error: {}", e))
})
.map_err(|e| DbError::log(DbErrorKind::UpdateError, e))
}
fn discard(&mut self, id: &str) -> Result<(), DbError> {
trace_db_call("update", "enhance_job_discard", |_span| {
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
sql_query("UPDATE enhance_jobs SET discarded = 1, updated_at = ? WHERE id = ?")
.bind::<BigInt, _>(chrono::Utc::now().timestamp())
.bind::<Text, _>(id)
.execute(&mut *conn)
.map(|_| ())
.map_err(|e| anyhow::anyhow!("Update error: {}", e))
})
.map_err(|e| DbError::log(DbErrorKind::UpdateError, e))
}
fn prune(&mut self, cutoff: i64) -> Result<usize, DbError> {
trace_db_call("delete", "enhance_job_prune", |_span| {
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
sql_query(
"DELETE FROM enhance_jobs WHERE updated_at < ? \
AND status NOT IN ('queued', 'running')",
)
.bind::<BigInt, _>(cutoff)
.execute(&mut *conn)
.map_err(|e| anyhow::anyhow!("Delete error: {}", e))
})
.map_err(|e| DbError::log(DbErrorKind::UpdateError, e))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::database::test::in_memory_db_connection;
fn dao() -> SqliteEnhanceJobDao {
SqliteEnhanceJobDao::from_shared(Arc::new(Mutex::new(in_memory_db_connection())))
}
fn job(id: &str, hash: &str, status: &str, created: i64) -> EnhanceJobRow {
EnhanceJobRow {
id: id.into(),
original_hash: hash.into(),
library_id: 1,
rel_path: "a.jpg".into(),
template: "restore".into(),
prompt: Some("warmer".into()),
reference_hash: None,
reference_path: None,
comfy_prompt_id: format!("p-{id}"),
status: status.into(),
error: None,
queue_position: None,
pending: "[]".into(),
saved: "[]".into(),
discarded: 0,
created_at: created,
updated_at: created,
}
}
#[test]
fn insert_get_list_update_prune() {
let mut d = dao();
d.insert(&job("a", "h1", "queued", 100)).unwrap();
d.insert(&job("b", "h1", "ready", 200)).unwrap();
d.insert(&job("c", "h2", "running", 300)).unwrap();
assert_eq!(d.get("a").unwrap().unwrap().comfy_prompt_id, "p-a");
assert!(d.get("zzz").unwrap().is_none());
let ids: Vec<String> = d
.list_for_hash("h1", 10)
.unwrap()
.into_iter()
.map(|j| j.id)
.collect();
assert_eq!(ids, ["b", "a"], "newest first");
assert_eq!(d.list_for_hash("h1", 1).unwrap().len(), 1);
let active: Vec<String> = d.list_active().unwrap().into_iter().map(|j| j.id).collect();
assert_eq!(active, ["a", "c"]);
d.update(
"a",
&JobUpdate {
status: "ready".into(),
error: None,
queue_position: None,
pending: Some("[{\"id\":\"x\"}]".into()),
saved: None,
},
)
.unwrap();
let a = d.get("a").unwrap().unwrap();
assert_eq!(a.status, "ready");
assert_eq!(a.pending, "[{\"id\":\"x\"}]");
assert_eq!(a.saved, "[]", "unset columns are left alone");
assert!(a.updated_at >= 100);
d.discard("b").unwrap();
assert_eq!(
d.list_for_hash("h1", 10).unwrap().len(),
1,
"discarded hidden"
);
// Prune drops only old, finished jobs (updated_at was just bumped
// for `a`; `c` is active).
assert_eq!(d.prune(i64::MAX).unwrap(), 2);
assert!(d.get("c").unwrap().is_some());
assert!(d.get("b").unwrap().is_none());
}
}
+2
View File
@@ -43,6 +43,7 @@ pub struct UtcBackfillCandidate {
pub mod calendar_dao;
pub mod daily_summary_dao;
pub mod enhance_dao;
pub mod insight_generation_job_dao;
pub mod insights_dao;
pub mod knowledge_dao;
@@ -59,6 +60,7 @@ pub mod version_dao;
pub use calendar_dao::{CalendarEventDao, SqliteCalendarEventDao};
pub use daily_summary_dao::{DailySummaryDao, InsertDailySummary, SqliteDailySummaryDao};
pub use enhance_dao::{EnhanceJobDao, SqliteEnhanceJobDao};
pub use insight_generation_job_dao::{InsightGenerationJobDao, SqliteInsightGenerationJobDao};
pub use insights_dao::{InsightDao, SqliteInsightDao};
pub use knowledge_dao::{
+23
View File
@@ -41,6 +41,28 @@ diesel::table! {
}
}
diesel::table! {
enhance_jobs (id) {
id -> Text,
original_hash -> Text,
library_id -> Integer,
rel_path -> Text,
template -> Text,
prompt -> Nullable<Text>,
reference_hash -> Nullable<Text>,
reference_path -> Nullable<Text>,
comfy_prompt_id -> Text,
status -> Text,
error -> Nullable<Text>,
queue_position -> Nullable<Integer>,
pending -> Text,
saved -> Text,
discarded -> Integer,
created_at -> BigInt,
updated_at -> BigInt,
}
}
diesel::table! {
entities (id) {
id -> Integer,
@@ -406,6 +428,7 @@ diesel::allow_tables_to_appear_in_same_query!(
calendar_events,
clip_failures,
daily_conversation_summaries,
enhance_jobs,
entities,
entity_facts,
entity_photo_links,
+428
View File
@@ -0,0 +1,428 @@
//! Thin client for the parts of ComfyUI's HTTP API Enhance uses (#192).
//!
//! | Step | Call |
//! |-----------|----------------------------------------|
//! | reachable | `GET /system_stats` |
//! | submit | `POST /prompt {prompt, client_id}` |
//! | status | `GET /queue`, `GET /history/{id}` |
//! | preview | `GET /view?filename=&subfolder=&type=` |
//! | approve | `POST /imageapi/approve {id}` |
//! | free VRAM | `POST /free` |
use std::time::Duration;
use bytes::Bytes;
use serde::Deserialize;
use serde_json::{Value, json};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ComfyError {
/// Connect/timeout failure: ComfyUI isn't running (a normal state).
Unreachable(String),
/// ComfyUI answered with a status we don't handle.
Http { status: u16, body: String },
/// ComfyUI answered with something we can't read.
BadResponse(String),
}
impl std::fmt::Display for ComfyError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ComfyError::Unreachable(e) => write!(f, "ComfyUI is unreachable: {e}"),
ComfyError::Http { status, body } => write!(f, "ComfyUI returned {status}: {body}"),
ComfyError::BadResponse(e) => write!(f, "unexpected ComfyUI response: {e}"),
}
}
}
impl std::error::Error for ComfyError {}
impl From<reqwest::Error> for ComfyError {
fn from(e: reqwest::Error) -> Self {
if e.is_connect() || e.is_timeout() || e.is_request() {
ComfyError::Unreachable(e.without_url().to_string())
} else {
ComfyError::BadResponse(e.without_url().to_string())
}
}
}
/// `POST /prompt` outcomes other than success.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SubmitError {
Comfy(ComfyError),
/// ComfyUI refused the prompt (HTTP 400): validation failed.
Rejected {
message: String,
/// `node_errors` was non-empty (a node is missing an input, a model
/// or a link target).
node_errors: bool,
},
}
impl std::fmt::Display for SubmitError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SubmitError::Comfy(e) => e.fmt(f),
SubmitError::Rejected { message, .. } => {
write!(f, "ComfyUI rejected the workflow: {message}")
}
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ApproveError {
Comfy(ComfyError),
/// The staged result is gone (ComfyUI restarted: its temp folder is cleared).
Gone,
/// ComfyUI refused the approval (bad id, write failure).
Refused(String),
}
/// Where `GET /view` finds a staged result.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, Deserialize)]
pub struct ViewRef {
pub filename: String,
#[serde(default)]
pub subfolder: String,
#[serde(default = "temp", rename = "type")]
pub kind: String,
}
fn temp() -> String {
"temp".into()
}
/// A result the save node staged: ComfyUI's pending id, where to preview it,
/// and the library path when it was already saved (never, with approval on).
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, Deserialize)]
pub struct Staged {
pub id: String,
pub view: ViewRef,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HistoryEntry {
pub completed: bool,
/// `success` / `error`.
pub status_str: String,
/// ComfyUI's exception message when the run failed.
pub error: Option<String>,
pub staged: Vec<Staged>,
/// Last event time, unix millis.
pub timestamp_ms: Option<i64>,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct QueueState {
pub running: Vec<String>,
pub pending: Vec<String>,
}
impl QueueState {
pub fn is_empty(&self) -> bool {
self.running.is_empty() && self.pending.is_empty()
}
/// Jobs ahead of `prompt_id`; `None` when it isn't queued.
pub fn position_of(&self, prompt_id: &str) -> Option<(bool, usize)> {
if self.running.iter().any(|p| p == prompt_id) {
return Some((true, 0));
}
self.pending
.iter()
.position(|p| p == prompt_id)
.map(|i| (false, self.running.len() + i))
}
}
#[derive(Clone)]
pub struct ComfyClient {
http: reqwest::Client,
base: String,
/// Client id sent with submissions.
client_id: String,
}
impl ComfyClient {
pub fn new(base_url: &str, timeout: Duration) -> Self {
let http = reqwest::Client::builder()
.timeout(timeout)
.connect_timeout(Duration::from_secs(3).min(timeout))
.build()
.expect("reqwest client");
ComfyClient {
http,
base: base_url.trim_end_matches('/').to_string(),
client_id: "imageapi".to_string(),
}
}
fn url(&self, path: &str) -> String {
format!("{}{}", self.base, path)
}
/// `GET /system_stats` with its own (short) timeout.
pub async fn probe(&self, timeout: Duration) -> bool {
self.http
.get(self.url("/system_stats"))
.timeout(timeout)
.send()
.await
.is_ok_and(|r| r.status().is_success())
}
pub async fn submit(&self, prompt: &Value) -> Result<String, SubmitError> {
let body = json!({ "prompt": prompt, "client_id": self.client_id });
let resp = self
.http
.post(self.url("/prompt"))
.json(&body)
.send()
.await
.map_err(|e| SubmitError::Comfy(e.into()))?;
let status = resp.status();
let text = resp
.text()
.await
.map_err(|e| SubmitError::Comfy(e.into()))?;
let json: Value = serde_json::from_str(&text).unwrap_or(Value::Null);
if status.is_success() {
// A 200 can still carry node_errors for outputs that were pruned;
// the prompt itself queued.
return json["prompt_id"]
.as_str()
.map(str::to_string)
.ok_or_else(|| {
SubmitError::Comfy(ComfyError::BadResponse(format!("no prompt_id in {text}")))
});
}
if status.as_u16() == 400 {
return Err(SubmitError::Rejected {
message: rejection_message(&json, &text),
node_errors: json["node_errors"]
.as_object()
.is_some_and(|m| !m.is_empty()),
});
}
Err(SubmitError::Comfy(ComfyError::Http {
status: status.as_u16(),
body: truncate(&text),
}))
}
pub async fn queue(&self) -> Result<QueueState, ComfyError> {
let json: Value = self.get_json("/queue").await?;
let ids = |key: &str| -> Vec<String> {
json[key]
.as_array()
.map(|a| {
a.iter()
.filter_map(|item| item.get(1).and_then(Value::as_str))
.map(str::to_string)
.collect()
})
.unwrap_or_default()
};
Ok(QueueState {
running: ids("queue_running"),
pending: ids("queue_pending"),
})
}
/// `GET /history/{id}`; `None` when ComfyUI doesn't know the id.
pub async fn history(&self, prompt_id: &str) -> Result<Option<HistoryEntry>, ComfyError> {
let json: Value = self
.get_json(&format!("/history/{}", urlencoding::encode(prompt_id)))
.await?;
Ok(json.get(prompt_id).map(parse_history_entry))
}
/// The newest finished run, for the idle check.
pub async fn newest_history(&self) -> Result<Option<HistoryEntry>, ComfyError> {
let json: Value = self.get_json("/history?max_items=1").await?;
Ok(json
.as_object()
.and_then(|m| m.values().next())
.map(parse_history_entry))
}
/// `GET /view`: the staged PNG. `Ok(None)` for 404.
pub async fn view(&self, r: &ViewRef) -> Result<Option<(String, Bytes)>, ComfyError> {
let resp = self
.http
.get(self.url("/view"))
.query(&[
("filename", r.filename.as_str()),
("subfolder", r.subfolder.as_str()),
("type", r.kind.as_str()),
])
.send()
.await?;
if resp.status().as_u16() == 404 {
return Ok(None);
}
if !resp.status().is_success() {
return Err(ComfyError::Http {
status: resp.status().as_u16(),
body: truncate(&resp.text().await.unwrap_or_default()),
});
}
let content_type = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.unwrap_or("image/png")
.to_string();
Ok(Some((content_type, resp.bytes().await?)))
}
/// `POST /imageapi/approve`: ComfyUI writes the file; returns its path
/// as ComfyUI sees it.
pub async fn approve(&self, pending_id: &str) -> Result<String, ApproveError> {
let resp = self
.http
.post(self.url("/imageapi/approve"))
.json(&json!({ "id": pending_id }))
.send()
.await
.map_err(|e| ApproveError::Comfy(e.into()))?;
let status = resp.status();
let text = resp
.text()
.await
.map_err(|e| ApproveError::Comfy(e.into()))?;
let json: Value = serde_json::from_str(&text).unwrap_or(Value::Null);
if status.is_success() {
return json["path"].as_str().map(str::to_string).ok_or_else(|| {
ApproveError::Comfy(ComfyError::BadResponse(format!("no path in {text}")))
});
}
let message = json["error"].as_str().unwrap_or(&text).to_string();
match status.as_u16() {
404 => Err(ApproveError::Gone),
400 | 500 => Err(ApproveError::Refused(truncate(&message))),
s => Err(ApproveError::Comfy(ComfyError::Http {
status: s,
body: truncate(&text),
})),
}
}
/// `POST /free`: unload models and release cached memory.
pub async fn free(&self) -> Result<(), ComfyError> {
let resp = self
.http
.post(self.url("/free"))
.json(&json!({ "unload_models": true, "free_memory": true }))
.send()
.await?;
if resp.status().is_success() {
Ok(())
} else {
Err(ComfyError::Http {
status: resp.status().as_u16(),
body: truncate(&resp.text().await.unwrap_or_default()),
})
}
}
async fn get_json(&self, path: &str) -> Result<Value, ComfyError> {
let resp = self.http.get(self.url(path)).send().await?;
if !resp.status().is_success() {
return Err(ComfyError::Http {
status: resp.status().as_u16(),
body: truncate(&resp.text().await.unwrap_or_default()),
});
}
Ok(resp.json().await?)
}
}
fn truncate(s: &str) -> String {
s.chars().take(300).collect()
}
fn rejection_message(json: &Value, raw: &str) -> String {
let err = &json["error"];
let head = err["message"].as_str().or_else(|| err.as_str());
let detail = err["details"].as_str().filter(|d| !d.is_empty());
// Per-node errors say which input is wrong.
let nodes: Vec<String> = json["node_errors"]
.as_object()
.map(|m| {
m.values()
.filter_map(|n| {
let class = n["class_type"].as_str().unwrap_or("node");
let first = n["errors"].as_array()?.first()?;
Some(format!(
"{class}: {}",
first["message"].as_str().unwrap_or("invalid")
))
})
.collect()
})
.unwrap_or_default();
let mut parts: Vec<String> = Vec::new();
parts.extend(head.map(str::to_string));
parts.extend(detail.map(str::to_string));
parts.extend(nodes);
if parts.is_empty() {
truncate(raw)
} else {
truncate(&parts.join("; "))
}
}
pub(crate) fn parse_history_entry(entry: &Value) -> HistoryEntry {
let status = &entry["status"];
let status_str = status["status_str"]
.as_str()
.unwrap_or_default()
.to_string();
let completed = status["completed"].as_bool().unwrap_or(false);
let messages = status["messages"].as_array().cloned().unwrap_or_default();
let error = messages.iter().find_map(|m| {
let (kind, data) = (m.get(0)?.as_str()?, m.get(1)?);
(kind == "execution_error").then(|| {
data["exception_message"]
.as_str()
.unwrap_or("execution error")
.trim()
.to_string()
})
});
let timestamp_ms = messages
.iter()
.filter_map(|m| m.get(1)?.get("timestamp")?.as_i64())
.max();
let mut staged = Vec::new();
if let Some(outputs) = entry["outputs"].as_object() {
for out in outputs.values() {
let Some(results) = out["imageapi_versions"].as_array() else {
continue;
};
let images = out["images"].as_array().cloned().unwrap_or_default();
for (i, r) in results.iter().enumerate() {
let (Some(id), Some(image)) = (r["id"].as_str(), images.get(i)) else {
continue;
};
if let Ok(view) = serde_json::from_value::<ViewRef>(image.clone()) {
staged.push(Staged {
id: id.to_string(),
view,
});
}
}
}
}
HistoryEntry {
completed,
status_str,
error,
staged,
timestamp_ms,
}
}
+263
View File
@@ -0,0 +1,263 @@
//! Enhance configuration (#192): one `enhance.toml`, next to `llm.toml`.
//!
//! Path from `ENHANCE_CONFIG_PATH`, else `./enhance.toml`. A missing file at
//! the default path means Enhance is simply off (no `url`); a missing
//! explicit path is an error. Nothing here is a secret: ComfyUI has no auth
//! and stays LAN-only. `enhance.example.toml` is the annotated schema (a
//! test parses it).
use std::path::{Path, PathBuf};
use anyhow::{Context as _, Result, bail};
use serde::Deserialize;
use super::path_map::PathMap;
const DEFAULT_PATH: &str = "enhance.toml";
/// How the app's text combines with the template's own prompt.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, serde::Serialize, Default)]
#[serde(rename_all = "lowercase")]
pub enum PromptMode {
/// The template's prompt is a full base instruction; the app's text is
/// appended after `prompt_join`. Empty text leaves the prompt as is.
#[default]
Append,
/// The app's text *is* the prompt (and is required).
Replace,
}
#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct EnhanceConfig {
/// ComfyUI base URL, e.g. `http://192.168.10.218:8188`. Unset = Enhance
/// is disabled (`/enhance/templates` reports `available: false`).
pub url: Option<String>,
/// Per-request timeout for ComfyUI calls.
pub timeout_secs: u64,
/// How long the `available` probe result is cached.
pub probe_ttl_secs: u64,
/// Job poll interval.
pub poll_secs: u64,
/// Idle time (empty queue, newest history entry this old) before
/// ImageApi calls `POST /free` once. 0 disables it.
pub idle_free_secs: u64,
/// An active job whose ComfyUI stays unreachable this long becomes
/// `unavailable`.
pub unreachable_grace_secs: u64,
/// Template preselected by clients; default = the first one.
pub default_template: Option<String>,
/// Clients warn above this many megapixels in the original.
pub warn_above_megapixels: f64,
/// Library path prefix <-> ComfyUI path prefix. Longest prefix wins.
pub path_map: Vec<PathMapping>,
#[serde(rename = "template")]
pub templates: Vec<TemplateConfig>,
}
impl Default for EnhanceConfig {
fn default() -> Self {
Self {
url: None,
timeout_secs: 15,
probe_ttl_secs: 10,
poll_secs: 2,
idle_free_secs: 300,
unreachable_grace_secs: 120,
default_template: None,
warn_above_megapixels: 1.5,
path_map: Vec::new(),
templates: Vec::new(),
}
}
}
#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct PathMapping {
/// Prefix as ImageApi sees it, e.g. `/media/nas/`.
pub from: String,
/// Prefix as ComfyUI sees it, e.g. `Z:\` or `\\nas\`.
pub to: String,
}
#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct TemplateConfig {
/// Stable id clients send (`POST /enhance { template }`).
pub name: String,
pub label: String,
#[serde(default)]
pub description: String,
/// API-format workflow JSON (ComfyUI -> Workflow -> Export (API)).
/// Relative paths resolve against the config file's directory.
pub file: PathBuf,
#[serde(default)]
pub prompt_mode: PromptMode,
/// Text between the template's prompt and the app's text (append mode).
#[serde(default = "default_join")]
pub prompt_join: String,
/// Approximate output size, for the client's large-original warning.
#[serde(default = "default_output_mp")]
pub output_megapixels: f64,
}
fn default_join() -> String {
" ".into()
}
fn default_output_mp() -> f64 {
1.0
}
impl EnhanceConfig {
pub fn validate(&self) -> Result<()> {
if let Some(url) = &self.url
&& !(url.starts_with("http://") || url.starts_with("https://"))
{
bail!("url must start with http:// or https://, got {url:?}");
}
for (key, v) in [
("timeout_secs", self.timeout_secs),
("probe_ttl_secs", self.probe_ttl_secs),
("poll_secs", self.poll_secs),
("unreachable_grace_secs", self.unreachable_grace_secs),
] {
if v == 0 {
bail!("{key} must be greater than 0");
}
}
if self.warn_above_megapixels.is_nan() || self.warn_above_megapixels <= 0.0 {
bail!("warn_above_megapixels must be greater than 0");
}
let mut seen = std::collections::HashSet::new();
for t in &self.templates {
let ok = !t.name.is_empty()
&& t.name
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-');
if !ok {
bail!(
"template name {:?} must be non-empty and use only letters, digits, '_' and '-'",
t.name
);
}
if !seen.insert(t.name.as_str()) {
bail!("duplicate template name {:?}", t.name);
}
if t.label.trim().is_empty() {
bail!("template {:?} needs a label", t.name);
}
if t.output_megapixels.is_nan() || t.output_megapixels <= 0.0 {
bail!("template {:?}: output_megapixels must be > 0", t.name);
}
}
if let Some(d) = &self.default_template
&& !self.templates.iter().any(|t| &t.name == d)
{
bail!("default_template {d:?} is not a configured template");
}
// Parses every mapping (rejects empty prefixes).
PathMap::new(&self.path_map)?;
Ok(())
}
/// The preselected template's name.
pub fn default_name(&self) -> Option<&str> {
self.default_template
.as_deref()
.or_else(|| self.templates.first().map(|t| t.name.as_str()))
}
}
/// Where the config was read from, and what it said.
#[derive(Debug)]
pub struct Loaded {
pub config: EnhanceConfig,
/// Directory relative template `file`s resolve against.
pub base_dir: PathBuf,
/// The file, when one was read.
pub path: Option<PathBuf>,
}
/// Read `enhance.toml` (see module docs). A bad file is an error so a typo
/// fails startup instead of silently disabling the feature.
pub fn load() -> Result<Loaded> {
let (path, explicit) = match std::env::var("ENHANCE_CONFIG_PATH") {
Ok(p) if !p.trim().is_empty() => (PathBuf::from(p.trim()), true),
_ => (PathBuf::from(DEFAULT_PATH), false),
};
match std::fs::read_to_string(&path) {
Ok(text) => {
let config = parse(&text).with_context(|| format!("in {}", path.display()))?;
let base_dir = path
.parent()
.filter(|p| !p.as_os_str().is_empty())
.map(Path::to_path_buf)
.unwrap_or_else(|| PathBuf::from("."));
Ok(Loaded {
config,
base_dir,
path: Some(path),
})
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound && !explicit => Ok(Loaded {
config: EnhanceConfig::default(),
base_dir: PathBuf::from("."),
path: None,
}),
Err(e) => Err(e).with_context(|| format!("reading enhance config {}", path.display())),
}
}
pub fn parse(text: &str) -> Result<EnhanceConfig> {
let config: EnhanceConfig = toml::from_str(text)?;
config.validate()?;
Ok(config)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn example_file_parses() {
let cfg = parse(include_str!("../../enhance.example.toml")).unwrap();
assert!(cfg.url.is_some());
assert_eq!(cfg.templates[0].name, "restore");
assert_eq!(cfg.default_name(), Some("restore"));
assert!(
cfg.templates
.iter()
.any(|t| t.prompt_mode == PromptMode::Replace)
);
}
#[test]
fn empty_file_is_disabled_defaults() {
let cfg = parse("").unwrap();
assert_eq!(cfg, EnhanceConfig::default());
assert!(cfg.url.is_none());
}
#[test]
fn rejects_unknown_keys_and_bad_values() {
assert!(parse("nope = 1").is_err());
assert!(parse("url = \"ftp://x\"").is_err());
assert!(parse("poll_secs = 0").is_err());
assert!(parse("default_template = \"x\"").is_err());
let dup = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n[[template]]\nname='a'\nlabel='B'\nfile='b.json'";
assert!(parse(dup).is_err());
assert!(parse("[[template]]\nname='a b'\nlabel='A'\nfile='a.json'").is_err());
assert!(parse("[[path_map]]\nfrom=''\nto='Z:\\\\'").is_err());
}
#[test]
fn template_defaults() {
let cfg = parse("[[template]]\nname='a'\nlabel='A'\nfile='a.json'").unwrap();
let t = &cfg.templates[0];
assert_eq!(t.prompt_mode, PromptMode::Append);
assert_eq!(t.prompt_join, " ");
assert_eq!(t.output_megapixels, 1.0);
}
}
+523
View File
@@ -0,0 +1,523 @@
//! HTTP surface of Enhance (#192). ImageApi relays everything, previews
//! included: the apps never talk to ComfyUI.
//!
//! - `GET /enhance/templates`
//! - `POST /enhance`
//! - `GET /enhance/jobs?path=`
//! - `GET /enhance/jobs/{id}`
//! - `GET /enhance/jobs/{id}/preview/{n}`
//! - `POST /enhance/jobs/{id}/approve`
//! - `POST /enhance/jobs/{id}/discard`
use std::path::{Component, Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use actix_web::{HttpResponse, Responder, get, post, web};
use log::warn;
use serde::{Deserialize, Serialize};
use serde_json::json;
use super::config::PromptMode;
use super::service::{EnhanceError, EnhanceService, SourceFile};
use crate::data::Claims;
use crate::database::ExifDao;
use crate::database::enhance_dao::{EnhanceJobRow, STATUS_APPROVED, STATUS_READY};
use crate::libraries::{self, Library};
use crate::state::AppState;
use crate::utils::normalize_path;
use crate::versions::handlers::build_response;
type ExifData = web::Data<Mutex<Box<dyn ExifDao>>>;
fn error_response(e: &EnhanceError) -> HttpResponse {
let (mut builder, code) = match e {
EnhanceError::Disabled => (HttpResponse::ServiceUnavailable(), "disabled"),
EnhanceError::Unavailable => (HttpResponse::ServiceUnavailable(), "unavailable"),
EnhanceError::UnknownTemplate(_) => (HttpResponse::NotFound(), "unknown_template"),
EnhanceError::BadRequest(_) => (HttpResponse::BadRequest(), "bad_request"),
EnhanceError::Rejected(_) => (HttpResponse::UnprocessableEntity(), "rejected"),
EnhanceError::Expired => (HttpResponse::Gone(), "expired"),
EnhanceError::NotFound => (HttpResponse::NotFound(), "not_found"),
EnhanceError::Unplaceable(_) => (HttpResponse::BadGateway(), "unplaceable"),
EnhanceError::Upstream(_) => (HttpResponse::BadGateway(), "upstream"),
EnhanceError::Internal(_) => (HttpResponse::InternalServerError(), "internal"),
};
builder.json(json!({ "error": e.message(), "code": code }))
}
fn bad(msg: &str) -> HttpResponse {
error_response(&EnhanceError::BadRequest(msg.to_string()))
}
// ---- views --------------------------------------------------------------
#[derive(Debug, Serialize)]
struct PendingView {
/// Index for `preview/{n}` and `approve { n }`.
n: usize,
/// Path of the preview on this API.
preview: String,
approved: bool,
}
#[derive(Debug, Serialize)]
struct SavedView {
n: usize,
path: Option<String>,
library_id: Option<i32>,
}
#[derive(Debug, Serialize)]
pub struct JobView {
id: String,
status: String,
#[serde(skip_serializing_if = "Option::is_none")]
queue_position: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
path: String,
library_id: i32,
template: String,
#[serde(skip_serializing_if = "Option::is_none")]
prompt: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
reference_path: Option<String>,
created_at: i64,
updated_at: i64,
pending: Vec<PendingView>,
saved: Vec<SavedView>,
}
pub fn job_view(job: &EnhanceJobRow) -> JobView {
let saved = EnhanceService::saved_of(job);
// An expired job's staged files are gone; don't offer them.
let pending = if job.status == STATUS_READY || job.status == STATUS_APPROVED {
EnhanceService::pending_of(job)
.iter()
.enumerate()
.map(|(n, _)| PendingView {
n,
preview: format!("/enhance/jobs/{}/preview/{n}", job.id),
approved: saved.iter().any(|s| s.n == n),
})
.collect()
} else {
Vec::new()
};
let queued = matches!(job.status.as_str(), "queued" | "running");
JobView {
id: job.id.clone(),
status: job.status.clone(),
queue_position: job.queue_position.filter(|_| queued),
error: job.error.clone(),
path: job.rel_path.clone(),
library_id: job.library_id,
template: job.template.clone(),
prompt: job.prompt.clone(),
reference_path: job.reference_path.clone(),
created_at: job.created_at,
updated_at: job.updated_at,
pending,
saved: saved
.into_iter()
.map(|s| SavedView {
n: s.n,
path: s.rel_path,
library_id: s.library_id,
})
.collect(),
}
}
// ---- file resolution ----------------------------------------------------
fn safe_rel(rel: &str) -> bool {
!rel.is_empty()
&& Path::new(rel)
.components()
.all(|c| matches!(c, Component::Normal(_)))
}
/// An indexed, hashed file in a library.
fn resolve_file(
app_state: &AppState,
exif_dao: &Mutex<Box<dyn ExifDao>>,
path: &str,
library: Option<&str>,
) -> Result<SourceFile, HttpResponse> {
let library_id = match libraries::resolve_library_param_state(app_state, library) {
Ok(l) => l.map(|l| l.id),
Err(msg) => return Err(bad(&msg)),
};
let rel = normalize_path(path);
if !safe_rel(&rel) {
return Err(bad("invalid path"));
}
let mut rows = exif_dao
.lock()
.expect("Unable to lock ExifDao")
.get_exif_batch(library_id, std::slice::from_ref(&rel))
.map_err(|_| error_response(&EnhanceError::Internal("lookup failed".into())))?;
rows.sort_by_key(|r| r.id);
let Some(row) = rows.iter().find(|r| r.content_hash.is_some()) else {
return Err(error_response(&EnhanceError::NotFound));
};
let Some(lib) = app_state.library_by_id(row.library_id) else {
return Err(error_response(&EnhanceError::NotFound));
};
let abs = Path::new(&lib.root_path).join(&rel);
if !abs.is_file() {
return Err(error_response(&EnhanceError::NotFound));
}
Ok(SourceFile {
library_id: row.library_id,
rel_path: rel,
abs,
hash: row.content_hash.clone().unwrap_or_default(),
})
}
/// Library and relative path of an absolute file ComfyUI wrote. The job's own
/// library wins when roots overlap.
fn place_in_libraries(libs: &[Library], prefer: i32, abs: &Path) -> Option<(i32, String)> {
let mut ordered: Vec<&Library> = libs.iter().collect();
ordered.sort_by_key(|l| l.id != prefer);
for lib in ordered {
if let Ok(rest) = abs.strip_prefix(&lib.root_path) {
let rel = rest.to_str()?.replace('\\', "/");
if safe_rel(&rel) {
return Some((lib.id, rel));
}
}
}
None
}
// ---- templates ----------------------------------------------------------
#[derive(Serialize)]
struct TemplateView {
name: String,
label: String,
description: String,
accepts_prompt: bool,
prompt_mode: PromptMode,
/// Replace-mode templates need non-empty text (empty is a 400).
prompt_required: bool,
accepts_reference: bool,
default: bool,
/// Approximate output size; clients warn when the original is larger.
output_megapixels: f64,
warn_above_megapixels: f64,
}
#[get("/enhance/templates")]
pub async fn get_templates(_: Claims, svc: web::Data<EnhanceService>) -> impl Responder {
let available = svc.available().await;
let reason = if !svc.is_configured() {
Some(super::service::MSG_DISABLED)
} else if !available {
Some(super::service::MSG_UNAVAILABLE)
} else {
None
};
let default = svc.default_template().map(str::to_string);
let templates: Vec<TemplateView> = svc
.templates()
.iter()
.map(|t| TemplateView {
name: t.cfg.name.clone(),
label: t.cfg.label.clone(),
description: t.cfg.description.clone(),
accepts_prompt: t.accepts_prompt(),
prompt_mode: t.cfg.prompt_mode,
prompt_required: t.prompt_required(),
accepts_reference: t.accepts_reference(),
default: default.as_deref() == Some(t.cfg.name.as_str()),
output_megapixels: t.cfg.output_megapixels,
warn_above_megapixels: svc.config.warn_above_megapixels,
})
.collect();
let mut body = json!({
"available": available,
"default": default,
"warn_above_megapixels": svc.config.warn_above_megapixels,
"templates": templates,
});
if let Some(r) = reason {
body["reason"] = json!(r);
}
HttpResponse::Ok().json(body)
}
// ---- jobs ---------------------------------------------------------------
#[derive(Debug, Deserialize)]
pub struct LibraryQuery {
pub library: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct SubmitBody {
pub path: String,
pub template: String,
#[serde(default)]
pub prompt: Option<String>,
#[serde(default)]
pub reference_path: Option<String>,
#[serde(default)]
pub library: Option<String>,
}
#[post("/enhance")]
#[tracing::instrument(name = "enhance_submit", skip_all)]
pub async fn submit(
_: Claims,
query: web::Query<LibraryQuery>,
body: web::Json<SubmitBody>,
svc: web::Data<EnhanceService>,
app_state: web::Data<AppState>,
exif_dao: ExifData,
) -> impl Responder {
if !svc.is_configured() {
return error_response(&EnhanceError::Disabled);
}
let library = body.library.as_deref().or(query.library.as_deref());
let photo = match resolve_file(&app_state, &exif_dao, &body.path, library) {
Ok(f) => f,
Err(r) => return r,
};
let reference = match body
.reference_path
.as_deref()
.filter(|p| !p.trim().is_empty())
{
Some(p) => {
// The reference lives in the photo's own library.
let lib = app_state
.library_by_id(photo.library_id)
.map(|l| l.name.clone());
match resolve_file(&app_state, &exif_dao, p, lib.as_deref()) {
Ok(f) => Some(f),
Err(r) => return r,
}
}
None => None,
};
match svc
.submit(
&photo,
&body.template,
body.prompt.as_deref(),
reference.as_ref(),
)
.await
{
Ok(job) => HttpResponse::Ok().json(job_view(&job)),
Err(e) => error_response(&e),
}
}
#[get("/enhance/jobs/{id}")]
pub async fn get_job(
_: Claims,
id: web::Path<String>,
svc: web::Data<EnhanceService>,
) -> impl Responder {
match svc.get_job(&id) {
Ok(job) => HttpResponse::Ok().json(job_view(&job)),
Err(e) => error_response(&e),
}
}
#[derive(Debug, Deserialize)]
pub struct JobsQuery {
pub path: String,
pub library: Option<String>,
}
/// A photo's recent jobs, newest first, so the viewer can pick up an enhance
/// in progress after reopening. Discarded jobs are not listed.
#[get("/enhance/jobs")]
pub async fn list_jobs(
_: Claims,
q: web::Query<JobsQuery>,
svc: web::Data<EnhanceService>,
app_state: web::Data<AppState>,
exif_dao: ExifData,
) -> impl Responder {
let file = match resolve_file(&app_state, &exif_dao, &q.path, q.library.as_deref()) {
Ok(f) => f,
Err(r) => return r,
};
match svc.jobs_for_hash(&file.hash) {
Ok(jobs) => {
let views: Vec<JobView> = jobs.iter().map(job_view).collect();
HttpResponse::Ok().json(json!({ "jobs": views }))
}
Err(e) => error_response(&e),
}
}
#[get("/enhance/jobs/{id}/preview/{n}")]
pub async fn preview(
_: Claims,
path: web::Path<(String, usize)>,
svc: web::Data<EnhanceService>,
) -> impl Responder {
let (id, n) = path.into_inner();
let job = match svc.get_job(&id) {
Ok(j) => j,
Err(e) => return error_response(&e),
};
match svc.preview(&job, n).await {
Ok((content_type, bytes)) => HttpResponse::Ok()
.content_type(content_type)
.insert_header(("Cache-Control", "private, max-age=3600"))
.body(bytes),
Err(e) => error_response(&e),
}
}
#[post("/enhance/jobs/{id}/discard")]
pub async fn discard(
_: Claims,
id: web::Path<String>,
svc: web::Data<EnhanceService>,
) -> impl Responder {
match svc.discard(&id) {
Ok(job) => HttpResponse::Ok().json(job_view(&job)),
Err(e) => error_response(&e),
}
}
#[derive(Debug, Deserialize)]
pub struct ApproveBody {
#[serde(default)]
pub n: usize,
}
/// Ask ComfyUI to write result `n` into the library, ingest that file now,
/// and answer with the version (an `/image/versions` `versions[]` entry).
#[post("/enhance/jobs/{id}/approve")]
#[tracing::instrument(name = "enhance_approve", skip_all)]
pub async fn approve(
_: Claims,
id: web::Path<String>,
body: Option<web::Json<ApproveBody>>,
svc: web::Data<EnhanceService>,
app_state: web::Data<AppState>,
exif_dao: ExifData,
) -> impl Responder {
let n = body.map(|b| b.n).unwrap_or(0);
let job = match svc.get_job(&id) {
Ok(j) => j,
Err(e) => return error_response(&e),
};
let libs = app_state.libraries.clone();
let prefer = job.library_id;
let place = move |abs: &Path| place_in_libraries(&libs, prefer, abs);
let (job, saved) = match svc.approve(&id, n, &place).await {
Ok(v) => v,
Err(e) => return error_response(&e),
};
let (Some(library_id), Some(rel)) = (saved.library_id, saved.rel_path.clone()) else {
return error_response(&EnhanceError::Unplaceable(
"the saved file isn't in a library".into(),
));
};
let Some(library) = app_state.library_by_id(library_id).cloned() else {
return error_response(&EnhanceError::Unplaceable("unknown library".into()));
};
let abs = PathBuf::from(&library.root_path).join(&rel);
// The file was written by another machine over a share; give a lagging
// mount a moment to show it.
for _ in 0..25 {
if abs.is_file() {
break;
}
tokio::time::sleep(Duration::from_millis(200)).await;
}
if !abs.is_file() {
return error_response(&EnhanceError::Unplaceable(format!(
"ComfyUI saved {rel} but it isn't visible to ImageApi yet; it will appear after the next scan"
)));
}
let version_dao = app_state.version_dao.clone();
let exif_arc: Arc<Mutex<Box<dyn ExifDao>>> = exif_dao.clone().into_inner();
let (lib2, rel2, abs2) = (library.clone(), rel.clone(), abs.clone());
let ingested =
web::block(move || ingest_saved(&lib2, &rel2, &abs2, &exif_arc, &version_dao)).await;
let hash = match ingested {
Ok(Some(h)) => h,
_ => {
warn!("enhance: ingest of {rel} failed");
return HttpResponse::Accepted().json(json!({
"linked": false,
"path": rel,
"library_id": library_id,
"job": job_view(&job),
}));
}
};
// Which version row did that become?
let entry = {
let mut dao = app_state
.version_dao
.lock()
.expect("Unable to lock VersionDao");
let original = match dao.get_link(&hash) {
Ok(Some(link)) => link.original_hash,
_ => job.original_hash.clone(),
};
build_response(dao.as_mut(), &original, job.library_id)
.ok()
.and_then(|r| r.versions.into_iter().find(|v| v.hash == hash))
};
match entry {
Some(v) => HttpResponse::Ok().json(v),
None => HttpResponse::Accepted().json(json!({
"linked": false,
"path": rel,
"library_id": library_id,
"job": job_view(&job),
})),
}
}
/// Register the saved file (if the watcher hasn't yet) and link it as a
/// version. Idempotent. Returns the file's content hash.
pub fn ingest_saved(
library: &Library,
rel: &str,
abs: &Path,
exif_dao: &Arc<Mutex<Box<dyn ExifDao>>>,
version_dao: &Arc<Mutex<Box<dyn crate::database::VersionDao>>>,
) -> Option<String> {
let rel_s = rel.to_string();
let existing = exif_dao
.lock()
.expect("Unable to lock ExifDao")
.get_exif_batch(Some(library.id), std::slice::from_ref(&rel_s))
.ok()?;
let file = vec![(abs.to_path_buf(), rel_s.clone())];
if existing.is_empty() {
crate::watcher::register_new_files(library, file.clone(), exif_dao);
}
let linked = crate::versions::link_new_files(library, &file, exif_dao, version_dao)
+ crate::versions::sweep_name_links(library, exif_dao, version_dao);
if linked > 0 {
crate::duplicates::invalidate_caches();
}
exif_dao
.lock()
.expect("Unable to lock ExifDao")
.get_exif_batch(Some(library.id), std::slice::from_ref(&rel_s))
.ok()?
.into_iter()
.find_map(|r| r.content_hash)
}
+985
View File
@@ -0,0 +1,985 @@
//! End-to-end checks for Enhance (#192) against a mock ComfyUI: nothing here
//! touches a real ComfyUI, a real library or the prod DB.
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use actix_web::test as awtest;
use actix_web::web::Data;
use actix_web::{App, HttpResponse, HttpServer, web};
use diesel::prelude::*;
use diesel::sql_query;
use diesel::sql_types::Text;
use diesel::sqlite::SqliteConnection;
use serde_json::{Value, json};
use super::config::{EnhanceConfig, PathMapping, PromptMode, TemplateConfig};
use super::service::EnhanceService;
use super::template::Template;
use crate::data::Claims;
use crate::database::enhance_dao::{EnhanceJobDao, SqliteEnhanceJobDao};
use crate::database::test::in_memory_db_connection;
use crate::database::{ExifDao, SqliteExifDao, SqliteVersionDao};
use crate::state::AppState;
use crate::versions::xmp::test_png::{producer_xmp, write_png};
const RESTORE: &str = include_str!("../../tests/fixtures/enhance/restore_turbo.json");
const EDIT: &str = include_str!("../../tests/fixtures/enhance/edit_turbo.json");
// ---- mock ComfyUI ---------------------------------------------------------
#[derive(Default)]
struct Mock {
submitted: Mutex<Vec<Value>>,
queue_running: Mutex<Vec<String>>,
queue_pending: Mutex<Vec<String>>,
/// prompt_id -> history entry
history: Mutex<HashMap<String, Value>>,
newest: Mutex<Option<Value>>,
approve_calls: Mutex<Vec<String>>,
approve_path: Mutex<String>,
approve_gone: Mutex<bool>,
free_calls: AtomicUsize,
reject_prompts: Mutex<Option<Value>>,
view_404: Mutex<bool>,
}
async fn start_mock() -> (Arc<Mock>, String) {
let mock = Arc::new(Mock::default());
let m = mock.clone();
let server = HttpServer::new(move || {
App::new()
.app_data(Data::from(m.clone()))
.route("/system_stats", web::get().to(|| async { HttpResponse::Ok().json(json!({})) }))
.route(
"/prompt",
web::post().to(|m: Data<Mock>, body: web::Json<Value>| async move {
if let Some(err) = m.reject_prompts.lock().unwrap().clone() {
return HttpResponse::BadRequest().json(err);
}
let mut s = m.submitted.lock().unwrap();
s.push(body.into_inner());
let id = format!("prompt-{}", s.len());
m.queue_pending.lock().unwrap().push(id.clone());
HttpResponse::Ok().json(json!({"prompt_id": id, "number": s.len(), "node_errors": {}}))
}),
)
.route(
"/queue",
web::get().to(|m: Data<Mock>| async move {
let item = |id: &String| json!([1, id, {}, {}, []]);
HttpResponse::Ok().json(json!({
"queue_running": m.queue_running.lock().unwrap().iter().map(item).collect::<Vec<_>>(),
"queue_pending": m.queue_pending.lock().unwrap().iter().map(item).collect::<Vec<_>>(),
}))
}),
)
.route(
"/history",
web::get().to(|m: Data<Mock>| async move {
match m.newest.lock().unwrap().clone() {
Some(e) => HttpResponse::Ok().json(json!({"newest": e})),
None => HttpResponse::Ok().json(json!({})),
}
}),
)
.route(
"/history/{id}",
web::get().to(|m: Data<Mock>, id: web::Path<String>| async move {
match m.history.lock().unwrap().get(id.as_str()) {
Some(e) => HttpResponse::Ok().json(json!({ id.as_str(): e })),
None => HttpResponse::Ok().json(json!({})),
}
}),
)
.route(
"/view",
web::get().to(|m: Data<Mock>, q: web::Query<HashMap<String, String>>| async move {
if *m.view_404.lock().unwrap() {
return HttpResponse::NotFound().finish();
}
assert_eq!(q.get("type").map(String::as_str), Some("temp"));
HttpResponse::Ok()
.content_type("image/png")
.body(format!("PNG:{}", q["filename"]))
}),
)
.route(
"/imageapi/approve",
web::post().to(|m: Data<Mock>, body: web::Json<Value>| async move {
if *m.approve_gone.lock().unwrap() {
return HttpResponse::NotFound().json(json!({"error": "no longer pending"}));
}
m.approve_calls
.lock()
.unwrap()
.push(body["id"].as_str().unwrap().to_string());
HttpResponse::Ok().json(json!({"path": m.approve_path.lock().unwrap().clone()}))
}),
)
.route(
"/free",
web::post().to(|m: Data<Mock>, body: web::Json<Value>| async move {
assert_eq!(body["unload_models"], true);
m.free_calls.fetch_add(1, Ordering::SeqCst);
HttpResponse::Ok().finish()
}),
)
})
.workers(1)
.bind(("127.0.0.1", 0))
.unwrap();
let addr = server.addrs()[0];
actix_rt::spawn(server.run());
(mock, format!("http://{addr}"))
}
fn finished_entry(pending_id: &str, ts: i64) -> Value {
json!({
"status": {"status_str": "success", "completed": true,
"messages": [["execution_start", {"timestamp": ts - 3000}], ["execution_success", {"timestamp": ts}]]},
"outputs": {"36": {
"images": [{"filename": format!("imageapi_pending_{pending_id}.png"), "subfolder": "", "type": "temp"}],
"imageapi_versions": [{"id": pending_id, "saved_path": null}],
}},
})
}
// ---- fixture --------------------------------------------------------------
struct Fx {
mock: Arc<Mock>,
svc: Data<EnhanceService>,
state: Data<AppState>,
exif: Data<Mutex<Box<dyn ExifDao>>>,
conn: Arc<Mutex<SqliteConnection>>,
root: PathBuf,
}
fn template(name: &str, mode: PromptMode, json: &str) -> Template {
Template::from_value(
TemplateConfig {
name: name.into(),
label: name.to_uppercase(),
description: format!("the {name} template"),
file: "x.json".into(),
prompt_mode: mode,
prompt_join: " ".into(),
output_megapixels: 1.0,
},
serde_json::from_str(json).unwrap(),
)
.unwrap()
}
async fn fx_with(url: Option<String>, tweak: impl FnOnce(&mut EnhanceConfig)) -> Fx {
let (mock, mock_url) = start_mock().await;
let mut state = AppState::test_state();
let conn = Arc::new(Mutex::new(in_memory_db_connection()));
state.version_dao = Arc::new(Mutex::new(Box::new(SqliteVersionDao::from_shared(
conn.clone(),
))));
let library = state.primary_library().clone();
let root = PathBuf::from(&library.root_path);
sql_query("UPDATE libraries SET root_path = ? WHERE id = 1")
.bind::<Text, _>(&library.root_path)
.execute(&mut *conn.lock().unwrap())
.unwrap();
let mut cfg = EnhanceConfig {
url: Some(url.unwrap_or(mock_url)),
path_map: vec![PathMapping {
from: format!("{}/", library.root_path),
to: "Z:\\".into(),
}],
..EnhanceConfig::default()
};
tweak(&mut cfg);
let dao: Arc<Mutex<Box<dyn EnhanceJobDao>>> = Arc::new(Mutex::new(Box::new(
SqliteEnhanceJobDao::from_shared(conn.clone()),
)));
let svc = EnhanceService::new(
cfg,
vec![
template("restore", PromptMode::Append, RESTORE),
template("edit", PromptMode::Replace, EDIT),
],
dao,
)
.unwrap();
Fx {
mock,
svc: Data::new(svc),
exif: Data::new(Mutex::new(
Box::new(SqliteExifDao::from_shared(conn.clone())) as Box<dyn ExifDao>,
)),
state: Data::new(state),
conn,
root,
}
}
async fn fx() -> Fx {
fx_with(None, |_| {}).await
}
impl Fx {
fn put(&self, rel: &str, bytes: &[u8]) -> String {
let abs = self.root.join(rel);
std::fs::create_dir_all(abs.parent().unwrap()).unwrap();
std::fs::write(&abs, bytes).unwrap();
self.register(rel)
}
fn register(&self, rel: &str) -> String {
let hash = crate::content_hash::compute(&self.root.join(rel))
.unwrap()
.content_hash;
sql_query(
"INSERT INTO image_exif (library_id, rel_path, created_time, last_modified, content_hash) \
VALUES (1, ?, 0, 0, ?)",
)
.bind::<Text, _>(rel)
.bind::<Text, _>(&hash)
.execute(&mut *self.conn.lock().unwrap())
.unwrap();
hash
}
fn token() -> String {
jsonwebtoken::encode(
&jsonwebtoken::Header::default(),
&Claims::valid_user("1".to_string()),
&jsonwebtoken::EncodingKey::from_secret(b"test_key"),
)
.unwrap()
}
async fn call(&self, method: &str, uri: &str, body: Option<Value>) -> (u16, Value, Vec<u8>) {
let app = awtest::init_service(
App::new()
.service(super::get_templates)
.service(super::submit)
.service(super::list_jobs)
.service(super::get_job)
.service(super::preview)
.service(super::approve)
.service(super::discard)
.service(crate::versions::get_versions)
.app_data(self.svc.clone())
.app_data(self.state.clone())
.app_data(self.exif.clone()),
)
.await;
let mut req = match method {
"POST" => awtest::TestRequest::post(),
_ => awtest::TestRequest::get(),
}
.uri(uri)
.insert_header(("Authorization", format!("Bearer {}", Fx::token())));
if let Some(b) = body {
req = req.set_json(b);
}
let resp = awtest::call_service(&app, req.to_request()).await;
let status = resp.status().as_u16();
let bytes = awtest::read_body(resp).await.to_vec();
let json = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
(status, json, bytes)
}
}
async fn call(fx: &Fx, method: &str, uri: &str, body: Option<Value>) -> (u16, Value, Vec<u8>) {
fx.call(method, uri, body).await
}
const PHOTO: &str = "2009_200/019_19.JPG";
// ---- tests ----------------------------------------------------------------
#[actix_rt::test]
async fn templates_report_availability_and_shape() {
let fx = fx().await;
let (st, body, _) = call(&fx, "GET", "/enhance/templates", None).await;
assert_eq!(st, 200);
assert_eq!(body["available"], true);
assert!(body.get("reason").is_none());
assert_eq!(body["default"], "restore");
let t = &body["templates"];
assert_eq!(t[0]["name"], "restore");
assert_eq!(t[0]["label"], "RESTORE");
assert_eq!(t[0]["accepts_prompt"], true);
assert_eq!(t[0]["prompt_mode"], "append");
assert_eq!(t[0]["prompt_required"], false);
assert_eq!(t[0]["accepts_reference"], false);
assert_eq!(t[0]["default"], true);
assert_eq!(t[0]["output_megapixels"], 1.0);
assert_eq!(t[0]["warn_above_megapixels"], 1.5);
assert_eq!(t[1]["name"], "edit");
assert_eq!(t[1]["prompt_mode"], "replace");
assert_eq!(t[1]["prompt_required"], true);
assert_eq!(t[1]["accepts_reference"], true);
assert_eq!(t[1]["default"], false);
}
#[actix_rt::test]
async fn comfy_down_means_unavailable_and_a_clear_submit_error() {
// Nothing listens on port 9 (discard) of loopback.
let fx = fx_with(Some("http://127.0.0.1:9".into()), |_| {}).await;
fx.put(PHOTO, b"photo");
let (st, body, _) = call(&fx, "GET", "/enhance/templates", None).await;
assert_eq!(st, 200);
assert_eq!(body["available"], false);
assert_eq!(body["reason"], "ComfyUI isn't running");
assert_eq!(body["templates"].as_array().unwrap().len(), 2);
let (st, body, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
assert_eq!(st, 503);
assert_eq!(body["error"], "ComfyUI isn't running");
assert_eq!(body["code"], "unavailable");
}
#[actix_rt::test]
async fn unconfigured_enhance_is_off_not_an_error() {
let fx = fx_with(None, |c| c.url = None).await;
let (st, body, _) = call(&fx, "GET", "/enhance/templates", None).await;
assert_eq!((st, &body["available"]), (200, &json!(false)));
let (st, _, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
assert_eq!(st, 503);
}
#[actix_rt::test]
async fn submit_maps_the_path_and_uses_a_fresh_seed_each_time() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
let body = json!({"path": PHOTO, "template": "restore", "prompt": " warmer colours "});
let (st, job1, _) = call(&fx, "POST", "/enhance", Some(body.clone())).await;
assert_eq!(st, 200, "{job1}");
let (_, job2, _) = call(&fx, "POST", "/enhance", Some(body)).await;
assert_eq!(job1["status"], "queued");
assert_eq!(job1["path"], PHOTO);
assert_eq!(job1["template"], "restore");
assert_eq!(job1["prompt"], "warmer colours");
assert_eq!(job1["pending"], json!([]));
assert!(job1["created_at"].as_i64().unwrap() > 0);
assert_ne!(job1["id"], job2["id"]);
let sent = fx.mock.submitted.lock().unwrap();
assert_eq!(sent.len(), 2);
let p1 = &sent[0]["prompt"];
let p2 = &sent[1]["prompt"];
assert_eq!(p1["37"]["inputs"]["path"], "Z:\\2009_200\\019_19.JPG");
assert_eq!(p1["36"]["inputs"]["require_approval"], true);
assert!(
p1["13"]["inputs"]["value"]
.as_str()
.unwrap()
.ends_with("overlays. warmer colours")
);
assert_ne!(p1["5"]["inputs"]["seed"], p2["5"]["inputs"]["seed"]);
assert_ne!(
p1["14"]["inputs"]["sampling_mode.seed"],
p2["14"]["inputs"]["sampling_mode.seed"]
);
assert_eq!(sent[0]["client_id"], "imageapi");
}
#[actix_rt::test]
async fn submit_validation() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
// Edit needs an instruction.
let (st, body, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "edit"})),
)
.await;
assert_eq!(st, 400, "{body}");
assert!(body["error"].as_str().unwrap().contains("instruction"));
let (st, _, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "edit", "prompt": " "})),
)
.await;
assert_eq!(st, 400);
// Unknown template / photo / traversal / reference on a template without one.
let (st, _, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "nope"})),
)
.await;
assert_eq!(st, 404);
let (st, _, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": "missing.jpg", "template": "restore"})),
)
.await;
assert_eq!(st, 404);
let (st, _, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": "../etc/passwd", "template": "restore"})),
)
.await;
assert_eq!(st, 400);
fx.put("ref.jpg", b"other");
let (st, _, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore", "reference_path": "ref.jpg"})),
)
.await;
assert_eq!(st, 400);
assert!(
fx.mock.submitted.lock().unwrap().is_empty(),
"nothing reached ComfyUI"
);
}
#[actix_rt::test]
async fn edit_without_reference_drops_the_node_and_with_one_keeps_it() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
fx.put("ref/other.jpg", b"other view");
let (st, _, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "edit", "prompt": "make it b&w"})),
)
.await;
assert_eq!(st, 200);
let (st, job, _) = call(
&fx, "POST", "/enhance",
Some(json!({"path": PHOTO, "template": "edit", "prompt": "x", "reference_path": "ref/other.jpg"})),
)
.await;
assert_eq!(st, 200);
assert_eq!(job["reference_path"], "ref/other.jpg");
let sent = fx.mock.submitted.lock().unwrap();
let without = &sent[0]["prompt"];
assert!(without.get("38").is_none());
assert!(!without.to_string().contains("[\"38\""));
assert_eq!(without["13"]["inputs"]["value"], "make it b&w");
let with = &sent[1]["prompt"];
assert_eq!(with["38"]["inputs"]["path"], "Z:\\ref\\other.jpg");
}
#[actix_rt::test]
async fn comfy_rejection_is_reported_and_a_missing_reference_is_named() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
*fx.mock.reject_prompts.lock().unwrap() = Some(json!({
"error": {"type": "prompt_outputs_failed_validation", "message": "Prompt outputs failed validation", "details": ""},
"node_errors": {"36": {"class_type": "ImageApiSaveImageVersion", "errors": [{"message": "Required input is missing"}]}},
}));
let (st, body, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "edit", "prompt": "x"})),
)
.await;
assert_eq!(st, 422);
assert!(
body["error"]
.as_str()
.unwrap()
.contains("needs a reference photo"),
"{body}"
);
// Restore never dropped a reference, so the message is ComfyUI's own.
let (st, body, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
assert_eq!(st, 422);
assert!(!body["error"].as_str().unwrap().contains("reference"));
assert!(
body["error"]
.as_str()
.unwrap()
.contains("Required input is missing")
);
}
#[actix_rt::test]
async fn job_moves_queued_running_ready_and_the_preview_is_relayed() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
let (_, job, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
let id = job["id"].as_str().unwrap().to_string();
// Behind someone else's run: queued with a position.
fx.mock.queue_running.lock().unwrap().push("other".into());
fx.svc.poll_once().await;
let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await;
assert_eq!(
(j["status"].as_str().unwrap(), j["queue_position"].as_i64()),
("queued", Some(1))
);
// Now it runs.
*fx.mock.queue_running.lock().unwrap() = vec!["prompt-1".into()];
fx.mock.queue_pending.lock().unwrap().clear();
fx.svc.poll_once().await;
let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await;
assert_eq!(j["status"], "running");
// And finishes.
fx.mock.queue_running.lock().unwrap().clear();
fx.mock
.history
.lock()
.unwrap()
.insert("prompt-1".into(), finished_entry("abc", 10_000));
fx.svc.poll_once().await;
let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await;
assert_eq!(j["status"], "ready");
assert!(j.get("queue_position").is_none());
assert_eq!(j["pending"].as_array().unwrap().len(), 1);
assert_eq!(j["pending"][0]["n"], 0);
assert_eq!(j["pending"][0]["approved"], false);
// The staged file name is ComfyUI-internal and not exposed.
assert!(!j.to_string().contains("imageapi_pending"));
let uri = j["pending"][0]["preview"].as_str().unwrap().to_string();
let (st, _, bytes) = call(&fx, "GET", &uri, None).await;
assert_eq!(st, 200);
assert_eq!(bytes, b"PNG:imageapi_pending_abc.png");
let (st, _, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}/preview/5"), None).await;
assert_eq!(st, 404);
// The photo's job list shows it (and a failed one too).
let (_, list, _) = call(&fx, "GET", &format!("/enhance/jobs?path={PHOTO}"), None).await;
assert_eq!(list["jobs"].as_array().unwrap().len(), 1);
assert_eq!(list["jobs"][0]["id"], id.as_str());
}
#[actix_rt::test]
async fn failed_run_reports_comfys_message() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
let (_, job, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
fx.mock.queue_pending.lock().unwrap().clear();
fx.mock.history.lock().unwrap().insert(
"prompt-1".into(),
json!({"status": {"status_str": "error", "completed": false, "messages": [
["execution_error", {"exception_message": "CUDA out of memory\n", "timestamp": 5}]]}, "outputs": {}}),
);
fx.svc.poll_once().await;
let (_, j, _) = call(
&fx,
"GET",
&format!("/enhance/jobs/{}", job["id"].as_str().unwrap()),
None,
)
.await;
assert_eq!(j["status"], "failed");
assert_eq!(j["error"], "CUDA out of memory");
}
#[actix_rt::test]
async fn a_job_comfy_forgot_expires() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
let (_, job, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
// ComfyUI restarted: empty queue, empty history.
fx.mock.queue_pending.lock().unwrap().clear();
fx.svc.poll_once().await;
let (_, j, _) = call(
&fx,
"GET",
&format!("/enhance/jobs/{}", job["id"].as_str().unwrap()),
None,
)
.await;
assert_eq!(j["status"], "expired");
assert!(j["error"].as_str().unwrap().contains("restarted"));
assert_eq!(j["pending"], json!([]));
}
#[actix_rt::test]
async fn ready_job_expires_when_the_preview_is_gone_and_approve_says_so() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
let (_, job, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
let id = job["id"].as_str().unwrap().to_string();
fx.mock.queue_pending.lock().unwrap().clear();
fx.mock
.history
.lock()
.unwrap()
.insert("prompt-1".into(), finished_entry("abc", 1));
fx.svc.poll_once().await;
*fx.mock.view_404.lock().unwrap() = true;
let (st, body, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}/preview/0"), None).await;
assert_eq!((st, body["code"].as_str()), (410, Some("expired")));
let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await;
assert_eq!(j["status"], "expired");
let (st, _, _) = call(
&fx,
"POST",
&format!("/enhance/jobs/{id}/approve"),
Some(json!({"n": 0})),
)
.await;
assert_eq!(st, 410);
assert!(fx.mock.approve_calls.lock().unwrap().is_empty());
}
#[actix_rt::test]
async fn approve_gone_on_comfy_side_expires_the_job() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
let (_, job, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
let id = job["id"].as_str().unwrap().to_string();
fx.mock.queue_pending.lock().unwrap().clear();
fx.mock
.history
.lock()
.unwrap()
.insert("prompt-1".into(), finished_entry("abc", 1));
fx.svc.poll_once().await;
*fx.mock.approve_gone.lock().unwrap() = true;
let (st, _, _) = call(
&fx,
"POST",
&format!("/enhance/jobs/{id}/approve"),
Some(json!({"n": 0})),
)
.await;
assert_eq!(st, 410);
let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await;
assert_eq!(j["status"], "expired");
}
#[actix_rt::test]
async fn approve_ingests_links_and_is_idempotent() {
let fx = fx().await;
let original = fx.put(PHOTO, b"the original photo");
let (_, job, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
let id = job["id"].as_str().unwrap().to_string();
fx.mock.queue_pending.lock().unwrap().clear();
fx.mock
.history
.lock()
.unwrap()
.insert("prompt-1".into(), finished_entry("abc", 1));
fx.svc.poll_once().await;
// What ComfyUI's approve would have written: a PNG with the XMP link.
let saved_rel = "2009_200/aiedits/019_19.JPG.restored.png";
let xmp = producer_xmp(&original, "restored", &[]);
write_png(&fx.root.join(saved_rel), Some(&xmp), None, 7);
*fx.mock.approve_path.lock().unwrap() = "Z:\\2009_200\\aiedits\\019_19.JPG.restored.png".into();
let (st, v, _) = call(
&fx,
"POST",
&format!("/enhance/jobs/{id}/approve"),
Some(json!({"n": 0})),
)
.await;
assert_eq!(st, 200, "{v}");
assert_eq!(v["path"], saved_rel);
assert_eq!(v["library_id"], 1);
assert_eq!(v["kind"], "restored");
assert_eq!(v["source"], "xmp");
for k in ["hash", "created_at", "references", "width", "height"] {
assert!(v.get(k).is_some(), "missing {k}: {v}");
}
assert_eq!(fx.mock.approve_calls.lock().unwrap().as_slice(), ["abc"]);
// Visible in /image/versions right away.
let (_, vs, _) = call(&fx, "GET", &format!("/image/versions?path={PHOTO}"), None).await;
assert_eq!(vs["versions"].as_array().unwrap().len(), 1);
assert_eq!(vs["versions"][0]["path"], saved_rel);
// Approving again: same version, ComfyUI not asked, no duplicate row.
let (st, v2, _) = call(
&fx,
"POST",
&format!("/enhance/jobs/{id}/approve"),
Some(json!({"n": 0})),
)
.await;
assert_eq!(st, 200);
assert_eq!(v2["hash"], v["hash"]);
assert_eq!(fx.mock.approve_calls.lock().unwrap().len(), 1);
let (_, vs, _) = call(&fx, "GET", &format!("/image/versions?path={PHOTO}"), None).await;
assert_eq!(vs["versions"].as_array().unwrap().len(), 1);
let rows: i64 = {
#[derive(QueryableByName)]
struct N {
#[diesel(sql_type = diesel::sql_types::BigInt)]
n: i64,
}
sql_query("SELECT COUNT(*) AS n FROM image_exif WHERE rel_path = ?")
.bind::<Text, _>(saved_rel)
.get_result::<N>(&mut *fx.conn.lock().unwrap())
.unwrap()
.n
};
assert_eq!(rows, 1);
let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await;
assert_eq!(j["status"], "approved");
assert_eq!(j["pending"][0]["approved"], true);
assert_eq!(j["saved"][0]["path"], saved_rel);
}
#[actix_rt::test]
async fn approve_of_a_path_outside_every_library_is_reported_not_lost() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
let (_, job, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
let id = job["id"].as_str().unwrap().to_string();
fx.mock.queue_pending.lock().unwrap().clear();
fx.mock
.history
.lock()
.unwrap()
.insert("prompt-1".into(), finished_entry("abc", 1));
fx.svc.poll_once().await;
*fx.mock.approve_path.lock().unwrap() = "C:\\somewhere\\else.png".into();
let (st, body, _) = call(
&fx,
"POST",
&format!("/enhance/jobs/{id}/approve"),
Some(json!({"n": 0})),
)
.await;
assert_eq!(st, 502);
assert!(
body["error"]
.as_str()
.unwrap()
.contains("C:\\somewhere\\else.png")
);
}
#[actix_rt::test]
async fn approve_before_ready_is_a_400() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
let (_, job, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
let (st, _, _) = call(
&fx,
"POST",
&format!("/enhance/jobs/{}/approve", job["id"].as_str().unwrap()),
Some(json!({"n": 0})),
)
.await;
assert_eq!(st, 400);
let (st, _, _) = call(
&fx,
"POST",
"/enhance/jobs/nope/approve",
Some(json!({"n": 0})),
)
.await;
assert_eq!(st, 404);
}
#[actix_rt::test]
async fn discard_hides_the_job_from_the_photo_list_only() {
let fx = fx().await;
fx.put(PHOTO, b"photo");
let (_, a, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
let (_, b, _) = call(
&fx,
"POST",
"/enhance",
Some(json!({"path": PHOTO, "template": "restore"})),
)
.await;
let (st, _, _) = call(
&fx,
"POST",
&format!("/enhance/jobs/{}/discard", a["id"].as_str().unwrap()),
None,
)
.await;
assert_eq!(st, 200);
let (_, list, _) = call(&fx, "GET", &format!("/enhance/jobs?path={PHOTO}"), None).await;
let ids: Vec<&str> = list["jobs"]
.as_array()
.unwrap()
.iter()
.map(|j| j["id"].as_str().unwrap())
.collect();
assert_eq!(ids, [b["id"].as_str().unwrap()]);
// Still fetchable by id.
let (st, _, _) = call(
&fx,
"GET",
&format!("/enhance/jobs/{}", a["id"].as_str().unwrap()),
None,
)
.await;
assert_eq!(st, 200);
let (st, _, _) = call(&fx, "POST", "/enhance/jobs/nope/discard", None).await;
assert_eq!(st, 404);
}
#[actix_rt::test]
async fn unreachable_active_jobs_become_unavailable_after_the_grace_period() {
// Submit against the mock, then point nothing at it: use a service whose
// client targets a dead port but whose DB holds an active job.
let fx = fx_with(Some("http://127.0.0.1:9".into()), |c| {
c.unreachable_grace_secs = 1
})
.await;
let now = chrono::Utc::now().timestamp();
let dao_job = crate::database::enhance_dao::EnhanceJobRow {
id: "j1".into(),
original_hash: "h".into(),
library_id: 1,
rel_path: "a.jpg".into(),
template: "restore".into(),
prompt: None,
reference_hash: None,
reference_path: None,
comfy_prompt_id: "p".into(),
status: "running".into(),
error: None,
queue_position: None,
pending: "[]".into(),
saved: "[]".into(),
discarded: 0,
created_at: now,
updated_at: now,
};
SqliteEnhanceJobDao::from_shared(fx.conn.clone())
.insert(&dao_job)
.unwrap();
fx.svc.poll_once().await;
assert_eq!(
fx.svc.get_job("j1").unwrap().status,
"running",
"within grace"
);
tokio::time::sleep(std::time::Duration::from_millis(1100)).await;
fx.svc.poll_once().await;
let j = fx.svc.get_job("j1").unwrap();
assert_eq!(j.status, "unavailable");
assert_eq!(j.error.as_deref(), Some("ComfyUI isn't running"));
}
#[actix_rt::test]
async fn idle_comfy_is_freed_once_and_activity_rearms_it() {
let fx = fx_with(None, |c| c.idle_free_secs = 1).await;
*fx.mock.newest.lock().unwrap() = Some(finished_entry("x", 100));
fx.svc.idle_check().await; // baseline
assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 0);
tokio::time::sleep(std::time::Duration::from_millis(1100)).await;
fx.svc.idle_check().await;
fx.svc.idle_check().await;
assert_eq!(
fx.mock.free_calls.load(Ordering::SeqCst),
1,
"freed exactly once"
);
// Someone runs a workflow in ComfyUI's own page: newest history changes.
*fx.mock.newest.lock().unwrap() = Some(finished_entry("y", 200));
fx.svc.idle_check().await;
assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 1);
tokio::time::sleep(std::time::Duration::from_millis(1100)).await;
fx.svc.idle_check().await;
assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 2);
}
#[actix_rt::test]
async fn busy_queue_blocks_the_free() {
let fx = fx_with(None, |c| c.idle_free_secs = 1).await;
fx.mock.queue_running.lock().unwrap().push("someone".into());
fx.svc.idle_check().await;
tokio::time::sleep(std::time::Duration::from_millis(1100)).await;
fx.svc.idle_check().await;
assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 0);
}
+25
View File
@@ -0,0 +1,25 @@
//! Enhance on demand (#192): run a ComfyUI restore/edit workflow on a photo
//! from the apps, approve the result in-app, and have it land in the library
//! as a version (#191).
//!
//! - [`config`]: `enhance.toml` (templates, path map, timings).
//! - [`template`]: workflow loading/validation and per-photo filling.
//! - [`comfy`]: the ComfyUI HTTP calls used.
//! - [`service`]: submit, poll, approve, idle `/free`.
//! - [`handlers`]: the HTTP endpoints.
//!
//! ImageApi never writes under a library root: ComfyUI writes the approved
//! file, ImageApi only ingests it.
pub mod comfy;
pub mod config;
pub mod handlers;
pub mod path_map;
pub mod service;
pub mod template;
#[cfg(test)]
mod integration_tests;
pub use handlers::{approve, discard, get_job, get_templates, list_jobs, preview, submit};
pub use service::EnhanceService;
+175
View File
@@ -0,0 +1,175 @@
//! Library path <-> ComfyUI path mapping (#192).
//!
//! ImageApi sees `/media/nas/Media/Pictures/a.jpg`; ComfyUI (on Windows)
//! sees `Z:\Media\Pictures\a.jpg` or `\\nas\Media\Pictures\a.jpg`. Each
//! mapping is a prefix pair; the longest matching prefix wins and the
//! separators are converted. Used in both directions: library -> ComfyUI on
//! submit, and ComfyUI -> library to recognise the path `approve` returns.
//! With no mappings configured, paths pass through unchanged.
use std::path::{Path, PathBuf};
use anyhow::{Result, bail};
use super::config::PathMapping;
#[derive(Debug, Clone)]
struct Entry {
/// Library side, `/`-separated, always ending in `/`.
from: String,
/// ComfyUI side without a trailing separator.
to: String,
/// Separator ComfyUI's side uses.
to_sep: char,
}
#[derive(Debug, Clone, Default)]
pub struct PathMap {
entries: Vec<Entry>,
}
impl PathMap {
pub fn new(mappings: &[PathMapping]) -> Result<Self> {
let mut entries = Vec::new();
for m in mappings {
let from = m.from.replace('\\', "/");
let to = m.to.trim();
if from.trim_matches('/').is_empty() || to.is_empty() {
bail!("path_map entries need non-empty from and to ({m:?})");
}
let to_sep = if to.contains('\\') { '\\' } else { '/' };
entries.push(Entry {
from: format!("{}/", from.trim_end_matches('/')),
to: to.trim_end_matches(['/', '\\']).to_string(),
to_sep,
});
}
Ok(PathMap { entries })
}
pub fn is_identity(&self) -> bool {
self.entries.is_empty()
}
/// The path ComfyUI should open for a library file. `None` when mappings
/// exist but none covers the file.
pub fn to_comfy(&self, abs: &Path) -> Option<String> {
let text = abs.to_str()?.replace('\\', "/");
if self.entries.is_empty() {
return Some(abs.to_string_lossy().into_owned());
}
let best = self
.entries
.iter()
.filter(|e| text.starts_with(e.from.as_str()))
.max_by_key(|e| e.from.len())?;
let rest = &text[best.from.len()..];
let sep = best.to_sep.to_string();
Some(format!("{}{}{}", best.to, sep, rest.replace('/', &sep)))
}
/// The library-side absolute path of a path ComfyUI reported. Matching
/// is case-insensitive (Windows). `None` when no mapping covers it or it
/// tries to climb out with `..`.
pub fn from_comfy(&self, comfy: &str) -> Option<PathBuf> {
let text = comfy.trim().replace('\\', "/");
if self.entries.is_empty() {
return Some(PathBuf::from(text));
}
let best = self
.entries
.iter()
.filter_map(|e| {
let prefix = format!("{}/", e.to.replace('\\', "/"));
let head = text.get(..prefix.len())?;
head.eq_ignore_ascii_case(&prefix)
.then_some((e, prefix.len()))
})
.max_by_key(|(_, len)| *len)?;
let rest = &text[best.1..];
if rest.split('/').any(|c| c == "..") {
return None;
}
Some(PathBuf::from(format!("{}{}", best.0.from, rest)))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn map(pairs: &[(&str, &str)]) -> PathMap {
let m: Vec<PathMapping> = pairs
.iter()
.map(|(f, t)| PathMapping {
from: f.to_string(),
to: t.to_string(),
})
.collect();
PathMap::new(&m).unwrap()
}
#[test]
fn drive_letter_round_trip() {
let m = map(&[("/media/nas/", "Z:\\")]);
let c = m
.to_comfy(Path::new("/media/nas/Media/Pictures/a b.jpg"))
.unwrap();
assert_eq!(c, "Z:\\Media\\Pictures\\a b.jpg");
assert_eq!(
m.from_comfy(&c).unwrap(),
PathBuf::from("/media/nas/Media/Pictures/a b.jpg")
);
// Drive letters are case-insensitive.
assert_eq!(
m.from_comfy("z:\\x\\y.png").unwrap(),
PathBuf::from("/media/nas/x/y.png")
);
}
#[test]
fn unc_target_and_missing_trailing_slash() {
let m = map(&[("/media/nas", "\\\\nas\\")]);
assert_eq!(
m.to_comfy(Path::new("/media/nas/Media/a.jpg")).unwrap(),
"\\\\nas\\Media\\a.jpg"
);
assert_eq!(
m.from_comfy("\\\\nas\\Media\\a.jpg").unwrap(),
PathBuf::from("/media/nas/Media/a.jpg")
);
}
#[test]
fn longest_prefix_wins_both_ways() {
let m = map(&[("/media/nas/", "Z:\\"), ("/media/nas/Media/", "Y:\\")]);
assert_eq!(
m.to_comfy(Path::new("/media/nas/Media/a.jpg")).unwrap(),
"Y:\\a.jpg"
);
assert_eq!(
m.to_comfy(Path::new("/media/nas/Other/a.jpg")).unwrap(),
"Z:\\Other\\a.jpg"
);
assert_eq!(
m.from_comfy("Y:\\a.jpg").unwrap(),
PathBuf::from("/media/nas/Media/a.jpg")
);
}
#[test]
fn uncovered_and_escaping_paths_are_none() {
let m = map(&[("/media/nas/", "Z:\\")]);
assert!(m.to_comfy(Path::new("/elsewhere/a.jpg")).is_none());
assert!(m.to_comfy(Path::new("/media/nasty/a.jpg")).is_none());
assert!(m.from_comfy("C:\\a.png").is_none());
assert!(m.from_comfy("Z:\\..\\etc\\passwd").is_none());
}
#[test]
fn no_mappings_is_identity() {
let m = PathMap::default();
assert_eq!(m.to_comfy(Path::new("/a/b.jpg")).unwrap(), "/a/b.jpg");
assert_eq!(m.from_comfy("/a/b.png").unwrap(), PathBuf::from("/a/b.png"));
}
}
+742
View File
@@ -0,0 +1,742 @@
//! Enhance orchestration (#192): submit, poll, approve, idle `/free`.
//!
//! The job row (`enhance_jobs`) is the only state that matters. ComfyUI is
//! allowed to be down (a normal state) and to forget everything when it
//! restarts (unapproved results are cleared), so every transition tolerates
//! both.
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use log::{info, warn};
use serde::{Deserialize, Serialize};
use super::comfy::{ApproveError, ComfyClient, ComfyError, HistoryEntry, Staged, SubmitError};
use super::config::{EnhanceConfig, Loaded, PromptMode};
use super::path_map::PathMap;
use super::template::Template;
use crate::database::enhance_dao::{
EnhanceJobDao, EnhanceJobRow, JobUpdate, STATUS_APPROVED, STATUS_EXPIRED, STATUS_FAILED,
STATUS_QUEUED, STATUS_READY, STATUS_RUNNING, STATUS_UNAVAILABLE,
};
/// Short timeout for the `available` probe, so the templates call stays snappy
/// when ComfyUI's host is off.
const PROBE_TIMEOUT: Duration = Duration::from_secs(2);
/// How often the idle check looks at ComfyUI.
const IDLE_CHECK_EVERY: Duration = Duration::from_secs(30);
/// Finished jobs are kept this long.
const JOB_RETENTION_SECS: i64 = 14 * 24 * 3600;
/// Jobs shown per photo.
pub const JOBS_PER_PHOTO: i64 = 20;
pub const MSG_UNAVAILABLE: &str = "ComfyUI isn't running";
pub const MSG_EXPIRED: &str = "This result expired (ComfyUI was restarted). Run it again.";
pub const MSG_DISABLED: &str = "Enhance isn't configured";
/// A library file a request names.
#[derive(Debug, Clone)]
pub struct SourceFile {
pub library_id: i32,
pub rel_path: String,
pub abs: PathBuf,
pub hash: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EnhanceError {
/// No `url` configured.
Disabled,
/// ComfyUI isn't reachable.
Unavailable,
UnknownTemplate(String),
BadRequest(String),
/// ComfyUI refused the workflow.
Rejected(String),
/// The result is gone (ComfyUI restarted).
Expired,
NotFound,
/// Approved, but ImageApi can't place the file in a library.
Unplaceable(String),
Upstream(String),
Internal(String),
}
impl EnhanceError {
pub fn message(&self) -> String {
match self {
EnhanceError::Disabled => MSG_DISABLED.into(),
EnhanceError::Unavailable => MSG_UNAVAILABLE.into(),
EnhanceError::UnknownTemplate(t) => format!("unknown template {t:?}"),
EnhanceError::BadRequest(m)
| EnhanceError::Rejected(m)
| EnhanceError::Unplaceable(m)
| EnhanceError::Upstream(m)
| EnhanceError::Internal(m) => m.clone(),
EnhanceError::Expired => MSG_EXPIRED.into(),
EnhanceError::NotFound => "not found".into(),
}
}
}
fn comfy_err(e: ComfyError) -> EnhanceError {
match e {
ComfyError::Unreachable(_) => EnhanceError::Unavailable,
other => EnhanceError::Upstream(other.to_string()),
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SavedEntry {
pub n: usize,
pub comfy_path: String,
pub library_id: Option<i32>,
pub rel_path: Option<String>,
}
/// What the idle check decided.
#[derive(Debug, Default)]
pub struct IdleState {
initialised: bool,
marker: Option<i64>,
seen_at_ms: i64,
freed: bool,
}
impl IdleState {
/// One observation of a reachable ComfyUI. `marker` identifies the
/// newest history entry (its last event time): when it changes, or the
/// queue isn't empty, or one of our jobs is active, ComfyUI is in use.
/// Returns true once per quiet spell when it has been idle for
/// `idle_ms` and `/free` should be called.
///
/// The marker's *value* is ComfyUI's clock; only changes in it matter,
/// and "how long ago" is measured on ImageApi's clock, so clock skew
/// between the hosts can't cause an early unload. After ImageApi
/// (re)starts the first observation counts as activity.
pub fn observe(
&mut self,
now_ms: i64,
queue_empty: bool,
active_jobs: usize,
marker: Option<i64>,
idle_ms: i64,
) -> bool {
if !self.initialised || marker != self.marker || !queue_empty || active_jobs > 0 {
self.initialised = true;
self.marker = marker;
self.seen_at_ms = now_ms;
self.freed = false;
return false;
}
if !self.freed && now_ms - self.seen_at_ms >= idle_ms {
self.freed = true;
return true;
}
false
}
/// ComfyUI went away: the next observation starts a fresh baseline.
pub fn reset(&mut self) {
*self = IdleState::default();
}
/// `/free` failed; try again on the next check.
pub fn free_failed(&mut self) {
self.freed = false;
}
}
pub struct EnhanceService {
pub config: EnhanceConfig,
templates: Vec<Template>,
client: Option<ComfyClient>,
pub path_map: PathMap,
dao: Arc<Mutex<Box<dyn EnhanceJobDao>>>,
probe: Mutex<Option<(Instant, bool)>>,
unreachable_since: Mutex<HashMap<String, Instant>>,
idle: Mutex<IdleState>,
/// Serialises approvals (rare; keeps "approve twice" simple).
approve_lock: tokio::sync::Mutex<()>,
}
impl EnhanceService {
/// Load the templates named in `loaded`, logging (not failing on) bad
/// ones: one broken template shouldn't take the rest down.
pub fn from_loaded(
loaded: Loaded,
dao: Arc<Mutex<Box<dyn EnhanceJobDao>>>,
) -> anyhow::Result<Self> {
let mut templates = Vec::new();
for cfg in &loaded.config.templates {
match Template::load(cfg, &loaded.base_dir) {
Ok(t) => {
for w in &t.warnings {
warn!("enhance: template '{}': {w}", cfg.name);
}
info!(
"enhance: template '{}' loaded (prompt: {}{}, reference: {})",
cfg.name,
if t.accepts_prompt() { "yes" } else { "no" },
if t.accepts_prompt() {
match cfg.prompt_mode {
PromptMode::Append => " append",
PromptMode::Replace => " replace",
}
} else {
""
},
if t.accepts_reference() { "yes" } else { "no" },
);
templates.push(t);
}
Err(e) => warn!("enhance: template '{}' rejected: {e:#}", cfg.name),
}
}
match &loaded.path {
Some(p) if loaded.config.url.is_some() => info!(
"enhance: config {} ({} of {} templates usable)",
p.display(),
templates.len(),
loaded.config.templates.len()
),
Some(p) => info!("enhance: config {} has no url; Enhance is off", p.display()),
None => info!("enhance: no enhance.toml; Enhance is off"),
}
Self::new(loaded.config, templates, dao)
}
pub fn new(
config: EnhanceConfig,
templates: Vec<Template>,
dao: Arc<Mutex<Box<dyn EnhanceJobDao>>>,
) -> anyhow::Result<Self> {
let client = config
.url
.as_deref()
.map(|u| ComfyClient::new(u, Duration::from_secs(config.timeout_secs)));
let path_map = PathMap::new(&config.path_map)?;
Ok(EnhanceService {
config,
templates,
client,
path_map,
dao,
probe: Mutex::new(None),
unreachable_since: Mutex::new(HashMap::new()),
idle: Mutex::new(IdleState::default()),
approve_lock: tokio::sync::Mutex::new(()),
})
}
pub fn is_configured(&self) -> bool {
self.client.is_some()
}
pub fn templates(&self) -> &[Template] {
&self.templates
}
pub fn template(&self, name: &str) -> Option<&Template> {
self.templates.iter().find(|t| t.cfg.name == name)
}
pub fn default_template(&self) -> Option<&str> {
self.config
.default_template
.as_deref()
.and_then(|w| self.template(w))
.or_else(|| self.templates.first())
.map(|t| t.cfg.name.as_str())
}
/// Cached `GET /system_stats` probe.
pub async fn available(&self) -> bool {
let Some(client) = &self.client else {
return false;
};
let ttl = Duration::from_secs(self.config.probe_ttl_secs);
if let Some((at, up)) = *self.probe.lock().unwrap()
&& at.elapsed() < ttl
{
return up;
}
let up = client.probe(PROBE_TIMEOUT).await;
*self.probe.lock().unwrap() = Some((Instant::now(), up));
up
}
/// Fresh probe (ignores and refreshes the cache): a submit shouldn't be
/// refused because ComfyUI came up since the last probe.
async fn probe_now(&self) -> bool {
let Some(client) = &self.client else {
return false;
};
let up = client.probe(PROBE_TIMEOUT).await;
*self.probe.lock().unwrap() = Some((Instant::now(), up));
up
}
// ---- jobs -----------------------------------------------------------
pub async fn submit(
&self,
photo: &SourceFile,
template_name: &str,
prompt: Option<&str>,
reference: Option<&SourceFile>,
) -> Result<EnhanceJobRow, EnhanceError> {
let Some(client) = &self.client else {
return Err(EnhanceError::Disabled);
};
let template = self
.template(template_name)
.ok_or_else(|| EnhanceError::UnknownTemplate(template_name.to_string()))?;
let unmapped = |p: &Path| {
EnhanceError::Unplaceable(format!(
"{} isn't under any path_map entry, so ComfyUI can't open it",
p.display()
))
};
let source = self
.path_map
.to_comfy(&photo.abs)
.ok_or_else(|| unmapped(&photo.abs))?;
let reference_path = match reference {
Some(r) => Some(
self.path_map
.to_comfy(&r.abs)
.ok_or_else(|| unmapped(&r.abs))?,
),
None => None,
};
let filled = template
.fill(&source, prompt, reference_path.as_deref(), &mut || {
rand::random::<u64>() >> 12
})
.map_err(|e| EnhanceError::BadRequest(e.0))?;
if !self.probe_now().await {
return Err(EnhanceError::Unavailable);
}
let prompt_id = match client.submit(&filled.prompt).await {
Ok(id) => id,
Err(SubmitError::Comfy(e)) => return Err(comfy_err(e)),
Err(SubmitError::Rejected {
message,
node_errors,
}) => {
// A reference-less run of a template that wants one fails
// validation on the dropped node's consumers.
return Err(if filled.dropped_reference && node_errors {
EnhanceError::Rejected(format!(
"this template needs a reference photo ({message})"
))
} else {
EnhanceError::Rejected(message)
});
}
};
let now = chrono::Utc::now().timestamp();
let user_prompt = prompt.map(str::trim).filter(|p| !p.is_empty());
let row = EnhanceJobRow {
id: uuid::Uuid::new_v4().to_string(),
original_hash: photo.hash.clone(),
library_id: photo.library_id,
rel_path: photo.rel_path.clone(),
template: template.cfg.name.clone(),
prompt: user_prompt.map(str::to_string),
reference_hash: reference.map(|r| r.hash.clone()),
reference_path: reference.map(|r| r.rel_path.clone()),
comfy_prompt_id: prompt_id,
status: STATUS_QUEUED.to_string(),
error: None,
queue_position: None,
pending: "[]".to_string(),
saved: "[]".to_string(),
discarded: 0,
created_at: now,
updated_at: now,
};
self.dao
.lock()
.unwrap()
.insert(&row)
.map_err(|e| EnhanceError::Internal(format!("saving the job failed: {e}")))?;
info!(
"enhance: job {} submitted ({} for {}, prompt {})",
row.id, row.template, row.rel_path, row.comfy_prompt_id
);
Ok(row)
}
pub fn get_job(&self, id: &str) -> Result<EnhanceJobRow, EnhanceError> {
self.dao
.lock()
.unwrap()
.get(id)
.map_err(|e| EnhanceError::Internal(e.to_string()))?
.ok_or(EnhanceError::NotFound)
}
pub fn jobs_for_hash(&self, hash: &str) -> Result<Vec<EnhanceJobRow>, EnhanceError> {
self.dao
.lock()
.unwrap()
.list_for_hash(hash, JOBS_PER_PHOTO)
.map_err(|e| EnhanceError::Internal(e.to_string()))
}
/// Hide a job from the per-photo list. Nothing is deleted on ComfyUI:
/// staged results vanish on its restart anyway.
pub fn discard(&self, id: &str) -> Result<EnhanceJobRow, EnhanceError> {
self.get_job(id)?;
self.dao
.lock()
.unwrap()
.discard(id)
.map_err(|e| EnhanceError::Internal(e.to_string()))?;
self.get_job(id)
}
fn update_job(&self, id: &str, u: &JobUpdate) {
if let Err(e) = self.dao.lock().unwrap().update(id, u) {
warn!("enhance: updating job {id} failed: {e}");
}
}
/// The staged result `n` of a job, for the preview relay.
pub fn pending_of(job: &EnhanceJobRow) -> Vec<Staged> {
serde_json::from_str(&job.pending).unwrap_or_default()
}
pub fn saved_of(job: &EnhanceJobRow) -> Vec<SavedEntry> {
serde_json::from_str(&job.saved).unwrap_or_default()
}
/// Relay the staged PNG `n` of a job. A 404 from ComfyUI means it was
/// restarted: the job becomes `expired`.
pub async fn preview(
&self,
job: &EnhanceJobRow,
n: usize,
) -> Result<(String, bytes::Bytes), EnhanceError> {
let client = self.client.as_ref().ok_or(EnhanceError::Disabled)?;
let staged = Self::pending_of(job);
let item = staged.get(n).ok_or(EnhanceError::NotFound)?;
match client.view(&item.view).await {
Ok(Some(found)) => Ok(found),
Ok(None) => {
self.expire(job);
Err(EnhanceError::Expired)
}
Err(e) => Err(comfy_err(e)),
}
}
/// Mark a job whose results are gone `expired` (unless it was approved:
/// the saved file is in the library regardless).
fn expire(&self, job: &EnhanceJobRow) {
if job.status == STATUS_READY {
let mut u = JobUpdate::status(STATUS_EXPIRED);
u.error = Some(MSG_EXPIRED.to_string());
self.update_job(&job.id, &u);
}
}
// ---- poller ---------------------------------------------------------
/// Advance every active job once. Public so tests (and the loop) drive it.
pub async fn poll_once(&self) {
let Some(client) = &self.client else { return };
let active = match self.dao.lock().unwrap().list_active() {
Ok(a) => a,
Err(e) => {
warn!("enhance: listing active jobs failed: {e}");
return;
}
};
if active.is_empty() {
self.unreachable_since.lock().unwrap().clear();
return;
}
// Queue first, history second: a job that finishes between the two
// calls is then found in history, never misread as forgotten.
let queue = match client.queue().await {
Ok(q) => q,
Err(ComfyError::Unreachable(_)) => {
self.note_unreachable(&active);
return;
}
Err(e) => {
warn!("enhance: /queue failed: {e}");
return;
}
};
for job in active {
self.unreachable_since.lock().unwrap().remove(&job.id);
if let Some((running, ahead)) = queue.position_of(&job.comfy_prompt_id) {
let status = if running {
STATUS_RUNNING
} else {
STATUS_QUEUED
};
let mut u = JobUpdate::status(status);
u.queue_position = Some(ahead as i32);
if job.status != status || job.queue_position != u.queue_position {
self.update_job(&job.id, &u);
}
continue;
}
match client.history(&job.comfy_prompt_id).await {
Ok(Some(entry)) if entry.completed || entry.status_str == "error" => {
self.finish(&job, entry)
}
// Known but not finished (a race with the queue): keep waiting.
Ok(Some(_)) => {}
Ok(None) => {
info!("enhance: job {} expired (ComfyUI forgot it)", job.id);
let mut u = JobUpdate::status(STATUS_EXPIRED);
u.error = Some(MSG_EXPIRED.to_string());
self.update_job(&job.id, &u);
}
Err(ComfyError::Unreachable(_)) => {
self.note_unreachable(std::slice::from_ref(&job))
}
Err(e) => warn!("enhance: /history for job {} failed: {e}", job.id),
}
}
}
fn finish(&self, job: &EnhanceJobRow, entry: HistoryEntry) {
if entry.status_str == "error" || (!entry.completed && entry.error.is_some()) {
let mut u = JobUpdate::status(STATUS_FAILED);
u.error = Some(
entry
.error
.unwrap_or_else(|| "ComfyUI reported an error".to_string()),
);
self.update_job(&job.id, &u);
return;
}
if entry.staged.is_empty() {
let mut u = JobUpdate::status(STATUS_FAILED);
u.error = Some("ComfyUI finished without producing a result".to_string());
self.update_job(&job.id, &u);
return;
}
let mut u = JobUpdate::status(STATUS_READY);
u.pending = serde_json::to_string(&entry.staged).ok();
self.update_job(&job.id, &u);
info!(
"enhance: job {} ready ({} result(s))",
job.id,
entry.staged.len()
);
}
fn note_unreachable(&self, jobs: &[EnhanceJobRow]) {
let grace = Duration::from_secs(self.config.unreachable_grace_secs);
let mut seen = self.unreachable_since.lock().unwrap();
for job in jobs {
let since = *seen.entry(job.id.clone()).or_insert_with(Instant::now);
if since.elapsed() >= grace {
seen.remove(&job.id);
let mut u = JobUpdate::status(STATUS_UNAVAILABLE);
u.error = Some(MSG_UNAVAILABLE.to_string());
self.update_job(&job.id, &u);
}
}
}
/// Release ComfyUI's VRAM once it has been idle. The idle test uses
/// ComfyUI's own history, so a session you're working in by hand keeps
/// it from unloading.
pub async fn idle_check(&self) {
let (Some(client), true) = (&self.client, self.config.idle_free_secs > 0) else {
return;
};
let reachable = async {
let queue = client.queue().await?;
let newest = client.newest_history().await?;
Ok::<_, ComfyError>((queue, newest))
}
.await;
let (queue, newest) = match reachable {
Ok(v) => v,
Err(_) => {
self.idle.lock().unwrap().reset();
return;
}
};
let active = self
.dao
.lock()
.unwrap()
.list_active()
.map(|a| a.len())
.unwrap_or(1);
let marker = newest.and_then(|h| h.timestamp_ms);
let should_free = self.idle.lock().unwrap().observe(
chrono::Utc::now().timestamp_millis(),
queue.is_empty(),
active,
marker,
self.config.idle_free_secs as i64 * 1000,
);
if should_free {
match client.free().await {
Ok(()) => info!(
"enhance: ComfyUI idle for {}s, released its models",
self.config.idle_free_secs
),
Err(e) => {
warn!("enhance: POST /free failed: {e}");
self.idle.lock().unwrap().free_failed();
}
}
}
}
/// The background loop: poll active jobs every `poll_secs`, run the idle
/// check every 30 s, prune old jobs once at start.
pub fn spawn(self: &Arc<Self>) {
if !self.is_configured() {
return;
}
let svc = Arc::clone(self);
tokio::spawn(async move {
let cutoff = chrono::Utc::now().timestamp() - JOB_RETENTION_SECS;
if let Err(e) = svc.dao.lock().unwrap().prune(cutoff) {
warn!("enhance: pruning old jobs failed: {e}");
}
let mut last_idle = Instant::now() - IDLE_CHECK_EVERY;
loop {
svc.poll_once().await;
if last_idle.elapsed() >= IDLE_CHECK_EVERY {
last_idle = Instant::now();
svc.idle_check().await;
}
tokio::time::sleep(Duration::from_secs(svc.config.poll_secs)).await;
}
});
}
// ---- approve --------------------------------------------------------
/// Ask ComfyUI to write result `n` into the library. Idempotent: a result
/// that is already saved returns the stored entry without calling ComfyUI.
/// The caller ingests the returned file.
pub async fn approve(
&self,
job_id: &str,
n: usize,
place: &dyn Fn(&Path) -> Option<(i32, String)>,
) -> Result<(EnhanceJobRow, SavedEntry), EnhanceError> {
let client = self.client.as_ref().ok_or(EnhanceError::Disabled)?;
let _guard = self.approve_lock.lock().await;
let job = self.get_job(job_id)?;
let mut saved = Self::saved_of(&job);
if let Some(done) = saved.iter().find(|s| s.n == n) {
return Ok((job.clone(), done.clone()));
}
match job.status.as_str() {
STATUS_READY | STATUS_APPROVED => {}
STATUS_EXPIRED => return Err(EnhanceError::Expired),
other => {
return Err(EnhanceError::BadRequest(format!(
"the job is {other}, there is nothing to approve"
)));
}
}
let staged = Self::pending_of(&job);
let item = staged
.get(n)
.ok_or_else(|| EnhanceError::BadRequest(format!("this job has no result {n}")))?;
let comfy_path = match client.approve(&item.id).await {
Ok(p) => p,
Err(ApproveError::Gone) => {
self.expire(&job);
return Err(EnhanceError::Expired);
}
Err(ApproveError::Refused(m)) => return Err(EnhanceError::Upstream(m)),
Err(ApproveError::Comfy(e)) => return Err(comfy_err(e)),
};
let located = self
.path_map
.from_comfy(&comfy_path)
.and_then(|abs| place(&abs));
let entry = SavedEntry {
n,
comfy_path: comfy_path.clone(),
library_id: located.as_ref().map(|(l, _)| *l),
rel_path: located.map(|(_, r)| r),
};
// Record it even when it can't be placed: the file exists now.
saved.push(entry.clone());
let mut u = JobUpdate::status(STATUS_APPROVED);
u.error = None;
u.saved = serde_json::to_string(&saved).ok();
self.update_job(&job.id, &u);
let job = self.get_job(job_id)?;
if entry.rel_path.is_none() {
return Err(EnhanceError::Unplaceable(format!(
"ComfyUI saved {comfy_path}, but it isn't under any library or path_map entry"
)));
}
Ok((job, entry))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn idle_frees_once_after_the_quiet_period() {
let mut s = IdleState::default();
let idle = 300_000;
// First sight is activity.
assert!(!s.observe(0, true, 0, Some(5), idle));
assert!(!s.observe(299_000, true, 0, Some(5), idle));
assert!(s.observe(300_000, true, 0, Some(5), idle));
// Only once.
assert!(!s.observe(400_000, true, 0, Some(5), idle));
assert!(!s.observe(9_000_000, true, 0, Some(5), idle));
// A new history entry (e.g. an interactive run) restarts the clock
// and re-arms the free.
assert!(!s.observe(9_100_000, true, 0, Some(6), idle));
assert!(!s.observe(9_399_000, true, 0, Some(6), idle));
assert!(s.observe(9_400_000, true, 0, Some(6), idle));
}
#[test]
fn idle_never_frees_while_busy() {
let mut s = IdleState::default();
assert!(!s.observe(0, true, 0, Some(1), 1000));
// Queue not empty: activity, clock restarts each time.
assert!(!s.observe(5000, false, 0, Some(1), 1000));
assert!(!s.observe(5500, true, 0, Some(1), 1000));
// One of our jobs is active.
assert!(!s.observe(9000, true, 1, Some(1), 1000));
assert!(!s.observe(9500, true, 0, Some(1), 1000));
assert!(s.observe(10_000, true, 0, Some(1), 1000));
}
#[test]
fn idle_restarts_after_comfy_goes_away_and_retries_failures() {
let mut s = IdleState::default();
assert!(!s.observe(0, true, 0, None, 1000));
assert!(s.observe(1000, true, 0, None, 1000));
s.free_failed();
assert!(s.observe(1500, true, 0, None, 1000), "retried");
s.reset();
assert!(!s.observe(10_000, true, 0, None, 1000), "fresh baseline");
assert!(s.observe(11_000, true, 0, None, 1000));
}
}
+502
View File
@@ -0,0 +1,502 @@
//! Workflow templates (#192): load, validate, and fill in an API-format
//! ComfyUI workflow for one photo.
//!
//! Nodes are found by `_meta.title`, never by id, so re-exporting after
//! editing the graph keeps working.
use std::path::Path;
use anyhow::{Context as _, Result, bail};
use serde_json::{Map, Value, json};
use super::config::{PromptMode, TemplateConfig};
pub const TITLE_SOURCE: &str = "imageapi:source";
pub const TITLE_SAVE: &str = "imageapi:save";
pub const TITLE_PROMPT: &str = "imageapi:prompt";
pub const TITLE_REFERENCE: &str = "imageapi:reference";
/// Prompt-node inputs, in the order they are tried.
const PROMPT_INPUTS: &[&str] = &["text", "prompt", "value"];
/// Output nodes that write into ComfyUI's output folder on every API run.
const STRAY_OUTPUT_CLASSES: &[&str] = &[
"SaveImage",
"SaveAnimatedWEBP",
"SaveAnimatedPNG",
"SaveWEBM",
"SaveVideo",
"SaveGLB",
];
#[derive(Debug, Clone)]
pub struct Template {
pub cfg: TemplateConfig,
workflow: Map<String, Value>,
source_id: String,
save_id: String,
/// `(node id, input name)` of the prompt text.
prompt: Option<(String, &'static str)>,
reference_id: Option<String>,
/// Startup warnings (stray output nodes).
pub warnings: Vec<String>,
}
/// Why a request can't be turned into a ComfyUI prompt. The message is for
/// the caller, so it names the problem in the user's terms.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FillError(pub String);
impl std::fmt::Display for FillError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
/// A filled-in prompt plus whether the reference node had to be removed
/// (so a ComfyUI rejection can be reported as "needs a reference photo").
#[derive(Debug)]
pub struct Filled {
pub prompt: Value,
pub dropped_reference: bool,
}
impl Template {
pub fn load(cfg: &TemplateConfig, base_dir: &Path) -> Result<Template> {
let path = if cfg.file.is_absolute() {
cfg.file.clone()
} else {
base_dir.join(&cfg.file)
};
let text = std::fs::read_to_string(&path)
.with_context(|| format!("reading workflow {}", path.display()))?;
let value: Value = serde_json::from_str(&text)
.with_context(|| format!("parsing workflow {}", path.display()))?;
Self::from_value(cfg.clone(), value)
}
pub fn from_value(cfg: TemplateConfig, value: Value) -> Result<Template> {
let Value::Object(workflow) = value else {
bail!("workflow is not a JSON object (is it the API-format export?)");
};
for (id, node) in &workflow {
if node.get("class_type").and_then(Value::as_str).is_none()
|| !node.get("inputs").is_some_and(Value::is_object)
{
bail!(
"node {id} has no class_type/inputs: this looks like a UI-format workflow, \
export with Workflow -> Export (API)"
);
}
}
let find = |title: &str| -> Result<Option<String>> {
let ids: Vec<&String> = workflow
.iter()
.filter(|(_, n)| title_of(n) == Some(title))
.map(|(id, _)| id)
.collect();
match ids.len() {
0 => Ok(None),
1 => Ok(Some(ids[0].clone())),
_ => bail!("more than one node is titled {title:?}"),
}
};
let Some(source_id) = find(TITLE_SOURCE)? else {
bail!("no node titled {TITLE_SOURCE:?} (the Load Image From Path that gets the photo)");
};
let Some(save_id) = find(TITLE_SAVE)? else {
bail!("no node titled {TITLE_SAVE:?} (the Save Image Version node)");
};
require_input(&workflow, &source_id, "path", TITLE_SOURCE)?;
let prompt = match find(TITLE_PROMPT)? {
None => None,
Some(id) => {
let inputs = inputs_of(&workflow, &id);
let Some(name) = PROMPT_INPUTS
.iter()
.copied()
.find(|n| inputs.get(*n).is_some_and(Value::is_string))
else {
bail!(
"node {TITLE_PROMPT:?} has none of the inputs {PROMPT_INPUTS:?} \
as a plain string (a linked input can't take the user's text)"
);
};
Some((id, name))
}
};
if cfg.prompt_mode == PromptMode::Replace && prompt.is_none() {
bail!("prompt_mode = \"replace\" needs a node titled {TITLE_PROMPT:?}");
}
let reference_id = find(TITLE_REFERENCE)?;
if let Some(id) = &reference_id {
require_input(&workflow, id, "path", TITLE_REFERENCE)?;
}
let mut warnings = Vec::new();
for (id, node) in &workflow {
let class = node["class_type"].as_str().unwrap_or_default();
if STRAY_OUTPUT_CLASSES.contains(&class) {
warnings.push(format!(
"node {id} ({class}) also writes into ComfyUI's output folder on every run"
));
}
}
Ok(Template {
cfg,
workflow,
source_id,
save_id,
prompt,
reference_id,
warnings,
})
}
pub fn accepts_prompt(&self) -> bool {
self.prompt.is_some()
}
pub fn accepts_reference(&self) -> bool {
self.reference_id.is_some()
}
pub fn prompt_required(&self) -> bool {
self.cfg.prompt_mode == PromptMode::Replace
}
/// Fill the workflow for one photo.
///
/// - `source` / `reference` are paths as ComfyUI sees them.
/// - Without a reference the reference node and every link to it are
/// removed: its leftover export path must never be used.
/// - Every `seed` / `noise_seed` / `*.seed` input gets a fresh value
/// (ComfyUI's "randomize after generate" only exists in its page), or a
/// retry would return the cached result.
pub fn fill(
&self,
source: &str,
prompt: Option<&str>,
reference: Option<&str>,
next_seed: &mut dyn FnMut() -> u64,
) -> Result<Filled, FillError> {
let user_text = prompt.map(str::trim).filter(|t| !t.is_empty());
match (&self.prompt, user_text) {
(None, Some(_)) => {
return Err(FillError("this template doesn't take a prompt".to_string()));
}
(Some(_), None) if self.prompt_required() => {
return Err(FillError(
"this template needs an instruction (the prompt can't be empty)".to_string(),
));
}
_ => {}
}
if reference.is_some() && self.reference_id.is_none() {
return Err(FillError(
"this template doesn't take a reference photo".to_string(),
));
}
let mut wf = self.workflow.clone();
set_input(&mut wf, &self.source_id, "path", json!(source));
set_input(&mut wf, &self.save_id, "require_approval", json!(true));
if let (Some((id, input)), Some(text)) = (&self.prompt, user_text) {
let base = inputs_of(&wf, id)
.get(*input)
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
let combined = match self.cfg.prompt_mode {
PromptMode::Replace => text.to_string(),
PromptMode::Append if base.trim().is_empty() => text.to_string(),
PromptMode::Append => format!("{base}{}{text}", self.cfg.prompt_join),
};
set_input(&mut wf, id, input, json!(combined));
}
let mut dropped_reference = false;
if let Some(ref_id) = &self.reference_id {
match reference {
Some(path) => set_input(&mut wf, ref_id, "path", json!(path)),
None => {
wf.remove(ref_id);
for node in wf.values_mut() {
if let Some(inputs) = node.get_mut("inputs").and_then(Value::as_object_mut)
{
inputs.retain(|_, v| !links_to(v, ref_id));
}
}
dropped_reference = true;
}
}
}
for node in wf.values_mut() {
let Some(inputs) = node.get_mut("inputs").and_then(Value::as_object_mut) else {
continue;
};
for (key, value) in inputs.iter_mut() {
if is_seed_input(key) && value.is_number() {
*value = json!(next_seed());
}
}
}
Ok(Filled {
prompt: Value::Object(wf),
dropped_reference,
})
}
}
fn title_of(node: &Value) -> Option<&str> {
node.get("_meta")?.get("title")?.as_str()
}
fn inputs_of<'a>(wf: &'a Map<String, Value>, id: &str) -> &'a Map<String, Value> {
static EMPTY: std::sync::LazyLock<Map<String, Value>> = std::sync::LazyLock::new(Map::new);
wf.get(id)
.and_then(|n| n.get("inputs"))
.and_then(Value::as_object)
.unwrap_or(&EMPTY)
}
fn require_input(wf: &Map<String, Value>, id: &str, input: &str, title: &str) -> Result<()> {
if !inputs_of(wf, id).get(input).is_some_and(Value::is_string) {
bail!("node {title:?} has no string input {input:?}");
}
Ok(())
}
fn set_input(wf: &mut Map<String, Value>, id: &str, input: &str, value: Value) {
if let Some(inputs) = wf
.get_mut(id)
.and_then(|n| n.get_mut("inputs"))
.and_then(Value::as_object_mut)
{
inputs.insert(input.to_string(), value);
}
}
/// `["<id>", n]`: a link to output `n` of node `id`.
fn links_to(v: &Value, id: &str) -> bool {
matches!(v.as_array().map(Vec::as_slice), Some([Value::String(s), Value::Number(_)]) if s == id)
}
fn is_seed_input(key: &str) -> bool {
key == "seed" || key == "noise_seed" || key.ends_with(".seed")
}
#[cfg(test)]
mod tests {
use super::*;
pub const RESTORE: &str = include_str!("../../tests/fixtures/enhance/restore_turbo.json");
pub const EDIT: &str = include_str!("../../tests/fixtures/enhance/edit_turbo.json");
pub const DEV: &str = include_str!("../../tests/fixtures/enhance/restore_dev.json");
fn cfg(name: &str, mode: PromptMode) -> TemplateConfig {
TemplateConfig {
name: name.into(),
label: name.into(),
description: String::new(),
file: "x.json".into(),
prompt_mode: mode,
prompt_join: " ".into(),
output_megapixels: 1.0,
}
}
fn load(json: &str, mode: PromptMode) -> Result<Template> {
Template::from_value(cfg("t", mode), serde_json::from_str(json).unwrap())
}
fn seeds() -> impl FnMut() -> u64 {
let mut n = 1000;
move || {
n += 1;
n
}
}
fn input<'a>(p: &'a Value, id: &str, name: &str) -> &'a Value {
&p[id]["inputs"][name]
}
#[test]
fn real_templates_load() {
let restore = load(RESTORE, PromptMode::Append).unwrap();
assert!(restore.accepts_prompt());
assert!(!restore.accepts_reference());
assert!(!restore.prompt_required());
let edit = load(EDIT, PromptMode::Replace).unwrap();
assert!(edit.accepts_prompt() && edit.accepts_reference() && edit.prompt_required());
assert!(load(DEV, PromptMode::Append).is_ok());
// The prompt node is a String Multiline primitive: input `value`.
assert_eq!(restore.prompt.as_ref().unwrap().1, "value");
assert!(restore.warnings.is_empty());
}
#[test]
fn missing_required_titles_are_rejected_with_a_clear_message() {
let no_source = RESTORE.replace("imageapi:source", "something else");
let err = load(&no_source, PromptMode::Append)
.unwrap_err()
.to_string();
assert!(err.contains("imageapi:source"), "{err}");
let no_save = RESTORE.replace("imageapi:save", "something else");
let err = load(&no_save, PromptMode::Append).unwrap_err().to_string();
assert!(err.contains("imageapi:save"), "{err}");
}
#[test]
fn rejects_ui_format_duplicates_and_bad_prompt_nodes() {
let ui = r#"{"nodes": [], "links": []}"#;
assert!(load(ui, PromptMode::Append).is_err());
let dup = RESTORE.replace("\"Load CLIP\"", "\"imageapi:source\"");
assert!(load(&dup, PromptMode::Append).is_err());
// Prompt node with none of text / prompt / value.
let mut v: Value = serde_json::from_str(RESTORE).unwrap();
v["13"]["inputs"] = json!({"other": "x"});
assert!(Template::from_value(cfg("t", PromptMode::Append), v).is_err());
// replace without a prompt node
let no_prompt = RESTORE.replace("imageapi:prompt", "nope");
assert!(load(&no_prompt, PromptMode::Replace).is_err());
assert!(load(&no_prompt, PromptMode::Append).is_ok());
}
#[test]
fn prompt_node_accepts_text_or_prompt_inputs() {
let mut v: Value = serde_json::from_str(RESTORE).unwrap();
v["13"]["inputs"] = json!({"prompt": "base"});
let t = Template::from_value(cfg("t", PromptMode::Append), v).unwrap();
let f = t
.fill("Z:\\a.jpg", Some("more"), None, &mut seeds())
.unwrap();
assert_eq!(input(&f.prompt, "13", "prompt"), "base more");
}
#[test]
fn stray_save_image_is_a_warning_not_an_error() {
let mut v: Value = serde_json::from_str(RESTORE).unwrap();
v["99"] = json!({"class_type": "SaveImage", "inputs": {"images": ["6", 0]}});
let t = Template::from_value(cfg("t", PromptMode::Append), v).unwrap();
assert_eq!(t.warnings.len(), 1);
assert!(t.warnings[0].contains("SaveImage"));
}
#[test]
fn fill_sets_path_forces_approval_and_randomises_every_seed() {
let t = load(RESTORE, PromptMode::Append).unwrap();
let f = t.fill("Z:\\p\\a.jpg", None, None, &mut seeds()).unwrap();
assert_eq!(input(&f.prompt, "37", "path"), "Z:\\p\\a.jpg");
assert_eq!(input(&f.prompt, "36", "require_approval"), true);
// KSampler.seed and TextGenerate's sampling_mode.seed, both fresh.
let a = input(&f.prompt, "5", "seed").as_u64().unwrap();
let b = input(&f.prompt, "14", "sampling_mode.seed")
.as_u64()
.unwrap();
assert!(a > 1000 && b > 1000 && a != b);
assert_ne!(a, 774744764437710);
// Two fills with a real RNG differ.
let mut rng = || rand::random::<u64>() >> 12;
let x = t.fill("p", None, None, &mut rng).unwrap();
let y = t.fill("p", None, None, &mut rng).unwrap();
assert_ne!(input(&x.prompt, "5", "seed"), input(&y.prompt, "5", "seed"));
}
#[test]
fn force_approval_even_if_the_export_had_it_off() {
let off = RESTORE.replace("\"require_approval\": true", "\"require_approval\": false");
let t = load(&off, PromptMode::Append).unwrap();
let f = t.fill("p", None, None, &mut seeds()).unwrap();
assert_eq!(input(&f.prompt, "36", "require_approval"), true);
}
#[test]
fn append_uses_the_join_text_and_empty_leaves_the_base() {
let base = "Increase the clarity";
let mut c = cfg("t", PromptMode::Append);
c.prompt_join = "\n\nAdditional instructions: ".into();
let mut v: Value = serde_json::from_str(RESTORE).unwrap();
v["13"]["inputs"]["value"] = json!(base);
let t = Template::from_value(c, v).unwrap();
let f = t.fill("p", Some(" warmer "), None, &mut seeds()).unwrap();
assert_eq!(
input(&f.prompt, "13", "value"),
"Increase the clarity\n\nAdditional instructions: warmer"
);
let f = t.fill("p", Some(" "), None, &mut seeds()).unwrap();
assert_eq!(input(&f.prompt, "13", "value"), base);
let f = t.fill("p", None, None, &mut seeds()).unwrap();
assert_eq!(input(&f.prompt, "13", "value"), base);
}
#[test]
fn default_join_is_a_single_space_sentence() {
let t = load(RESTORE, PromptMode::Append).unwrap();
let f = t
.fill("p", Some("Make it warmer."), None, &mut seeds())
.unwrap();
let text = input(&f.prompt, "13", "value").as_str().unwrap();
assert!(text.ends_with("Remove any phone interface overlays. Make it warmer."));
}
#[test]
fn replace_mode_requires_and_uses_the_instruction() {
let t = load(EDIT, PromptMode::Replace).unwrap();
let err = t.fill("p", None, None, &mut seeds()).unwrap_err();
assert!(err.0.contains("instruction"));
assert!(t.fill("p", Some(" "), None, &mut seeds()).is_err());
let f = t
.fill("p", Some("make it black and white"), None, &mut seeds())
.unwrap();
assert_eq!(input(&f.prompt, "13", "value"), "make it black and white");
}
#[test]
fn prompt_or_reference_the_template_cant_take_are_refused() {
let mut v: Value = serde_json::from_str(RESTORE).unwrap();
v.as_object_mut().unwrap().remove("13");
// 14 links to 13; drop that link so the graph stays plausible.
v["14"]["inputs"].as_object_mut().unwrap().remove("prompt");
let t = Template::from_value(cfg("t", PromptMode::Append), v).unwrap();
assert!(t.fill("p", Some("x"), None, &mut seeds()).is_err());
assert!(t.fill("p", None, None, &mut seeds()).is_ok());
let restore = load(RESTORE, PromptMode::Append).unwrap();
assert!(restore.fill("p", None, Some("r"), &mut seeds()).is_err());
}
#[test]
fn reference_given_sets_its_path_and_keeps_links() {
let t = load(EDIT, PromptMode::Replace).unwrap();
let f = t
.fill("src", Some("x"), Some("Z:\\ref.jpg"), &mut seeds())
.unwrap();
assert!(!f.dropped_reference);
assert_eq!(input(&f.prompt, "38", "path"), "Z:\\ref.jpg");
assert_eq!(input(&f.prompt, "36", "reference_1"), &json!(["38", 1]));
}
#[test]
fn no_reference_removes_the_node_and_every_link_to_it() {
let t = load(EDIT, PromptMode::Replace).unwrap();
let f = t.fill("src", Some("x"), None, &mut seeds()).unwrap();
assert!(f.dropped_reference);
let wf = f.prompt.as_object().unwrap();
assert!(!wf.contains_key("38"));
// No input anywhere still links to node 38, and the leftover export
// path can't leak.
assert!(
!serde_json::to_string(&f.prompt)
.unwrap()
.contains("[\"38\"")
);
assert!(wf["36"]["inputs"].get("reference_1").is_none());
// Links to other nodes survive.
assert_eq!(input(&f.prompt, "36", "source"), &json!(["37", 1]));
}
}
+1
View File
@@ -17,6 +17,7 @@ pub mod database;
pub mod date_resolver;
pub mod date_view;
pub mod duplicates;
pub mod enhance;
pub mod error;
pub mod exif;
pub mod face_watch;
+23
View File
@@ -234,6 +234,21 @@ fn main() -> std::io::Result<()> {
}
}
// Enhance on demand (#192): templates from enhance.toml, the job
// poller and the idle /free. A bad enhance.toml fails startup; a
// missing one just leaves Enhance off.
let enhance_service = {
let loaded = image_api::enhance::config::load().expect("invalid enhance config");
let dao: Arc<Mutex<Box<dyn database::EnhanceJobDao>>> =
Arc::new(Mutex::new(Box::new(database::SqliteEnhanceJobDao::new())));
let svc = Arc::new(
image_api::enhance::EnhanceService::from_loaded(loaded, dao)
.expect("invalid enhance config"),
);
svc.spawn();
Data::from(svc)
};
// Spawn the nightly pre-generation scheduler (Section D).
reels::spawn_pregen_scheduler(app_state.clone()).await;
// Spawn the on-disk reel-cache sweeper (bounds pre-gen + on-demand reels).
@@ -326,6 +341,13 @@ fn main() -> std::io::Result<()> {
.service(handlers::favorites::favorites)
.service(handlers::favorites::put_add_favorite)
.service(handlers::favorites::delete_favorite)
.service(enhance::get_templates)
.service(enhance::submit)
.service(enhance::list_jobs)
.service(enhance::get_job)
.service(enhance::preview)
.service(enhance::approve)
.service(enhance::discard)
.service(versions::get_versions)
.service(versions::put_display)
.service(handlers::image::get_file_metadata)
@@ -392,6 +414,7 @@ fn main() -> std::io::Result<()> {
.add_feature(faces::add_face_services::<_, faces::SqliteFaceDao>)
.add_feature(duplicates::add_duplicate_services)
.app_data(app_data.clone())
.app_data(enhance_service.clone())
.app_data::<Data<RealFileSystem>>(Data::new(RealFileSystem::new(
app_data.base_path.clone(),
)))
+130 -113
View File
@@ -740,6 +740,135 @@ pub fn playlist_needs_generation(video_path: &Path, playlist_path: &Path) -> boo
true
}
/// Register files in `image_exif` (hash, perceptual hashes, EXIF, resolved
/// date). Returns the `(absolute, rel_path)` pairs whose row was written.
/// Shared by the watcher tick and the targeted ingest after an Enhance
/// approval (#192).
pub fn register_new_files(
library: &libraries::Library,
files_needing_row: Vec<(PathBuf, String)>,
exif_dao: &Arc<Mutex<Box<dyn ExifDao>>>,
) -> Vec<(PathBuf, String)> {
let mut newly_registered: Vec<(PathBuf, String)> = Vec::new();
if files_needing_row.is_empty() {
return newly_registered;
}
info!(
"Registering {} new files in image_exif",
files_needing_row.len()
);
for (file_path, relative_path) in files_needing_row {
let registered_file = (file_path.clone(), relative_path.clone());
let timestamp = Utc::now().timestamp();
// Hash + size from filesystem metadata — always attempted so
// every file gets a content_hash, even when EXIF is absent.
let (content_hash, size_bytes) = match content_hash::compute(&file_path) {
Ok(id) => (Some(id.content_hash), Some(id.size_bytes)),
Err(e) => {
warn!("Failed to hash {}: {:?}", file_path.display(), e);
(None, None)
}
};
// Perceptual hashes (pHash + dHash). Best-effort — None for
// videos and decode failures. Drives near-duplicate detection
// in the Apollo duplicates surface; failure here is non-fatal
// and never blocks indexing.
let perceptual = perceptual_hash::compute(&file_path);
// EXIF is best-effort enrichment. When extraction fails (or the
// file type doesn't support EXIF) we still store a row with all
// EXIF fields NULL; the file remains visible to sort-by-date
// and tag queries via its rel_path and filesystem timestamps.
let exif_fields = if exif::supports_exif(&file_path) {
match exif::extract_exif_from_path(&file_path) {
Ok(data) => Some(data),
Err(e) => {
debug!(
"No EXIF or parse error for {}: {:?}",
file_path.display(),
e
);
None
}
}
} else {
None
};
// Canonical date_taken via the waterfall — kamadak-exif (already
// computed above) → exiftool fallback for videos / MakerNote /
// QuickTime → filename regex → earliest_fs_time. Source is
// recorded so the per-tick backfill drain can re-run weak
// resolutions later. The EXIF branch additionally runs the
// camera-offset cascade to produce the new-contract columns
// (date_utc / camera_offset_minutes / date_taken_tz_source,
// ADR-0001) while the legacy date_taken stays byte-identical.
let resolved_date = date_resolver::resolve_date_taken(
&file_path,
exif_fields
.as_ref()
.and_then(date_resolver::exif_prior_from)
.as_ref(),
);
let insert_exif = InsertImageExif {
library_id: library.id,
file_path: relative_path.clone(),
camera_make: exif_fields.as_ref().and_then(|e| e.camera_make.clone()),
camera_model: exif_fields.as_ref().and_then(|e| e.camera_model.clone()),
lens_model: exif_fields.as_ref().and_then(|e| e.lens_model.clone()),
width: exif_fields.as_ref().and_then(|e| e.width),
height: exif_fields.as_ref().and_then(|e| e.height),
orientation: exif_fields.as_ref().and_then(|e| e.orientation),
gps_latitude: exif_fields
.as_ref()
.and_then(|e| e.gps_latitude.map(|v| v as f32)),
gps_longitude: exif_fields
.as_ref()
.and_then(|e| e.gps_longitude.map(|v| v as f32)),
gps_altitude: exif_fields
.as_ref()
.and_then(|e| e.gps_altitude.map(|v| v as f32)),
focal_length: exif_fields
.as_ref()
.and_then(|e| e.focal_length.map(|v| v as f32)),
aperture: exif_fields
.as_ref()
.and_then(|e| e.aperture.map(|v| v as f32)),
shutter_speed: exif_fields.as_ref().and_then(|e| e.shutter_speed.clone()),
iso: exif_fields.as_ref().and_then(|e| e.iso),
date_taken: resolved_date.map(|r| r.timestamp),
created_time: timestamp,
last_modified: timestamp,
content_hash,
size_bytes,
phash_64: perceptual.map(|h| h.phash_64),
dhash_64: perceptual.map(|h| h.dhash_64),
date_taken_source: resolved_date.map(|r| r.source.as_str().to_string()),
date_utc: resolved_date.and_then(|r| r.date_utc),
camera_offset_minutes: resolved_date.and_then(|r| r.camera_offset_minutes),
date_taken_tz_source: resolved_date
.and_then(|r| r.tz_source)
.map(|t| t.as_str().to_string()),
};
let mut dao = exif_dao.lock().expect("Unable to lock ExifDao");
if let Err(e) = dao.store_exif(insert_exif) {
error!(
"Failed to register {} in image_exif: {:?}",
relative_path, e
);
} else {
debug!("Registered {} in image_exif", relative_path);
newly_registered.push(registered_file);
}
}
newly_registered
}
pub fn process_new_files(
library: &libraries::Library,
exif_dao: Arc<Mutex<Box<dyn ExifDao>>>,
@@ -834,119 +963,7 @@ pub fn process_new_files(
}
if !files_needing_row.is_empty() {
info!(
"Registering {} new files in image_exif",
files_needing_row.len()
);
for (file_path, relative_path) in files_needing_row {
let registered_file = (file_path.clone(), relative_path.clone());
let timestamp = Utc::now().timestamp();
// Hash + size from filesystem metadata — always attempted so
// every file gets a content_hash, even when EXIF is absent.
let (content_hash, size_bytes) = match content_hash::compute(&file_path) {
Ok(id) => (Some(id.content_hash), Some(id.size_bytes)),
Err(e) => {
warn!("Failed to hash {}: {:?}", file_path.display(), e);
(None, None)
}
};
// Perceptual hashes (pHash + dHash). Best-effort — None for
// videos and decode failures. Drives near-duplicate detection
// in the Apollo duplicates surface; failure here is non-fatal
// and never blocks indexing.
let perceptual = perceptual_hash::compute(&file_path);
// EXIF is best-effort enrichment. When extraction fails (or the
// file type doesn't support EXIF) we still store a row with all
// EXIF fields NULL; the file remains visible to sort-by-date
// and tag queries via its rel_path and filesystem timestamps.
let exif_fields = if exif::supports_exif(&file_path) {
match exif::extract_exif_from_path(&file_path) {
Ok(data) => Some(data),
Err(e) => {
debug!(
"No EXIF or parse error for {}: {:?}",
file_path.display(),
e
);
None
}
}
} else {
None
};
// Canonical date_taken via the waterfall — kamadak-exif (already
// computed above) → exiftool fallback for videos / MakerNote /
// QuickTime → filename regex → earliest_fs_time. Source is
// recorded so the per-tick backfill drain can re-run weak
// resolutions later. The EXIF branch additionally runs the
// camera-offset cascade to produce the new-contract columns
// (date_utc / camera_offset_minutes / date_taken_tz_source,
// ADR-0001) while the legacy date_taken stays byte-identical.
let resolved_date = date_resolver::resolve_date_taken(
&file_path,
exif_fields
.as_ref()
.and_then(date_resolver::exif_prior_from)
.as_ref(),
);
let insert_exif = InsertImageExif {
library_id: library.id,
file_path: relative_path.clone(),
camera_make: exif_fields.as_ref().and_then(|e| e.camera_make.clone()),
camera_model: exif_fields.as_ref().and_then(|e| e.camera_model.clone()),
lens_model: exif_fields.as_ref().and_then(|e| e.lens_model.clone()),
width: exif_fields.as_ref().and_then(|e| e.width),
height: exif_fields.as_ref().and_then(|e| e.height),
orientation: exif_fields.as_ref().and_then(|e| e.orientation),
gps_latitude: exif_fields
.as_ref()
.and_then(|e| e.gps_latitude.map(|v| v as f32)),
gps_longitude: exif_fields
.as_ref()
.and_then(|e| e.gps_longitude.map(|v| v as f32)),
gps_altitude: exif_fields
.as_ref()
.and_then(|e| e.gps_altitude.map(|v| v as f32)),
focal_length: exif_fields
.as_ref()
.and_then(|e| e.focal_length.map(|v| v as f32)),
aperture: exif_fields
.as_ref()
.and_then(|e| e.aperture.map(|v| v as f32)),
shutter_speed: exif_fields.as_ref().and_then(|e| e.shutter_speed.clone()),
iso: exif_fields.as_ref().and_then(|e| e.iso),
date_taken: resolved_date.map(|r| r.timestamp),
created_time: timestamp,
last_modified: timestamp,
content_hash,
size_bytes,
phash_64: perceptual.map(|h| h.phash_64),
dhash_64: perceptual.map(|h| h.dhash_64),
date_taken_source: resolved_date.map(|r| r.source.as_str().to_string()),
date_utc: resolved_date.and_then(|r| r.date_utc),
camera_offset_minutes: resolved_date.and_then(|r| r.camera_offset_minutes),
date_taken_tz_source: resolved_date
.and_then(|r| r.tz_source)
.map(|t| t.as_str().to_string()),
};
let mut dao = exif_dao.lock().expect("Unable to lock ExifDao");
if let Err(e) = dao.store_exif(insert_exif) {
error!(
"Failed to register {} in image_exif: {:?}",
relative_path, e
);
} else {
debug!("Registered {} in image_exif", relative_path);
newly_registered.push(registered_file);
}
}
newly_registered = register_new_files(library, files_needing_row, &exif_dao);
}
// ── Photo versions (#191) ──────────────────────────────────────────
+283
View File
@@ -0,0 +1,283 @@
{
"1": {
"inputs": {
"unet_name": "fibo2_edit_turbo_int8_convrot.safetensors",
"weight_dtype": "default"
},
"class_type": "UNETLoader",
"_meta": {
"title": "Load Diffusion Model"
}
},
"2": {
"inputs": {
"clip_name": "qwen3vl_4b_int8_convrot.safetensors",
"type": "fibo2",
"device": "default"
},
"class_type": "CLIPLoader",
"_meta": {
"title": "Load CLIP"
}
},
"3": {
"inputs": {
"vae_name": "flux2-vae.safetensors"
},
"class_type": "VAELoader",
"_meta": {
"title": "Load VAE"
}
},
"4": {
"inputs": {
"prompt": [
"25",
0
],
"negative_prompt": "{\"short_description\":\"\",\"objects\":[{\"description\":\"the main subject\",\"texture\":\"smooth, featureless, plastic-like, airbrushed, no fine surface detail\",\"appearance_details\":\"soft blurred edges, simplified shapes, no fine detail\",\"skin_tone_and_texture\":\"waxy airbrushed skin, poreless, plastic smooth\"}],\"background_setting\":\"a plain seamless studio backdrop, perfectly smooth gradient, blurred and featureless, no visible texture\",\"lighting\":{},\"aesthetics\":{\"artifacts\":[\"oversaturated colours\",\"neon colour cast\",\"overexposed highlights\",\"lens flare\",\"slight blur\",\"motion blur\",\"haze\"]},\"photographic_characteristics\":{\"depth_of_field\":\"shallow, heavy bokeh obscuring most of the frame\",\"focus\":\"soft, slightly out of focus, lacking fine texture\"},\"style_medium\":\"\",\"context\":\"\",\"artistic_style\":\"\"}",
"width": 0,
"height": 0,
"batch_size": 1,
"clip": [
"2",
0
],
"images.image_1": [
"37",
0
],
"images.image_2": [
"38",
0
],
"vae": [
"3",
0
]
},
"class_type": "TextEncodeFibo2",
"_meta": {
"title": "TextEncodeFibo2"
}
},
"5": {
"inputs": {
"seed": 774744764437710,
"steps": 4,
"cfg": 1,
"sampler_name": "euler",
"scheduler": "simple",
"denoise": 1,
"model": [
"1",
0
],
"positive": [
"4",
0
],
"negative": [
"4",
1
],
"latent_image": [
"4",
2
]
},
"class_type": "KSampler",
"_meta": {
"title": "KSampler"
}
},
"6": {
"inputs": {
"samples": [
"5",
0
],
"vae": [
"3",
0
]
},
"class_type": "VAEDecode",
"_meta": {
"title": "VAEDecode"
}
},
"12": {
"inputs": {
"value": "You are a highly skilled Creative Director for Visual Adaptation at a leading Generative AI company.\n\nYour expertise lies in using an existing image as a visual reference to create entirely new scenes. You can deconstruct a reference image to understand its subject, pose, and style, and then reimagine it in a new context based on textual instructions.\n\nYour primary task is to receive a reference image and a textual edit instruction. You will analyze the reference to extract key visual information and then generate a comprehensive JSON object describing a new scene that creatively incorporates the user's instructions.\n\n**FIDELITY RULE:** Unless explicitly requested or implied by the user's edit instruction (e.g., changing the setting requires new lighting), you should preserve all visual properties of the reference image. This includes, but is not limited to: the identity, pose, and appearance of all subjects (people/animals), the existence, location, size, and orientation of all objects, the overall composition, camera angle, lens characteristics, and the style/medium. Only change what is strictly required to fulfill the user's edit.\nE.g. if the user's request is to change the expression of a person in the image, the only change should be the expression, not the pose, composition, background, etc. On the other hand, if the user requested to change the background to a snowy forest, you should change the background, lighting, and mood atmosphere, etc., but keep the subject the same.\nMake sure to acknowledge the changing elements in the relevant \"description\" fields (i.e. \"short_description\", object \"description\", etc.). e.g. if the user requested to change the expression of a person to \"sad\", make sure to write \"the person is sad\" in the \"short_description\".\n\n**Describe the FINAL image only, never the edit (STRICT):** describe changed elements by their resulting STATE, not the action \u2014 e.g. `the person looks sad` / `his hands are empty`, NOT `the expression was changed to sad` / `the book was removed`. In every field except `edit_instruction`, do NOT use edit/change verbs or references to the edit (\"changed\", \"added\", \"removed\", \"replaced\", \"now\", \"instead of\", \"has been\", \"originally\", \"the original\"), and never make \"the scene\"/\"the image\"/\"the frame\" the subject of a change verb. Only `edit_instruction` states the change.\n\n**CONTROL-MAP / CONDITIONING INPUTS (act differently):** Sometimes the input image is NOT a normal photograph or illustration but a *conditioning / control map* \u2014 e.g. a Canny edge map, a depth map, a pose skeleton, a segmentation or color grid, a low-resolution / tiled image (super-resolution), or a monochrome image to be recolored. You can tell from the instruction (e.g. \"match this edge map\", \"use this depth map to generate ...\", \"use this pose guidance to generate ...\", \"recolor this monochrome image into ...\", \"increase the clarity and resolution\"). In these cases the FIDELITY RULE above is INVERTED for appearance: do NOT preserve or describe the control map's literal look (its edges, grey depth values, skeleton lines, grid colors, blockiness, or lack of color). Instead treat the input purely as STRUCTURAL guidance \u2014 outline/geometry (edges), spatial layout and relative distance (depth), body position and framing (pose), regional layout (segmentation/color grid), or coarse composition (low-res) \u2014 and describe the NEW, fully-rendered, realistic image the instruction asks for. Take the subject, colors, materials, textures, and style from the instruction's target description (defaulting to a realistic photograph unless it says otherwise), while keeping the structure/geometry/pose that the control map dictates. Every field (`objects`, `background_setting`, `style_medium`, `lighting`, etc.) must describe that intended final image \u2014 never the control map itself.\n\nAdhere strictly to the following structure and guidelines:\n\n1. **Inputs:** You will receive a reference image and a textual edit instruction. You will NOT receive a starting JSON.\n2. **Core Logic (Analyze and Synthesize):**\n * **Analyze:** First, deeply analyze the provided reference image. Identify its primary subject(s), their specific poses, expressions, and appearance. Also note the overall composition, lighting style, and artistic medium.\n * **Synthesize:** Next, interpret the textual edit instruction to understand what elements to keep from the reference and what to change. You will then construct a brand new JSON object from scratch that describes the desired final scene. For example, if the instruction is \"the same dog and pose, but at the beach,\" you must describe the dog from the reference image in the `objects` array but create a new `background_setting` for a beach, with appropriate `lighting` and `mood_atmosphere`. **The description of every unchanged element (objects, people, composition, camera settiLine truncated
},
"class_type": "PrimitiveStringMultiline",
"_meta": {
"title": "String (System Prompt VGL)"
}
},
"13": {
"inputs": {
"value": ""
},
"class_type": "PrimitiveStringMultiline",
"_meta": {
"title": "imageapi:prompt"
}
},
"14": {
"inputs": {
"prompt": [
"13",
0
],
"max_length": 3000,
"sampling_mode": "on",
"sampling_mode.temperature": 0.3,
"sampling_mode.top_k": 20,
"sampling_mode.top_p": 0.95,
"sampling_mode.min_p": 0.05,
"sampling_mode.repetition_penalty": 1.05,
"sampling_mode.seed": 1235,
"sampling_mode.presence_penalty": 0,
"thinking": false,
"use_default_template": true,
"mtp": "off",
"clip": [
"27",
0
],
"image": [
"26",
0
],
"system_prompt": [
"12",
0
]
},
"class_type": "TextGenerate",
"_meta": {
"title": "Generate Text (PE rewrite)"
}
},
"15": {
"inputs": {
"switch": false,
"on_false": [
"13",
0
],
"on_true": [
"14",
0
]
},
"class_type": "ComfySwitchNode",
"_meta": {
"title": "If/Else Switch (PE on/off)"
}
},
"24": {
"inputs": {
"compare_view": {
"__value__": [
"",
""
]
},
"image_a": [
"37",
0
],
"image_b": [
"6",
0
]
},
"class_type": "ImageCompare",
"_meta": {
"title": "Compare Images"
}
},
"25": {
"inputs": {
"source": [
"15",
0
]
},
"class_type": "PreviewAny",
"_meta": {
"title": "Preview as Text"
}
},
"26": {
"inputs": {
"images.image0": [
"37",
0
],
"images.image1": [
"38",
0
]
},
"class_type": "BatchImagesNode",
"_meta": {
"title": "Batch Images"
}
},
"27": {
"inputs": {
"clip_name": "qwen3vl_8b_int8_convrot.safetensors",
"type": "fibo2",
"device": "default"
},
"class_type": "CLIPLoader",
"_meta": {
"title": "Load CLIP (PE)"
}
},
"36": {
"inputs": {
"kind": "edit",
"subfolder": "aiedits",
"require_approval": true,
"images": [
"6",
0
],
"source": [
"37",
1
],
"reference_1": [
"38",
1
]
},
"class_type": "ImageApiSaveImageVersion",
"_meta": {
"title": "imageapi:save"
}
},
"37": {
"inputs": {
"path": ""
},
"class_type": "ImageApiLoadImageFromPath",
"_meta": {
"title": "imageapi:source"
}
},
"38": {
"inputs": {
"path": ""
},
"class_type": "ImageApiLoadImageFromPath",
"_meta": {
"title": "imageapi:reference"
}
}
}
+262
View File
@@ -0,0 +1,262 @@
{
"1": {
"inputs": {
"unet_name": "fibo2_dev_int8_convrot.safetensors",
"weight_dtype": "default"
},
"class_type": "UNETLoader",
"_meta": {
"title": "Load Diffusion Model"
}
},
"2": {
"inputs": {
"clip_name": "qwen3vl_4b_int8_convrot.safetensors",
"type": "fibo2",
"device": "default"
},
"class_type": "CLIPLoader",
"_meta": {
"title": "Load CLIP"
}
},
"3": {
"inputs": {
"vae_name": "flux2-vae.safetensors"
},
"class_type": "VAELoader",
"_meta": {
"title": "Load VAE"
}
},
"4": {
"inputs": {
"prompt": [
"25",
0
],
"negative_prompt": "{\"short_description\":\"\",\"objects\":[{\"description\":\"the main subject\",\"texture\":\"smooth, featureless, plastic-like, airbrushed, no fine surface detail\",\"appearance_details\":\"soft blurred edges, simplified shapes, no fine detail\",\"skin_tone_and_texture\":\"waxy airbrushed skin, poreless, plastic smooth\"}],\"background_setting\":\"a plain seamless studio backdrop, perfectly smooth gradient, blurred and featureless, no visible texture\",\"lighting\":{},\"aesthetics\":{\"artifacts\":[\"oversaturated colours\",\"neon colour cast\",\"overexposed highlights\",\"lens flare\",\"slight blur\",\"motion blur\",\"haze\"]},\"photographic_characteristics\":{\"depth_of_field\":\"shallow, heavy bokeh obscuring most of the frame\",\"focus\":\"soft, slightly out of focus, lacking fine texture\"},\"style_medium\":\"\",\"context\":\"\",\"artistic_style\":\"\"}",
"width": 0,
"height": 0,
"batch_size": 1,
"clip": [
"2",
0
],
"images.image_1": [
"37",
0
],
"vae": [
"3",
0
]
},
"class_type": "TextEncodeFibo2",
"_meta": {
"title": "TextEncodeFibo2"
}
},
"5": {
"inputs": {
"seed": 774744764437710,
"steps": 50,
"cfg": 5,
"sampler_name": "euler",
"scheduler": "simple",
"denoise": 1,
"model": [
"1",
0
],
"positive": [
"4",
0
],
"negative": [
"4",
1
],
"latent_image": [
"4",
2
]
},
"class_type": "KSampler",
"_meta": {
"title": "KSampler"
}
},
"6": {
"inputs": {
"samples": [
"5",
0
],
"vae": [
"3",
0
]
},
"class_type": "VAEDecode",
"_meta": {
"title": "VAEDecode"
}
},
"12": {
"inputs": {
"value": "You are a highly skilled Creative Director for Visual Adaptation at a leading Generative AI company.\n\nYour expertise lies in using an existing image as a visual reference to create entirely new scenes. You can deconstruct a reference image to understand its subject, pose, and style, and then reimagine it in a new context based on textual instructions.\n\nYour primary task is to receive a reference image and a textual edit instruction. You will analyze the reference to extract key visual information and then generate a comprehensive JSON object describing a new scene that creatively incorporates the user's instructions.\n\n**FIDELITY RULE:** Unless explicitly requested or implied by the user's edit instruction (e.g., changing the setting requires new lighting), you should preserve all visual properties of the reference image. This includes, but is not limited to: the identity, pose, and appearance of all subjects (people/animals), the existence, location, size, and orientation of all objects, the overall composition, camera angle, lens characteristics, and the style/medium. Only change what is strictly required to fulfill the user's edit.\nE.g. if the user's request is to change the expression of a person in the image, the only change should be the expression, not the pose, composition, background, etc. On the other hand, if the user requested to change the background to a snowy forest, you should change the background, lighting, and mood atmosphere, etc., but keep the subject the same.\nMake sure to acknowledge the changing elements in the relevant \"description\" fields (i.e. \"short_description\", object \"description\", etc.). e.g. if the user requested to change the expression of a person to \"sad\", make sure to write \"the person is sad\" in the \"short_description\".\n\n**Describe the FINAL image only, never the edit (STRICT):** describe changed elements by their resulting STATE, not the action \u2014 e.g. `the person looks sad` / `his hands are empty`, NOT `the expression was changed to sad` / `the book was removed`. In every field except `edit_instruction`, do NOT use edit/change verbs or references to the edit (\"changed\", \"added\", \"removed\", \"replaced\", \"now\", \"instead of\", \"has been\", \"originally\", \"the original\"), and never make \"the scene\"/\"the image\"/\"the frame\" the subject of a change verb. Only `edit_instruction` states the change.\n\n**CONTROL-MAP / CONDITIONING INPUTS (act differently):** Sometimes the input image is NOT a normal photograph or illustration but a *conditioning / control map* \u2014 e.g. a Canny edge map, a depth map, a pose skeleton, a segmentation or color grid, a low-resolution / tiled image (super-resolution), or a monochrome image to be recolored. You can tell from the instruction (e.g. \"match this edge map\", \"use this depth map to generate ...\", \"use this pose guidance to generate ...\", \"recolor this monochrome image into ...\", \"increase the clarity and resolution\"). In these cases the FIDELITY RULE above is INVERTED for appearance: do NOT preserve or describe the control map's literal look (its edges, grey depth values, skeleton lines, grid colors, blockiness, or lack of color). Instead treat the input purely as STRUCTURAL guidance \u2014 outline/geometry (edges), spatial layout and relative distance (depth), body position and framing (pose), regional layout (segmentation/color grid), or coarse composition (low-res) \u2014 and describe the NEW, fully-rendered, realistic image the instruction asks for. Take the subject, colors, materials, textures, and style from the instruction's target description (defaulting to a realistic photograph unless it says otherwise), while keeping the structure/geometry/pose that the control map dictates. Every field (`objects`, `background_setting`, `style_medium`, `lighting`, etc.) must describe that intended final image \u2014 never the control map itself.\n\nAdhere strictly to the following structure and guidelines:\n\n1. **Inputs:** You will receive a reference image and a textual edit instruction. You will NOT receive a starting JSON.\n2. **Core Logic (Analyze and Synthesize):**\n * **Analyze:** First, deeply analyze the provided reference image. Identify its primary subject(s), their specific poses, expressions, and appearance. Also note the overall composition, lighting style, and artistic medium.\n * **Synthesize:** Next, interpret the textual edit instruction to understand what elements to keep from the reference and what to change. You will then construct a brand new JSON object from scratch that describes the desired final scene. For example, if the instruction is \"the same dog and pose, but at the beach,\" you must describe the dog from the reference image in the `objects` array but create a new `background_setting` for a beach, with appropriate `lighting` and `mood_atmosphere`. **The description of every unchanged element (objects, people, composition, camera settiLine truncated
},
"class_type": "PrimitiveStringMultiline",
"_meta": {
"title": "String (System Prompt VGL)"
}
},
"13": {
"inputs": {
"value": "Increase the clarity and resolution of the image, and refine it to obtain sharp, finely detailed, natural textures as if shot on a modern mirrorless camera, keeping the subject, pose, composition, colours and lighting the same. Remove any phone interface overlays."
},
"class_type": "PrimitiveStringMultiline",
"_meta": {
"title": "imageapi:prompt"
}
},
"14": {
"inputs": {
"prompt": [
"13",
0
],
"max_length": 3000,
"sampling_mode": "on",
"sampling_mode.temperature": 0.3,
"sampling_mode.top_k": 20,
"sampling_mode.top_p": 0.95,
"sampling_mode.min_p": 0.05,
"sampling_mode.repetition_penalty": 1.05,
"sampling_mode.seed": 1235,
"sampling_mode.presence_penalty": 0,
"thinking": false,
"use_default_template": true,
"mtp": "off",
"clip": [
"27",
0
],
"image": [
"26",
0
],
"system_prompt": [
"12",
0
]
},
"class_type": "TextGenerate",
"_meta": {
"title": "Generate Text (PE rewrite)"
}
},
"15": {
"inputs": {
"switch": false,
"on_false": [
"13",
0
],
"on_true": [
"14",
0
]
},
"class_type": "ComfySwitchNode",
"_meta": {
"title": "If/Else Switch (PE on/off)"
}
},
"24": {
"inputs": {
"compare_view": {
"__value__": [
"",
""
]
},
"image_a": [
"37",
0
],
"image_b": [
"6",
0
]
},
"class_type": "ImageCompare",
"_meta": {
"title": "Compare Images"
}
},
"25": {
"inputs": {
"source": [
"15",
0
]
},
"class_type": "PreviewAny",
"_meta": {
"title": "Preview as Text"
}
},
"26": {
"inputs": {
"images.image0": [
"37",
0
]
},
"class_type": "BatchImagesNode",
"_meta": {
"title": "Batch Images"
}
},
"27": {
"inputs": {
"clip_name": "qwen3vl_8b_int8_convrot.safetensors",
"type": "fibo2",
"device": "default"
},
"class_type": "CLIPLoader",
"_meta": {
"title": "Load CLIP (PE)"
}
},
"36": {
"inputs": {
"kind": "restored",
"subfolder": "aiedits",
"require_approval": true,
"images": [
"6",
0
],
"source": [
"37",
1
]
},
"class_type": "ImageApiSaveImageVersion",
"_meta": {
"title": "imageapi:save"
}
},
"37": {
"inputs": {
"path": ""
},
"class_type": "ImageApiLoadImageFromPath",
"_meta": {
"title": "imageapi:source"
}
}
}
+262
View File
@@ -0,0 +1,262 @@
{
"1": {
"inputs": {
"unet_name": "fibo2_edit_turbo_int8_convrot.safetensors",
"weight_dtype": "default"
},
"class_type": "UNETLoader",
"_meta": {
"title": "Load Diffusion Model"
}
},
"2": {
"inputs": {
"clip_name": "qwen3vl_4b_int8_convrot.safetensors",
"type": "fibo2",
"device": "default"
},
"class_type": "CLIPLoader",
"_meta": {
"title": "Load CLIP"
}
},
"3": {
"inputs": {
"vae_name": "flux2-vae.safetensors"
},
"class_type": "VAELoader",
"_meta": {
"title": "Load VAE"
}
},
"4": {
"inputs": {
"prompt": [
"25",
0
],
"negative_prompt": "{\"short_description\":\"\",\"objects\":[{\"description\":\"the main subject\",\"texture\":\"smooth, featureless, plastic-like, airbrushed, no fine surface detail\",\"appearance_details\":\"soft blurred edges, simplified shapes, no fine detail\",\"skin_tone_and_texture\":\"waxy airbrushed skin, poreless, plastic smooth\"}],\"background_setting\":\"a plain seamless studio backdrop, perfectly smooth gradient, blurred and featureless, no visible texture\",\"lighting\":{},\"aesthetics\":{\"artifacts\":[\"oversaturated colours\",\"neon colour cast\",\"overexposed highlights\",\"lens flare\",\"slight blur\",\"motion blur\",\"haze\"]},\"photographic_characteristics\":{\"depth_of_field\":\"shallow, heavy bokeh obscuring most of the frame\",\"focus\":\"soft, slightly out of focus, lacking fine texture\"},\"style_medium\":\"\",\"context\":\"\",\"artistic_style\":\"\"}",
"width": 0,
"height": 0,
"batch_size": 1,
"clip": [
"2",
0
],
"images.image_1": [
"37",
0
],
"vae": [
"3",
0
]
},
"class_type": "TextEncodeFibo2",
"_meta": {
"title": "TextEncodeFibo2"
}
},
"5": {
"inputs": {
"seed": 774744764437710,
"steps": 4,
"cfg": 1,
"sampler_name": "euler",
"scheduler": "simple",
"denoise": 1,
"model": [
"1",
0
],
"positive": [
"4",
0
],
"negative": [
"4",
1
],
"latent_image": [
"4",
2
]
},
"class_type": "KSampler",
"_meta": {
"title": "KSampler"
}
},
"6": {
"inputs": {
"samples": [
"5",
0
],
"vae": [
"3",
0
]
},
"class_type": "VAEDecode",
"_meta": {
"title": "VAEDecode"
}
},
"12": {
"inputs": {
"value": "You are a highly skilled Creative Director for Visual Adaptation at a leading Generative AI company.\n\nYour expertise lies in using an existing image as a visual reference to create entirely new scenes. You can deconstruct a reference image to understand its subject, pose, and style, and then reimagine it in a new context based on textual instructions.\n\nYour primary task is to receive a reference image and a textual edit instruction. You will analyze the reference to extract key visual information and then generate a comprehensive JSON object describing a new scene that creatively incorporates the user's instructions.\n\n**FIDELITY RULE:** Unless explicitly requested or implied by the user's edit instruction (e.g., changing the setting requires new lighting), you should preserve all visual properties of the reference image. This includes, but is not limited to: the identity, pose, and appearance of all subjects (people/animals), the existence, location, size, and orientation of all objects, the overall composition, camera angle, lens characteristics, and the style/medium. Only change what is strictly required to fulfill the user's edit.\nE.g. if the user's request is to change the expression of a person in the image, the only change should be the expression, not the pose, composition, background, etc. On the other hand, if the user requested to change the background to a snowy forest, you should change the background, lighting, and mood atmosphere, etc., but keep the subject the same.\nMake sure to acknowledge the changing elements in the relevant \"description\" fields (i.e. \"short_description\", object \"description\", etc.). e.g. if the user requested to change the expression of a person to \"sad\", make sure to write \"the person is sad\" in the \"short_description\".\n\n**Describe the FINAL image only, never the edit (STRICT):** describe changed elements by their resulting STATE, not the action \u2014 e.g. `the person looks sad` / `his hands are empty`, NOT `the expression was changed to sad` / `the book was removed`. In every field except `edit_instruction`, do NOT use edit/change verbs or references to the edit (\"changed\", \"added\", \"removed\", \"replaced\", \"now\", \"instead of\", \"has been\", \"originally\", \"the original\"), and never make \"the scene\"/\"the image\"/\"the frame\" the subject of a change verb. Only `edit_instruction` states the change.\n\n**CONTROL-MAP / CONDITIONING INPUTS (act differently):** Sometimes the input image is NOT a normal photograph or illustration but a *conditioning / control map* \u2014 e.g. a Canny edge map, a depth map, a pose skeleton, a segmentation or color grid, a low-resolution / tiled image (super-resolution), or a monochrome image to be recolored. You can tell from the instruction (e.g. \"match this edge map\", \"use this depth map to generate ...\", \"use this pose guidance to generate ...\", \"recolor this monochrome image into ...\", \"increase the clarity and resolution\"). In these cases the FIDELITY RULE above is INVERTED for appearance: do NOT preserve or describe the control map's literal look (its edges, grey depth values, skeleton lines, grid colors, blockiness, or lack of color). Instead treat the input purely as STRUCTURAL guidance \u2014 outline/geometry (edges), spatial layout and relative distance (depth), body position and framing (pose), regional layout (segmentation/color grid), or coarse composition (low-res) \u2014 and describe the NEW, fully-rendered, realistic image the instruction asks for. Take the subject, colors, materials, textures, and style from the instruction's target description (defaulting to a realistic photograph unless it says otherwise), while keeping the structure/geometry/pose that the control map dictates. Every field (`objects`, `background_setting`, `style_medium`, `lighting`, etc.) must describe that intended final image \u2014 never the control map itself.\n\nAdhere strictly to the following structure and guidelines:\n\n1. **Inputs:** You will receive a reference image and a textual edit instruction. You will NOT receive a starting JSON.\n2. **Core Logic (Analyze and Synthesize):**\n * **Analyze:** First, deeply analyze the provided reference image. Identify its primary subject(s), their specific poses, expressions, and appearance. Also note the overall composition, lighting style, and artistic medium.\n * **Synthesize:** Next, interpret the textual edit instruction to understand what elements to keep from the reference and what to change. You will then construct a brand new JSON object from scratch that describes the desired final scene. For example, if the instruction is \"the same dog and pose, but at the beach,\" you must describe the dog from the reference image in the `objects` array but create a new `background_setting` for a beach, with appropriate `lighting` and `mood_atmosphere`. **The description of every unchanged element (objects, people, composition, camera settiLine truncated
},
"class_type": "PrimitiveStringMultiline",
"_meta": {
"title": "String (System Prompt VGL)"
}
},
"13": {
"inputs": {
"value": "Increase the clarity and resolution of the image, and refine it to obtain sharp, finely detailed, natural textures as if shot on a modern mirrorless camera, keeping the subject, pose, composition, colours and lighting the same. Remove any phone interface overlays."
},
"class_type": "PrimitiveStringMultiline",
"_meta": {
"title": "imageapi:prompt"
}
},
"14": {
"inputs": {
"prompt": [
"13",
0
],
"max_length": 3000,
"sampling_mode": "on",
"sampling_mode.temperature": 0.3,
"sampling_mode.top_k": 20,
"sampling_mode.top_p": 0.95,
"sampling_mode.min_p": 0.05,
"sampling_mode.repetition_penalty": 1.05,
"sampling_mode.seed": 1235,
"sampling_mode.presence_penalty": 0,
"thinking": false,
"use_default_template": true,
"mtp": "off",
"clip": [
"27",
0
],
"image": [
"26",
0
],
"system_prompt": [
"12",
0
]
},
"class_type": "TextGenerate",
"_meta": {
"title": "Generate Text (PE rewrite)"
}
},
"15": {
"inputs": {
"switch": false,
"on_false": [
"13",
0
],
"on_true": [
"14",
0
]
},
"class_type": "ComfySwitchNode",
"_meta": {
"title": "If/Else Switch (PE on/off)"
}
},
"24": {
"inputs": {
"compare_view": {
"__value__": [
"",
""
]
},
"image_a": [
"37",
0
],
"image_b": [
"6",
0
]
},
"class_type": "ImageCompare",
"_meta": {
"title": "Compare Images"
}
},
"25": {
"inputs": {
"source": [
"15",
0
]
},
"class_type": "PreviewAny",
"_meta": {
"title": "Preview as Text"
}
},
"26": {
"inputs": {
"images.image0": [
"37",
0
]
},
"class_type": "BatchImagesNode",
"_meta": {
"title": "Batch Images"
}
},
"27": {
"inputs": {
"clip_name": "qwen3vl_8b_int8_convrot.safetensors",
"type": "fibo2",
"device": "default"
},
"class_type": "CLIPLoader",
"_meta": {
"title": "Load CLIP (PE)"
}
},
"36": {
"inputs": {
"kind": "restored",
"subfolder": "aiedits",
"require_approval": true,
"images": [
"6",
0
],
"source": [
"37",
1
]
},
"class_type": "ImageApiSaveImageVersion",
"_meta": {
"title": "imageapi:save"
}
},
"37": {
"inputs": {
"path": ""
},
"class_type": "ImageApiLoadImageFromPath",
"_meta": {
"title": "imageapi:source"
}
}
}