Cameron CordesandClaude Opus 5.5
e2cf18678f
Templates with `imageapi:lora-slot:<n>:high|low` nodes take LoRAs picked
per job from a catalog built live from ComfyUI.
- Catalog: ComfyUI's LoraLoaderModelOnly `lora_name` choices (cached
60 s, re-read when a submit names an unknown id), paired by High/Low
name tokens (StyleHigh/StyleLow, x-HIGH/x-LOW, x_high_noise/...). Files
the workflow already loads outside the slots (lightx2v, switches) and
`[template.catalog] exclude` substrings never show; ids derive from the
names so they stay stable.
- POST /enhance `slots: [{lora, scale?}]` fills the slots in order: each
pick sets the slot nodes' lora_name and scales their exported strength;
unused slots run at 0, so the graph is never rewired. More picks than
slots, duplicates or unknown ids are 400s.
- GET /enhance/templates gains `slots` and `catalog: {loras, strength,
min, max}`; jobs store and return the picks (new enhance_jobs.lora_slots
JSON column).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
916 lines
32 KiB
Rust
916 lines
32 KiB
Rust
//! 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::lora_catalog::{self, CatalogEntry};
|
|
use super::path_map::PathMap;
|
|
use super::template::{SlotPick, Template, random_seed};
|
|
use crate::database::enhance_dao::{
|
|
EnhanceJobDao, EnhanceJobRow, JobUpdate, RecentPrompt, 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;
|
|
/// The jobs manager lists finished-but-unhandled jobs (ready, failed,
|
|
/// expired) for this long; active ones always.
|
|
pub const OPEN_JOBS_WINDOW_SECS: i64 = 3 * 24 * 3600;
|
|
/// Most jobs the jobs manager lists.
|
|
pub const OPEN_JOBS_LIMIT: i64 = 100;
|
|
|
|
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<()>,
|
|
/// ComfyUI's LoRA files and when they were read (#220).
|
|
lora_names: Mutex<Option<(Instant, Vec<String>)>>,
|
|
}
|
|
|
|
/// How long ComfyUI's LoRA list is reused; a new download shows up after it.
|
|
const LORA_LIST_TTL: Duration = Duration::from_secs(60);
|
|
|
|
/// One `POST /enhance { slots }` entry: a catalog id and an optional scale.
|
|
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
|
|
pub struct SlotRequest {
|
|
pub lora: String,
|
|
#[serde(default)]
|
|
pub scale: Option<f64>,
|
|
}
|
|
|
|
/// A resolved pick, as stored on the job and returned in its view.
|
|
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
|
pub struct SlotChoice {
|
|
pub lora: String,
|
|
pub label: String,
|
|
pub scale: f64,
|
|
}
|
|
|
|
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(()),
|
|
lora_names: Mutex::new(None),
|
|
})
|
|
}
|
|
|
|
/// ComfyUI's LoRA files, cached for [`LORA_LIST_TTL`]. `fresh` re-reads
|
|
/// them (a submit naming an id the cache doesn't know). An unreachable
|
|
/// ComfyUI keeps the last list.
|
|
async fn lora_names(&self, fresh: bool) -> Vec<String> {
|
|
let Some(client) = &self.client else {
|
|
return Vec::new();
|
|
};
|
|
if !fresh
|
|
&& let Some((at, names)) = &*self.lora_names.lock().unwrap()
|
|
&& at.elapsed() < LORA_LIST_TTL
|
|
{
|
|
return names.clone();
|
|
}
|
|
match client.lora_names().await {
|
|
Ok(names) => {
|
|
*self.lora_names.lock().unwrap() = Some((Instant::now(), names.clone()));
|
|
names
|
|
}
|
|
Err(e) => {
|
|
warn!("enhance: couldn't list ComfyUI's LoRAs: {e}");
|
|
self.lora_names
|
|
.lock()
|
|
.unwrap()
|
|
.as_ref()
|
|
.map(|(_, n)| n.clone())
|
|
.unwrap_or_default()
|
|
}
|
|
}
|
|
}
|
|
|
|
/// The LoRAs a template's slots can take; empty without slots.
|
|
pub async fn catalog(&self, template: &Template) -> Vec<CatalogEntry> {
|
|
if template.slot_count() == 0 {
|
|
return Vec::new();
|
|
}
|
|
let names = self.lora_names(false).await;
|
|
lora_catalog::build(
|
|
&names,
|
|
template.fixed_loras(),
|
|
&template.cfg.catalog.exclude,
|
|
)
|
|
}
|
|
|
|
/// Turn `slots` into picks for the template's slots, checking the ids
|
|
/// against the catalog (re-read once if one is unknown).
|
|
async fn resolve_slots(
|
|
&self,
|
|
template: &Template,
|
|
slots: &[SlotRequest],
|
|
) -> Result<(Vec<SlotPick>, Vec<SlotChoice>), EnhanceError> {
|
|
if slots.is_empty() {
|
|
return Ok((Vec::new(), Vec::new()));
|
|
}
|
|
if slots.len() > template.slot_count() {
|
|
return Err(EnhanceError::BadRequest(if template.slot_count() == 0 {
|
|
"this template doesn't take picked LoRAs".to_string()
|
|
} else {
|
|
format!(
|
|
"this template takes at most {} picked LoRAs",
|
|
template.slot_count()
|
|
)
|
|
}));
|
|
}
|
|
let mut seen = std::collections::HashSet::new();
|
|
if let Some(dup) = slots.iter().find(|s| !seen.insert(s.lora.as_str())) {
|
|
return Err(EnhanceError::BadRequest(format!(
|
|
"LoRA {:?} is picked twice",
|
|
dup.lora
|
|
)));
|
|
}
|
|
let mut catalog = self.catalog(template).await;
|
|
if slots
|
|
.iter()
|
|
.any(|s| !catalog.iter().any(|e| e.id == s.lora))
|
|
{
|
|
let names = self.lora_names(true).await;
|
|
catalog = lora_catalog::build(
|
|
&names,
|
|
template.fixed_loras(),
|
|
&template.cfg.catalog.exclude,
|
|
);
|
|
}
|
|
let mut picks = Vec::new();
|
|
let mut choices = Vec::new();
|
|
for s in slots {
|
|
let Some(entry) = catalog.iter().find(|e| e.id == s.lora) else {
|
|
return Err(EnhanceError::BadRequest(format!(
|
|
"no LoRA {:?} (it may have been removed from ComfyUI)",
|
|
s.lora
|
|
)));
|
|
};
|
|
let scale = template.cfg.catalog.resolve(s.scale);
|
|
picks.push(SlotPick {
|
|
high: entry.high.clone(),
|
|
low: entry.low.clone(),
|
|
scale,
|
|
});
|
|
choices.push(SlotChoice {
|
|
lora: entry.id.clone(),
|
|
label: entry.label.clone(),
|
|
scale,
|
|
});
|
|
}
|
|
Ok((picks, choices))
|
|
}
|
|
|
|
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>,
|
|
loras: &std::collections::BTreeMap<String, super::config::LoraRequest>,
|
|
slots: &[SlotRequest],
|
|
) -> 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 (picks, choices) = self.resolve_slots(template, slots).await?;
|
|
let filled = template
|
|
.fill(
|
|
&source,
|
|
prompt,
|
|
reference_path.as_deref(),
|
|
loras,
|
|
&picks,
|
|
&mut random_seed,
|
|
)
|
|
.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,
|
|
run_ms: None,
|
|
loras: (!filled.loras.is_empty())
|
|
.then(|| serde_json::to_string(&filled.loras).unwrap_or_default()),
|
|
lora_slots: (!choices.is_empty())
|
|
.then(|| serde_json::to_string(&choices).unwrap_or_default()),
|
|
};
|
|
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()))
|
|
}
|
|
|
|
/// Every photo's open jobs for the jobs manager, newest first. DB only:
|
|
/// the poller keeps the rows current, so this never calls ComfyUI.
|
|
pub fn open_jobs(&self) -> Result<Vec<EnhanceJobRow>, EnhanceError> {
|
|
let since = chrono::Utc::now().timestamp() - OPEN_JOBS_WINDOW_SECS;
|
|
self.dao
|
|
.lock()
|
|
.unwrap()
|
|
.list_open(since, OPEN_JOBS_LIMIT)
|
|
.map_err(|e| EnhanceError::Internal(e.to_string()))
|
|
}
|
|
|
|
/// Distinct recent prompts; DB only, so it works with ComfyUI down.
|
|
pub fn recent_prompts(
|
|
&self,
|
|
templates: Option<&[String]>,
|
|
limit: i64,
|
|
) -> Result<Vec<RecentPrompt>, EnhanceError> {
|
|
self.dao
|
|
.lock()
|
|
.unwrap()
|
|
.recent_prompts(templates, limit)
|
|
.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()),
|
|
);
|
|
u.run_ms = entry.run_ms;
|
|
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();
|
|
u.run_ms = entry.run_ms;
|
|
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));
|
|
}
|
|
}
|