From 80eb294283572e46fec06a770418d7f05a370e57 Mon Sep 17 00:00:00 2001
From: cdoko <190060110+cdoko@users.noreply.github.com>
Date: Tue, 31 Dec 2024 04:57:53 -0400
Subject: [PATCH] Add device mapper

---
 .../src/vision_models/llava/llava_inputs_processor.rs        | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/mistralrs-core/src/vision_models/llava/llava_inputs_processor.rs b/mistralrs-core/src/vision_models/llava/llava_inputs_processor.rs
index 5820ae736..e9a13bb5b 100644
--- a/mistralrs-core/src/vision_models/llava/llava_inputs_processor.rs
+++ b/mistralrs-core/src/vision_models/llava/llava_inputs_processor.rs
@@ -14,6 +14,7 @@ use tracing::warn;
 
 use super::llava15::LLaVAVisionSpecificArgs;
 use super::utils::{expand2square, LLaVAImageProcessor};
+use crate::device_map::DeviceMapper;
 use crate::pipeline::text_models_inputs_processor::{
     get_completion_input, get_prompt_input, PagedAttentionMeta,
 };
@@ -87,6 +88,7 @@ impl InputsProcessor for LLaVAInputProcessor {
         other_config: Option<Arc<dyn Any>>,
         mut paged_attn_metadata: Option<PagedAttentionMeta<'_>>,
         prompt_batchsize: Option<NonZeroUsize>,
+        _mapper: Option<&dyn DeviceMapper>,
     ) -> Box<dyn Iterator<Item = anyhow::Result<InputProcessorOutput>>> {
         if is_xlora {
             return Box::new(std::iter::once(Err(anyhow::Error::msg(
@@ -167,6 +169,7 @@ impl InputsProcessor for LLaVAInputProcessor {
                         other_config,
                         paged_attn_metadata,
                         None, // TODO
+                        None,
                     )
                     .map(|metadata| {
                         let InputProcessorOutput {
@@ -283,6 +286,7 @@ impl InputsProcessor for LLaVAInputProcessor {
                 return_raw_logits,
                 paged_attn_metadata.as_mut(),
                 None, // TODO: evaluate if it is possible to batch this
+                None,
             )
         } else {
             get_completion_input(
@@ -294,6 +298,7 @@ impl InputsProcessor for LLaVAInputProcessor {
                 return_raw_logits,
                 paged_attn_metadata.as_mut(),
                 None, // TODO: evaluate if it is possible to batch this
+                None,
             )
         };