feat(xla): admit image requests in CLI and continuous-batch serving

src/backend/session.rs

@@ -46,6 +46,98 @@ use mlxcel_core::session::{MlxInferenceSession, SessionCapabilities};
 #[cfg(feature = "xla-backend")]
 use mlxcel_xla::XlaInferenceSession;
 
+#[cfg(feature = "xla-backend")]
+pub(super) fn qualified_xla_capabilities(
+    mut engine_capabilities: SessionCapabilities,
+    has_image_preprocessor: bool,
+) -> SessionCapabilities {
+    engine_capabilities.multimodal = engine_capabilities.multimodal && has_image_preprocessor;
+    engine_capabilities
+}
+
+#[cfg(feature = "xla-backend")]
+pub(super) fn qualified_xla_supports_images(
+    engine_capabilities: SessionCapabilities,
+    has_image_preprocessor: bool,
+) -> bool {
+    qualified_xla_capabilities(engine_capabilities, has_image_preprocessor).multimodal
+}
+
+/// Root-crate OpenXLA session with the matching host image preprocessor.
+///
+/// The compiler runtime lives in `mlxcel-xla`, while the host preprocessor
+/// deliberately reuses this crate's qualified MLX vision tower and projector.
+/// Keeping both in one session makes capability reporting and image execution
+/// depend on the same loaded value.
+#[cfg(feature = "xla-backend")]
+pub struct XlaBackendSession {
+    engine: XlaInferenceSession,
+    image_preprocessor: Option<Box<dyn crate::HostMultimodalPreprocessor>>,
+}
+
+#[cfg(feature = "xla-backend")]
+impl XlaBackendSession {
+    #[must_use]
+    pub(crate) fn new(
+        engine: XlaInferenceSession,
+        image_preprocessor: Option<Box<dyn crate::HostMultimodalPreprocessor>>,
+    ) -> Self {
+        Self {
+            engine,
+            image_preprocessor,
+        }
+    }
+
+    /// Report only capabilities whose complete host/runtime path was loaded.
+    #[must_use]
+    pub fn capabilities(&self) -> SessionCapabilities {
+        qualified_xla_capabilities(
+            self.engine.capabilities(),
+            self.image_preprocessor.is_some(),
+        )
+    }
+
+    /// Whether this session can prepare and execute image-prefill requests.
+    #[must_use]
+    pub fn supports_images(&self) -> bool {
+        qualified_xla_supports_images(
+            self.engine.capabilities(),
+            self.image_preprocessor.is_some(),
+        )
+    }
+
+    /// Prepare decoded images through the exact preprocessor used by the
+    /// capability predicate.
+    pub fn prepare_images(
+        &self,
+        token_ids: &[i32],
+        images: &[image::DynamicImage],
+    ) -> Result<PreparedPrefill, String> {
+        let preprocessor = self.image_preprocessor.as_ref().ok_or_else(|| {
+            "the loaded OpenXLA model/runtime bundle does not support image input".to_string()
+        })?;
+        preprocessor
+            .prepare(token_ids, images)
+            .map_err(|error| error.to_string())
+    }
+}
+
+#[cfg(feature = "xla-backend")]
+impl std::ops::Deref for XlaBackendSession {
+    type Target = XlaInferenceSession;
+
+    fn deref(&self) -> &Self::Target {
+        &self.engine
+    }
+}
+
+#[cfg(feature = "xla-backend")]
+impl std::ops::DerefMut for XlaBackendSession {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.engine
+    }
+}
+
 /// The inference session selected for a load.
 ///
 /// Under default features this enum has a single variant, [`Session::Mlx`].
