car-inference 0.63.0

//! External Flux image backend — shells out to `mflux-generate`.
//!
//! **Not a temporary bridge any more.** This module was written as one, while
//! the native Rust port collapsed every prompt to near-white output; that is
//! fixed and `mlx_flux.rs` now reaches prompt-faithful parity, so the old
//! header describing this as a stopgap was stale and actively misleading — it
//! reads as "the native path is broken" long after it stopped being true.
//!
//! What this module is now is the **architecture escape hatch**, the image
//! counterpart of `vllm_pool` for text. `mlx_flux.rs` implements exactly one
//! architecture (Flux.1-lite-8B). mflux implements many — FLUX.2, Krea 2,
//! Qwen-Image, Z-Image, Fibo — and gains more without CAR writing any Rust.
//! [`native_backend_serves`] decides between them on what the requested model
//! actually needs.
//!
//! Dispatch is by capability, not by an env toggle. `CAR_IMAGE_BACKEND` used to
//! choose the backend and was removed: it existed to A/B the Rust port during
//! the parity migration, outlived that migration, and defaulted to `native`
//! unconditionally — so asking for a FLUX.2 checkpoint silently routed to a
//! backend that cannot load it.
//!
//! The subprocess call stays at the feature boundary (command → PNG path on
//! disk), never an internal runtime dependency.

use std::process::Command;

use crate::tasks::generate_image::{GenerateImageRequest, GenerateImageResult};
use crate::InferenceError;

const CLI_BINARY: &str = "mflux-generate";

/// Default Flux checkpoint to pass via `--model`. Same quantized weights
/// the native `mlx_flux` backend loads — so no additional download cost
/// when swapping backends.
const DEFAULT_MODEL: &str = "mlx-community/Flux-1.lite-8B-MLX-Q4";
/// mflux needs a `--base-model` hint for custom `--model` paths; the
/// Flux.1-lite checkpoint is a distilled variant of Flux.1-dev, so `dev`
/// is the right architecture family.
const DEFAULT_BASE_MODEL: &str = "dev";

/// Checkpoints the native Rust `mlx_flux` backend implements.
///
/// `mlx_flux.rs` is a hand-written implementation of the **Flux.1-lite-8B**
/// architecture against these exact quantized weights — it is not a general
/// Flux loader. Anything else (FLUX.2, Krea 2, Qwen-Image, Z-Image, Fibo) has
/// no Rust implementation and must go through mflux.
const NATIVE_SERVED_MODELS: &[&str] = &["mlx-community/Flux-1.lite-8B-MLX-Q4"];

/// Whether the native Rust backend can serve `model` (`None` = CAR's default).
///
/// This replaces the old `CAR_IMAGE_BACKEND` env toggle. The toggle existed to
/// A/B the Rust port against mflux while the port was reaching parity; parity
/// landed, and the toggle outlived it. Worse, it defaulted to `native`
/// unconditionally, so asking for a FLUX.2 checkpoint silently routed to a
/// backend that cannot load that architecture. Letting the model decide is the
/// same rule the text path uses in `backend::local::has_native_backend`.
pub fn native_backend_serves(model: Option<&str>) -> bool {
    match model {
        // The default is the checkpoint the native backend implements.
        None => true,
        Some(m) => NATIVE_SERVED_MODELS.contains(&m),
    }
}

/// mflux's own built-in model identifiers, from `mflux-generate --base-model`.
///
/// When the caller names one of these it *is* the base model: mflux resolves
/// the checkpoint itself and no `--model` path is needed. Pairing one with the
/// hardcoded `dev` hint describes a FLUX.2 or Krea checkpoint as Flux.1-dev,
/// which is why every architecture past Flux.1 was unreachable through this
/// bridge even though the installed mflux already supported them.
const BUILTIN_BASE_MODELS: &[&str] = &[
    "dev",
    "schnell",
    "krea-dev",
    "dev-krea",
    "qwen",
    "fibo",
    "fibo-lite",
    "fibo-edit",
    "fibo-edit-rmbg",
    "z-image",
    "z-image-turbo",
    "flux2-klein-4b",
    "flux2-klein-9b",
    "flux2-klein-base-4b",
    "flux2-klein-base-9b",
];

/// Whether `model` is one of mflux's built-in identifiers rather than a
/// HuggingFace repo path.
fn is_builtin_base_model(model: &str) -> bool {
    BUILTIN_BASE_MODELS.contains(&model)
}

