use std::path::Path;
use mold_core::{GenerateRequest, OutputFormat, OutputMetadata};
use mold_db::{GenerationRecord, MetadataDb, RecordSource};
pub fn install_config_db_hooks() {
mold_db::config_sync::install_config_post_load_hook();
}
pub fn handle() -> Option<&'static MetadataDb> {
mold_db::global_db()
}
fn backend_label() -> Option<String> {
if cfg!(feature = "cuda") {
Some("cuda".into())
} else if cfg!(feature = "metal") {
Some("metal".into())
} else {
Some("cpu".into())
}
}
pub fn record_local_save(
saved_path: &Path,
req: &GenerateRequest,
seed_used: u64,
generation_time_ms: u64,
format: OutputFormat,
) -> bool {
let Some(db) = handle() else {
return false;
};
let abs = std::fs::canonicalize(saved_path).unwrap_or_else(|_| saved_path.to_path_buf());
let Some(filename) = abs
.file_name()
.and_then(|f| f.to_str())
.map(|s| s.to_string())
else {
return false;
};
let Some(output_dir) = abs.parent() else {
return false;
};
let metadata = OutputMetadata::from_generate_request(
req,
seed_used,
None,
mold_core::build_info::version_string(),
);
let now_ms = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or(0);
let mut rec = GenerationRecord::from_save(
output_dir,
filename,
format,
metadata,
RecordSource::Cli,
now_ms,
);
rec.stat_from_disk(&abs);
rec.generation_time_ms = Some(generation_time_ms as i64);
rec.backend = backend_label();
rec.hostname = hostname::get().ok().and_then(|s| s.into_string().ok());
if let Err(e) = db.upsert(&rec) {
tracing::warn!("metadata DB upsert failed for {}: {e:#}", abs.display());
return false;
}
true
}
#[cfg(test)]
mod tests {
use super::*;
use mold_core::{GenerateRequest, OutputFormat};
fn req() -> GenerateRequest {
serde_json::from_str(
r#"{
"prompt":"a stoic owl",
"model":"flux-dev:q4",
"width":1024,
"height":1024,
"steps":20,
"guidance":4.0
}"#,
)
.unwrap()
}
#[test]
fn round_trip_constructs_record_from_request_and_seed() {
let dir = tempfile::tempdir().unwrap();
let saved = dir.path().join("mold-flux-dev-q4-1.png");
std::fs::write(&saved, b"fake-bytes").unwrap();
let db = MetadataDb::open(&dir.path().join("mold.db")).unwrap();
let metadata = OutputMetadata::from_generate_request(
&req(),
42,
None,
mold_core::build_info::version_string(),
);
let mut rec = GenerationRecord::from_save(
dir.path(),
"mold-flux-dev-q4-1.png",
OutputFormat::Png,
metadata,
RecordSource::Cli,
1_700_000_000_000,
);
rec.stat_from_disk(&saved);
rec.generation_time_ms = Some(2_500);
rec.backend = backend_label();
db.upsert(&rec).unwrap();
let got = db
.get(dir.path(), "mold-flux-dev-q4-1.png")
.unwrap()
.unwrap();
assert_eq!(got.metadata.prompt, "a stoic owl");
assert_eq!(got.metadata.seed, 42);
assert_eq!(got.source, RecordSource::Cli);
assert_eq!(got.generation_time_ms, Some(2_500));
assert_eq!(got.file_size_bytes, Some(b"fake-bytes".len() as i64));
}
#[test]
fn backend_label_is_non_empty() {
let label = backend_label().unwrap();
assert!(["cuda", "metal", "cpu"].contains(&label.as_str()));
}
}