dtourolle
06422a07db
AI Denoise's Apply switch becomes Method: Bilinear, Fast, Medium, Best, default Best, so an untouched raw writes nothing and develops through the mixture. `apply` is still read and never written: 0 is Bilinear, 1 keeps a network already chosen. - Best is the mixture of a flat and an edge expert with a learned gate; Medium and Fast are students distilled from it. 2.48 s, 0.79 s and 0.57 s for a 20 MP frame on TensorRT fp16. - Each network carries its own tile border (256 for the mixture, 192 for the students) through `dr_denoise::Shipped` and `TileNet::halo`. - The file is hashed once at open and each network keys its own cached result; Bilinear keeps the result in memory for the way back. - Each has an .a16w16 sibling for the Hexagon: 0.00 dB on the 6D gate, at most 0.11 dB with the noise scaled x0.5 to x4. - APK BUNDLED 19 -> 23; the PKGBUILD installs all three.
87 lines
3.0 KiB
Rust
87 lines
3.0 KiB
Rust
//! TRACES: FR-DEV-3g
|
||
//! The denoise network under the inference engine.
|
||
//!
|
||
//! The shipped export takes `mosaic` and `sigma`, `1×1×1408×1408`, and
|
||
//! returns `rgb`, `1×3×1408×1408` (darkroom-denoise `denoise/export.py`,
|
||
//! fixed shape because every model the engine runs is). The engine picks the
|
||
//! rung: fp16 on TensorRT and MIGraphX, which measured 0.00 dB from f32; f32
|
||
//! on CUDA and the CPU; on the Hexagon the `.a16w16.onnx` sibling, 16-bit
|
||
//! activations and weights, 0.00 dB from f32 on the tablet itself where int8
|
||
//! lost 5–9 dB (docs/dev/inference.md §1.5). That sibling is the same network
|
||
//! with the Bayer packing spelled `SpaceToDepth`, which QNN can hold and the
|
||
//! 6-D reshape it replaces it cannot.
|
||
|
||
use crate::tile::TileNet;
|
||
use crate::DenoiseError;
|
||
use dr_inference_engine::{Model, Role};
|
||
|
||
/// The edge of the tile the shipped export takes.
|
||
pub const TILE: usize = 1408;
|
||
|
||
pub struct OnnxNet {
|
||
model: Model,
|
||
tile: usize,
|
||
halo: usize,
|
||
}
|
||
|
||
impl OnnxNet {
|
||
/// The network at `path`, which needs `halo` photosites of context
|
||
/// ([`crate::Shipped::halo`]).
|
||
pub fn from_path(path: &std::path::Path, halo: usize) -> Result<Self, DenoiseError> {
|
||
let (path, form) = dr_inference_engine::resolve_model(Role::Denoiser, path);
|
||
let bytes = std::fs::read(&path)?;
|
||
Ok(OnnxNet {
|
||
model: dr_inference_engine::open(Role::Denoiser, form, &bytes)?,
|
||
tile: TILE,
|
||
halo,
|
||
})
|
||
}
|
||
|
||
/// Where it runs, for a status line.
|
||
pub fn rung(&self) -> Result<dr_inference_engine::Rung, DenoiseError> {
|
||
Ok(self.model.acquire()?.rung())
|
||
}
|
||
}
|
||
|
||
impl TileNet for OnnxNet {
|
||
fn tile(&self) -> usize {
|
||
self.tile
|
||
}
|
||
|
||
fn halo(&self) -> usize {
|
||
self.halo
|
||
}
|
||
|
||
fn run(
|
||
&mut self,
|
||
mosaic: Vec<f32>,
|
||
sigma: Vec<f32>,
|
||
write: &mut dyn FnMut(&[f32]),
|
||
) -> Result<(), DenoiseError> {
|
||
let n = self.tile;
|
||
let shape = ndarray::IxDyn(&[1, 1, n, n]);
|
||
// The vectors become the tensors: no copy on the way in.
|
||
let m = ort::value::Tensor::from_array(
|
||
ndarray::Array::from_shape_vec(shape.clone(), mosaic)
|
||
.map_err(|e| DenoiseError::Model(e.to_string()))?,
|
||
)?;
|
||
let s = ort::value::Tensor::from_array(
|
||
ndarray::Array::from_shape_vec(shape, sigma)
|
||
.map_err(|e| DenoiseError::Model(e.to_string()))?,
|
||
)?;
|
||
let acquired = self.model.acquire()?;
|
||
let mut session = acquired.lock();
|
||
let outputs = session.run(ort::inputs!["mosaic" => m, "sigma" => s])?;
|
||
let (shape, data) = outputs[0].try_extract_tensor::<f32>()?;
|
||
let dims: Vec<i64> = shape.iter().copied().collect();
|
||
if dims != [1, 3, n as i64, n as i64] {
|
||
return Err(DenoiseError::Model(format!(
|
||
"output is {dims:?}, expected [1, 3, {n}, {n}]"
|
||
)));
|
||
}
|
||
// And none on the way out: the frame is written from the runtime's buffer.
|
||
write(data);
|
||
Ok(())
|
||
}
|
||
}
|