From 84b8995d7b85c0eabcdedbfa86f5d61ff869993d Mon Sep 17 00:00:00 2001 From: akvaithi Date: Sat, 10 Oct 2026 15:21:43 -0500 Subject: [PATCH 1/2] LR-EDIT-DETAIL-DENOISE: run tiles on another computer's graphics card MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - `lightcraft_denoise::remote`: protocol v1 over TCP (JSON lines + byte-shuffled, deflated half-float tiles), a `Remote` client (connection pool, four slots, greeting checks the model's SHA-256 and tile size, gives up after three failed tiles) and `serve` (shared token compared in constant time, size caps including payloads that inflate past their tile, connection limit, timeouts). - Engine: with LIGHTCRAFT_DENOISE_URL / LIGHTCRAFT_DENOISE_TOKEN set, the loaded model greets the server on first use and adds its slots to this computer's tiles at once; a tile goes to a free server slot, else runs here, and a tile the server fails runs here too. `denoise.status` → `device.remote` says where. - `lightcraft-cli denoise-serve` serves the installed, selected model on this computer's card (checked against the CPU, as in the app) or CPU. - docs/denoise.md: setup, safety (unencrypted: VPN/LAN only) and the protocol. Tested: protocol round trips, refusals (token, model, tile size, hostile payloads), fallback and give-up; the engine model with a live, a mismatched and a dead server; the CLI binary serving a synthetic model end to end. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 1 + apps/lightcraft-cli/Cargo.toml | 1 + apps/lightcraft-cli/src/main.rs | 59 ++++ apps/lightcraft-cli/tests/cli.rs | 66 ++++ crates/denoise/src/lib.rs | 2 + crates/denoise/src/remote.rs | 562 +++++++++++++++++++++++++++++++ crates/engine/src/cmd/denoise.rs | 2 +- crates/engine/src/denoise.rs | 225 ++++++++++++- docs/denoise.md | 40 +++ docs/parity-checklist.md | 2 +- 10 files changed, 945 insertions(+), 15 deletions(-) create mode 100644 crates/denoise/src/remote.rs diff --git a/Cargo.lock b/Cargo.lock index 0ad4ab843..bb011db07 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3014,6 +3014,7 @@ version = "0.5.0" dependencies = [ "dhat", "lightcraft-codecs", + "lightcraft-denoise", "lightcraft-engine", "lightcraft-mcp", "lightcraft-merge", diff --git a/apps/lightcraft-cli/Cargo.toml b/apps/lightcraft-cli/Cargo.toml index 6471ab86d..b82aca1d5 100644 --- a/apps/lightcraft-cli/Cargo.toml +++ b/apps/lightcraft-cli/Cargo.toml @@ -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 } diff --git a/apps/lightcraft-cli/src/main.rs b/apps/lightcraft-cli/src/main.rs index 592a534a9..38c49ebe4 100644 --- a/apps/lightcraft-cli/src/main.rs +++ b/apps/lightcraft-cli/src/main.rs @@ -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=: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… @@ -145,6 +158,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")); @@ -897,3 +911,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 = None; + let mut dir: Option = 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()) +} diff --git a/apps/lightcraft-cli/tests/cli.rs b/apps/lightcraft-cli/tests/cli.rs index ecda1102a..9aafc9310 100644 --- a/apps/lightcraft-cli/tests/cli.rs +++ b/apps/lightcraft-cli/tests/cli.rs @@ -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); +} diff --git a/crates/denoise/src/lib.rs b/crates/denoise/src/lib.rs index e72860ba4..f9005c7d5 100644 --- a/crates/denoise/src/lib.rs +++ b/crates/denoise/src/lib.rs @@ -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; diff --git a/crates/denoise/src/remote.rs b/crates/denoise/src/remote.rs new file mode 100644 index 000000000..b40eacd6e --- /dev/null +++ b/crates/denoise/src/remote.rs @@ -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":"","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) -> std::io::Error { + std::io::Error::new(std::io::ErrorKind::InvalidData, m.into()) +} + +fn read_line(r: &mut impl BufRead) -> std::io::Result> { + 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 { + 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> { + 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> { + 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 { + 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 Result, 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, +) -> 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, 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::(&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>, + pool: Mutex, 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::().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 { + 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), 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) -> Result { + 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 { + 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, 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, 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, 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, Error> { + self.1.fetch_add(1, Ordering::SeqCst); + (self.0)(input) + } + } + + fn tile_input(seed: f32) -> Vec { + (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 = (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"); + } +} diff --git a/crates/engine/src/cmd/denoise.rs b/crates/engine/src/cmd/denoise.rs index 75c6c0a7a..81c9eb39d 100644 --- a/crates/engine/src/cmd/denoise.rs +++ b/crates/engine/src/cmd/denoise.rs @@ -653,7 +653,7 @@ pub fn specs() -> Vec { ), 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 | given up, device?, 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), ] diff --git a/crates/engine/src/denoise.rs b/crates/engine/src/denoise.rs index 30df142ca..8e939e611 100644 --- a/crates/engine/src/denoise.rs +++ b/crates/engine/src/denoise.rs @@ -157,6 +157,56 @@ 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, + /// 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, + greeted: OnceLock>, +} + +#[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 { + 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 { runner: lightcraft_denoise::remote::Remote::new(&address, &token, sha256, tile), greeted: OnceLock::new() }) + } + + /// 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.runner.hello().map_err(|e| e.to_string()))) } 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), + } + } + + 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)) => { + let mut v = json!({"address": info.address, "device": info.device, "state": if self.runner.healthy() { "used" } else { "given up" }, "failures": self.runner.failures()}); + 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 +388,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, RunError> { + if let Some(Ok(remote)) = self.remote.as_ref().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,9 +420,16 @@ 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); + match self.remote.as_ref().map(|r| r.usable(true)) { + // the server's slots come on top of this computer's own tiles at once + Some(Ok(remote)) => (self, local + remote.slots()), + _ if card => (self, local), + _ => (&self.cpu, local), } } @@ -403,6 +469,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 +481,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) + // 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) }) } #[cfg(not(feature = "denoise"))] @@ -429,6 +497,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) -> Result { + 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 { + 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>>>; @@ -1502,6 +1633,70 @@ 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 { + runner: lightcraft_denoise::remote::Remote::new(addr, "remote-test-token-0123", sha256, tile as usize), + greeted: OnceLock::new(), + }; + 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 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); @@ -1548,6 +1743,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(); @@ -1576,6 +1773,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); diff --git a/docs/denoise.md b/docs/denoise.md index be208a9b8..b5a171767 100644 --- a/docs/denoise.md +++ b/docs/denoise.md @@ -121,6 +121,46 @@ LC_DENOISE_MODEL= LC_DENOISE_RAW= cargo test --release - LC_DENOISE_MODEL= LC_DENOISE_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…' +``` + +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, so a server behind a slow link still adds speed instead of costing it. 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":"","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`, diff --git a/docs/parity-checklist.md b/docs/parity-checklist.md index 099411089..53282adb6 100644 --- a/docs/parity-checklist.md +++ b/docs/parity-checklist.md @@ -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 | 🚫 | | | From 88f424e7e70c22cbaf2914e2fb20ffcbcadc860f Mon Sep 17 00:00:00 2001 From: akvaithi Date: Sat, 10 Oct 2026 15:53:39 -0500 Subject: [PATCH 2/2] LR-EDIT-DETAIL-DENOISE: share tiles with a server only when it is faster MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measured with RawNIND on a 20.8 MP raw: an Apple M3 does a tile in 116 ms, an RTX 3060 Ti over Tailscale 345 ms there and back; sharing blindly made the export take 34 s instead of 6.4 s while the photo waited for the server's last tiles. The greeting now times a check tile there and back, and tiles are shared only when that beats this computer's own time per tile (the card's, or the CPU's with its threads counted as `card_is_faster` does). Otherwise `denoise.status` → `device.remote` says `slower` with both times. With the server configured but slower the export takes 7.2 s (greeting + check tile). Co-Authored-By: Claude Opus 5.5 --- crates/engine/src/cmd/denoise.rs | 2 +- crates/engine/src/denoise.rs | 97 +++++++++++++++++++++++++++----- docs/denoise.md | 11 +++- 3 files changed, 91 insertions(+), 19 deletions(-) diff --git a/crates/engine/src/cmd/denoise.rs b/crates/engine/src/cmd/denoise.rs index 81c9eb39d..8802460a1 100644 --- a/crates/engine/src/cmd/denoise.rs +++ b/crates/engine/src/cmd/denoise.rs @@ -653,7 +653,7 @@ pub fn specs() -> Vec { ), 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?, remote?: {address, state: pending | unavailable | used | given up, device?, reason?, failures?, lastError?}}, 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), ] diff --git a/crates/engine/src/denoise.rs b/crates/engine/src/denoise.rs index 8e939e611..33582aaed 100644 --- a/crates/engine/src/denoise.rs +++ b/crates/engine/src/denoise.rs @@ -168,7 +168,11 @@ struct Network { #[cfg(feature = "denoise")] struct RemoteSide { runner: lightcraft_denoise::remote::Remote, - greeted: OnceLock>, + /// The greeting, and the server's time for the check tile there and back (ms; the faster of two runs). + greeted: OnceLock>, + /// Whether the photo started last shares tiles with the server, and if not because it is slower, the two times. + share: AtomicBool, + slower: Mutex>, } #[cfg(feature = "denoise")] @@ -178,12 +182,33 @@ impl 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 { runner: lightcraft_denoise::remote::Remote::new(&address, &token, sha256, tile), greeted: OnceLock::new() }) + 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.runner.hello().map_err(|e| e.to_string()))) } else { self.greeted.get() }; + 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()), @@ -194,12 +219,28 @@ impl RemoteSide { } } + /// The server's time for a tile, once greeted. + fn tile_ms(&self) -> Option { + 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)) => { - let mut v = json!({"address": info.address, "device": info.device, "state": if self.runner.healthy() { "used" } else { "given up" }, "failures": self.runner.failures()}); + 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)); } @@ -393,7 +434,7 @@ impl Network { #[cfg(feature = "denoise")] impl TileRunner for Network { fn run(&self, input: &[f32]) -> Result, RunError> { - if let Some(Ok(remote)) = self.remote.as_ref().map(|r| r.usable(false)) + 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); @@ -425,12 +466,32 @@ impl Model for Network { Err(_) => (false, threads), }; self.local_card.store(card, Ordering::Relaxed); - match self.remote.as_ref().map(|r| r.usable(true)) { - // the server's slots come on top of this computer's own tiles at once - Some(Ok(remote)) => (self, local + remote.slots()), - _ if card => (self, local), - _ => (&self.cpu, local), + 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 { @@ -1660,10 +1721,7 @@ mod gpu_tests { overlap: 16, gain: Gain::MatchMean { nominal: 1.0, max_deviation: 0.05 }, }; - let remote = RemoteSide { - runner: lightcraft_denoise::remote::Remote::new(addr, "remote-test-token-0123", sha256, tile as usize), - greeted: OnceLock::new(), - }; + let remote = RemoteSide::new(lightcraft_denoise::remote::Remote::new(addr, "remote-test-token-0123", sha256, tile as usize)); network(path, &manifest, Some(remote)).unwrap() } @@ -1682,6 +1740,15 @@ mod gpu_tests { 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); diff --git a/docs/denoise.md b/docs/denoise.md index b5a171767..926b1bee7 100644 --- a/docs/denoise.md +++ b/docs/denoise.md @@ -138,9 +138,14 @@ export LIGHTCRAFT_DENOISE_URL=100.101.102.103:7990 export LIGHTCRAFT_DENOISE_TOKEN='…the same token…' ``` -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, so a server behind a slow link still adds speed instead of costing it. A -tile the server fails runs locally; after three failures the server is no longer used for that model. Both ends must +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.