diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 17a545146..d7027f087 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -42,6 +42,8 @@ class NPUEncoder { virtual bool is_available() const = 0; virtual std::vector get_input_shape() const = 0; virtual std::vector get_output_shape() const = 0; + virtual bool has_input(const std::string&) const { return false; } + virtual std::vector 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_source_encoder_; bool audio_encode_via_npu(const std::vector& audio_features); - bool vision_encode_via_npu(const std::vector& pixel_values); + bool vision_encode_via_npu(const std::vector& pixel_values, + const std::vector* pixel_position_ids = nullptr); bool source_encode_via_npu(const std::vector& tokens); std::map> media_features_; diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 333c5d4a0..53a3fbb28 100644 --- a/cactus-engine/src/model.cpp +++ b/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& 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"); diff --git a/cactus-engine/src/model_npu.cpp b/cactus-engine/src/model_npu.cpp index 9302dcfd2..9dd8b4348 100644 --- a/cactus-engine/src/model_npu.cpp +++ b/cactus-engine/src/model_npu.cpp @@ -194,11 +194,12 @@ bool Model::audio_encode_via_npu(const std::vector& audio_features) { return true; } -bool Model::vision_encode_via_npu(const std::vector& pixel_values) { +bool Model::vision_encode_via_npu(const std::vector& pixel_values, + const std::vector* pixel_position_ids) { if (!npu_vision_encoder_ || !npu_vision_encoder_->is_available() || !vision_encoder_) { return false; } - const std::vector input_shape = npu_vision_encoder_->get_input_shape(); + const std::vector 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& 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& 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 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(d); + } + if (pixel_position_ids->size() > pos_elems) return false; + + std::vector positions_i32(pos_elems, -1); + for (size_t i = 0; i < pixel_position_ids->size(); ++i) { + positions_i32[i] = static_cast((*pixel_position_ids)[i]); + } + + const std::vector 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() diff --git a/cactus-engine/src/npu_ane.h b/cactus-engine/src/npu_ane.h index d291a885f..c03dcc0cf 100644 --- a/cactus-engine/src/npu_ane.h +++ b/cactus-engine/src/npu_ane.h @@ -44,6 +44,10 @@ class ANEEncoder : public NPUEncoder { std::vector get_output_shape() const override; + bool has_input(const std::string& name) const override; + + std::vector 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 get_output_shape() const override { return {}; } + bool has_input(const std::string&) const override { return false; } + + std::vector 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; } diff --git a/cactus-engine/src/npu_ane.mm b/cactus-engine/src/npu_ane.mm index 116f91585..d7334067a 100644 --- a/cactus-engine/src/npu_ane.mm +++ b/cactus-engine/src/npu_ane.mm @@ -917,6 +917,29 @@ - (size_t)predictMultiInputWithDict:(NSMutableDictionary ANEEncoder::get_input_shape_for(const std::string& name) const { + std::vector 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_; diff --git a/cactus-engine/src/tokenizer.cpp b/cactus-engine/src/tokenizer.cpp index c79fa1cad..e0ef7cec0 100644 --- a/cactus-engine/src/tokenizer.cpp +++ b/cactus-engine/src/tokenizer.cpp @@ -598,9 +598,7 @@ std::string Tokenizer::format_gemma4_style(const std::vector& 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_; diff --git a/cactus-kernels/cactus_kernels.h b/cactus-kernels/cactus_kernels.h index f371cffbe..4260e9e2a 100644 --- a/cactus-kernels/cactus_kernels.h +++ b/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(); diff --git a/cactus-kernels/src/image.cpp b/cactus-kernels/src/image.cpp index db0c3421b..2473af9b5 100644 --- a/cactus-kernels/src/image.cpp +++ b/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); } diff --git a/python/cactus/transpile/component_pipeline.py b/python/cactus/transpile/component_pipeline.py index 332096a31..0f79bf5fa 100644 --- a/python/cactus/transpile/component_pipeline.py +++ b/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 diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index 3324cae97..a2ae0a36a 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/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( diff --git a/python/cactus/transpile/npu/pipeline.py b/python/cactus/transpile/npu/pipeline.py index d864fbf9e..e04429298 100644 --- a/python/cactus/transpile/npu/pipeline.py +++ b/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, ) diff --git a/python/cactus/transpile/npu/vision.py b/python/cactus/transpile/npu/vision.py index 30c49bd8e..710230d66 100644 --- a/python/cactus/transpile/npu/vision.py +++ b/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