@@ -61,7 +153,7 @@ pub enum Session {
     /// on-device, so it self-drives the engine-neutral prefill / decode_step
     /// contract and ignores the threaded MLX model and sampling config (greedy).
     #[cfg(feature = "xla-backend")]
-    Xla(XlaInferenceSession),
+    Xla(XlaBackendSession),
 }
 
 impl Session {
@@ -76,8 +168,11 @@ impl Session {
     #[cfg(feature = "xla-backend")]
     #[inline]
     #[must_use]
-    pub fn xla(session: XlaInferenceSession) -> Self {
-        Session::Xla(session)
+    pub(crate) fn xla(
+        session: XlaInferenceSession,
+        image_preprocessor: Option<Box<dyn crate::HostMultimodalPreprocessor>>,
+    ) -> Self {
+        Session::Xla(XlaBackendSession::new(session, image_preprocessor))
     }
 
     /// What this session can do, so the control plane can gate features.

src/backend/tests.rs

@@ -124,6 +124,29 @@ fn mlx_session_threads_the_token_bias_through() {
     }
 }
 
+#[cfg(feature = "xla-backend")]
+#[test]
+fn xla_image_capability_requires_runtime_and_host_preprocessor() {
+    use super::session::{qualified_xla_capabilities, qualified_xla_supports_images};
+    use mlxcel_core::session::SessionCapabilities;
+
+    let text_runtime = SessionCapabilities::single_sequence();
+    let multimodal_runtime = text_runtime.with_multimodal();
+
+    assert!(!qualified_xla_supports_images(text_runtime, false));
+    assert!(!qualified_xla_supports_images(text_runtime, true));
+    assert!(!qualified_xla_supports_images(multimodal_runtime, false));
+    assert!(qualified_xla_supports_images(multimodal_runtime, true));
+    assert!(
+        !qualified_xla_capabilities(multimodal_runtime, false).multimodal,
+        "missing required host artifact must keep capability false"
+    );
+    assert!(
+        qualified_xla_capabilities(multimodal_runtime, true).multimodal,
+        "only a qualified runtime plus loaded preprocessor advertises images"
+    );
+}
+
 /// The experimental scaffold has no engine wired in, so it cannot produce a
 /// session. Compiled only under the optional feature; default builds carry none
 /// of it.
@@ -161,13 +184,14 @@ fn xla_backend_creates_a_single_sequence_session_scaffold() {
         "the XLA backend drives generation through the session, not load_model"
     );
 
+    let model_dir = tempfile::tempdir().expect("temporary text-model fixture");
+    std::fs::write(
+        model_dir.path().join("config.json"),
+        r#"{"model_type":"llama"}"#,
+    )
+    .expect("write text-model config");
     let session = backend
-        .create_session(
-            std::path::Path::new("/tmp/model"),
-            16,
-            KVCacheMode::Fp16,
-            TokenBiasMap::new(),
-        )
+        .create_session(model_dir.path(), 16, KVCacheMode::Fp16, TokenBiasMap::new())
         .expect("xla session creation must succeed");
     let caps = session.capabilities();
     assert!(

src/backend/xla.rs

@@ -107,9 +107,11 @@ impl ComputeBackend for XlaBackend {
         // without it, the session loads but `prefill` / `decode_step` report that
         // the `iree` feature is off. KV mode and token bias do not apply: the
         // session owns its own KV and samples greedily on-device.
+        let image_preprocessor = crate::load_xla_image_preprocessor(model_path)
+            .map_err(|error| anyhow::anyhow!("OpenXLA image preprocessor load failed: {error}"))?;
         let session =
             XlaInferenceSession::load(model_path, num_layers).map_err(|e| anyhow::anyhow!(e))?;
-        Ok(Session::xla(session))
+        Ok(Session::xla(session, image_preprocessor))
     }
 
     fn supports_batched_serving(&self) -> bool {

src/commands/generate.rs

@@ -18,6 +18,8 @@
 //! prompt preparation, generation-mode selection, and terminal output helpers.
 
 use anyhow::{Result, anyhow, ensure};
+#[cfg(feature = "xla-backend")]
+use std::io::Read;
 use std::io::{self, IsTerminal, Write as IoWrite};
 use std::path::Path;
 use std::time::{Duration, Instant};
@@ -1031,31 +1033,108 @@ fn xla_num_layers(model_dir: &Path) -> usize {
         .map_or(0, |n| n as usize)
 }
 
-/// Self-contained text generation for the OpenXLA backend (issue #449 Phase 3).
+#[cfg(any(feature = "xla-backend", test))]
+fn validate_xla_output_audio(output_audio: Option<&Path>) -> Result<()> {
+    ensure!(
+        output_audio.is_none(),
+        "--output-audio is not supported by the OpenXLA backend"
+    );
+    Ok(())
+}
+
+#[cfg(any(feature = "xla-backend", test))]
+fn validate_xla_cli_image_cardinality(declared: usize, decoded: usize) -> Result<()> {
+    ensure!(
+        declared == decoded,
+        "OpenXLA image decode cardinality mismatch: {declared} image path(s) provided, \
+         {decoded} decoded; refusing partial image execution"
+    );
+    Ok(())
+}
+
+#[cfg(feature = "xla-backend")]
+fn decode_xla_cli_images(paths: &[std::path::PathBuf]) -> Result<Vec<image::DynamicImage>> {
+    let limits = mlxcel::current_image_input_limits();
+    ensure!(
+        paths.len() <= limits.max_images_per_request,
+        "OpenXLA image request contains {} image(s), exceeding the configured maximum of {}",
+        paths.len(),
+        limits.max_images_per_request
+    );
+    let read_limit = limits
+        .max_payload_bytes
+        .checked_add(1)
+        .ok_or_else(|| anyhow!("configured image payload limit overflowed"))?;
+    let mut payloads = Vec::with_capacity(paths.len());
+    for path in paths {
+        let mut bytes = Vec::new();
+        std::fs::File::open(path)
+            .map_err(|error| anyhow!("Failed to open image {path:?}: {error}"))?
+            .take(read_limit as u64)
+            .read_to_end(&mut bytes)
+            .map_err(|error| anyhow!("Failed to read image {path:?}: {error}"))?;
+        ensure!(
+            bytes.len() <= limits.max_payload_bytes,
+            "Image payload {path:?} is {} bytes, exceeding the configured maximum of {}",
+            bytes.len(),
+            limits.max_payload_bytes
+        );
+        payloads.push(bytes);
+    }
+    let images = mlxcel::decode_image_payloads_with_limits(&payloads, limits)?;
+    validate_xla_cli_image_cardinality(paths.len(), images.len())?;
+    Ok(images)
+}
+
+/// Self-contained generation for the OpenXLA backend (issue #449 Phase 3).
 ///
 /// The OpenXLA engine drives generation from its own session and has no MLX
 /// `LoadedModel`, so this path skips `load_model` (which the XLA backend rejects)
 /// and the generic model-threaded loop: it creates the session straight from the
-/// model directory and runs the session's own greedy loop, which seeds KV and
-/// samples on-device. Text-only and greedy by design (no VLM / draft / sampling
-/// knobs), so the caller routes here only when `MLXCEL_BACKEND=xla` is selected.
+/// model directory and runs the session's own greedy loop. Image requests use
+/// the same decoded-image security limits and owned `PreparedPrefill` contract
+/// as serving; audio and video remain explicit unsupported modalities.
 #[cfg(feature = "xla-backend")]
-fn generate_xla_text(
+fn generate_xla(
     model_path: &Path,
     num_layers: usize,
     prompt_tokens: &[i32],
+    image_paths: &[std::path::PathBuf],
+    has_audio: bool,
+    has_video: bool,
     max_tokens: usize,
     kv_cache_mode: KVCacheMode,
     token_bias: TokenBiasMap,
 ) -> Result<(Vec<i32>, GenerationStats)> {
+    ensure!(
+        !has_audio,
+        "the OpenXLA backend does not support audio input yet"
+    );
+    ensure!(
+        !has_video,
+        "the OpenXLA backend does not support video input yet"
+    );
     let mut session =
         select_backend().create_session(model_path, num_layers, kv_cache_mode, token_bias)?;
     let start_time = Instant::now();
     let tokens = match &mut session {
         mlxcel::Session::Xla(s) => {
             let eos = s.eos_token_ids().to_vec();
-            s.generate_greedy(prompt_tokens, max_tokens, &eos)
-                .map_err(|e| anyhow!("OpenXLA generation failed: {e}"))?
+            if image_paths.is_empty() {
+                s.generate_greedy(prompt_tokens, max_tokens, &eos)
+                    .map_err(|e| anyhow!("OpenXLA generation failed: {e}"))?
+            } else {
+                ensure!(
+                    s.supports_images(),
+                    "the loaded OpenXLA model/runtime bundle does not support image input"
+                );
+                let images = decode_xla_cli_images(image_paths)?;
+                let prepared = s
+                    .prepare_images(prompt_tokens, &images)
+                    .map_err(|error| anyhow!("OpenXLA image preprocessing failed: {error}"))?;
+                s.generate_prepared_greedy(&prepared, max_tokens, &eos)
+                    .map_err(|error| anyhow!("OpenXLA image generation failed: {error}"))?
+            }
         }
         // `select_backend()` returned the XLA backend, so its `create_session`
         // yields an XLA session; any other variant would be a wiring bug.
@@ -1817,12 +1896,16 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> {
     // `token_bias` is sound because this branch diverges with `return`.)
     #[cfg(feature = "xla-backend")]
     if select_backend().name() == "xla" {
+        validate_xla_output_audio(args.generation.output_audio.as_deref())?;
         let num_layers = xla_num_layers(&args.model.model);
         print_generation_preamble(&user_prompt)?;
-        let (generated_tokens, stats) = generate_xla_text(
+        let (generated_tokens, stats) = generate_xla(
             &args.model.model,
             num_layers,
             &prompt_tokens,
+            &args.generation.image,
+            args.generation.audio.is_some(),
+            !args.generation.video.is_empty(),
             args.generation.max_tokens,
             kv_cache_mode,
             token_bias,

src/commands/generate_tests.rs

@@ -17,14 +17,45 @@ use super::{
     estimate_delta_label_and_bytes, generated_suffix, generation_stats_from_duration,
     memory_preflight_ctx_len, resolve_cli_pipeline_assignments, resolve_cli_prompt,
     should_route_offline_mtp, strip_trailing_eos, validate_pipeline_parallel_args,
-    validate_tensor_parallel_args,
+    validate_tensor_parallel_args, validate_xla_cli_image_cardinality, validate_xla_output_audio,
 };
 use mlxcel::server::chat_template::ChatTemplateProcessor;
 use mlxcel_core::drafter::DrafterKind;
 use std::fs;
 use std::path::PathBuf;
 use std::time::Duration;
 
+#[test]
+fn xla_rejects_output_audio_before_generation() {
+    assert!(validate_xla_output_audio(None).is_ok());
+    let error = validate_xla_output_audio(Some(std::path::Path::new("out.wav"))).unwrap_err();
+    assert!(
+        error
+            .to_string()
+            .contains("not supported by the OpenXLA backend")
+    );
+}
+
+#[test]
+fn xla_cli_rejects_partial_multi_image_decode() {
+    let payloads = vec![
+        include_bytes!("../../tests/fixtures/test_image.png").to_vec(),
+        b"not-an-image".to_vec(),
+    ];
+    let images =
+        mlxcel::decode_image_payloads_with_limits(&payloads, mlxcel::ImageInputLimits::default())
+            .unwrap();
+    assert_eq!(images.len(), 1, "the shared decoder remains tolerant");
+
+    let error = validate_xla_cli_image_cardinality(payloads.len(), images.len())
+        .unwrap_err()
+        .to_string();
+    assert!(error.contains("2 image path(s) provided"));
+    assert!(error.contains("1 decoded"));
+    assert!(error.contains("refusing partial image execution"));
+    assert!(validate_xla_cli_image_cardinality(1, images.len()).is_ok());
+}
+
 // issue #166: the offline MTP loop is constructed only when the operator
 // explicitly passes `--draft-kind mtp`. An auto-detected MTP shape (no explicit
 // flag) keeps the classic SpeculativeGenerator path, and DFlash / InternalMtp