/// Absolute path to `mflux-generate`, preferring CAR's managed runtime — but
/// only a copy that actually runs.
///
/// Two failure modes, and fixing one naively causes the other:
///
/// - `Command::new("mflux-generate")` searches PATH only, so on a machine where
///   CAR provisioned the venv but the user never ran `uv tool install mflux`,
///   availability probed false and image generation fell back to the native
///   backend — which this bridge exists precisely because it is broken.
/// - Preferring the managed copy *unconditionally* is equally wrong: the
///   sibling `ltx-2-mlx` in this very runtime raises `ModuleNotFoundError: No
///   module named 'ltx_core_mlx'` while the PATH install works fine. A
///   dependency-broken venv entry point exists on disk and fails on import.
///
/// So preference is by *working*, not by existing: probe the managed copy and
/// fall back to PATH when it cannot run. Same best-available rule as the
/// vllm-mlx binary, keyed on functioning rather than version.
fn cli_path() -> std::path::PathBuf {
    if let Some(managed) = managed_cli() {
        if runs(&managed) {
            return managed;
        }
    }
    std::path::PathBuf::from(CLI_BINARY)
}

/// The managed runtime's copy of the CLI, if the file is present.
fn managed_cli() -> Option<std::path::PathBuf> {
    let home = std::env::var_os("HOME")?;
    let p = std::path::Path::new(&home)
        .join(".car")
        .join("visual-runtime")
        .join("bin")
        .join(CLI_BINARY);
    p.is_file().then_some(p)
}

/// Whether `bin` executes successfully. A venv entry point whose imports are
/// broken exits non-zero here despite existing on disk.
fn runs(bin: &std::path::Path) -> bool {
    Command::new(bin)
        .arg("--help")
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
}

pub fn is_available() -> bool {
    Command::new(cli_path())
        .arg("--help")
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
}

/// The `mflux-generate` command for one image, with the hub cache and token.
fn generate_command(
    program: &std::path::Path,
    req: &GenerateImageRequest,
    output_path: &str,
) -> Command {
    let mut cmd = Command::new(program);
    crate::hf_cache::apply_with_token_std(&mut cmd);
    let model = req.model.as_deref().unwrap_or(DEFAULT_MODEL);
    if is_builtin_base_model(model) {
        // The name IS the architecture; mflux resolves the checkpoint.
        cmd.arg("--base-model").arg(model);
    } else {
        // A custom HuggingFace checkpoint still needs an architecture hint.
        cmd.arg("--base-model")
            .arg(DEFAULT_BASE_MODEL)
            .arg("--model")
            .arg(model);
    }
    cmd.arg("--prompt")
        .arg(&req.prompt)
        .arg("--output")
        .arg(output_path);

    if let Some(w) = req.width {
        cmd.arg("--width").arg(w.to_string());
    }
    if let Some(h) = req.height {
        cmd.arg("--height").arg(h.to_string());
    }
    if let Some(s) = req.steps {
        cmd.arg("--steps").arg(s.to_string());
    }
    if let Some(g) = req.guidance {
        cmd.arg("--guidance").arg(g.to_string());
    }
    if let Some(seed) = req.seed {
        cmd.arg("--seed").arg(seed.to_string());
    }
    cmd
}

pub fn generate_image(req: &GenerateImageRequest) -> Result<GenerateImageResult, InferenceError> {
    let output_path = req
        .output_path
        .clone()
        .unwrap_or_else(|| "output.png".to_string());

    let mut cmd = generate_command(&cli_path(), req, &output_path);

    tracing::info!(prompt = %req.prompt, output = %output_path, "external mflux: invoking");
    let output = cmd.output().map_err(|e| {
        InferenceError::InferenceFailed(format!(
            "failed to spawn `{CLI_BINARY}`: {e}. \
             Install with `uv pip install mflux` and put its venv's bin on PATH."
        ))
    })?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(InferenceError::InferenceFailed(format!(
            "mflux-generate exited with status {}: stderr={stderr}",
            output.status
        )));
    }

    Ok(GenerateImageResult {
        image_path: output_path,
        media_type: "image/png".to_string(),
        model_used: Some(format!(
            "external:{}",
            req.model.as_deref().unwrap_or(DEFAULT_MODEL)
        )),
    })
}

#[cfg(test)]
mod model_selection_tests {
    use super::*;

    /// The bug: `--base-model` was hardcoded to `dev`, so asking for a FLUX.2
    /// or Krea checkpoint described it to mflux as Flux.1-dev. Every
    /// architecture past Flux.1 was unreachable through this bridge even though
    /// the installed mflux already listed them under `--base-model`.
    #[test]
    fn mflux_builtin_identifiers_are_recognized() {
        for m in [
            "flux2-klein-4b",
            "flux2-klein-9b",
            "flux2-klein-base-4b",
            "flux2-klein-base-9b",
            "krea-dev",
            "qwen",
            "z-image-turbo",
            "schnell",
            "dev",
        ] {
            assert!(
                is_builtin_base_model(m),
                "`{m}` is an mflux base model and must be passed as --base-model"
            );
        }
    }

