Add Jev-Omni MLX 4-bit Mac mini release

Apple Silicon MLX conversion of akhilaaa3/Jev-Omni with local M4/16GB benchmarks and inference code.

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE.md

@@ -0,0 +1,7 @@
+# Attribution
+
+This repository is an Apple Silicon MLX conversion of [`akhilaaa3/Jev-Omni`](https://huggingface.co/akhilaaa3/Jev-Omni), pinned to revision `c050d51354147985d13286cf4acf90f562f2c631`.
+
+The converted weights use the upstream unified checkpoint and decision head. The upstream model card identifies the model as Apache-2.0 and states that it is an independent open model, not affiliated with TypeSafe AI.
+
+The conversion and inference wrapper in `omni_mlx/` are provided under Apache-2.0. The base Gemma 4 model and any dataset terms remain subject to their respective upstream notices.

README.md

@@ -0,0 +1,87 @@
+---
+license: apache-2.0
+base_model: akhilaaa3/Jev-Omni
+library_name: mlx
+pipeline_tag: image-text-to-text
+tags:
+- mlx
+- apple-silicon
+- multimodal
+- text-classification
+- image-classification
+- typed-decision
+---
+
+# Jev-Omni-MLX-4bit
+
+An Apple Silicon MLX conversion of [`akhilaaa3/Jev-Omni`](https://huggingface.co/akhilaaa3/Jev-Omni) for local inference on a Mac mini with 16GB unified memory.
+
+This repository keeps the upstream Jev-Omni unified multimodal checkpoint and its trained 256-way decision head, then converts the language-model weights to **4-bit affine quantization with group size 64**. The vision weights remain BF16 and the decision head remains FP32.
+
+It is an independent conversion. It is not an official TypeSafe Jev release and does not claim to reproduce TypeSafe's proprietary system. It is also not a new fine-tune.
+
+## Hardware and speed
+
+Measured on:
+
+- Mac mini, Apple M4, 10 CPU cores, 16GB unified memory
+- macOS 26.5.1
+- Python 3.13.12
+- MLX 0.32.2
+- MLX-VLM 0.7.1
+- single request, batch size 1, no token generation
+
+Ten warm requests were measured after one warm-up request. The image test used the recommended 20 visual-token budget and a 3-option question.
+
+| Mode | Median | P95 | Peak Metal memory |
+|---|---:|---:|---:|
+| Text decision | ~963 ms | ~998 ms | ~7.0 GB |
+| Image decision, 20 visual tokens | ~994 ms | ~1,021 ms | ~7.0 GB |
+
+The first request includes MLX graph and memory warm-up. On the same machine, a 70-token image request is slower (roughly 1.8 seconds warm in an earlier run). Lowering visual tokens reduces latency but can lose small details; validate on your own game frames.
+
+The published Jev-Omni H200 numbers are not transferable to this Mac mini. This model card reports local measurements only.
+
+## Validation
+
+- Upstream unified verification cases: 4/4 argmax decisions matched after 4-bit conversion.
+- Six simple red/blue/green circle and square image checks: 6/6 color decisions matched.
+- Maximum absolute probability difference on the four upstream text cases: 0.244 in this small check.
+- JevBench and DecisionBench were **not** re-run for this conversion.
+- Quantization changes the probability distribution; probabilities are not recalibrated here and must not be treated as calibrated confidence.
+
+## Installation
+
+This release is intended for Apple Silicon. Download the repository and install the small MLX runtime:
+
+```bash
+hf download Ruiruiz30/Jev-Omni-MLX-4bit \
+  --local-dir Jev-Omni-MLX-4bit
+
+cd Jev-Omni-MLX-4bit
+python3.13 -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+```
+
+## Image decision
+
+```bash
+python -m omni_mlx.classifier \
+  --model . \
+  --image /path/to/frame.png \
+  --state "A kart is approaching a right turn." \
+  --question "Which steering action is best?" \
+  --options "Turn left" "Hold center" "Turn right" \
+  --image-tokens 20
+```
+
+The classifier returns candidate probabilities and the selected option. It does not generate a free-form explanation. Use `--image-tokens 70` when the scene contains small or dense visual details.
+
+## Local conversion code
+
+`omni_mlx/convert.py` contains the conversion path used for this release. The original unquantized checkpoint is not bundled here; it can be obtained from the upstream repository under its own license and terms.
+
+## License and attribution
+
+Apache-2.0. See [LICENSE](LICENSE) and [NOTICE.md](NOTICE.md). The upstream model card, Gemma 4 terms, and dataset rights remain authoritative for their respective components.

chat_template.jinja

@@ -0,0 +1,390 @@
+{#
+  Template: Google Gemma 4 Canonical Chat Template
+  Author: Google Gemma Engineering Team
+  Published: 2026-07-09
+  Context: Fixed tool-calling loops, turn closures, and thinking content-ordering.
+#}
+{%- macro format_parameters(properties, required, filter_keys=false) -%}
+    {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
+    {%- set ns = namespace(found_first=false) -%}
+    {%- for key, value in properties | dictsort -%}
+        {%- set add_comma = false -%}
+        {%- if not filter_keys or key not in standard_keys -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {{ key }}:{
+            {%- if value['description'] -%}
+                description:<|"|>{{ value['description'] }}<|"|>
+                {%- set add_comma = true -%}
+            {%- endif -%}
+            {%- if value['type'] | upper == 'STRING' -%}
+                {%- if value['enum'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    enum:{{ format_argument(value['enum']) }}
+                {%- endif -%}
+            {%- elif value['type'] | upper == 'ARRAY' -%}
+                {%- if value['items'] is mapping and value['items'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    items:{
+                    {%- set ns_items = namespace(found_first=false) -%}
+                    {%- for item_key, item_value in value['items'] | dictsort -%}
+                        {%- if item_value is not none -%}
+                            {%- if ns_items.found_first %},{% endif -%}
+                            {%- set ns_items.found_first = true -%}
+                            {%- if item_key == 'properties' -%}
+                                properties:{
+                                {%- if item_value is mapping -%}
+                                    {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
+                                {%- endif -%}
+                                }
+                            {%- elif item_key == 'required' -%}
+                                required:[
+                                {%- for req_item in item_value -%}
+                                    <|"|>{{- req_item -}}<|"|>
+                                    {%- if not loop.last %},{% endif -%}
+                                {%- endfor -%}
+                                ]
+                            {%- elif item_key == 'type' -%}
+                                {%- if item_value is string -%}
+                                    type:{{ format_argument(item_value | upper) }}
+                                {%- else -%}
+                                    type:{{ format_argument(item_value | map('upper') | list) }}
+                                {%- endif -%}
+                            {%- else -%}
+                                {{ item_key }}:{{ format_argument(item_value) }}
+                            {%- endif -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                    }
+                {%- endif -%}
+            {%- endif -%}
+            {%- if value['nullable'] %}
+                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                nullable:true
+            {%- endif -%}
+            {%- if value['type'] | upper == 'OBJECT' -%}
+                {%- if value['properties'] is defined and value['properties'] is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+                    }
+                {%- elif value is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
+                    }
+                {%- endif -%}
+                {%- if value['required'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    required:[
+                    {%- for item in value['required'] | default([]) -%}
+                        <|"|>{{- item -}}<|"|>
+                        {%- if not loop.last %},{% endif -%}
+                    {%- endfor -%}
+                    ]
+                {%- endif -%}
+            {%- endif -%}
+            {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+            type:<|"|>{{ value['type'] | upper }}<|"|>}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endmacro -%}
+{%- macro format_function_declaration(tool_data) -%}
+    declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
+    {%- set params = tool_data['function']['parameters'] -%}
+    {%- if params -%}
+        ,parameters:{
+        {%- if params['properties'] -%}
+            properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
+        {%- endif -%}
+        {%- if params['required'] -%}
+            required:[
+            {%- for item in params['required'] -%}
+                <|"|>{{- item -}}<|"|>
+                {{- ',' if not loop.last -}}
+            {%- endfor -%}
+            ],
+        {%- endif -%}
+        {%- if params['type'] -%}
+            type:<|"|>{{- params['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if 'response' in tool_data['function'] -%}
+        {%- set response_declaration = tool_data['function']['response'] -%}
+        ,response:{
+        {%- if response_declaration['description'] -%}
+            description:<|"|>{{- response_declaration['description'] -}}<|"|>,
+        {%- endif -%}
+        {%- if response_declaration['type'] | upper == 'OBJECT' -%}
+            type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    }
+{%- endmacro -%}
+{%- macro format_argument(argument, escape_keys=True) -%}
+    {%- if argument is none -%}
+        {{- 'null' -}}
+    {%- elif argument is string -%}
+        {{- '<|"|>' + argument + '<|"|>' -}}
+    {%- elif argument is boolean -%}
+        {{- 'true' if argument else 'false' -}}
+    {%- elif argument is mapping -%}
+        {{- '{' -}}
+        {%- set ns = namespace(found_first=false) -%}
+        {%- for key, value in argument | dictsort -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {%- if escape_keys -%}
+                {{- '<|"|>' + key + '<|"|>' -}}
+            {%- else -%}
+                {{- key -}}
+            {%- endif -%}
+            :{{- format_argument(value, escape_keys=escape_keys) -}}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- elif argument is sequence -%}
+        {{- '[' -}}
+        {%- for item in argument -%}
+            {{- format_argument(item, escape_keys=escape_keys) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- ']' -}}
+    {%- else -%}
+        {{- argument -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro strip_thinking(text) -%}
+    {%- set ns = namespace(result='') -%}
+    {%- for part in text.split('<channel|>') -%}
+        {%- if '<|channel>' in part -%}
+            {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
+        {%- else -%}
+            {%- set ns.result = ns.result + part -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- ns.result | trim -}}
+{%- endmacro -%}
+
+{%- macro format_tool_response_block(tool_name, response) -%}
+    {{- '<|tool_response>' -}}
+    {%- if response is mapping -%}
+        {{- 'response:' + tool_name + '{' -}}
+        {%- for key, value in response | dictsort -%}
+            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- else -%}
+        {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
+    {%- endif -%}
+    {{- '<tool_response|>' -}}
+{%- endmacro -%}
+
+{#- ===== SETUP ===== -#}
+{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%}
+{%- set loop_messages = messages -%}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set preserve_thinking = preserve_thinking | default(false) -%}
+{{- bos_token -}}
+{#- Handle System/Tool Definitions Block -#}
+{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%}
+    {{- '<|turn>system\n' -}}
+    {#- Inject Thinking token at the very top of the FIRST system turn -#}
+    {%- if enable_thinking -%}
+        {{- '<|think|>\n' -}}
+        {%- set ns.prev_message_type = 'think' -%}
+    {%- endif -%}
+    {%- if messages and messages[0]['role'] in ['system', 'developer'] -%}
+        {%- if messages[0]['content'] is string -%}
+            {{- messages[0]['content'] | trim -}}
+        {%- elif messages[0]['content'] is sequence -%}
+            {%- for item in messages[0]['content'] -%}
+                {{- item['text'] | trim + ' '-}}
+            {%- endfor -%}
+        {%- endif -%}
+        {%- set loop_messages = messages[1:] -%}
+    {%- endif -%}
+    {%- if tools -%}
+        {%- for tool in tools %}
+            {{- '<|tool>' -}}
+            {{- format_function_declaration(tool) | trim -}}
+            {{- '<tool|>' -}}
+        {%- endfor %}
+        {%- set ns.prev_message_type = 'tool' -%}
+    {%- endif -%}
+    {{- '<turn|>\n' -}}
+{%- endif %}
+
+{#- Pre-scan: find last user message index for reasoning guard -#}
+{%- set ns_turn = namespace(last_user_idx=-1) -%}
+{%- for i in range(loop_messages | length) -%}
+    {%- if loop_messages[i]['role'] == 'user' -%}
+        {%- set ns_turn.last_user_idx = i -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- Loop through messages -#}
+{%- for message in loop_messages -%}
+    {%- if message['role'] != 'tool' -%}
+    {%- set ns.prev_message_type = None -%}
+    {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
+    {#- Detect continuation using tracked state — O(1) instead of O(n) backward scan -#}
+    {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%}
+    {%- if not continue_same_model_turn -%}
+        {{- '<|turn>' + role + '\n' }}
+
+    {%- endif -%}
+
+    {#- Render reasoning/reasoning_content as thinking channel -#}
+    {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
+    {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or (preserve_thinking and message.get('tool_calls')) -%}
+    {%- if thinking_text and thinking_gate -%}
+        {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
+    {%- endif -%}
+
+            {%- if message.get('tool_calls') -%}
+                {%- for tool_call in message.get('tool_calls') -%}
+                    {%- set function = tool_call['function'] -%}
+                    {{- '<|tool_call>call:' + function['name'] + '{' -}}
+                    {%- if function['arguments'] is mapping -%}
+                        {%- set ns_args = namespace(found_first=false) -%}
+                        {%- for key, value in function['arguments'] | dictsort -%}
+                            {%- if ns_args.found_first %},{% endif -%}
+                            {%- set ns_args.found_first = true -%}
+                            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+                        {%- endfor -%}
+                    {%- elif function['arguments'] is none -%}
+                    {%- else -%}
+                        {{- raise_exception(
+                            "chat_template: tool_calls[].function.arguments must be a "
+                            "JSON object (mapping), not a string. Deserialize arguments "
+                            "before passing to the template."
+                        ) -}}
+                    {%- endif -%}
+                    {{- '}<tool_call|>' -}}
+                {%- endfor -%}
+                {%- set ns.prev_message_type = 'tool_call' -%}
+            {%- endif -%}
+
+            {%- set ns_tr_out = namespace(flag=false) -%}
+            {%- if message.get('tool_responses') -%}
+                {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
+                {%- for tool_response in message.get('tool_responses') -%}
+                    {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}}
+                    {%- set ns_tr_out.flag = true -%}
+                    {%- set ns.prev_message_type = 'tool_response' -%}
+                {%- endfor -%}
+            {%- elif message.get('tool_calls') -%}
+                {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
+                {%- set ns_tool_scan = namespace(stopped=false) -%}
+                {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
+                    {%- if ns_tool_scan.stopped -%}
+                    {%- elif loop_messages[k]['role'] != 'tool' -%}
+                        {%- set ns_tool_scan.stopped = true -%}
+                    {%- else -%}
+                        {%- set follow = loop_messages[k] -%}
+                        {#- Resolve tool_call_id to function name -#}
+                        {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%}
+                        {%- for tc in message.get('tool_calls') -%}
+                            {%- if tc.get('id') == follow.get('tool_call_id') -%}
+                                {%- set ns_tname.name = tc['function']['name'] -%}
+                            {%- endif -%}
+                        {%- endfor -%}
+                        {#- Handle content as string or content-parts array -#}
+                        {%- set tool_body = follow.get('content') -%}
+                        {%- if tool_body is string -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- elif tool_body is sequence and tool_body is not string -%}
+                            {%- set ns_txt = namespace(s='') -%}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'text' -%}
+                                    {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
+                                {%- endif -%}
+                            {%- endfor -%}
+                            {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') in ['image', 'image_url'] -%}
+                                    {{- '<|image|>' -}}
+                                {%- elif part.get('type') in ['audio', 'input_audio'] -%}
+                                    {{- '<|audio|>' -}}
+                                {%- elif part.get('type') == 'video' -%}
+                                    {{- '<|video|>' -}}
+                                {%- endif -%}
+                            {%- endfor -%}
+                        {%- else -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- endif -%}
+                        {%- set ns_tr_out.flag = true -%}
+                        {%- set ns.prev_message_type = 'tool_response' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+
+            {%- set captured_content -%}
+            {%- if message.get('content') is string -%}
+                {%- if role == 'model' -%}
+                    {{- strip_thinking(message['content']) -}}
+                {%- else -%}
+                    {{- message['content'] | trim -}}
+                {%- endif -%}
+            {%- elif message.get('content') is sequence -%}
+                {%- for item in message['content'] -%}
+                    {%- if item.get('type') == 'text' -%}
+                        {%- if role == 'model' -%}
+                            {{- strip_thinking(item['text']) -}}
+                        {%- else -%}
+                            {{- item['text'] | trim -}}
+                        {%- endif -%}
+                    {%- elif item.get('type') in ['image', 'image_url'] -%}
+                        {{- '<|image|>' -}}
+                    {%- elif item.get('type') in ['audio', 'input_audio'] -%}
+                        {{- '<|audio|>' -}}
+                    {%- elif item.get('type') == 'video' -%}
+                        {{- '<|video|>' -}}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+            {%- endset -%}
+
+            {{- captured_content -}}
+            {%- set has_content = captured_content | trim | length > 0 -%}
+
+        {#- Forward-scan: find next non-tool message role for continuation detection -#}
+        {%- set next_nt = namespace(role=None, found=false) -%}
+        {%- for j in range(loop.index0 + 1, loop_messages | length) -%}
+            {%- if not next_nt.found -%}
+                {%- if loop_messages[j]['role'] != 'tool' -%}
+                    {%- set next_nt.role = loop_messages[j]['role'] -%}
+                    {%- set next_nt.found = true -%}
+                {%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+
+        {%- set continues_into_next = (
+            role == 'model'
+            and next_nt.role == 'assistant'
+            and (not message.get('tool_calls') or ns_tr_out.flag)
+        ) -%}
+
+        {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
+            {{- '<|tool_response>' -}}
+        {%- elif continues_into_next -%}
+        {%- elif not (ns_tr_out.flag and not has_content and not next_nt.found) -%}
+            {{- '<turn|>\n' -}}
+        {%- endif -%}
+
+    {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#}
+    {%- set ns.prev_non_tool_role = message['role'] -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
+        {{- '<|turn>model\n' -}}
+        {%- if not enable_thinking -%}
+            {{- '<|channel>thought\n<channel|>' -}}
+        {%- endif -%}
+    {%- elif ns.prev_message_type == 'tool_response' and enable_thinking -%}
+        {{- '<|channel>thought\n' -}}
+    {%- endif -%}
+{%- endif -%}

config.json

@@ -0,0 +1,182 @@
+{
+  "architectures": [
+    "Gemma4UnifiedForConditionalGeneration"
+  ],
+  "audio_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "audio_embed_dim": 640,
+    "audio_samples_per_token": 640,
+    "chunk_size_feed_forward": 0,
+    "dtype": null,
+    "hidden_size": 640,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "model_type": "gemma4_unified_audio",
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_proj_dims": 640,
+    "problem_type": null,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06
+  },
+  "audio_token_id": 258881,
+  "boa_token_id": 256000,
+  "boi_token_id": 255999,
+  "dtype": "bfloat16",
+  "eoa_token_index": 258883,
+  "eoi_token_id": 258882,
+  "eos_token_id": [
+    1,
+    106
+  ],
+  "image_token_id": 258880,
+  "initializer_range": 0.02,
+  "model_type": "gemma4_unified",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attention_k_eq_v": true,
+    "bos_token_id": 2,
+    "enable_moe_block": false,
+    "eos_token_id": 1,
+    "final_logit_softcapping": 30.0,
+    "global_head_dim": 512,
+    "head_dim": 256,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 3840,
+    "hidden_size_per_layer_input": 0,
+    "initializer_range": 0.02,
+    "intermediate_size": 15360,
+    "layer_types": [
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 262144,
+    "model_type": "gemma4_unified_text",
+    "moe_intermediate_size": null,
+    "num_attention_heads": 16,
+    "num_experts": null,
+    "num_global_key_value_heads": 1,
+    "num_hidden_layers": 48,
+    "num_key_value_heads": 8,
+    "num_kv_shared_layers": 0,
+    "pad_token_id": 0,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "full_attention": {
+        "partial_rotary_factor": 0.25,
+        "rope_theta": 1000000.0,
+        "rope_type": "proportional"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 1024,
+    "tie_word_embeddings": true,
+    "top_k_experts": null,
+    "use_bidirectional_attention": "vision",
+    "use_cache": true,
+    "use_double_wide_mlp": false,
+    "vocab_size": 262144,
+    "vocab_size_per_layer_input": 262144
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.10.0.dev0",
+  "video_token_id": 258884,
+  "vision_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "chunk_size_feed_forward": 0,
+    "dtype": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "mm_embed_dim": 3840,
+    "mm_posemb_size": 1120,
+    "model_patch_size": 48,
+    "model_type": "gemma4_unified_vision",
+    "num_soft_tokens": 280,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_proj_dims": 3840,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "problem_type": null,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06
+  },
+  "quantization": {
+    "bits": 4,
+    "group_size": 64,
+    "mode": "affine"
+  },
+  "quantization_config": {
+    "bits": 4,
+    "group_size": 64,
+    "mode": "affine"
+  }
+}

conversion.json

@@ -0,0 +1,11 @@
+{
+  "source": "akhilaaa3/Jev-Omni",
+  "revision": "c050d51354147985d13286cf4acf90f562f2c631",
+  "source_checkpoint": "unified/model.safetensors",
+  "bits": 4,
+  "group_size": 64,
+  "vision_precision": "bfloat16",
+  "head_precision": "float32",
+  "calibrated_after_conversion": false,
+  "weight_bytes": 6803751008
+}

decision_head/weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21ee232647cd6aaa61277ed11b25aee0726956ac950ae6bfbdc49cf277daf59b
+size 3964224

generation_config.json

@@ -0,0 +1,18 @@
+{
+  "bos_token_id": 2,
+  "do_sample": true,
+  "eos_token_id": [
+    1,
+    106,
+    50
+  ],
+  "pad_token_id": 0,
+  "suppress_tokens": [
+    258883,
+    258882
+  ],
+  "temperature": 1.0,
+  "top_k": 64,
+  "top_p": 0.95,
+  "transformers_version": "5.10.0.dev0"
+}

omni_mlx/__init__.py

@@ -0,0 +1 @@
+"""Local experimental MLX port of the Jev-Omni decision classifier."""

omni_mlx/classifier.py

@@ -0,0 +1,120 @@
+"""Image/text classification with Jev's trained head; no vocabulary projection or generation."""
+import argparse
+import json
+import statistics
+import time
+from pathlib import Path
+
+import mlx.core as mx
+from mlx_vlm.models.gemma4_unified.processing_gemma4_unified import Gemma4UnifiedProcessor
+from mlx_vlm.utils import load_model, prepare_inputs
+from PIL import Image
+
+
+def prompt(state, question, options):
+    choices = "\n".join(f"{i+1}. {value}" for i, value in enumerate(options))
+    return (f"{state}\n\n---\n\nQUESTION: {question}\n\nOPTIONS:\n{choices}\n\n"
+            f"Reply with only the number of the correct option (1-{len(options)}).\n"
+            "Output a single number and nothing else.")
+
+
+def head_probabilities(hidden, head, count):
+    features = (hidden.astype(mx.float32) - head["mu"]) / head["sd"]
+    logits = features @ head["linear.weight"][:count].T + head["linear.bias"][:count]
+    return mx.softmax(logits, axis=-1)
+
+
+class Classifier:
+    def __init__(self, path=".models/Jev-Omni-MLX-4bit"):
+        self.path = Path(path).resolve()
+        mx.set_cache_limit(256 * 1024**2)
+        self.model = load_model(self.path, strict=True)
+        self.processor = Gemma4UnifiedProcessor.from_pretrained(self.path, trust_remote_code=False)
+        self.head = mx.load(str(self.path / "decision_head/weights.safetensors"))
+        mx.eval(self.head)
+        self.provenance = json.loads((self.path / "conversion.json").read_text())
+
+    def predict(self, state, question, options, image=None, image_tokens=280):
+        if not 2 <= len(options) <= 20 or len(set(options)) != len(options):
+            raise ValueError("Supply 2–20 distinct options")
+        if image_tokens not in (10, 20, 35, 70, 140, 280):
+            raise ValueError("image_tokens must be 10, 20, 35, 70, 140, or 280")
+        started = time.perf_counter()
+        self.processor.image_processor.max_soft_tokens = image_tokens
+        self.processor.image_seq_length = image_tokens
+        content = ([{"type": "image"}] if image is not None else [])
+        content.append({"type": "text", "text": prompt(state, question, options)})
+        formatted = self.processor.apply_chat_template(
+            [{"role": "user", "content": content}],
+            add_generation_prompt=True, tokenize=False, enable_thinking=False,
+        )
+        if image is not None and not isinstance(image, Image.Image):
+            with Image.open(image) as opened:
+                image = opened.convert("RGB")
+        inputs = prepare_inputs(self.processor, images=[image] if image is not None else None,
+                                prompts=formatted, add_special_tokens=False)
+        inputs = {k: v.astype(mx.bfloat16) if isinstance(v, mx.array) and mx.issubdtype(v.dtype, mx.floating) else v for k,v in inputs.items()}
+        ids = inputs["input_ids"]
+        if ids.shape[1] > 1024:
+            raise ValueError("Local prototype limited to 1024 input tokens to bound memory")
+        # Batch size 1, no padding. Let the decoder build its causal/sliding and vision masks.
+        extra = {k:v for k,v in inputs.items() if k not in {"input_ids", "attention_mask"}}
+        mx.eval(inputs)
+        preprocessing_ms = (time.perf_counter() - started) * 1000
+        vision_started = time.perf_counter()
+        embedded = self.model.get_input_embeddings(input_ids=ids, **extra)
+        mx.eval(embedded.inputs_embeds)
+        vision_ms = (time.perf_counter() - vision_started) * 1000
+        decoder_started = time.perf_counter()
+        hidden = self.model.language_model.model(
+            inputs_embeds=embedded.inputs_embeds,
+            per_layer_inputs=embedded.per_layer_inputs,
+            mm_token_type_ids=inputs.get("mm_token_type_ids"),
+        )
+        probabilities = head_probabilities(hidden[:, -1], self.head, len(options))[0]
+        mx.eval(probabilities)
+        decoder_ms = (time.perf_counter() - decoder_started) * 1000
+        values = probabilities.tolist()
+        if not all(0 <= value <= 1 for value in values):
+            raise RuntimeError("Non-finite classifier output")
+        elapsed_ms = (time.perf_counter() - started) * 1000
+        best = max(range(len(values)), key=values.__getitem__)
+        return {
+            "model": "akhilaaa3/Jev-Omni", "backend": "mlx", "quantization_bits": self.provenance["bits"],
+            "prediction": options[best], "prediction_index": best,
+            "probabilities": dict(zip(options, values)), "calibrated": False,
+            "metrics": {"elapsed_ms": elapsed_ms, "preprocessing_ms": preprocessing_ms,
+                        "vision_ms": vision_ms, "decoder_ms": decoder_ms,
+                        "input_tokens": ids.shape[1], "image_token_budget": image_tokens if image is not None else 0,
+                        "peak_metal_memory_gb": mx.get_peak_memory()/1e9, "generated_tokens": 0},
+        }
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", default=".models/Jev-Omni-MLX-4bit")
+    parser.add_argument("--image")
+    parser.add_argument("--state", default="")
+    parser.add_argument("--question", required=True)
+    parser.add_argument("--options", nargs="+", required=True)
+    parser.add_argument("--image-tokens", type=int, choices=[10, 20, 35, 70, 140, 280], default=70)
+    parser.add_argument("--repeat", type=int, default=1)
+    parser.add_argument("--output", type=Path)
+    args = parser.parse_args()
+    if args.repeat < 1:
+        parser.error("repeat must be positive")
+    start = time.perf_counter()
+    classifier = Classifier(args.model)
+    load_ms = (time.perf_counter() - start) * 1000
+    results = [classifier.predict(args.state, args.question, args.options, args.image, args.image_tokens) for _ in range(args.repeat)]
+    result = {"load_ms": load_ms, "requests": results}
+    if len(results) > 1:
+        result["warm_median_ms"] = statistics.median(r["metrics"]["elapsed_ms"] for r in results[1:])
+    rendered = json.dumps(result, ensure_ascii=False, indent=2)
+    if args.output:
+        args.output.write_text(rendered + "\n")
+    print(rendered)
+
+
+if __name__ == "__main__":
+    main()

omni_mlx/convert.py

@@ -0,0 +1,133 @@
+"""Convert the pinned Jev-Omni unified checkpoint without loading 24 GB into RAM."""
+import argparse
+import gc
+import hashlib
+import json
+import shutil
+import struct
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+
+
+def remap(key):
+    key = key.removeprefix("model.")
+    if key.startswith("language_model."):
+        key = key.replace("language_model.", "language_model.model.", 1)
+    return key
+
+
+def convert(source: Path, output: Path, bits=4):
+    if output.exists():
+        raise FileExistsError(f"Refusing to overwrite {output}")
+    unified = source / "unified"
+    checkpoint = unified / "model.safetensors"
+    if not checkpoint.is_file():
+        raise FileNotFoundError(checkpoint)
+    output.mkdir(parents=True)
+    mx.set_cache_limit(128 * 1024**2)
+    weights, index, total_bytes = {}, {}, 0
+    shard_number = 0
+
+    def flush():
+        nonlocal weights, shard_number, total_bytes
+        if not weights:
+            return
+        shard_number += 1
+        name = f"weights-{shard_number:03d}.safetensors"
+        mx.save_safetensors(str(output / name), weights, metadata={"format": "mlx"})
+        total_bytes += sum(v.nbytes for v in weights.values())
+        index.update({key: name for key in weights})
+        weights = {}
+        gc.collect()
+        mx.clear_cache()
+
+    with checkpoint.open("rb") as stream:
+        header_size = struct.unpack("<Q", stream.read(8))[0]
+        header = json.loads(stream.read(header_size))
+        data_start = header_size + 8
+        for number, (key, info) in enumerate(header.items()):
+            if key == "__metadata__":
+                continue
+            if info["dtype"] != "BF16":
+                raise ValueError(f"Unexpected dtype for {key}: {info['dtype']}")
+            shape = info["shape"]
+            mapped = remap(key)
+            quantize = (
+                mapped.startswith("language_model.")
+                and mapped.endswith(".weight")
+                and len(shape) == 2
+                and shape[-1] % 64 == 0
+            )
+            stream.seek(data_start + info["data_offsets"][0])
+            if quantize:
+                packed, scales, biases = [], [], []
+                for row in range(0, shape[0], 1024):
+                    rows = min(1024, shape[0] - row)
+                    data = stream.read(rows * shape[1] * 2)
+                    block = mx.array(np.frombuffer(data, dtype=np.uint16).copy()).view(mx.bfloat16)
+                    block = block.reshape(rows, shape[1])
+                    q, s, b = mx.quantize(block, group_size=64, bits=bits)
+                    mx.eval(q, s, b)
+                    packed.append(q)
+                    scales.append(s)
+                    biases.append(b)
+                prefix = mapped.removesuffix(".weight")
+                converted = {
+                    prefix + ".weight": mx.concatenate(packed),
+                    prefix + ".scales": mx.concatenate(scales),
+                    prefix + ".biases": mx.concatenate(biases),
+                }
+                mx.eval(converted)
+                del packed, scales, biases, block, q, s, b, data
+            else:
+                length = info["data_offsets"][1] - info["data_offsets"][0]
+                data = stream.read(length)
+                converted = {mapped: mx.array(np.frombuffer(data, dtype=np.uint16).copy()).view(mx.bfloat16).reshape(shape)}
+                mx.eval(converted)
+                del data
+            weights.update(converted)
+            del converted
+            if sum(v.nbytes for v in weights.values()) >= 384 * 1024**2:
+                flush()
+            if number % 40 == 0:
+                print(f"Converted {number}/{len(header)-1} tensors; MLX peak {mx.get_peak_memory()/1e9:.2f} GB", flush=True)
+    flush()
+    config = json.loads((unified / "config.json").read_text())
+    config["quantization"] = {"bits": bits, "group_size": 64, "mode": "affine"}
+    config["quantization_config"] = config["quantization"]
+    (output / "config.json").write_text(json.dumps(config, indent=2) + "\n")
+    (output / "model.safetensors.index.json").write_text(json.dumps({"metadata": {"total_size": total_bytes}, "weight_map": index}, indent=2) + "\n")
+    for file in unified.iterdir():
+        if file.suffix in {".json", ".jinja"} and file.name != "config.json":
+            shutil.copy2(file, output / file.name)
+
+    # Only read tensors from the original small classification head; never execute pickle objects.
+    import torch
+    expected = json.loads((source / "sha256.json").read_text())["head.pt"]
+    if hashlib.sha256((source / "head.pt").read_bytes()).hexdigest() != expected:
+        raise ValueError("head.pt hash differs from the source manifest")
+    state = torch.load(source / "head.pt", map_location="cpu", weights_only=True)
+    if set(state) != {"mu", "sd", "linear.weight", "linear.bias"}:
+        raise ValueError(f"Unexpected classification head keys: {list(state)}")
+    target = output / "decision_head"
+    target.mkdir()
+    mx.save_safetensors(str(target / "weights.safetensors"), {k: mx.array(v.float().numpy()) for k,v in state.items()})
+    revision = Path("artifacts/jev-omni/source-revision.txt").read_text().strip()
+    (output / "conversion.json").write_text(json.dumps({
+        "source": "akhilaaa3/Jev-Omni", "revision": revision,
+        "source_checkpoint": "unified/model.safetensors", "bits": bits,
+        "group_size": 64, "vision_precision": "bfloat16", "head_precision": "float32",
+        "calibrated_after_conversion": False, "weight_bytes": total_bytes,
+    }, indent=2) + "\n")
+    print(f"Saved {total_bytes/1e9:.2f} GB of weights to {output}", flush=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--source", type=Path, default=Path(".models/Jev-Omni-source"))
+    parser.add_argument("--output", type=Path, default=Path(".models/Jev-Omni-MLX-4bit"))
+    parser.add_argument("--bits", type=int, choices=[4, 6, 8], default=4)
+    args = parser.parse_args()
+    convert(args.source, args.output, args.bits)

processor_config.json

@@ -0,0 +1,62 @@
+{
+  "audio_ms_per_token": 40,
+  "audio_seq_length": 750,
+  "feature_extractor": {
+    "audio_samples_per_token": 640,
+    "feature_extractor_type": "Gemma4UnifiedAudioFeatureExtractor",
+    "feature_size": 640,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "return_attention_mask": true,
+    "sampling_rate": 16000
+  },
+  "image_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": false,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_processor_type": "Gemma4UnifiedImageProcessor",
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_soft_tokens": 280,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098
+  },
+  "image_seq_length": 280,
+  "processor_class": "Gemma4UnifiedProcessor",
+  "video_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "do_sample_frames": true,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_soft_tokens": 70,
+    "num_frames": 32,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "return_metadata": false,
+    "video_processor_type": "Gemma4UnifiedVideoProcessor"
+  }
+}

requirements.txt

@@ -0,0 +1,4 @@
+mlx==0.32.2
+mlx-vlm==0.7.1
+numpy>=2.0
+Pillow>=10.0

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f
+size 32169626

tokenizer_config.json

@@ -0,0 +1,142 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": false,
+  "local_files_only": true,
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>"
+  },
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "processor_class": "Gemma4UnifiedProcessor",
+  "response_schema": {
+    "properties": {
+      "content": {
+        "type": "string"
+      },
+      "role": {
+        "const": "assistant"
+      },
+      "thinking": {
+        "type": "string"
+      },
+      "tool_calls": {
+        "items": {
+          "properties": {
+            "function": {
+              "properties": {
+                "arguments": {
+                  "additionalProperties": {},
+                  "type": "object",
+                  "x-parser": "gemma4-tool-call"
+                },
+                "name": {
+                  "type": "string"
+                }
+              },
+              "type": "object",
+              "x-regex": "call\\:(?P<name>\\w+)(?P<arguments>\\{.*\\})"
+            },
+            "type": {
+              "const": "function"
+            }
+          },
+          "type": "object"
+        },
+        "type": "array",
+        "x-regex-iterator": "<\\|tool_call>(.*?)<tool_call\\|>"
+      }
+    },
+    "type": "object",
+    "x-regex": "(\\<\\|channel\\>thought\\n(?P<thinking>.*?)\\<channel\\|\\>)?(?P<tool_calls>\\<\\|tool_call\\>.*\\<tool_call\\|\\>)?(?P<content>(?:(?!\\<turn\\|\\>)(?!\\<\\|tool_response\\>).)+)?(?:\\<turn\\|\\>|\\<\\|tool_response\\>)?"
+  },
+  "response_template": {
+    "defaults": {
+      "role": "assistant"
+    },
+    "fields": {
+      "content": {
+        "close": [
+          "<turn|>",
+          "<|tool_response>",
+          "<eos>"
+        ],
+        "content": "text"
+      },
+      "thinking": {
+        "close": "<channel|>",
+        "content": "text",
+        "open": "<|channel>thought\n"
+      },
+      "tool_calls": {
+        "close": "<tool_call|>",
+        "content": "json",
+        "content_args": {
+          "string_delims": [
+            [
+              "<|\"|>",
+              "<|\"|>"
+            ]
+          ],
+          "unquoted_keys": true
+        },
+        "open_pattern": "<\\|tool_call>call:(?P<name>\\w+)",
+        "repeats": true,
+        "transform": {
+          "function": {
+            "arguments": "{content}",
+            "name": "{name}"
+          },
+          "type": "function"
+        }
+      }
+    },
+    "start_anchor": [
+      "<|turn>model\n",
+      "<tool_response|>"
+    ]
+  },
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>"
+}

verification_unified.json

@@ -0,0 +1,67 @@
+{
+ "cases": [
+  {
+   "state": "The meeting starts at 10 AM. It is now 9 AM.",
+   "question": "Has the meeting started?",
+   "options": [
+    "Yes",
+    "No"
+   ]
+  },
+  {
+   "state": "Customer: I was charged twice.\nAgent: I've refunded $29 to your card.\nCustomer: Got it. All sorted, thanks!",
+   "question": "Was the issue actually resolved?",
+   "options": [
+    "Yes",
+    "No"
+   ]
+  },
+  {
+   "state": "A fair six-sided die is rolled once.",
+   "question": "Which number comes up?",
+   "options": [
+    "1",
+    "2",
+    "3",
+    "4",
+    "5",
+    "6"
+   ]
+  },
+  {
+   "state": "An urn holds one blue, one yellow, one red and one green marble. One is drawn without looking.",
+   "question": "Which marble is drawn?",
+   "options": [
+    "Blue",
+    "Yellow",
+    "Red",
+    "Green"
+   ]
+  }
+ ],
+ "reference": [
+  {
+   "Yes": 0.00011235327838221565,
+   "No": 0.9998875856399536
+  },
+  {
+   "Yes": 0.9399133324623108,
+   "No": 0.06008664518594742
+  },
+  {
+   "1": 0.23407739400863647,
+   "2": 0.09292787313461304,
+   "3": 0.04372488334774971,
+   "4": 0.08901944011449814,
+   "5": 0.08105313032865524,
+   "6": 0.4591972529888153
+  },
+  {
+   "Blue": 0.41296717524528503,
+   "Yellow": 0.12619410455226898,
+   "Red": 0.0792793333530426,
+   "Green": 0.38155943155288696
+  }
+ ],
+ "worst_abs_diff": 0.01936584711074829
+}

weights-001.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fdb0abb73a0957abbf7498f16dac0eb4f9af50a40ca89b997ef8fa03a44b21ec
+size 600638052

weights-002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ce22e87811ec5522a9de8e686d5f3d70f9867a6917bf67d8dbb4adeab293374
+size 411516454

weights-003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89ea0ae6aec9e59f9b9b3f59b469024afd84a6264c47f6e017a08b05c9db1338
+size 421462508

weights-004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a176555c9df5d762a4013586c0b5e05a6fe953faa3543a238ca3402f6ba6b4e8
+size 411508567

weights-005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd17b1be4ab7559bde0cc43d0a719c14da4585a2cf20925303f5092233914619
+size 410428736

weights-006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0ea901604540e01e5b217f6b0d1973f2f8755652e2622c80cd2989b3ddc96a9
+size 415940536

weights-007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11c415ff671f0a0fd609f70aaa78140b7a08e45e6e853f82d73279a24a690bef
+size 421462504

weights-008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7891f453644091d28330466309ee95c26d59dacc3a9201e77f0b2925b78474f9
+size 411508583

weights-009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b04affc543c23a78ab8d4a30ffe3c33aa5cc40ecbb0a43816ace2ce1e2c87b3
+size 410428759

weights-010.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8dcbae180260c5a1bbdb1dcb9f5e20b95a5fac044b1bf46c6c3224e90590ccd
+size 415940519

weights-011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d7c18bd26c8781c66fde122dfd6969bd2d0a54b35f215c8518860b05756e092
+size 421462482

weights-012.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0c452c1abccbf8c6802b7b1f6495450d30ef6a902fa433ab066ee9fd6fd4032
+size 411508537

weights-013.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:114c9752669f7f672966054e1623a9da99e4aac9f8f20055f1dd3ce466b10310
+size 410428751

weights-014.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85134bc69f35b0d89af9a4caafb2d40ce414aa72dce67cb883b20db4f69fa9c5
+size 415940575

weights-015.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed1833b879184d16efd62836578019ce68c02c9cffa23207e1813b062117f633
+size 431416352

weights-016.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46438c76c4ed86356cfb4be6b04eaeed4747194c3d1c500c8afd74a9613ef3e9
+size 382328776