Fix Gemma 4 NPU vision encoder on non-square images

cactus-engine/src/engine.h

@@ -42,6 +42,8 @@ class NPUEncoder {
     virtual bool is_available() const = 0;
     virtual std::vector<int> get_input_shape() const = 0;
     virtual std::vector<int> get_output_shape() const = 0;
+    virtual bool has_input(const std::string&) const { return false; }
+    virtual std::vector<int> get_input_shape_for(const std::string&) const { return {}; }
     virtual __fp16* get_output_buffer() = 0;
     virtual size_t get_output_buffer_size() const = 0;
     virtual size_t encode_multimodal_input(
@@ -837,7 +839,8 @@ class Model {
     std::unique_ptr<npu::NPUEncoder> npu_source_encoder_;
 
     bool audio_encode_via_npu(const std::vector<float>& audio_features);
-    bool vision_encode_via_npu(const std::vector<float>& pixel_values);
+    bool vision_encode_via_npu(const std::vector<float>& pixel_values,
+                               const std::vector<int64_t>* pixel_position_ids = nullptr);
     bool source_encode_via_npu(const std::vector<uint32_t>& tokens);
 
     std::map<std::string, std::vector<uint8_t>> media_features_;

cactus-engine/src/model.cpp

@@ -1314,7 +1314,7 @@ void Model::run_vision_encoder(const std::string& image_path) {
                            Precision::FP32);
     } else {
         Gemma4ImagePreprocessed prep = preprocess_gemma4_image(image_path, config_);
-        if (has_npu_vision_encoder() && vision_encode_via_npu(prep.pixel_values)) {
+        if (has_npu_vision_encoder() && vision_encode_via_npu(prep.pixel_values, &prep.pixel_position_ids)) {
             return;
         }
         int pv_idx = input_index(*vision_encoder_, "pixel_values");
@@ -2155,7 +2155,7 @@ bool Model::run_chunk_prefill_path(const std::vector<uint32_t>& tokens,
                 qwen_images.push_back(std::move(prep));
             } else {
                 Gemma4ImagePreprocessed prep = preprocess_gemma4_image(path, config_);
-                if (has_npu_vision_encoder() && vision_encode_via_npu(prep.pixel_values)) {
+                if (has_npu_vision_encoder() && vision_encode_via_npu(prep.pixel_values, &prep.pixel_position_ids)) {
                     continue;
                 }
                 int pv_idx = input_index(*vision_encoder_, "pixel_values");

cactus-engine/src/model_npu.cpp

@@ -194,11 +194,12 @@ bool Model::audio_encode_via_npu(const std::vector<float>& audio_features) {
     return true;
 }
 
-bool Model::vision_encode_via_npu(const std::vector<float>& pixel_values) {
+bool Model::vision_encode_via_npu(const std::vector<float>& pixel_values,
+                                  const std::vector<int64_t>* pixel_position_ids) {
     if (!npu_vision_encoder_ || !npu_vision_encoder_->is_available() || !vision_encoder_) {
         return false;
     }
-    const std::vector<int> input_shape = npu_vision_encoder_->get_input_shape();
+    const std::vector<int> input_shape = npu_vision_encoder_->get_input_shape_for("x");
     if (input_shape.empty()) return false;
 
     size_t expected_elems = 1;
@@ -208,6 +209,13 @@ bool Model::vision_encode_via_npu(const std::vector<float>& pixel_values) {
     }
     if (pixel_values.size() > expected_elems) return false;
 
+    const bool package_takes_positions = npu_vision_encoder_->has_input("pixel_position_ids");
+    if (package_takes_positions != (pixel_position_ids != nullptr)) {
+        CACTUS_LOG_WARN("model", "NPU vision encoder and pixel_position_ids mismatch; "
+            "falling back to CPU vision encoder (re-transpile with --npu to fix)");
+        return false;
+    }
+
     std::vector<__fp16> input_fp16(expected_elems, __fp16(0));
     for (size_t i = 0; i < pixel_values.size(); ++i) {
         input_fp16[i] = static_cast<__fp16>(pixel_values[i]);
@@ -224,8 +232,33 @@ bool Model::vision_encode_via_npu(const std::vector<float>& pixel_values) {
     }
     std::vector<__fp16> output_fp16(output_elems, __fp16(0));
 
-    size_t written = npu_vision_encoder_->encode(
-        input_fp16.data(), output_fp16.data(), input_shape, "x", "encoded");
+    size_t written = 0;
+    if (package_takes_positions) {
+        const std::vector<int> pos_shape =
+            npu_vision_encoder_->get_input_shape_for("pixel_position_ids");
+        if (pos_shape.empty()) return false;
+        size_t pos_elems = 1;
+        for (int d : pos_shape) {
+            if (d <= 0) return false;
+            pos_elems *= static_cast<size_t>(d);
+        }
+        if (pixel_position_ids->size() > pos_elems) return false;
+
+        std::vector<int32_t> positions_i32(pos_elems, -1);
+        for (size_t i = 0; i < pixel_position_ids->size(); ++i) {
+            positions_i32[i] = static_cast<int32_t>((*pixel_position_ids)[i]);
+        }
+
+        const std::vector<npu::NPUNamedInput> inputs = {
+            {"x", input_fp16.data(), npu::NPUNamedInput::DataType::FP16, input_shape},
+            {"pixel_position_ids", positions_i32.data(), npu::NPUNamedInput::DataType::INT32, pos_shape},
+        };
+        written = npu_vision_encoder_->encode_multimodal_input(
+            inputs, output_fp16.data(), "encoded");
+    } else {
+        written = npu_vision_encoder_->encode(
+            input_fp16.data(), output_fp16.data(), input_shape, "x", "encoded");
+    }
     if (written == 0) return false;
 
     for (size_t i = 0; i < vision_encoder_->output_node_ids.size()

cactus-engine/src/npu_ane.h

@@ -44,6 +44,10 @@ class ANEEncoder : public NPUEncoder {
 
     std::vector<int> get_output_shape() const override;
 
+    bool has_input(const std::string& name) const override;
+
+    std::vector<int> get_input_shape_for(const std::string& name) const override;
+
     __fp16* get_output_buffer() override;
 
     size_t get_output_buffer_size() const override;
@@ -82,6 +86,10 @@ class ANEEncoder : public NPUEncoder {
 
     std::vector<int> get_output_shape() const override { return {}; }
 
+    bool has_input(const std::string&) const override { return false; }
+
+    std::vector<int> get_input_shape_for(const std::string&) const override { return {}; }
+
     __fp16* get_output_buffer() override { return nullptr; }
 
     size_t get_output_buffer_size() const override { return 0; }

cactus-engine/src/npu_ane.mm

@@ -917,6 +917,29 @@ - (size_t)predictMultiInputWithDict:(NSMutableDictionary<NSString*, MLFeatureVal
     return result;
 }
 
+bool ANEEncoder::has_input(const std::string& name) const {
+    if (!impl_) return false;
+    CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_;
+    if (!impl.modelDescription) return false;
+    NSString* nsName = [NSString stringWithUTF8String:name.c_str()];
+    return impl.modelDescription.inputDescriptionsByName[nsName] != nil;
+}
+
+std::vector<int> ANEEncoder::get_input_shape_for(const std::string& name) const {
+    std::vector<int> result;
+    if (!impl_) return result;
+    CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_;
+    if (!impl.modelDescription) return result;
+    NSString* nsName = [NSString stringWithUTF8String:name.c_str()];
+    MLFeatureDescription* desc = impl.modelDescription.inputDescriptionsByName[nsName];
+    if (desc && desc.type == MLFeatureTypeMultiArray) {
+        for (NSNumber* dim in desc.multiArrayConstraint.shape) {
+            result.push_back([dim intValue]);
+        }
+    }
+    return result;
+}
+
 __fp16* ANEEncoder::get_output_buffer() {
     if (!impl_) return nullptr;
     CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_;

cactus-engine/src/tokenizer.cpp

@@ -598,9 +598,7 @@ std::string Tokenizer::format_gemma4_style(const std::vector<ChatMessage>& messa
         if (image_soft_token_count_ > 0) return image_soft_token_count_;
 
         int w = 0, h = 0, c = 0;
-        unsigned char* data = cactus_image_load(image_path.c_str(), &w, &h, &c, 3);
-        if (!data) return 0;
-        cactus_image_free(data);
+        if (!cactus_image_info(image_path.c_str(), &w, &h, &c)) return 0;
 
         uint32_t p = vision_patch_size_;
         uint32_t k = vision_pooling_kernel_size_;

cactus-kernels/cactus_kernels.h

@@ -795,6 +795,7 @@ void cactus_compute_spectrogram_f32(
     bool remove_dc_offset);
 
 unsigned char* cactus_image_load(const char* path, int* width, int* height, int* channels, int desired_channels);
+int cactus_image_info(const char* path, int* width, int* height, int* channels);
 void cactus_image_free(unsigned char* data);
 const char* cactus_image_failure_reason();
 

cactus-kernels/src/image.cpp

@@ -38,6 +38,10 @@ unsigned char* cactus_image_load(const char* path, int* width, int* height, int*
     return stbi_load(path, width, height, channels, desired_channels);
 }
 
+int cactus_image_info(const char* path, int* width, int* height, int* channels) {
+    return stbi_info(path, width, height, channels);
+}
+
 void cactus_image_free(unsigned char* data) {
     stbi_image_free(data);
 }

python/cactus/transpile/component_pipeline.py

@@ -28,6 +28,8 @@ class ComponentModuleSpec:
     output_keys: tuple[str, ...]
     graph_meta: dict[str, object] = field(default_factory=dict)
     metadata: dict[str, object] = field(default_factory=dict)
+    npu_module: torch.nn.Module | None = None
+    npu_runtime_input_count: int = 1
 
 
 @dataclass

python/cactus/transpile/model_adapters.py

@@ -798,9 +798,10 @@ def _gemma4_pool_vision_hidden_native_like(
     pooling_kernel_size = int(_module_or_config_attr(vision_tower, "pooling_kernel_size", 3) or 3)
     padding_positions = (pixel_position_ids == -1).all(dim=-1)
     pooled_batches: list[torch.Tensor] = []
-    for row_idx, (hidden_row, position_row, padding_row) in enumerate(
-        zip(hidden_states, pixel_position_ids, padding_positions, strict=True)
-    ):
+    for row_idx in range(int(hidden_states.shape[0])):
+        hidden_row = hidden_states[row_idx]
+        position_row = pixel_position_ids[row_idx]
+        padding_row = padding_positions[row_idx]
         if image_pool_shapes is not None and row_idx < len(image_pool_shapes):
             grid_h, grid_w, pooled_count = image_pool_shapes[row_idx]
             valid_patch_count = int(grid_h) * int(grid_w)
@@ -3458,6 +3459,12 @@ def _require(component: str) -> None:
         native_image_soft_token_counts=native_image_soft_token_counts,
         native_image_pool_shapes=native_image_pool_shapes,
     ).eval()
+    vision_encoder_npu = Gemma4VisionEncoderAdapter(
+        model,
+        weights_dir=weights_dir,
+        native_image_soft_token_counts=native_image_soft_token_counts,
+        native_image_pool_shapes=None,
+    ).eval()
     audio_encoder = Gemma4AudioEncoderAdapter(model, weights_dir=weights_dir).eval()
 
     image_features: torch.Tensor | None = None
@@ -3585,6 +3592,8 @@ def _require(component: str) -> None:
             output_keys=("image_features",),
             graph_meta={**common_graph_meta, "component": "vision_encoder"},
             metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"},
+            npu_module=vision_encoder_npu,
+            npu_runtime_input_count=2,
         ))
     if "audio_encoder" in expanded_components:
         specs.append(ComponentModuleSpec(

python/cactus/transpile/npu/pipeline.py

@@ -51,26 +51,36 @@ def _resolve(per_component: int | None, default: int | None) -> int | None:
         if not example_inputs:
             print(f"npu.pipeline: {component} spec has no example inputs; skipping")
             continue
-        primary = example_inputs[0]
-        baked = example_inputs[1:]
+        module = spec.npu_module or spec.module
+        n_runtime = min(max(1, spec.npu_runtime_input_count), len(example_inputs))
         qbits = component_quants[component]
         qdesc = f"int{qbits}" if qbits else "fp16"
         print(f"npu.pipeline: emitting {component} quant={qdesc}")
         if component == "source_encoder":
             emitted = emit_fn(
-                spec.module,
+                module,
                 bundle_root,
                 example_inputs=example_inputs,
                 input_names=tuple(getattr(spec, "input_keys", ()) or ()),
                 filename=filename,
                 quantize_bits=qbits,
             )
+        elif component == "vision_encoder":
+            names = ("x",) + tuple(spec.input_keys[1:n_runtime])
+            emitted = emit_fn(
+                module,
+                bundle_root,
+                runtime_inputs=tuple(zip(names, example_inputs[:n_runtime])),
+                baked_inputs=example_inputs[n_runtime:],
+                filename=filename,
+                quantize_bits=qbits,
+            )
         else:
             emitted = emit_fn(
-                spec.module,
+                module,
                 bundle_root,
-                example_input=primary,
-                baked_inputs=baked,
+                example_input=example_inputs[0],
+                baked_inputs=example_inputs[1:],
                 filename=filename,
                 quantize_bits=qbits,
             )

python/cactus/transpile/npu/vision.py

@@ -18,9 +18,13 @@ def __init__(
         for idx, tensor in enumerate(baked_inputs):
             self.register_buffer(f"_baked_{idx}", tensor, persistent=False)
 
-    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+    def forward(self, *runtime_inputs: torch.Tensor) -> torch.Tensor:
+        coerced = tuple(
+            t if torch.is_floating_point(t) else t.to(torch.long)
+            for t in runtime_inputs
+        )
         extra = tuple(getattr(self, f"_baked_{i}") for i in range(self._n_baked))
-        return self.vision(pixel_values, *extra)
+        return self.vision(*coerced, *extra)
 
 
 def _import_coremltools() -> Any:
@@ -56,22 +60,22 @@ def emit_vision_encoder_mlpackage(
     vision_module: torch.nn.Module,
     bundle_dir: Path,
     *,
-    example_input: torch.Tensor,
+    runtime_inputs: tuple[tuple[str, torch.Tensor], ...],
     baked_inputs: tuple[torch.Tensor, ...] = (),
     filename: str = "vision_encoder.mlpackage",
-    input_name: str = "x",
     output_name: str = "encoded",
     minimum_deployment_target: str = "iOS18",
     quantize_bits: int | None = None,
 ) -> str | None:
     ct = _import_coremltools()
 
-    wrapper = VisionEncoderWrapper(vision_module, baked_inputs)
+    example_tensors = tuple(tensor for _, tensor in runtime_inputs)
+    wrapper = VisionEncoderWrapper(vision_module, baked_inputs=baked_inputs)
     wrapper.eval()
 
     try:
         with torch.no_grad():
-            exported = torch.export.export(wrapper, (example_input,))
+            exported = torch.export.export(wrapper, example_tensors)
             exported = exported.run_decompositions({})
     except Exception as exc:
         print(f"npu.vision: torch.export failed ({type(exc).__name__}: {exc}); skipping mlpackage emit")
@@ -83,11 +87,19 @@ def emit_vision_encoder_mlpackage(
 
     target_attr = getattr(ct.target, minimum_deployment_target, None) or ct.target.iOS17
 
+    import numpy as np
+    ct_inputs = [
+        ct.TensorType(name=name, shape=tuple(tensor.shape))
+        if torch.is_floating_point(tensor)
+        else ct.TensorType(name=name, shape=tuple(tensor.shape), dtype=np.int32)
+        for name, tensor in runtime_inputs
+    ]
+
     from .coremltools_patches import build_cactus_pass_pipeline
     try:
         mlmodel = ct.convert(
             exported,
-            inputs=[ct.TensorType(name=input_name, shape=tuple(example_input.shape))],
+            inputs=ct_inputs,
             outputs=[ct.TensorType(name=output_name)],
             compute_precision=ct.precision.FLOAT16,
             convert_to="mlprogram",
@@ -112,5 +124,6 @@ def emit_vision_encoder_mlpackage(
         print(f"npu.vision: mlpackage save failed ({type(exc).__name__}: {exc})")
         return None
 
-    print(f"npu.vision: wrote {out_path} (input_shape={tuple(example_input.shape)})")
+    shapes = ", ".join(f"{name}={tuple(t.shape)}" for name, t in runtime_inputs)
+    print(f"npu.vision: wrote {out_path} (inputs: {shapes})")
     return filename