av-denoise 0.3.1

use clap::Parser;
use tracing_subscriber::EnvFilter;

mod cli;
mod file_mode;
mod ingest;
mod progress;
mod stream_mode;

use cli::{Args, Command, InputSource};
use ingest::CliOptions;

#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;

/// Scene workers used when `--workers` is not given.
const DEFAULT_WORKERS: usize = 2;

/// Routes an input source to the pipeline that fits it.
///
/// A path is opened with ffms2 and split across scenes. Anything piped
/// streams y4m frame by frame with no scene detection.
fn run_input(opts: &CliOptions, input: &InputSource, workers: Option<usize>) -> Result<(), anyhow::Error> {
    match input {
        InputSource::File(path) => file_mode::run_file(opts, path, workers.unwrap_or(DEFAULT_WORKERS)),
        stream @ (InputSource::Stdin | InputSource::Fd(_)) => {
            if workers.is_some() {
                tracing::warn!("--workers is ignored for piped input, which cannot be split by scene");
            }

            tracing::info!(input = %stream, "reading a y4m stream");

            stream_mode::run_stream(opts, stream.open_reader()?)
        },
    }
}

fn main() -> anyhow::Result<()> {
    // cubecl spawns its per-device worker thread with no explicit stack
    // size (uses Rust's default 2 MiB). GPU kernel codegen runs on that
    // thread; at large --search-radius the (2R+1)^2 unrolled body in
    // the windowed NLM kernels in src/nlmeans/kernels/fused.rs
    // overflows the default stack. RUST_MIN_STACK is cached on first
    // read, so set it here before any GPU thread spawns.
    if std::env::var_os("RUST_MIN_STACK").is_none() {
        // SAFETY: still single-threaded, no other thread can race the env mutation.
        unsafe { std::env::set_var("RUST_MIN_STACK", "16777216") };
    }

    let args = Args::parse();

    if std::env::var("RUST_LOG").is_err() {
        unsafe { std::env::set_var("RUST_LOG", "info") };
    }

    tracing_subscriber::fmt()
        .with_env_filter(EnvFilter::from_default_env())
        .with_writer(progress::tracing_writer())
        .init();

    // Honor AV_DENOISE_COMPILATION_CACHE. Must run before Denoiser::create
    // since the first CubeCL client lazily locks the global config.
    match av_denoise::apply_compilation_cache_env() {
        Ok(Some(path)) => {
            tracing::info!(?path, "AV_DENOISE_COMPILATION_CACHE override active")
        },
        Ok(None) => {},
        Err(_) => anyhow::bail!("unable to apply AV_DENOISE_COMPILATION_CACHE, this is a bug."),
    }

    let Command::Nlmeans(nlm) = &args.command;
    let opts = nlm.build_options(&args)?;

    run_input(&opts, &nlm.input, nlm.workers)
}