DarkRoom/core/dr-pano/src/migan.rs

//! TRACES: FR-MRG-4
//! MI-GAN, the border filler, under the inference engine.
//!
//! Sargsyan et al., ICCV 2023 (Picsart AI Research): inpainting built for
//! phones — about six million parameters of plain convolutions, no FFT and
//! no attention, so it quantises and runs on a DSP. MIT, code and weights
//! (`models/LICENCE.md`). The bare 512 generator is what ships, exported
//! at a fixed shape by `tools/export-migan.sh`; its six operator types load
//! on every rung, and what they cost is the whole story of whether a fill
//! is interactive: 7.4 s a tile under tract, 0.4 s under ONNX Runtime's
//! CPU pool, 23 ms in fp16 and 13 ms in int8 on a laptop's TensorRT
//! (2026-09-19, docs/dev/panorama.md §12).
//!
//! The model's contract, from the reference `export_inference_model.py`:
//! input `1×4×512×512` float — channel 0 is `mask − 0.5` with 1 where the
//! picture is known, channels 1–3 the RGB in −1..1 with the unknown pixels
//! zeroed; output `1×3×512×512` in −1..1, of which the caller keeps the
//! unknown pixels. That is [`crate::fill::Inpainter`], and the rest —
//! which tiles, what context, how to blend — is `fill.rs`.

use crate::fill::Inpainter;
use crate::PanoError;

/// The tile the shipped export takes.
pub const TILE: usize = 512;

pub struct MiGan {
    model: dr_inference_engine::Model,
}

impl MiGan {
    /// From the model file, in whichever form the engine's rung wants
    /// (`resolve_model` picks the `.a16w16.onnx` sibling on the Hexagon:
    /// int8 moved the fill 16 dB from f32's, 16-bit about 41).
    pub fn from_path(path: &std::path::Path) -> Result<Self, PanoError> {
        use dr_inference_engine::{resolve_model, Role};
        let (path, form) = resolve_model(Role::Inpainter, path);
        let bytes = std::fs::read(&path).map_err(PanoError::ModelRead)?;
        Self::from_bytes(&bytes, form)
    }

    pub fn from_bytes(bytes: &[u8], form: dr_inference_engine::Form) -> Result<Self, PanoError> {
        use dr_inference_engine::Role;
        Ok(MiGan {
            model: dr_inference_engine::open(Role::Inpainter, form, bytes)?,
        })
    }

    /// Where the fill runs, for a status line.
    pub fn rung(&self) -> Result<dr_inference_engine::Rung, PanoError> {
        Ok(self.model.acquire()?.rung())
    }
}

impl Inpainter for MiGan {
    fn tile(&self) -> usize {
        TILE
    }

    fn fill(&mut self, rgb: &[f32], known: &[bool]) -> Result<Vec<f32>, PanoError> {
        let n = TILE * TILE;
        if rgb.len() != n * 3 || known.len() != n {
            return Err(PanoError::Input(format!(
                "MI-GAN takes a {TILE}×{TILE} tile; given {} values and {} mask entries",
                rgb.len(),
                known.len()
            )));
        }
        // NCHW: the mask plane, then the three masked colour planes.
        let mut input = vec![0.0f32; 4 * n];
        for i in 0..n {
            let m = if known[i] { 1.0 } else { 0.0 };
            input[i] = m - 0.5;
            for c in 0..3 {
                input[(c + 1) * n + i] = (rgb[i * 3 + c] * 2.0 - 1.0) * m;
            }
        }
        let tensor = ort::value::Tensor::from_array(
            ndarray::Array::from_shape_vec(ndarray::IxDyn(&[1, 4, TILE, TILE]), input)
                .expect("shape matches by construction"),
        )?;
        let started = std::time::Instant::now();
        let acquired = self.model.acquire()?;
        let acquired_at = started.elapsed();
        let mut session = acquired.lock();
        let outputs = session.run(ort::inputs![tensor])?;
        log::trace!(
            "migan: tile on {} — acquire {:.1} ms, run {:.1} ms",
            acquired.rung().label(),
            acquired_at.as_secs_f64() * 1e3,
            (started.elapsed() - acquired_at).as_secs_f64() * 1e3
        );
        let (shape, data) = outputs[0].try_extract_tensor::<f32>()?;
        let dims: Vec<i64> = shape.iter().copied().collect();
        if dims != [1, 3, TILE as i64, TILE as i64] {
            return Err(PanoError::Model(format!(
                "MI-GAN output is {dims:?}, expected [1, 3, {TILE}, {TILE}]"
            )));
        }
        let mut out = vec![0.0f32; n * 3];
        for i in 0..n {
            for c in 0..3 {
                out[i * 3 + c] = (data[c * n + i] * 0.5 + 0.5).clamp(0.0, 1.0);
            }
        }
        Ok(out)
    }
}