mirror of
https://github.com/storytold/lightcraft.git
synced 2026-10-14 19:27:01 -04:00
Merge PR #813
This commit is contained in:
10 files changed
+1017
-15
No files matched your search
Generated
+1
@@ -3106,6 +3106,7 @@ version = "0.6.0"
|
||||
dependencies = [
|
||||
"dhat",
|
||||
"lightcraft-codecs",
|
||||
"lightcraft-denoise",
|
||||
"lightcraft-engine",
|
||||
"lightcraft-mcp",
|
||||
"lightcraft-merge",
|
||||
|
||||
@@ -15,6 +15,7 @@ workspace = true
|
||||
|
||||
[dependencies]
|
||||
lightcraft-engine = { workspace = true }
|
||||
lightcraft-denoise = { workspace = true }
|
||||
lightcraft-sysmem = { workspace = true }
|
||||
lightcraft-mcp = { workspace = true }
|
||||
lightcraft-ui-egui = { workspace = true }
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
//! lightcraft-cli snapshot [--library DIR | --demo] [--script FILE.jsonl] [-o OUT.png] [--size WxH] [--scale S] [FILES…]
|
||||
//! lightcraft-cli merge hdr|panorama|hdr-panorama [OPTIONS] FILES…
|
||||
//! lightcraft-cli synth-merge hdr|panorama -o DIR
|
||||
//! lightcraft-cli denoise-serve [--model ID] [--models DIR] [--listen ADDR] [--token-env VAR] [--run-on auto|gpu|cpu]
|
||||
//! lightcraft-cli commands [--json]
|
||||
//! lightcraft-cli controls [--json]
|
||||
//! lightcraft-cli calibrate [--max N] [--out DIR] FOLDERS/FILES…
|
||||
@@ -85,6 +86,18 @@ USAGE:
|
||||
--preview OUT.png only render a ≤ 1024 px preview (nothing written next to the files)
|
||||
lightcraft-cli synth-merge hdr|panorama -o DIR
|
||||
Write synthetic merge inputs (procedural scene; bracketed DNGs or overlapping PNG views).
|
||||
lightcraft-cli denoise-serve [OPTIONS]
|
||||
Run AI Denoise tiles for LightCraft on other computers (LIGHTCRAFT_DENOISE_URL=<this host>:7990 and
|
||||
LIGHTCRAFT_DENOISE_TOKEN there; see docs/denoise.md). Serves the installed model selected in Settings,
|
||||
on this computer's graphics card when it can run it, else the CPU. Plain TCP: listen on a VPN/LAN
|
||||
address (e.g. the Tailscale IP) or firewall the port to it. Options:
|
||||
--model ID an installed model other than the selected one
|
||||
--models DIR the models folder (default: LIGHTCRAFT_DENOISE_MODELS, else this user's)
|
||||
--listen ADDR address:port to listen on (default 127.0.0.1:7990)
|
||||
--token-env VAR environment variable holding the shared token, ≥ 16 characters
|
||||
(default LIGHTCRAFT_DENOISE_TOKEN)
|
||||
--run-on WHERE auto (default), gpu or cpu
|
||||
--max-connections N concurrent clients (default 8)
|
||||
lightcraft-cli commands [--json] list every command id with its parameters
|
||||
lightcraft-cli controls [--json] list every develop control id with its range
|
||||
lightcraft-cli calibrate [--max N] [--out DIR] FOLDERS/FILES…
|
||||
@@ -148,6 +161,7 @@ fn main() -> ExitCode {
|
||||
Some("merge") => merge(&args[1..]),
|
||||
Some("synth-merge") => synth_merge(&args[1..]),
|
||||
Some("controls") => controls(&args[1..]),
|
||||
Some("denoise-serve") => denoise_serve(&args[1..]),
|
||||
Some("calibrate") => calibrate(&args[1..]),
|
||||
Some("--version" | "-V" | "version") => {
|
||||
println!("lightcraft-cli {}", env!("CARGO_PKG_VERSION"));
|
||||
@@ -900,3 +914,48 @@ fn controls(args: &[String]) -> Result<(), String> {
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `denoise-serve`: the remote half of AI Denoise (`lightcraft_denoise::remote`).
|
||||
#[cfg(feature = "denoise")]
|
||||
fn denoise_serve(args: &[String]) -> Result<(), String> {
|
||||
use lightcraft_denoise::remote::{DEFAULT_PORT, ServerOptions, serve};
|
||||
use lightcraft_engine::denoise::RunOn;
|
||||
let mut model: Option<String> = None;
|
||||
let mut dir: Option<String> = None;
|
||||
let mut listen = format!("127.0.0.1:{DEFAULT_PORT}");
|
||||
let mut token_env = "LIGHTCRAFT_DENOISE_TOKEN".to_string();
|
||||
let mut run_on = RunOn::Auto;
|
||||
let mut max_connections = 8usize;
|
||||
let mut i = 0;
|
||||
while i < args.len() {
|
||||
match args[i].as_str() {
|
||||
"--model" => model = Some(take_value(args, &mut i, "--model")?.to_string()),
|
||||
"--models" => dir = Some(take_value(args, &mut i, "--models")?.to_string()),
|
||||
"--listen" => listen = take_value(args, &mut i, "--listen")?.to_string(),
|
||||
"--token-env" => token_env = take_value(args, &mut i, "--token-env")?.to_string(),
|
||||
"--run-on" => {
|
||||
let v = take_value(args, &mut i, "--run-on")?;
|
||||
run_on = RunOn::parse(v).ok_or_else(|| format!("--run-on expects auto, gpu or cpu, not `{v}`"))?;
|
||||
}
|
||||
"--max-connections" => {
|
||||
max_connections = take_value(args, &mut i, "--max-connections")?.parse().map_err(|_| "--max-connections expects a number")?
|
||||
}
|
||||
other => return Err(format!("denoise-serve: unknown option `{other}`")),
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
let token = std::env::var(&token_env).map_err(|_| format!("denoise-serve: set the shared token in ${token_env}"))?;
|
||||
let s = lightcraft_engine::denoise::serving(dir.as_deref().map(Path::new), model.as_deref(), run_on)?;
|
||||
let listener = std::net::TcpListener::bind(&listen).map_err(|e| format!("denoise-serve: {listen}: {e}"))?;
|
||||
eprintln!("lightcraft-cli denoise-serve: {} {} ({}) on {listen} via {}", s.model.id, s.model.version, s.model.sha256, s.device);
|
||||
if listen.starts_with("0.0.0.0") || listen.starts_with("[::]") {
|
||||
eprintln!("lightcraft-cli denoise-serve: listening on every interface; firewall the port to your VPN/LAN (the protocol is not encrypted)");
|
||||
}
|
||||
let opts = ServerOptions { token, max_connections, ..Default::default() };
|
||||
serve(listener, s.model, s.tile, s.device, s.run, opts, std::sync::Arc::new(|m: &str| eprintln!("lightcraft-cli denoise-serve: {m}")))
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "denoise"))]
|
||||
fn denoise_serve(_: &[String]) -> Result<(), String> {
|
||||
Err("denoise-serve: this build has no denoise support (build with the `denoise` feature)".into())
|
||||
}
|
||||
@@ -577,3 +577,69 @@ fn snapshot_leaves_desktop_gpu_preferences_untouched() {
|
||||
assert!(image.exists());
|
||||
}
|
||||
}
|
||||
|
||||
/// `denoise-serve` serves the selected installed model; a client greets it and gets the CPU's answer for a tile
|
||||
/// (within half-float rounding on the wire), and a wrong token or a missing token variable is refused.
|
||||
#[cfg(feature = "denoise")]
|
||||
#[test]
|
||||
fn denoise_serve_runs_tiles_for_another_computer() {
|
||||
use lightcraft_denoise::manifest::{DenoiserManifest, Domain, Gain};
|
||||
use lightcraft_denoise::remote::Remote;
|
||||
use lightcraft_denoise::run::TileRunner;
|
||||
let models = tmp("denoise-models");
|
||||
let _ = std::fs::remove_dir_all(&models);
|
||||
let dir = models.join("synthetic");
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let manifest = DenoiserManifest {
|
||||
id: "synthetic".into(),
|
||||
name: "Synthetic".into(),
|
||||
version: "1".into(),
|
||||
licence: Default::default(),
|
||||
source: None,
|
||||
sha256: None,
|
||||
size_bytes: None,
|
||||
provenance: "made for a test".into(),
|
||||
domain: Domain::BayerToRgb,
|
||||
tile: 64,
|
||||
overlap: 16,
|
||||
gain: Gain::None,
|
||||
};
|
||||
let onnx = dir.join("model.onnx");
|
||||
std::fs::write(&onnx, lightcraft_denoise::synthetic::unet_onnx(64, 8, 2, 7)).unwrap();
|
||||
std::fs::write(dir.join("denoise-model.json"), serde_json::to_vec(&manifest).unwrap()).unwrap();
|
||||
std::fs::write(models.join("settings.json"), r#"{"model": "synthetic"}"#).unwrap();
|
||||
let sha = lightcraft_denoise::hash::sha256_file(&onnx).unwrap();
|
||||
|
||||
// without its token the server doesn't start
|
||||
let o = Command::new(BIN).args(["denoise-serve", "--models", models.to_str().unwrap(), "--token-env", "LC_TEST_NO_SUCH_VAR"]).output().unwrap();
|
||||
assert!(!o.status.success() && String::from_utf8_lossy(&o.stderr).contains("LC_TEST_NO_SUCH_VAR"));
|
||||
|
||||
let port = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port();
|
||||
let addr = format!("127.0.0.1:{port}");
|
||||
let token = "cli-test-token-0123456789";
|
||||
let mut child = Command::new(BIN)
|
||||
.args(["denoise-serve", "--models", models.to_str().unwrap(), "--listen", &addr, "--run-on", "cpu"])
|
||||
.env("LIGHTCRAFT_DENOISE_TOKEN", token)
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.unwrap();
|
||||
// wait for the banner
|
||||
let mut err = BufReader::new(child.stderr.take().unwrap());
|
||||
let mut banner = String::new();
|
||||
err.read_line(&mut banner).unwrap();
|
||||
assert!(banner.contains("synthetic 1") && banner.contains(&addr), "{banner}");
|
||||
|
||||
let remote = Remote::new(&addr, token, &sha, 64);
|
||||
let info = remote.hello().unwrap();
|
||||
assert!(info.device.starts_with("CPU"), "{info:?}");
|
||||
let (input, _) = lightcraft_denoise::runtime::test_tile(64);
|
||||
let got = remote.try_run(&input).unwrap().unwrap();
|
||||
let want = lightcraft_denoise::runtime::CpuRunner::load(&onnx, &manifest).unwrap().run(&input).unwrap();
|
||||
let scale = want.iter().fold(1e-9f32, |m, v| m.max(v.abs()));
|
||||
let worst = want.iter().zip(&got).fold(0f32, |m, (a, b)| m.max((a - b).abs())) / scale;
|
||||
assert!(worst < 5e-3, "remote vs local: {worst}");
|
||||
assert!(Remote::new(&addr, "a-wrong-token-000000", &sha, 64).hello().is_err());
|
||||
let _ = child.kill();
|
||||
let _ = child.wait();
|
||||
let _ = std::fs::remove_dir_all(&models);
|
||||
}
|
||||
@@ -12,6 +12,7 @@
|
||||
//! blend with the plain demosaic that the Amount slider controls.
|
||||
//! - [`cpu`] and [`runtime`]: a single-thread pure-Rust CPU runner and its installation self-test.
|
||||
//! - [`net`] and [`onnx`]: checked plain network data and bounded ONNX decoding, shared by CPU and GPU.
|
||||
//! - [`remote`]: running tiles on another computer's graphics card over TCP (a [`run::TileRunner`] and its server).
|
||||
//!
|
||||
//! Everything here treats model files, manifests and cache files as hostile input: sizes are capped, numbers must
|
||||
//! be finite, and a malformed file is an error, never a panic.
|
||||
@@ -26,6 +27,7 @@ pub mod onnx;
|
||||
mod onnx_proto;
|
||||
|
||||
pub mod reference;
|
||||
pub mod remote;
|
||||
|
||||
#[cfg(feature = "runtime")]
|
||||
pub mod runtime;
|
||||
|
||||
@@ -0,0 +1,562 @@
|
||||
//! Remote denoise: run a model's tiles on another computer's graphics card over TCP. Meant for a LAN or a VPN such as
|
||||
//! Tailscale or WireGuard: the connection itself is not encrypted.
|
||||
//!
|
||||
//! [`Remote`] is a [`TileRunner`] that sends each tile to a server; [`serve`] runs one (`lightcraft-cli
|
||||
//! denoise-serve`). Both ends must run the same model, checked by its identity ([`ModelId`]) when they greet.
|
||||
//!
|
||||
//! Protocol (version 1). Each message is one JSON line (`\n`-terminated, ≤ 4 KiB), optionally followed by a payload of
|
||||
//! `bytes` bytes: the tile's samples as little-endian half floats, byte-shuffled (every sample's low byte, then every
|
||||
//! high byte) and deflated, which roughly halves what crosses a slow link.
|
||||
//!
|
||||
//! ```text
|
||||
//! → {"v":1,"token":"…","op":"hello"}
|
||||
//! ← {"ok":true,"v":1,"device":"GPU: …","model":{"id":"…","version":"…","sha256":"…"}}
|
||||
//! → {"v":1,"token":"…","op":"run","model":"<sha256>","tile":T,"bytes":N} + N bytes (4 planes of T × T)
|
||||
//! ← {"ok":true,"tile":T,"bytes":N} + N bytes (3 planes of 2T × 2T)
|
||||
//! ← {"ok":false,"error":"…"} (any failure)
|
||||
//! ```
|
||||
//!
|
||||
//! A connection serves any number of requests. The server refuses a wrong token, another model than its own, tiles of
|
||||
//! another size, oversized headers and payloads (also ones that inflate past their tile), and caps concurrent
|
||||
//! connections; the client gives up on a server that keeps failing, so the computer's own runner does the rest.
|
||||
|
||||
use std::io::{BufRead, BufReader, Read, Write};
|
||||
use std::net::{TcpListener, TcpStream, ToSocketAddrs};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex, PoisonError};
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::Serialize;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::product::{from_f16, to_f16};
|
||||
use crate::run::{Error, TileRunner};
|
||||
|
||||
pub const PROTOCOL: u64 = 1;
|
||||
pub const DEFAULT_PORT: u16 = 7990;
|
||||
const MAX_LINE: usize = 4096;
|
||||
/// Largest tile side in cells (a 1024-cell tile is a 2048 × 2048 pixel answer).
|
||||
pub const MAX_TILE: usize = 1024;
|
||||
/// Shortest token a server accepts.
|
||||
pub const MIN_TOKEN: usize = 16;
|
||||
|
||||
/// Which model a server runs: both ends must agree before tiles cross.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
|
||||
pub struct ModelId {
|
||||
pub id: String,
|
||||
pub version: String,
|
||||
/// SHA-256 of the model file (hex).
|
||||
pub sha256: String,
|
||||
}
|
||||
|
||||
fn bad(m: impl Into<String>) -> std::io::Error {
|
||||
std::io::Error::new(std::io::ErrorKind::InvalidData, m.into())
|
||||
}
|
||||
|
||||
fn read_line(r: &mut impl BufRead) -> std::io::Result<Option<String>> {
|
||||
let mut buf = Vec::new();
|
||||
let n = r.by_ref().take(MAX_LINE as u64 + 1).read_until(b'\n', &mut buf)?;
|
||||
if n == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
if buf.last() != Some(&b'\n') {
|
||||
return Err(bad("message header too long or truncated"));
|
||||
}
|
||||
String::from_utf8(buf).map(Some).map_err(|_| bad("header is not UTF-8"))
|
||||
}
|
||||
|
||||
/// Samples → half floats → byte-shuffled → deflated.
|
||||
pub fn pack(data: &[f32]) -> Vec<u8> {
|
||||
let n = data.len();
|
||||
let mut b = vec![0u8; n * 2];
|
||||
let (lo, hi) = b.split_at_mut(n);
|
||||
for ((v, l), h) in data.iter().zip(lo.iter_mut()).zip(hi.iter_mut()) {
|
||||
[*l, *h] = to_f16(*v).to_le_bytes();
|
||||
}
|
||||
miniz_oxide::deflate::compress_to_vec(&b, 1)
|
||||
}
|
||||
|
||||
/// The inverse of [`pack`] for exactly `n` samples; refuses anything that inflates to more.
|
||||
pub fn unpack(bytes: &[u8], n: usize) -> std::io::Result<Vec<f32>> {
|
||||
let b = miniz_oxide::inflate::decompress_to_vec_with_limit(bytes, n * 2).map_err(|_| bad("payload doesn't inflate to the tile size"))?;
|
||||
if b.len() != n * 2 {
|
||||
return Err(bad("payload size doesn't match the tile"));
|
||||
}
|
||||
let (lo, hi) = b.split_at(n);
|
||||
Ok(lo.iter().zip(hi).map(|(l, h)| from_f16(u16::from_le_bytes([*l, *h]))).collect())
|
||||
}
|
||||
|
||||
fn write_payload(w: &mut impl Write, mut header: Value, data: &[f32]) -> std::io::Result<()> {
|
||||
let packed = pack(data);
|
||||
header["bytes"] = json!(packed.len());
|
||||
writeln!(w, "{header}")?;
|
||||
w.write_all(&packed)
|
||||
}
|
||||
|
||||
/// Read a `bytes`-byte payload holding `n` samples (`bytes` is capped by the tile, so a hostile header can't make us
|
||||
/// allocate more than an uncompressed tile plus slack).
|
||||
fn read_payload(r: &mut impl Read, bytes: usize, n: usize) -> std::io::Result<Vec<f32>> {
|
||||
if bytes > n * 2 + 4096 {
|
||||
return Err(bad("payload larger than its tile"));
|
||||
}
|
||||
let mut b = vec![0u8; bytes];
|
||||
r.read_exact(&mut b)?;
|
||||
unpack(&b, n)
|
||||
}
|
||||
|
||||
fn field(v: &Value, k: &str) -> Option<usize> {
|
||||
v.get(k).and_then(Value::as_u64).and_then(|x| usize::try_from(x).ok())
|
||||
}
|
||||
|
||||
/// Constant-time comparison (the token must not leak through timing).
|
||||
fn same(a: &str, b: &str) -> bool {
|
||||
let (a, b) = (a.as_bytes(), b.as_bytes());
|
||||
let mut d = a.len() ^ b.len();
|
||||
for i in 0..a.len().max(b.len()) {
|
||||
d |= (a.get(i).copied().unwrap_or(0) ^ b.get(i).copied().unwrap_or(0)) as usize;
|
||||
}
|
||||
d == 0
|
||||
}
|
||||
|
||||
/// What a server runs a tile with: four planes of `tile × tile` cells in, three planes of `2·tile × 2·tile` out.
|
||||
pub type RunTile = Arc<dyn Fn(&[f32]) -> Result<Vec<f32>, Error> + Send + Sync>;
|
||||
|
||||
/// A server's settings.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ServerOptions {
|
||||
/// Shared secret every request must carry (at least [`MIN_TOKEN`] characters).
|
||||
pub token: String,
|
||||
pub max_connections: usize,
|
||||
/// Idle and transfer limit per connection.
|
||||
pub timeout: Duration,
|
||||
}
|
||||
|
||||
impl Default for ServerOptions {
|
||||
fn default() -> Self {
|
||||
ServerOptions { token: String::new(), max_connections: 8, timeout: Duration::from_secs(120) }
|
||||
}
|
||||
}
|
||||
|
||||
/// Serve `model` (tiles of `tile` cells, run by `run` on `device`) until the listener fails. `log` gets one line per event.
|
||||
pub fn serve(
|
||||
listener: TcpListener,
|
||||
model: ModelId,
|
||||
tile: usize,
|
||||
device: String,
|
||||
run: RunTile,
|
||||
opts: ServerOptions,
|
||||
log: Arc<dyn Fn(&str) + Send + Sync>,
|
||||
) -> Result<(), String> {
|
||||
if opts.token.trim().chars().count() < MIN_TOKEN {
|
||||
return Err(format!("the token must be at least {MIN_TOKEN} characters"));
|
||||
}
|
||||
if tile == 0 || tile > MAX_TILE {
|
||||
return Err(format!("tile size {tile} is outside 1..={MAX_TILE}"));
|
||||
}
|
||||
let active = Arc::new(AtomicUsize::new(0));
|
||||
let shared = Arc::new((model, device));
|
||||
for conn in listener.incoming() {
|
||||
let Ok(stream) = conn else { continue };
|
||||
let peer = stream.peer_addr().map(|a| a.to_string()).unwrap_or_else(|_| "?".into());
|
||||
if active.load(Ordering::SeqCst) >= opts.max_connections.max(1) {
|
||||
let mut s = stream;
|
||||
let _ = writeln!(s, "{}", json!({"ok": false, "error": "server busy"}));
|
||||
log(&format!("{peer}: refused (busy)"));
|
||||
continue;
|
||||
}
|
||||
active.fetch_add(1, Ordering::SeqCst);
|
||||
let (sh, r, o, l, n) = (shared.clone(), run.clone(), opts.clone(), log.clone(), active.clone());
|
||||
let spawned = std::thread::Builder::new().name(format!("denoise {peer}")).spawn(move || {
|
||||
if let Err(e) = handle(stream, &sh.0, tile, &sh.1, &*r, &o, &*l) {
|
||||
l(&format!("{peer}: {e}"));
|
||||
}
|
||||
n.fetch_sub(1, Ordering::SeqCst);
|
||||
});
|
||||
if spawned.is_err() {
|
||||
active.fetch_sub(1, Ordering::SeqCst);
|
||||
log("could not start a connection thread");
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle(
|
||||
stream: TcpStream,
|
||||
model: &ModelId,
|
||||
tile: usize,
|
||||
device: &str,
|
||||
run: &dyn Fn(&[f32]) -> Result<Vec<f32>, Error>,
|
||||
opts: &ServerOptions,
|
||||
log: &dyn Fn(&str),
|
||||
) -> std::io::Result<()> {
|
||||
stream.set_read_timeout(Some(opts.timeout))?;
|
||||
stream.set_write_timeout(Some(opts.timeout))?;
|
||||
stream.set_nodelay(true)?;
|
||||
let peer = stream.peer_addr().map(|a| a.to_string()).unwrap_or_default();
|
||||
let mut r = BufReader::new(stream.try_clone()?);
|
||||
let mut w = stream;
|
||||
let mut tiles = 0usize;
|
||||
let fail = |w: &mut TcpStream, msg: &str| writeln!(w, "{}", json!({"ok": false, "error": msg}));
|
||||
while let Some(line) = read_line(&mut r)? {
|
||||
let Ok(req) = serde_json::from_str::<Value>(&line) else {
|
||||
fail(&mut w, "malformed request")?;
|
||||
break;
|
||||
};
|
||||
if req.get("v").and_then(Value::as_u64) != Some(PROTOCOL) {
|
||||
fail(&mut w, &format!("protocol version {PROTOCOL} only"))?;
|
||||
break;
|
||||
}
|
||||
if !same(req.get("token").and_then(Value::as_str).unwrap_or(""), &opts.token) {
|
||||
fail(&mut w, "bad token")?;
|
||||
log(&format!("{peer}: bad token"));
|
||||
break;
|
||||
}
|
||||
match req.get("op").and_then(Value::as_str) {
|
||||
Some("hello") => writeln!(w, "{}", json!({"ok": true, "v": PROTOCOL, "device": device, "model": model, "tile": tile}))?,
|
||||
Some("run") => {
|
||||
let (Some(t), Some(bytes)) = (field(&req, "tile"), field(&req, "bytes")) else {
|
||||
fail(&mut w, "missing tile size or payload size")?;
|
||||
break;
|
||||
};
|
||||
if t != tile {
|
||||
fail(&mut w, &format!("this server runs {tile}-cell tiles, not {t}"))?;
|
||||
break;
|
||||
}
|
||||
// the payload is read before any refusal below, so the stream stays in step
|
||||
let input = match read_payload(&mut r, bytes, 4 * t * t) {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
fail(&mut w, &e.to_string())?;
|
||||
break;
|
||||
}
|
||||
};
|
||||
if req.get("model").and_then(Value::as_str) != Some(model.sha256.as_str()) {
|
||||
fail(&mut w, &format!("model mismatch: this server runs {} {}", model.id, model.version))?;
|
||||
continue;
|
||||
}
|
||||
match run(&input) {
|
||||
Ok(out) if out.len() == 3 * 4 * t * t && out.iter().all(|v| v.is_finite()) => {
|
||||
write_payload(&mut w, json!({"ok": true, "tile": t}), &out)?;
|
||||
tiles += 1;
|
||||
}
|
||||
Ok(_) => fail(&mut w, "the model's answer has the wrong size or numbers that are not finite")?,
|
||||
Err(e) => fail(&mut w, &e.to_string())?,
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
fail(&mut w, "unknown op")?;
|
||||
break;
|
||||
}
|
||||
}
|
||||
w.flush()?;
|
||||
}
|
||||
if tiles > 0 {
|
||||
log(&format!("{peer}: {tiles} tiles"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// What a server said in its greeting.
|
||||
#[derive(Clone, Debug, PartialEq, Serialize)]
|
||||
pub struct ServerInfo {
|
||||
pub address: String,
|
||||
pub device: String,
|
||||
pub model: ModelId,
|
||||
pub tile: usize,
|
||||
}
|
||||
|
||||
/// Tiles the server may get wrong (an error, a broken connection) before it stops being used.
|
||||
pub const MAX_FAILURES: usize = 3;
|
||||
|
||||
/// A server as a [`TileRunner`], for one model: keeps up to `slots` connections open and gives up on the server after
|
||||
/// [`MAX_FAILURES`] failed tiles (the caller runs those tiles itself).
|
||||
pub struct Remote {
|
||||
address: String,
|
||||
token: String,
|
||||
model_sha256: String,
|
||||
tile: usize,
|
||||
slots: usize,
|
||||
timeout: Duration,
|
||||
busy: AtomicUsize,
|
||||
failures: AtomicUsize,
|
||||
last_error: Mutex<Option<String>>,
|
||||
pool: Mutex<Vec<(BufReader<TcpStream>, TcpStream)>>,
|
||||
}
|
||||
|
||||
impl Remote {
|
||||
/// `address` as `host[:port]` (default port [`DEFAULT_PORT`]; `tcp://` and a trailing `/` are ignored).
|
||||
pub fn new(address: &str, token: &str, model_sha256: &str, tile: usize) -> Remote {
|
||||
let a = address.trim().trim_start_matches("tcp://").trim_end_matches('/');
|
||||
let has_port = a.rsplit_once(':').is_some_and(|(h, p)| !h.is_empty() && !h.ends_with(':') && p.parse::<u16>().is_ok()) && !a.ends_with(']');
|
||||
let address = if has_port { a.to_string() } else { format!("{a}:{DEFAULT_PORT}") };
|
||||
Remote {
|
||||
address,
|
||||
token: token.into(),
|
||||
model_sha256: model_sha256.to_ascii_lowercase(),
|
||||
tile,
|
||||
slots: 4,
|
||||
timeout: Duration::from_secs(60),
|
||||
busy: AtomicUsize::new(0),
|
||||
failures: AtomicUsize::new(0),
|
||||
last_error: Mutex::new(None),
|
||||
pool: Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_slots(mut self, n: usize) -> Remote {
|
||||
self.slots = n.clamp(1, 16);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn address(&self) -> &str {
|
||||
&self.address
|
||||
}
|
||||
|
||||
/// Tiles it can take at once.
|
||||
pub fn slots(&self) -> usize {
|
||||
self.slots
|
||||
}
|
||||
|
||||
/// Still used (fewer than [`MAX_FAILURES`] failed tiles).
|
||||
pub fn healthy(&self) -> bool {
|
||||
self.failures.load(Ordering::Relaxed) < MAX_FAILURES
|
||||
}
|
||||
|
||||
pub fn failures(&self) -> usize {
|
||||
self.failures.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Why the last tile failed.
|
||||
pub fn last_error(&self) -> Option<String> {
|
||||
self.last_error.lock().unwrap_or_else(PoisonError::into_inner).clone()
|
||||
}
|
||||
|
||||
fn err(&self, reason: impl std::fmt::Display) -> Error {
|
||||
Error::Runtime(format!("remote {}: {reason}", self.address))
|
||||
}
|
||||
|
||||
fn connect(&self) -> Result<(BufReader<TcpStream>, TcpStream), Error> {
|
||||
let addr = self.address.to_socket_addrs().map_err(|e| self.err(e))?.next().ok_or_else(|| self.err("the address doesn't resolve"))?;
|
||||
let s = TcpStream::connect_timeout(&addr, Duration::from_secs(5)).map_err(|e| self.err(e))?;
|
||||
s.set_read_timeout(Some(self.timeout)).map_err(|e| self.err(e))?;
|
||||
s.set_write_timeout(Some(self.timeout)).map_err(|e| self.err(e))?;
|
||||
s.set_nodelay(true).map_err(|e| self.err(e))?;
|
||||
let r = BufReader::new(s.try_clone().map_err(|e| self.err(e))?);
|
||||
Ok((r, s))
|
||||
}
|
||||
|
||||
fn reply(&self, r: &mut BufReader<TcpStream>) -> Result<Value, Error> {
|
||||
let line = read_line(r).map_err(|e| self.err(e))?.ok_or_else(|| self.err("the connection closed"))?;
|
||||
let v: Value = serde_json::from_str(&line).map_err(|_| self.err("malformed reply"))?;
|
||||
if v.get("ok").and_then(Value::as_bool) != Some(true) {
|
||||
return Err(self.err(v.get("error").and_then(Value::as_str).unwrap_or("request failed")));
|
||||
}
|
||||
Ok(v)
|
||||
}
|
||||
|
||||
/// Greet the server: its device and model. Fails when it runs another model or tile size than this one.
|
||||
pub fn hello(&self) -> Result<ServerInfo, Error> {
|
||||
let (mut r, mut s) = self.connect()?;
|
||||
writeln!(s, "{}", json!({"v": PROTOCOL, "token": self.token, "op": "hello"})).map_err(|e| self.err(e))?;
|
||||
let v = self.reply(&mut r)?;
|
||||
let m = v.get("model").cloned().unwrap_or_default();
|
||||
let text = |o: &Value, k: &str| o.get(k).and_then(Value::as_str).unwrap_or("").to_string();
|
||||
let info = ServerInfo {
|
||||
address: self.address.clone(),
|
||||
device: text(&v, "device"),
|
||||
model: ModelId { id: text(&m, "id"), version: text(&m, "version"), sha256: text(&m, "sha256") },
|
||||
tile: field(&v, "tile").unwrap_or(0),
|
||||
};
|
||||
if !info.model.sha256.eq_ignore_ascii_case(&self.model_sha256) {
|
||||
return Err(self.err(format!("it runs another model ({} {})", info.model.id, info.model.version)));
|
||||
}
|
||||
if info.tile != self.tile {
|
||||
return Err(self.err(format!("it runs {}-cell tiles, this model {}", info.tile, self.tile)));
|
||||
}
|
||||
self.pool.lock().unwrap_or_else(PoisonError::into_inner).push((r, s));
|
||||
Ok(info)
|
||||
}
|
||||
|
||||
/// Run a tile if a slot is free right now; `None` when all are busy or the server was given up (the caller runs
|
||||
/// the tile elsewhere). A failure counts towards giving up on the server.
|
||||
pub fn try_run(&self, input: &[f32]) -> Option<Result<Vec<f32>, Error>> {
|
||||
if !self.healthy() {
|
||||
return None;
|
||||
}
|
||||
if self.busy.fetch_add(1, Ordering::SeqCst) >= self.slots {
|
||||
self.busy.fetch_sub(1, Ordering::SeqCst);
|
||||
return None;
|
||||
}
|
||||
let r = self.run_now(input);
|
||||
self.busy.fetch_sub(1, Ordering::SeqCst);
|
||||
if let Err(e) = &r {
|
||||
self.failures.fetch_add(1, Ordering::Relaxed);
|
||||
*self.last_error.lock().unwrap_or_else(PoisonError::into_inner) = Some(e.to_string());
|
||||
}
|
||||
Some(r)
|
||||
}
|
||||
|
||||
fn run_now(&self, input: &[f32]) -> Result<Vec<f32>, Error> {
|
||||
let t = self.tile;
|
||||
if input.len() != 4 * t * t {
|
||||
return Err(Error::Input(format!("a tile of {} samples for {t}-cell tiles", input.len())));
|
||||
}
|
||||
let pooled = self.pool.lock().unwrap_or_else(PoisonError::into_inner).pop();
|
||||
let (mut r, mut s) = match pooled {
|
||||
Some(c) => c,
|
||||
None => self.connect()?,
|
||||
};
|
||||
let req = json!({"v": PROTOCOL, "token": self.token, "op": "run", "model": self.model_sha256, "tile": t});
|
||||
write_payload(&mut s, req, input).map_err(|e| self.err(e))?;
|
||||
s.flush().map_err(|e| self.err(e))?;
|
||||
let v = self.reply(&mut r)?;
|
||||
if field(&v, "tile") != Some(t) {
|
||||
return Err(self.err("the answer is for another tile size"));
|
||||
}
|
||||
let bytes = field(&v, "bytes").ok_or_else(|| self.err("an answer without its size"))?;
|
||||
let out = read_payload(&mut r, bytes, 3 * 4 * t * t).map_err(|e| self.err(e))?;
|
||||
self.pool.lock().unwrap_or_else(PoisonError::into_inner).push((r, s));
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
/// Tiles on the server when one of its slots is free, else (and for any tile it fails) on `local`: the server adds
|
||||
/// speed to this computer's runner instead of replacing it.
|
||||
pub struct Shared<'a> {
|
||||
pub remote: &'a Remote,
|
||||
pub local: &'a dyn TileRunner,
|
||||
}
|
||||
|
||||
impl TileRunner for Shared<'_> {
|
||||
fn run(&self, input: &[f32]) -> Result<Vec<f32>, Error> {
|
||||
match self.remote.try_run(input) {
|
||||
Some(Ok(out)) => Ok(out),
|
||||
Some(Err(_)) | None => self.local.run(input),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const TOKEN: &str = "0123456789abcdef-test";
|
||||
const T: usize = 8;
|
||||
|
||||
/// A stand-in model: each output pixel is its cell's first plane, times `k`.
|
||||
fn model(k: f32) -> RunTile {
|
||||
Arc::new(move |input: &[f32]| {
|
||||
let t = T;
|
||||
let mut out = vec![0f32; 3 * 4 * t * t];
|
||||
for c in 0..3 {
|
||||
for y in 0..2 * t {
|
||||
for x in 0..2 * t {
|
||||
out[(c * 2 * t + y) * 2 * t + x] = input[(y / 2) * t + x / 2] * k;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
})
|
||||
}
|
||||
|
||||
fn id(sha: &str) -> ModelId {
|
||||
ModelId { id: "test".into(), version: "1".into(), sha256: sha.into() }
|
||||
}
|
||||
|
||||
fn start(sha: &str, run: RunTile) -> String {
|
||||
let l = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let addr = l.local_addr().unwrap().to_string();
|
||||
let opts = ServerOptions { token: TOKEN.into(), ..Default::default() };
|
||||
let m = id(sha);
|
||||
std::thread::spawn(move || serve(l, m, T, "CPU (test)".into(), run, opts, Arc::new(|_| {})));
|
||||
addr
|
||||
}
|
||||
|
||||
struct Local(RunTile, AtomicUsize);
|
||||
impl TileRunner for Local {
|
||||
fn run(&self, input: &[f32]) -> Result<Vec<f32>, Error> {
|
||||
self.1.fetch_add(1, Ordering::SeqCst);
|
||||
(self.0)(input)
|
||||
}
|
||||
}
|
||||
|
||||
fn tile_input(seed: f32) -> Vec<f32> {
|
||||
(0..4 * T * T).map(|i| (i as f32 * 0.37 + seed).sin() * 0.5 + 0.5).collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remote_tiles_match_local_ones_within_half_float_rounding() {
|
||||
let addr = start("aa11", model(2.0));
|
||||
let remote = Remote::new(&addr, TOKEN, "AA11", T);
|
||||
let info = remote.hello().unwrap();
|
||||
assert_eq!((info.device.as_str(), info.tile), ("CPU (test)", T));
|
||||
let input = tile_input(0.3);
|
||||
let got = remote.try_run(&input).unwrap().unwrap();
|
||||
let want = model(2.0)(&input).unwrap();
|
||||
assert!(got.iter().zip(&want).all(|(a, b)| (a - b).abs() < 2e-3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shared_runner_falls_back_and_gives_up_on_a_failing_server() {
|
||||
let local = Local(model(2.0), AtomicUsize::new(0));
|
||||
// nobody listening: every tile runs locally, and after a few the server is given up
|
||||
let dead = Remote::new("127.0.0.1:1", TOKEN, "aa11", T);
|
||||
let shared = Shared { remote: &dead, local: &local };
|
||||
for i in 0..6 {
|
||||
assert!(shared.run(&tile_input(i as f32)).is_ok());
|
||||
}
|
||||
assert_eq!(local.1.load(Ordering::SeqCst), 6);
|
||||
assert!(!dead.healthy() && dead.last_error().is_some());
|
||||
// a server with another model is refused at the greeting and per tile
|
||||
let addr = start("bb22", model(2.0));
|
||||
let other = Remote::new(&addr, TOKEN, "aa11", T);
|
||||
assert!(other.hello().unwrap_err().to_string().contains("another model"));
|
||||
assert!(other.try_run(&tile_input(0.0)).unwrap().unwrap_err().to_string().contains("model mismatch"));
|
||||
// a wrong token
|
||||
let addr = start("aa11", model(2.0));
|
||||
assert!(Remote::new(&addr, "wrong-token-0000000", "aa11", T).hello().unwrap_err().to_string().contains("bad token"));
|
||||
// all slots busy: the tile runs locally without waiting
|
||||
let busy = Remote::new(&addr, TOKEN, "aa11", T).with_slots(1);
|
||||
busy.busy.store(1, Ordering::SeqCst);
|
||||
assert!(busy.try_run(&tile_input(0.0)).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hostile_requests_are_refused_and_the_server_keeps_working() {
|
||||
let addr = start("aa11", model(1.0));
|
||||
let talk = |bytes: &[u8]| -> String {
|
||||
let mut s = TcpStream::connect(&addr).unwrap();
|
||||
s.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
|
||||
s.write_all(bytes).unwrap();
|
||||
let mut line = String::new();
|
||||
let _ = BufReader::new(s).read_line(&mut line);
|
||||
line
|
||||
};
|
||||
assert!(talk(b"not json\n").contains("malformed"));
|
||||
let v2 = format!("{}\n", json!({"v": 2, "token": TOKEN, "op": "hello"}));
|
||||
assert!(talk(v2.as_bytes()).contains("protocol"));
|
||||
let big = format!("{}\n", json!({"v": 1, "token": TOKEN, "op": "run", "model": "aa11", "tile": 4096, "bytes": 10}));
|
||||
assert!(talk(big.as_bytes()).contains("8-cell"));
|
||||
// a payload that inflates past its tile
|
||||
let bomb = miniz_oxide::deflate::compress_to_vec(&vec![0u8; 1 << 20], 9);
|
||||
let hdr = format!("{}\n", json!({"v": 1, "token": TOKEN, "op": "run", "model": "aa11", "tile": T, "bytes": bomb.len()}));
|
||||
assert!(talk(&[hdr.as_bytes(), &bomb].concat()).contains("false"));
|
||||
assert!(Remote::new(&addr, TOKEN, "aa11", T).hello().is_ok());
|
||||
assert!(same("abc", "abc") && !same("abc", "abd") && !same("abc", "abcd"));
|
||||
// a short token is refused when the server starts
|
||||
let l = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let opts = ServerOptions { token: "short".into(), ..Default::default() };
|
||||
assert!(serve(l, id("aa11"), T, "x".into(), model(1.0), opts, Arc::new(|_| {})).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packing_round_trips_and_addresses_are_normalised() {
|
||||
let v: Vec<f32> = (0..4096).map(|i| 0.5 + ((i * 7919 % 101) as f32 - 50.0) * 0.002).collect();
|
||||
let back = unpack(&pack(&v), v.len()).unwrap();
|
||||
assert!(v.iter().zip(&back).all(|(a, b)| (a - b).abs() < 1e-3));
|
||||
assert!(unpack(&pack(&v), v.len() + 1).is_err());
|
||||
assert_eq!(Remote::new("tcp://100.1.2.3/", TOKEN, "a", T).address(), "100.1.2.3:7990");
|
||||
assert_eq!(Remote::new("box:8000", TOKEN, "a", T).address(), "box:8000");
|
||||
assert_eq!(Remote::new("[::1]", TOKEN, "a", T).address(), "[::1]:7990");
|
||||
}
|
||||
}
|
||||
@@ -653,7 +653,7 @@ pub fn specs() -> Vec<CommandSpec> {
|
||||
),
|
||||
cmd!(query "denoise.queue", "Denoise Photos", [], None, "{ids?: [photoId], scope?: selected | visible | withAmount, retry?: false} → {added, …status} — make the denoised pictures of raw photos that have none, in the background. `retry` tries photos that failed before", always, queue),
|
||||
cmd!(query "denoise.cancel", "Cancel Denoise", [], None, "{id?, all?} → {cancelled} — stop the photo being made and/or forget the ones waiting (all of them when `id` is omitted)", always, cancel),
|
||||
cmd!(query "denoise.status", "Denoise Status", [], None, "{id?} → {enabled, model, runtime, productsDir, queued, running: {photo, done, total, seconds} | null, ready, made, auto, runOn, device: {kind: gpu | cpu | pending | none, adapter?, reason?, cardMs?, cpuMs?, fellBack?}, cache: {files, bytes, limitBytes}, generation, photo?: {state: notApplicable | noModel | ready | queued | running | failed | idle, …}}", always, status),
|
||||
cmd!(query "denoise.status", "Denoise Status", [], None, "{id?} → {enabled, model, runtime, productsDir, queued, running: {photo, done, total, seconds} | null, ready, made, auto, runOn, device: {kind: gpu | cpu | pending | none, adapter?, reason?, cardMs?, cpuMs?, fellBack?, remote?: {address, state: pending | unavailable | used | idle | slower | given up, device?, remoteMs?, localMs?, reason?, failures?, lastError?}}, cache: {files, bytes, limitBytes}, generation, photo?: {state: notApplicable | noModel | ready | queued | running | failed | idle, …}}", always, status),
|
||||
cmd!(query "denoise.pump", "Denoise Pump", [], None, "{pace?: pause | light | normal | full} → {active, running, queued, ready, generation, downloading} — called every frame by the app: takes in finished work and keeps one photo going", always, pump),
|
||||
cmd!(query "denoise.clear", "Clear Denoise Cache", [], None, "{} → {deleted, bytes} — delete every cached denoised picture of this library (they are made again when wanted)", always, clear),
|
||||
]
|
||||
|
||||
+279
-13
@@ -157,6 +157,97 @@ struct Network {
|
||||
fallbacks: AtomicUsize,
|
||||
/// Tiles at once of the photo started last (0: none yet), to say where the work runs.
|
||||
last_threads: AtomicUsize,
|
||||
/// Another computer's graphics card (`LIGHTCRAFT_DENOISE_URL` + `LIGHTCRAFT_DENOISE_TOKEN`; see
|
||||
/// `docs/denoise.md` → Another computer's graphics card), greeted on first use: it takes tiles while it has a free
|
||||
/// slot, on top of this computer's runner.
|
||||
remote: Option<RemoteSide>,
|
||||
/// Whether this computer's share runs on the card (else the CPU), as chosen for the photo started last.
|
||||
local_card: AtomicBool,
|
||||
}
|
||||
|
||||
#[cfg(feature = "denoise")]
|
||||
struct RemoteSide {
|
||||
runner: lightcraft_denoise::remote::Remote,
|
||||
/// The greeting, and the server's time for the check tile there and back (ms; the faster of two runs).
|
||||
greeted: OnceLock<Result<(lightcraft_denoise::remote::ServerInfo, f64), String>>,
|
||||
/// Whether the photo started last shares tiles with the server, and if not because it is slower, the two times.
|
||||
share: AtomicBool,
|
||||
slower: Mutex<Option<(f64, f64)>>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "denoise")]
|
||||
impl RemoteSide {
|
||||
/// From the environment, for the model whose file has `sha256` and tiles of `tile` cells.
|
||||
fn from_env(sha256: &str, tile: usize) -> Option<RemoteSide> {
|
||||
let var = |k: &str| std::env::var(k).ok().map(|v| v.trim().to_string()).filter(|v| !v.is_empty());
|
||||
let address = var("LIGHTCRAFT_DENOISE_URL")?;
|
||||
let token = var("LIGHTCRAFT_DENOISE_TOKEN").unwrap_or_default();
|
||||
Some(RemoteSide::new(lightcraft_denoise::remote::Remote::new(&address, &token, sha256, tile)))
|
||||
}
|
||||
|
||||
fn new(runner: lightcraft_denoise::remote::Remote) -> RemoteSide {
|
||||
RemoteSide { runner, greeted: OnceLock::new(), share: AtomicBool::new(false), slower: Mutex::new(None) }
|
||||
}
|
||||
|
||||
/// Greet the server and time a check tile there and back: a server behind a slow link can take seconds a tile, and
|
||||
/// handing it tiles a fast local card would finish sooner only makes the photo wait for it.
|
||||
fn greet(&self) -> Result<(lightcraft_denoise::remote::ServerInfo, f64), String> {
|
||||
let info = self.runner.hello().map_err(|e| e.to_string())?;
|
||||
let input = lightcraft_denoise::runtime::check_tile(info.tile);
|
||||
let mut best = f64::INFINITY;
|
||||
for _ in 0..2 {
|
||||
let started = web_time::Instant::now();
|
||||
match self.runner.try_run(&input) {
|
||||
Some(Ok(_)) => best = best.min(started.elapsed().as_secs_f64() * 1000.0),
|
||||
Some(Err(e)) => return Err(format!("the check tile failed: {e}")),
|
||||
None => return Err("no free slot for the check tile".into()),
|
||||
}
|
||||
}
|
||||
Ok((info, best))
|
||||
}
|
||||
|
||||
/// The server when it can take tiles now (greeting it the first time), else why not.
|
||||
fn usable(&self, start: bool) -> Result<&lightcraft_denoise::remote::Remote, String> {
|
||||
let greeted = if start { Some(self.greeted.get_or_init(|| self.greet())) } else { self.greeted.get() };
|
||||
match greeted {
|
||||
None => Err("not contacted yet".into()),
|
||||
Some(Err(why)) => Err(why.clone()),
|
||||
Some(Ok(_)) if !self.runner.healthy() => {
|
||||
Err(format!("it failed {} tiles ({})", self.runner.failures(), self.runner.last_error().unwrap_or_default()))
|
||||
}
|
||||
Some(Ok(_)) => Ok(&self.runner),
|
||||
}
|
||||
}
|
||||
|
||||
/// The server's time for a tile, once greeted.
|
||||
fn tile_ms(&self) -> Option<f64> {
|
||||
self.greeted.get().and_then(|g| g.as_ref().ok()).map(|g| g.1)
|
||||
}
|
||||
|
||||
fn describe(&self) -> Value {
|
||||
match self.greeted.get() {
|
||||
None => json!({"address": self.runner.address(), "state": "pending"}),
|
||||
Some(Err(why)) => json!({"address": self.runner.address(), "state": "unavailable", "reason": why}),
|
||||
Some(Ok((info, ms))) => {
|
||||
let slower = *self.slower.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
let state = match () {
|
||||
_ if !self.runner.healthy() => "given up",
|
||||
_ if slower.is_some() => "slower",
|
||||
_ if self.share.load(Ordering::Relaxed) => "used",
|
||||
_ => "idle",
|
||||
};
|
||||
let mut v =
|
||||
json!({"address": info.address, "device": info.device, "state": state, "remoteMs": ms, "failures": self.runner.failures()});
|
||||
if let (Some((_, local)), Some(o)) = (slower, v.as_object_mut()) {
|
||||
o.insert("localMs".into(), json!(local));
|
||||
}
|
||||
if let (Some(e), Some(o)) = (self.runner.last_error(), v.as_object_mut()) {
|
||||
o.insert("lastError".into(), json!(e));
|
||||
}
|
||||
v
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "denoise")]
|
||||
@@ -338,10 +429,19 @@ impl Network {
|
||||
}
|
||||
}
|
||||
|
||||
/// The card, with the CPU running any tile it gets wrong. Handed out only for photos that run on the card.
|
||||
/// A remote server's free slot first (when there is one), then the card when this photo runs on it, with the CPU
|
||||
/// running any tile either gets wrong. Handed out for photos that run on the card or have a server to share.
|
||||
#[cfg(feature = "denoise")]
|
||||
impl TileRunner for Network {
|
||||
fn run(&self, input: &[f32]) -> Result<Vec<f32>, RunError> {
|
||||
if let Some(Ok(remote)) = self.remote.as_ref().filter(|r| r.share.load(Ordering::Relaxed)).map(|r| r.usable(false))
|
||||
&& let Some(Ok(out)) = remote.try_run(input)
|
||||
{
|
||||
return Ok(out);
|
||||
}
|
||||
if !self.local_card.load(Ordering::Relaxed) {
|
||||
return self.cpu.run(input);
|
||||
}
|
||||
if let Ok(g) = self.card(false) {
|
||||
match g.runner.run(input) {
|
||||
Ok(out) if out.iter().all(|v| v.is_finite()) => return Ok(out),
|
||||
@@ -361,10 +461,37 @@ impl Model for Network {
|
||||
fn runner(&self, run_on: RunOn, threads: usize) -> (&dyn TileRunner, usize) {
|
||||
let threads = threads.max(1);
|
||||
self.last_threads.store(threads, Ordering::Relaxed);
|
||||
match self.choose(run_on, threads, true) {
|
||||
Ok(_) => (self, threads.min(GPU_PARALLEL)),
|
||||
Err(_) => (&self.cpu, threads),
|
||||
let (card, local) = match self.choose(run_on, threads, true) {
|
||||
Ok(_) => (true, threads.min(GPU_PARALLEL)),
|
||||
Err(_) => (false, threads),
|
||||
};
|
||||
self.local_card.store(card, Ordering::Relaxed);
|
||||
if let Some(r) = &self.remote {
|
||||
// this computer's time for a tile: the card's, or the CPU's with `threads` tiles sharing it (as
|
||||
// `card_is_faster` counts them); unknown when the card was never set up (then the server is used)
|
||||
let local_ms = match self.gpu.get() {
|
||||
Some(Ok(g)) if card => g.card_ms,
|
||||
Some(Ok(g)) => g.cpu_ms / (1.0 + 0.5 * threads.saturating_sub(1) as f64),
|
||||
_ => f64::INFINITY,
|
||||
};
|
||||
let share = match (r.usable(true), r.tile_ms()) {
|
||||
(Ok(_), Some(ms)) if ms < local_ms => {
|
||||
*r.slower.lock().unwrap_or_else(PoisonError::into_inner) = None;
|
||||
true
|
||||
}
|
||||
(Ok(_), Some(ms)) => {
|
||||
*r.slower.lock().unwrap_or_else(PoisonError::into_inner) = Some((ms, local_ms));
|
||||
false
|
||||
}
|
||||
_ => false,
|
||||
};
|
||||
r.share.store(share, Ordering::Relaxed);
|
||||
if let (true, Ok(remote)) = (share, r.usable(false)) {
|
||||
// the server's slots come on top of this computer's own tiles at once
|
||||
return (self, local + remote.slots());
|
||||
}
|
||||
}
|
||||
if card { (self, local) } else { (&self.cpu, local) }
|
||||
}
|
||||
|
||||
fn self_test(&self, run_on: RunOn) -> Result<Value, String> {
|
||||
@@ -403,6 +530,9 @@ impl Model for Network {
|
||||
o.insert("cpuMs".into(), json!(g.cpu_ms));
|
||||
o.insert("fellBack".into(), json!(self.fallbacks.load(Ordering::Relaxed)));
|
||||
}
|
||||
if let (Some(r), Some(o)) = (&self.remote, v.as_object_mut()) {
|
||||
o.insert("remote".into(), r.describe());
|
||||
}
|
||||
v
|
||||
}
|
||||
}
|
||||
@@ -412,15 +542,14 @@ pub(crate) fn default_loader() -> Loader {
|
||||
#[cfg(feature = "denoise")]
|
||||
{
|
||||
Arc::new(|path, manifest| {
|
||||
let cpu = lightcraft_denoise::runtime::CpuRunner::load(path, manifest).map_err(|e| e.to_string())?;
|
||||
Ok(Arc::new(Network {
|
||||
cpu,
|
||||
path: path.to_path_buf(),
|
||||
tile: manifest.tile as usize,
|
||||
gpu: OnceLock::new(),
|
||||
fallbacks: AtomicUsize::new(0),
|
||||
last_threads: AtomicUsize::new(0),
|
||||
}) as Arc<dyn Model>)
|
||||
// the server must run this very file: identified by the manifest's hash, else the file's
|
||||
let sha256 = match &manifest.sha256 {
|
||||
Some(h) => Some(h.clone()),
|
||||
None if std::env::var_os("LIGHTCRAFT_DENOISE_URL").is_some() => lightcraft_denoise::hash::sha256_file(path).ok(),
|
||||
None => None,
|
||||
};
|
||||
let remote = sha256.and_then(|h| RemoteSide::from_env(&h, manifest.tile as usize));
|
||||
Ok(Arc::new(network(path, manifest, remote)?) as Arc<dyn Model>)
|
||||
})
|
||||
}
|
||||
#[cfg(not(feature = "denoise"))]
|
||||
@@ -429,6 +558,69 @@ pub(crate) fn default_loader() -> Loader {
|
||||
}
|
||||
}
|
||||
|
||||
/// The model at `path` on this computer, sharing tiles with `remote` when there is one.
|
||||
#[cfg(feature = "denoise")]
|
||||
fn network(path: &Path, manifest: &DenoiserManifest, remote: Option<RemoteSide>) -> Result<Network, String> {
|
||||
let cpu = lightcraft_denoise::runtime::CpuRunner::load(path, manifest).map_err(|e| e.to_string())?;
|
||||
Ok(Network {
|
||||
remote,
|
||||
local_card: AtomicBool::new(true),
|
||||
cpu,
|
||||
path: path.to_path_buf(),
|
||||
tile: manifest.tile as usize,
|
||||
gpu: OnceLock::new(),
|
||||
fallbacks: AtomicUsize::new(0),
|
||||
last_threads: AtomicUsize::new(0),
|
||||
})
|
||||
}
|
||||
|
||||
/// A model ready to be served to other computers (`lightcraft-cli denoise-serve`).
|
||||
#[cfg(feature = "denoise")]
|
||||
pub struct Serving {
|
||||
pub model: lightcraft_denoise::remote::ModelId,
|
||||
pub tile: usize,
|
||||
/// Where its tiles run here: the graphics card's name, or why the CPU does.
|
||||
pub device: String,
|
||||
pub run: lightcraft_denoise::remote::RunTile,
|
||||
}
|
||||
|
||||
/// The installed model `id` (default: the one selected in Settings) in the models folder `dir` (default: this user's),
|
||||
/// set up to run on `run_on` as the app would: the card checked against the CPU first, the CPU when the card can't.
|
||||
#[cfg(feature = "denoise")]
|
||||
pub fn serving(dir: Option<&Path>, id: Option<&str>, run_on: RunOn) -> Result<Serving, String> {
|
||||
let dir = match dir {
|
||||
Some(d) => d.to_path_buf(),
|
||||
None => crate::config::default_denoise_models_dir().ok_or("no denoise models folder (set LIGHTCRAFT_DENOISE_MODELS)")?,
|
||||
};
|
||||
let want = match id {
|
||||
Some(i) => i.to_string(),
|
||||
None => read_settings(&dir).model.ok_or_else(|| {
|
||||
format!("no denoise model is selected in {} (install one in Settings ▸ AI Denoise, or name it with --model)", dir.display())
|
||||
})?,
|
||||
};
|
||||
let installed = installed_models(&dir)
|
||||
.into_iter()
|
||||
.find(|m| m.manifest.id == want)
|
||||
.ok_or_else(|| format!("no installed denoise model `{want}` in {}", dir.display()))?;
|
||||
let m = &installed.manifest;
|
||||
// the same identity a client derives (see `default_loader`)
|
||||
let sha256 = match &m.sha256 {
|
||||
Some(h) => h.clone(),
|
||||
None => lightcraft_denoise::hash::sha256_file(&installed.onnx).map_err(|e| format!("{}: {e}", installed.onnx.display()))?,
|
||||
};
|
||||
let net = Arc::new(network(&installed.onnx, m, None)?);
|
||||
let tile = m.tile as usize;
|
||||
// set the card up now (and say where the tiles run) rather than with the first client's tile
|
||||
let _ = net.runner(run_on, GPU_PARALLEL);
|
||||
let d = net.device(run_on);
|
||||
let device = match d["kind"].as_str() {
|
||||
Some("gpu") => format!("GPU: {}", d["adapter"].as_str().unwrap_or("?")),
|
||||
_ => format!("CPU ({})", d["reason"].as_str().unwrap_or("chosen")),
|
||||
};
|
||||
let run: lightcraft_denoise::remote::RunTile = Arc::new(move |input: &[f32]| net.runner(run_on, GPU_PARALLEL).0.run(input));
|
||||
Ok(Serving { model: lightcraft_denoise::remote::ModelId { id: m.id.clone(), version: m.version.clone(), sha256 }, tile, device, run })
|
||||
}
|
||||
|
||||
/// A model that is loaded on first use and kept, shared by the background job and the exports.
|
||||
pub(crate) type ModelSlot = Arc<Mutex<Option<Arc<dyn Model>>>>;
|
||||
|
||||
@@ -1517,6 +1709,76 @@ mod gpu_tests {
|
||||
let _ = std::fs::remove_dir_all(path.parent().unwrap());
|
||||
}
|
||||
|
||||
/// A server on this computer running `cpu` for the model with `sha256`.
|
||||
fn server(cpu: lightcraft_denoise::runtime::CpuRunner, sha256: &str, tile: usize) -> String {
|
||||
use lightcraft_denoise::remote::{ModelId, ServerOptions, serve};
|
||||
let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let addr = l.local_addr().unwrap().to_string();
|
||||
let id = ModelId { id: "synthetic".into(), version: "1".into(), sha256: sha256.into() };
|
||||
let run: lightcraft_denoise::remote::RunTile = Arc::new(move |input: &[f32]| cpu.run(input));
|
||||
let opts = ServerOptions { token: "remote-test-token-0123".into(), ..Default::default() };
|
||||
std::thread::spawn(move || serve(l, id, tile, "CPU (test server)".into(), run, opts, Arc::new(|_| {})));
|
||||
addr
|
||||
}
|
||||
|
||||
fn with_remote(path: &Path, tile: u64, addr: &str, sha256: &str) -> Network {
|
||||
let manifest = DenoiserManifest {
|
||||
id: "synthetic".into(),
|
||||
name: "Synthetic".into(),
|
||||
version: "1".into(),
|
||||
licence: Default::default(),
|
||||
source: None,
|
||||
sha256: None,
|
||||
size_bytes: None,
|
||||
provenance: String::new(),
|
||||
domain: Domain::BayerToRgb,
|
||||
tile: tile as u32,
|
||||
overlap: 16,
|
||||
gain: Gain::MatchMean { nominal: 1.0, max_deviation: 0.05 },
|
||||
};
|
||||
let remote = RemoteSide::new(lightcraft_denoise::remote::Remote::new(addr, "remote-test-token-0123", sha256, tile as usize));
|
||||
network(path, &manifest, Some(remote)).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_remote_server_shares_the_tiles_and_drops_out_when_it_cannot() {
|
||||
let (_, cpu, path) = model("remote", 64);
|
||||
let sha = lightcraft_denoise::hash::sha256_file(&path).unwrap();
|
||||
let (input, _) = lightcraft_denoise::runtime::test_tile(64);
|
||||
let want = cpu.run(&input).unwrap();
|
||||
// the same model on the server: its slots come on top of this computer's threads, answers within f16 rounding
|
||||
let addr = server(lightcraft_denoise::runtime::CpuRunner::load(&path, &cpu.manifest().clone()).unwrap(), &sha, 64);
|
||||
let net = with_remote(&path, 64, &addr, &sha);
|
||||
let (runner, at_once) = net.runner(RunOn::Cpu, 2);
|
||||
assert_eq!(at_once, 2 + 4, "four server slots on top of two local threads");
|
||||
assert!(worst(&want, &runner.run(&input).unwrap()) < 5e-3);
|
||||
let d = net.device(RunOn::Cpu);
|
||||
assert_eq!(d["remote"]["state"], "used", "{d}");
|
||||
assert_eq!(d["remote"]["device"], "CPU (test server)");
|
||||
// a local card faster than the server there and back (here a stand-in that takes no time): nothing is shared,
|
||||
// the photo never waits for the slower server, and the status says why
|
||||
let mut net = with_remote(&path, 64, &addr, &sha);
|
||||
net.gpu = OnceLock::from(Ok(GpuSide { runner: Box::new(cpu.clone()), adapter: "Fast card".into(), card_ms: 1e-6, cpu_ms: 1000.0 }));
|
||||
let (_, at_once) = net.runner(RunOn::Gpu, 8);
|
||||
assert_eq!(at_once, GPU_PARALLEL, "no server slots");
|
||||
let d = net.device(RunOn::Gpu);
|
||||
assert_eq!(d["remote"]["state"], "slower", "{d}");
|
||||
assert!(d["remote"]["remoteMs"].as_f64().is_some_and(|ms| ms > 0.0) && d["remote"]["localMs"].as_f64().is_some(), "{d}");
|
||||
// a server running another model is refused: everything runs here
|
||||
let other = server(lightcraft_denoise::runtime::CpuRunner::load(&path, &cpu.manifest().clone()).unwrap(), "0000", 64);
|
||||
let net = with_remote(&path, 64, &other, &sha);
|
||||
let (runner, at_once) = net.runner(RunOn::Cpu, 2);
|
||||
assert_eq!(at_once, 2);
|
||||
assert_eq!(runner.run(&input).unwrap(), want, "the CPU path, bit for bit");
|
||||
let d = net.device(RunOn::Cpu);
|
||||
assert!(d["remote"]["reason"].as_str().is_some_and(|r| r.contains("another model")), "{d}");
|
||||
// nobody listening: the same
|
||||
let net = with_remote(&path, 64, "127.0.0.1:1", &sha);
|
||||
assert_eq!(net.runner(RunOn::Cpu, 3).1, 3);
|
||||
assert_eq!(net.device(RunOn::Cpu)["remote"]["state"], "unavailable");
|
||||
let _ = std::fs::remove_dir_all(path.parent().unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_the_processor_chosen_everything_runs_on_the_cpu() {
|
||||
let (model, cpu, path) = model("cpu", 64);
|
||||
@@ -1563,6 +1825,8 @@ mod gpu_tests {
|
||||
gpu: OnceLock::from(Ok(card)),
|
||||
fallbacks: AtomicUsize::new(0),
|
||||
last_threads: AtomicUsize::new(0),
|
||||
remote: None,
|
||||
local_card: AtomicBool::new(true),
|
||||
};
|
||||
let (input, _) = lightcraft_denoise::runtime::test_tile(64);
|
||||
let want = cpu.run(&input).unwrap();
|
||||
@@ -1591,6 +1855,8 @@ mod gpu_tests {
|
||||
gpu: OnceLock::from(Ok(slow)),
|
||||
fallbacks: AtomicUsize::new(0),
|
||||
last_threads: AtomicUsize::new(0),
|
||||
remote: None,
|
||||
local_card: AtomicBool::new(true),
|
||||
};
|
||||
assert_eq!(model.runner(RunOn::Auto, 16).1, 16, "sixteen threads beat it");
|
||||
let device = model.device(RunOn::Auto);
|
||||
|
||||
@@ -121,6 +121,51 @@ LC_DENOISE_MODEL=<model.onnx> LC_DENOISE_RAW=<corpus/raw> cargo test --release -
|
||||
LC_DENOISE_MODEL=<model.onnx> LC_DENOISE_RAW=<corpus/raw> cargo test --release -p lightcraft-engine --features rawnind-model --test denoise_real -- --ignored --nocapture
|
||||
```
|
||||
|
||||
## Another computer's graphics card
|
||||
|
||||
A laptop can hand tiles to a desktop with a bigger graphics card over a LAN or a VPN. The server runs the installed,
|
||||
selected model the way the app does (the card when it can run it, checked against the CPU first; else the CPU):
|
||||
|
||||
```text
|
||||
export LIGHTCRAFT_DENOISE_TOKEN='a long random string, at least 16 characters'
|
||||
lightcraft-cli denoise-serve --listen 100.101.102.103:7990 # e.g. the machine's Tailscale address
|
||||
```
|
||||
|
||||
On the laptop, before LightCraft starts:
|
||||
|
||||
```text
|
||||
export LIGHTCRAFT_DENOISE_URL=100.101.102.103:7990
|
||||
export LIGHTCRAFT_DENOISE_TOKEN='…the same token…'
|
||||
```
|
||||
|
||||
When the laptop first uses the model it greets the server and times a check tile there and back. If that is faster
|
||||
than a tile on the laptop (its graphics card, or its processor with the tiles it runs at once), the server's slots
|
||||
(four tiles in flight) are added to the laptop's own: a tile goes to the server when one of its slots is free and runs
|
||||
locally otherwise. If it is slower, nothing is sent, so a photo never waits for the slower machine:
|
||||
`denoise.status` → `device.remote.state` is `slower`, with both times (`remoteMs`, `localMs`). Measured with
|
||||
RawNIND: an Apple M3 laptop does a tile in 116 ms, an RTX 3060 Ti reached over Tailscale 345 ms there and back,
|
||||
so that laptop keeps its tiles; one without a fast card hands them over. A tile the server fails runs locally; after
|
||||
three failures the server is no longer used for that model. Both ends must
|
||||
run the same model file: the client compares the manifest's SHA-256 (or the file's) when it first greets the server,
|
||||
and the server checks it again for every tile. `denoise.status` → `device.remote` reports the server's address, its
|
||||
device and whether it is used, or why not.
|
||||
|
||||
**The protocol is not encrypted.** Listen on a VPN address (Tailscale, WireGuard) or a LAN, or bind every interface
|
||||
behind a firewall rule that admits only the VPN; never expose the port to the internet. Requests carry the shared
|
||||
token (compared in constant time); the server refuses other models, other tile sizes, oversized headers and payloads
|
||||
(also ones that inflate past their tile) and caps concurrent connections (`--max-connections`, default 8).
|
||||
|
||||
Protocol version 1 (`crates/denoise/src/remote.rs`): one JSON line per message, then a payload of `bytes` bytes —
|
||||
the tile's samples as little-endian half floats, byte-shuffled (low bytes, then high bytes) and deflated:
|
||||
|
||||
```text
|
||||
→ {"v":1,"token":"…","op":"hello"}
|
||||
← {"ok":true,"v":1,"device":"GPU: …","model":{"id":"…","version":"…","sha256":"…"},"tile":T}
|
||||
→ {"v":1,"token":"…","op":"run","model":"<sha256>","tile":T,"bytes":N} + N bytes (4 planes of T × T cells)
|
||||
← {"ok":true,"tile":T,"bytes":N} + N bytes (3 planes of 2T × 2T pixels)
|
||||
← {"ok":false,"error":"…"}
|
||||
```
|
||||
|
||||
## Commands and limits
|
||||
|
||||
UI, CLI, control channel and MCP dispatch `denoise.toggle`, `denoise.settings`,
|
||||
|
||||
@@ -205,7 +205,7 @@ differences and missing bindings are in [ui-parity.md](ui-parity.md#shortcuts).
|
||||
| LR-EDIT-DETAIL-SHARPEN | Sharpening | P0 | ✅ | `ctl:detail.sharpenAmount`, `ctl:detail.sharpenRadius`, `ctl:detail.sharpenDetail`, `ctl:detail.sharpenMasking` | no Alt-drag mask preview |
|
||||
| LR-EDIT-DETAIL-NR | Luminance noise reduction | P0 | ✅ | `ctl:detail.nrLuminance`, `ctl:detail.nrDetail`, `ctl:detail.nrContrast` | |
|
||||
| LR-EDIT-DETAIL-CNR | Colour noise reduction | P0 | ✅ | `ctl:detail.nrColor`, `ctl:detail.nrColorDetail`, `ctl:detail.nrColorSmoothness` | |
|
||||
| LR-EDIT-DETAIL-DENOISE | AI denoise | P2 | 🟡 | cmd:denoise.toggle, cmd:denoise.models.install, cmd:denoise.queue, ctl:enhance.denoise, `crates/denoise/`, `docs/denoise.md` | non-destructive Bayer RAW, opt-in weights, pure-Rust CPU/GPU; GPU convolution uses independently written shared scalars (Metal equivalence regression, issue #479); X-Trans/linear RGB, model policy and Lightroom fidelity remain |
|
||||
| LR-EDIT-DETAIL-DENOISE | AI denoise | P2 | 🟡 | cmd:denoise.toggle, cmd:denoise.models.install, cmd:denoise.queue, ctl:enhance.denoise, `crates/denoise/`, `docs/denoise.md` | non-destructive Bayer RAW, opt-in weights, pure-Rust CPU/GPU; GPU convolution uses independently written shared scalars (Metal equivalence regression, issue #479); tiles can also run on another computer's graphics card (`lightcraft-cli denoise-serve`, `crates/denoise/src/remote.rs`); X-Trans/linear RGB, model policy and Lightroom fidelity remain |
|
||||
| LR-EDIT-DETAIL-RAWDETAILS | Improved demosaic toggle | P2 | ⬜ | | |
|
||||
| LR-EDIT-DETAIL-SUPERRES | Super resolution | P2 | ⬜ | | |
|
||||
| LR-EDIT-DETAIL-AISHARPEN | AI sharpen | OOS | 🚫 | | |
|
||||
|
||||
Reference in new issue
Block a user