    /// A HuggingFace repo path is not a base model — it still needs the
    /// architecture hint alongside `--model`.
    #[test]
    fn huggingface_checkpoints_are_not_builtins() {
        for m in [
            "mlx-community/Flux-1.lite-8B-MLX-Q4",
            "AITRADER/FLUX2-klein-4B-mlx-4bit",
            "some-org/some-model",
        ] {
            assert!(!is_builtin_base_model(m), "`{m}` is a checkpoint path");
        }
        // The default CAR ships with is a checkpoint path, so the hint stays.
        assert!(!is_builtin_base_model(DEFAULT_MODEL));
    }

    /// Resolution must prefer CAR's managed runtime over PATH, so a machine
    /// where CAR provisioned the venv but the user never ran `uv tool install
    /// mflux` does not silently fall back to the broken native backend.
    #[test]
    fn cli_path_prefers_the_managed_runtime_when_present() {
        let p = cli_path();
        if p.is_absolute() {
            assert!(
                p.ends_with(std::path::Path::new("visual-runtime/bin").join(CLI_BINARY)),
                "an absolute resolution must be the managed runtime: {}",
                p.display()
            );
        } else {
            // No managed runtime on this machine — falls back to a PATH lookup.
            assert_eq!(p, std::path::PathBuf::from(CLI_BINARY));
        }
    }
}

#[cfg(test)]
mod dispatch_tests {
    use super::*;

    /// Dispatch must follow the model, not an env var. The removed
    /// `CAR_IMAGE_BACKEND` toggle defaulted to `native` unconditionally, so a
    /// FLUX.2 request silently routed to a backend that cannot load it.
    #[test]
    fn architectures_without_a_rust_backend_route_external() {
        for m in [
            "flux2-klein-4b",
            "flux2-klein-9b",
            "krea-dev",
            "qwen",
            "z-image-turbo",
            "AITRADER/FLUX2-klein-4B-mlx-4bit",
        ] {
            assert!(
                !native_backend_serves(Some(m)),
                "`{m}` has no Rust implementation and must route to mflux"
            );
        }
    }

    /// The one checkpoint `mlx_flux.rs` actually implements stays native — it
    /// reached prompt-faithful parity, and a subprocess hop would be pure cost.
    #[test]
    fn the_implemented_checkpoint_stays_native() {
        assert!(native_backend_serves(Some(DEFAULT_MODEL)));
        assert!(
            native_backend_serves(None),
            "CAR's default model is the one the native backend implements"
        );
    }

    /// No environment variable may change the answer — that is the whole point
    /// of removing the toggle.
    #[test]
    fn dispatch_is_independent_of_the_environment() {
        let before = (
            native_backend_serves(None),
            native_backend_serves(Some("flux2-klein-4b")),
        );
        for key in ["CAR_IMAGE_BACKEND", "CAR_VIDEO_BACKEND"] {
            // SAFETY: single-threaded test process; restored immediately below.
            unsafe { std::env::set_var(key, "external") };
        }
        let after = (
            native_backend_serves(None),
            native_backend_serves(Some("flux2-klein-4b")),
        );
        for key in ["CAR_IMAGE_BACKEND", "CAR_VIDEO_BACKEND"] {
            unsafe { std::env::remove_var(key) };
        }
        assert_eq!(
            before, after,
            "backend choice must not read the environment"
        );
    }
}

#[cfg(all(test, unix))]
mod hf_token_tests {
    use super::*;
    use crate::hf_cache::token_probe::{Probe, WITH_TOKEN};

    fn recorded_for(credentials: &[(&str, &str)]) -> String {
        let probe = Probe::new();
        let program = probe.dir().join(CLI_BINARY);
        probe.fake_program(&program);
        let req = GenerateImageRequest {
            prompt: "p".into(),
            ..Default::default()
        };
        let mut cmd = car_secrets::with_child_credentials(credentials, || {
            generate_command(&program, &req, "out.png")
        });
        cmd.output().unwrap();
        probe.recorded()
    }

    #[test]
    fn mflux_is_handed_the_token() {
        assert_eq!(recorded_for(WITH_TOKEN), "yes");
    }

    #[test]
    fn mflux_gets_no_token_when_there_is_none() {
        assert_eq!(recorded_for(&[]), "no");
    }
}