FEFER AI 🦖 — VisionPsy-Nano-460M for pegd.fun / Stable 988

.gitattributes

@@ -33,3 +33,6 @@ 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
+assets/category_performance.png filter=lfs diff=lfs merge=lfs -text
+assets/category_summary.png filter=lfs diff=lfs merge=lfs -text
+fefer-512.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,201 @@
+                                 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.

README.md

@@ -0,0 +1,106 @@
+---
+license: apache-2.0
+base_model: qvac/VisionPsy-Nano-460M
+pipeline_tag: image-text-to-text
+library_name: transformers
+tags:
+  - vision-language-model
+  - nanovlm
+  - chart-understanding
+  - ocr
+  - crypto
+  - launchpad
+  - stable-mainnet
+  - fefer
+  - pegd-fun
+language:
+  - en
+---
+
+<p align="center">
+  <img src="fefer-512.png" width="180" alt="FEFER AI" />
+</p>
+
+# FEFER-AI-460M 🦖
+
+**FEFER** is the resident chart analyst of [pegd.fun](https://pegd.fun) — the stablecoin
+launchpad on **Stable mainnet (chain id 988)**. This model is a FEFER-tuned build of
+[`qvac/VisionPsy-Nano-460M`](https://huggingface.co/qvac/VisionPsy-Nano-460M)
+(Tether AI Research), a ~460M-parameter vision-language model (SigLIP2 vision encoder +
+SmolLM2-360M LM, nanoVLM architecture) that punches far above its size on chart
+understanding, OCR, and visual instruction following.
+
+FEFER reads **one thing and one thing only**: terminal-style candlestick charts rendered by
+the pegd.fun indexer for Stable-988 launchpad tokens. Each chart has verified indexer facts
+(market cap, 24h volume, graduation progress) printed directly on the image, so the model
+cross-checks pixels against ground truth instead of inventing numbers.
+
+## What it's for
+
+- Reading pegd.fun 1024×512 candle charts: trend, momentum, volume behavior, drawdowns
+- Answering holder questions in the FEFER persona (plain language, risk note always included)
+- Explaining launchpad mechanics it was trained on: fixed 1B supply, permanently locked
+  Uniswap V3 1% USDT0 pools, 3,000 USDT0 opening FDV, 9,000 USDT0 graduation, 80/20
+  creator/platform fee split
+
+## What it's NOT for
+
+- Price prediction or financial advice — FEFER always says so, by training and by prompt
+- Multi-image reasoning (single-image by architecture)
+- Anything outside Stable mainnet 988 — other chains are out of scope on purpose
+
+## Quickstart (transformers)
+
+```python
+from transformers import AutoModelForImageTextToText, AutoProcessor
+from PIL import Image
+
+repo = "feferai/FEFER-AI-460M"  # replace after publishing
+model = AutoModelForImageTextToText.from_pretrained(repo, trust_remote_code=True)
+processor = AutoProcessor.from_pretrained(repo, trust_remote_code=True)
+
+chart = Image.open("fefer_chart.png")  # a pegd.fun /api/fefer/chart/<token>.png render
+messages = [{"role": "user", "content": [
+    {"type": "image", "image": chart},
+    {"type": "text", "text": "Read this chart. Trend, volume, and graduation status?"},
+]}]
+inputs = processor.apply_chat_template(messages, add_generation_prompt=True,
+                                       tokenize=True, return_dict=True, return_tensors="pt")
+out = model.generate(**inputs, max_new_tokens=320)
+print(processor.batch_decode(out[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0])
+```
+
+## Serving for pegd.fun
+
+Any OpenAI-compatible endpoint works. Then point the launchpad at it:
+
+```bash
+# GPU
+vllm serve feferai/FEFER-AI-460M --trust-remote-code --port 8000
+# or CPU (no GPU needed at 460M)
+python fefer-ai/serve_cpu.py --model feferai/FEFER-AI-460M --port 8008
+
+# pegd.fun server env
+FEFER_AI_URL=http://127.0.0.1:8000/v1
+FEFER_AI_MODEL=feferai/FEFER-AI-460M
+```
+
+## Training
+
+LoRA fine-tune on the language-model attention projections (vision tower frozen), on
+chart/answer pairs harvested from live pegd.fun launches by `fefer-ai/build_dataset.mjs`
+— every sample is a real indexer-rendered chart paired with facts computed from on-chain
+trade history, so supervision is grounded, not synthetic guesswork. Merged with
+`merge_and_unload` for standalone serving.
+
+## Limitations & bias
+
+460M parameters is small: FEFER is a sharp chart *reader*, not an oracle. It can misread
+dense wicks, and it inherits any biases of the base model. Launchpad tokens are volatile;
+nothing this model outputs is financial advice, and it is trained to say exactly that.
+
+## License & attribution
+
+Apache 2.0, same as the base model. Built on **VisionPsy-Nano-460M by Tether AI Research
+(qvac)** — all credit for the base architecture and pretraining to them. FEFER branding,
+chart pipeline, dataset tooling, and fine-tune by the pegd.fun team.

chat_template.jinja

@@ -0,0 +1,4 @@
+{% for message in messages %}{{'<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>' + '
+'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}

config.json

@@ -0,0 +1,125 @@
+{
+  "architectures": [
+    "VisionPsyNanoForConditionalGeneration"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_visionpsynano.VisionPsyNanoConfig",
+    "AutoModel": "modeling_visionpsynano.VisionPsyNanoForConditionalGeneration",
+    "AutoModelForImageTextToText": "modeling_visionpsynano.VisionPsyNanoForConditionalGeneration",
+    "AutoProcessor": "processing_visionpsynano.VisionPsyNanoProcessor"
+  },
+  "compile_inference": true,
+  "compile_inference_mode": "default",
+  "cuda_graphs_cache_quantum": 128,
+  "eos_check_interval": 16,
+  "extra_token_amount": 66,
+  "hf_repo_name": "qvac/VisionPsy-Nano-460M",
+  "inference_max_img_size": null,
+  "is_flash": false,
+  "lm_attn_scaling": 1.0,
+  "lm_base_vocab_size": 49152,
+  "lm_chat_template": "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
+  "lm_dropout": 0.0,
+  "lm_hidden_dim": 960,
+  "lm_inter_dim": 2560,
+  "lm_max_length": 8192,
+  "lm_max_position_embeddings": 8192,
+  "lm_model_type": "HuggingFaceTB/SmolLM2-360M-Instruct",
+  "lm_n_blocks": 32,
+  "lm_n_heads": 15,
+  "lm_n_kv_heads": 5,
+  "lm_re_base": 100000,
+  "lm_rms_eps": 1e-05,
+  "lm_tie_weights": true,
+  "lm_tokenizer": "HuggingFaceTB/SmolLM2-360M-Instruct",
+  "lm_use_tokens": false,
+  "lm_vocab_size": 49218,
+  "max_img_size": 2048,
+  "model_type": "visionpsynano",
+  "mp_image_token_length": 64,
+  "mp_pixel_shuffle_factor": 4,
+  "resize_min_side_len": null,
+  "resize_to_max_side_len": true,
+  "transformers_version": "5.14.1",
+  "vit_cls_flag": false,
+  "vit_dropout": 0.0,
+  "vit_hidden_dim": 768,
+  "vit_img_size": 512,
+  "vit_inter_dim": 3072,
+  "vit_ln_eps": 1e-06,
+  "vit_model_type": "google/siglip2-base-patch16-512",
+  "vit_n_blocks": 12,
+  "vit_n_heads": 12,
+  "vit_patch_size": 16,
+  "vlm_checkpoint_path": null,
+  "vlm_extra_tokens": {
+    "global_image_token": "<|global_image|>",
+    "image_token": "<|image|>",
+    "r1c1": "<row_1_col_1>",
+    "r1c2": "<row_1_col_2>",
+    "r1c3": "<row_1_col_3>",
+    "r1c4": "<row_1_col_4>",
+    "r1c5": "<row_1_col_5>",
+    "r1c6": "<row_1_col_6>",
+    "r1c7": "<row_1_col_7>",
+    "r1c8": "<row_1_col_8>",
+    "r2c1": "<row_2_col_1>",
+    "r2c2": "<row_2_col_2>",
+    "r2c3": "<row_2_col_3>",
+    "r2c4": "<row_2_col_4>",
+    "r2c5": "<row_2_col_5>",
+    "r2c6": "<row_2_col_6>",
+    "r2c7": "<row_2_col_7>",
+    "r2c8": "<row_2_col_8>",
+    "r3c1": "<row_3_col_1>",
+    "r3c2": "<row_3_col_2>",
+    "r3c3": "<row_3_col_3>",
+    "r3c4": "<row_3_col_4>",
+    "r3c5": "<row_3_col_5>",
+    "r3c6": "<row_3_col_6>",
+    "r3c7": "<row_3_col_7>",
+    "r3c8": "<row_3_col_8>",
+    "r4c1": "<row_4_col_1>",
+    "r4c2": "<row_4_col_2>",
+    "r4c3": "<row_4_col_3>",
+    "r4c4": "<row_4_col_4>",
+    "r4c5": "<row_4_col_5>",
+    "r4c6": "<row_4_col_6>",
+    "r4c7": "<row_4_col_7>",
+    "r4c8": "<row_4_col_8>",
+    "r5c1": "<row_5_col_1>",
+    "r5c2": "<row_5_col_2>",
+    "r5c3": "<row_5_col_3>",
+    "r5c4": "<row_5_col_4>",
+    "r5c5": "<row_5_col_5>",
+    "r5c6": "<row_5_col_6>",
+    "r5c7": "<row_5_col_7>",
+    "r5c8": "<row_5_col_8>",
+    "r6c1": "<row_6_col_1>",
+    "r6c2": "<row_6_col_2>",
+    "r6c3": "<row_6_col_3>",
+    "r6c4": "<row_6_col_4>",
+    "r6c5": "<row_6_col_5>",
+    "r6c6": "<row_6_col_6>",
+    "r6c7": "<row_6_col_7>",
+    "r6c8": "<row_6_col_8>",
+    "r7c1": "<row_7_col_1>",
+    "r7c2": "<row_7_col_2>",
+    "r7c3": "<row_7_col_3>",
+    "r7c4": "<row_7_col_4>",
+    "r7c5": "<row_7_col_5>",
+    "r7c6": "<row_7_col_6>",
+    "r7c7": "<row_7_col_7>",
+    "r7c8": "<row_7_col_8>",
+    "r8c1": "<row_8_col_1>",
+    "r8c2": "<row_8_col_2>",
+    "r8c3": "<row_8_col_3>",
+    "r8c4": "<row_8_col_4>",
+    "r8c5": "<row_8_col_5>",
+    "r8c6": "<row_8_col_6>",
+    "r8c7": "<row_8_col_7>",
+    "r8c8": "<row_8_col_8>"
+  },
+  "vlm_load_backbone_weights": false,
+  "torch_dtype": "float32"
+}

configuration_visionpsynano.py

@@ -0,0 +1,317 @@
+"""HuggingFace PretrainedConfig for VisionPsyNano."""
+from __future__ import annotations
+
+import os
+from typing import Any, Optional
+
+from transformers import PretrainedConfig
+
+
+FLASH_MIN_SIDE_LEN = 512
+
+
+def resolve_is_flash(
+    *,
+    is_flash: Optional[bool] = None,
+    variant: Optional[str] = None,
+    resize_to_max_side_len: Optional[bool] = None,
+) -> bool:
+    """Resolve Flash preprocess mode from explicit flag, legacy variant, or resize policy."""
+    if is_flash is not None:
+        return bool(is_flash)
+    if variant is not None:
+        v = str(variant).lower().strip()
+        if v in ("flash", "nano-flash", "visionpsy-nano-flash"):
+            return True
+        if v in ("nano", "plain", "nano-plain", "visionpsy-nano"):
+            return False
+        raise ValueError(
+            "legacy variant must be 'nano' or 'flash' "
+            f"(or aliases); got {variant!r}. Prefer is_flash=True/False."
+        )
+    if resize_to_max_side_len is not None:
+        return not bool(resize_to_max_side_len)
+    return False
+
+
+def apply_flash_preprocess(
+    *,
+    is_flash: bool,
+    resize_to_max_side_len: Optional[bool] = None,
+    resize_min_side_len: Optional[int] = None,
+) -> tuple[bool, Optional[int]]:
+    """Return (resize_to_max_side_len, resize_min_side_len) for the given Flash mode."""
+    if resize_to_max_side_len is None:
+        resize_to_max_side_len = not is_flash
+    resize_to_max_side_len = bool(resize_to_max_side_len)
+    if is_flash:
+        resize_min_side_len = max(int(resize_min_side_len or 0), FLASH_MIN_SIDE_LEN)
+    elif resize_to_max_side_len:
+        resize_min_side_len = None
+    return resize_to_max_side_len, resize_min_side_len
+
+
+_DEFAULT_CHAT_TEMPLATE = (
+    "{% for message in messages %}"
+    "{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}"
+    "{% endfor %}"
+    "{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}"
+)
+
+_DEFAULT_EXTRA_TOKENS = {
+    "image_token": "<|image|>",
+    "global_image_token": "<|global_image|>",
+    "r1c1": "<row_1_col_1>",
+    "r1c2": "<row_1_col_2>",
+    "r1c3": "<row_1_col_3>",
+    "r1c4": "<row_1_col_4>",
+    "r1c5": "<row_1_col_5>",
+    "r1c6": "<row_1_col_6>",
+    "r1c7": "<row_1_col_7>",
+    "r1c8": "<row_1_col_8>",
+    "r2c1": "<row_2_col_1>",
+    "r2c2": "<row_2_col_2>",
+    "r2c3": "<row_2_col_3>",
+    "r2c4": "<row_2_col_4>",
+    "r2c5": "<row_2_col_5>",
+    "r2c6": "<row_2_col_6>",
+    "r2c7": "<row_2_col_7>",
+    "r2c8": "<row_2_col_8>",
+    "r3c1": "<row_3_col_1>",
+    "r3c2": "<row_3_col_2>",
+    "r3c3": "<row_3_col_3>",
+    "r3c4": "<row_3_col_4>",
+    "r3c5": "<row_3_col_5>",
+    "r3c6": "<row_3_col_6>",
+    "r3c7": "<row_3_col_7>",
+    "r3c8": "<row_3_col_8>",
+    "r4c1": "<row_4_col_1>",
+    "r4c2": "<row_4_col_2>",
+    "r4c3": "<row_4_col_3>",
+    "r4c4": "<row_4_col_4>",
+    "r4c5": "<row_4_col_5>",
+    "r4c6": "<row_4_col_6>",
+    "r4c7": "<row_4_col_7>",
+    "r4c8": "<row_4_col_8>",
+    "r5c1": "<row_5_col_1>",
+    "r5c2": "<row_5_col_2>",
+    "r5c3": "<row_5_col_3>",
+    "r5c4": "<row_5_col_4>",
+    "r5c5": "<row_5_col_5>",
+    "r5c6": "<row_5_col_6>",
+    "r5c7": "<row_5_col_7>",
+    "r5c8": "<row_5_col_8>",
+    "r6c1": "<row_6_col_1>",
+    "r6c2": "<row_6_col_2>",
+    "r6c3": "<row_6_col_3>",
+    "r6c4": "<row_6_col_4>",
+    "r6c5": "<row_6_col_5>",
+    "r6c6": "<row_6_col_6>",
+    "r6c7": "<row_6_col_7>",
+    "r6c8": "<row_6_col_8>",
+    "r7c1": "<row_7_col_1>",
+    "r7c2": "<row_7_col_2>",
+    "r7c3": "<row_7_col_3>",
+    "r7c4": "<row_7_col_4>",
+    "r7c5": "<row_7_col_5>",
+    "r7c6": "<row_7_col_6>",
+    "r7c7": "<row_7_col_7>",
+    "r7c8": "<row_7_col_8>",
+    "r8c1": "<row_8_col_1>",
+    "r8c2": "<row_8_col_2>",
+    "r8c3": "<row_8_col_3>",
+    "r8c4": "<row_8_col_4>",
+    "r8c5": "<row_8_col_5>",
+    "r8c6": "<row_8_col_6>",
+    "r8c7": "<row_8_col_7>",
+    "r8c8": "<row_8_col_8>",
+}
+
+
+class VisionPsyNanoConfig(PretrainedConfig):
+    """Config for VisionPsyNano (``is_flash`` selects preprocess mode)."""
+
+    model_type = "visionpsynano"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        is_flash: Optional[bool] = None,
+        variant: Optional[str] = None,
+        vit_hidden_dim: int = 768,
+        vit_inter_dim: int = 3072,
+        vit_patch_size: int = 16,
+        vit_img_size: int = 512,
+        vit_n_heads: int = 12,
+        vit_dropout: float = 0.0,
+        vit_n_blocks: int = 12,
+        vit_ln_eps: float = 1e-6,
+        vit_cls_flag: bool = False,
+        vit_model_type: str = "google/siglip2-base-patch16-512",
+        lm_hidden_dim: int = 960,
+        lm_inter_dim: int = 2560,
+        lm_rms_eps: float = 1e-5,
+        lm_re_base: int = 100000,
+        lm_max_position_embeddings: int = 8192,
+        lm_base_vocab_size: int = 49152,
+        extra_token_amount: int = 66,
+        lm_vocab_size: Optional[int] = None,
+        lm_n_heads: int = 15,
+        lm_n_kv_heads: int = 5,
+        lm_dropout: float = 0.0,
+        lm_n_blocks: int = 32,
+        lm_attn_scaling: float = 1.0,
+        lm_max_length: int = 4096,
+        lm_use_tokens: bool = False,
+        lm_tie_weights: bool = True,
+        lm_model_type: str = "HuggingFaceTB/SmolLM2-360M-Instruct",
+        lm_tokenizer: str = "HuggingFaceTB/SmolLM2-360M-Instruct",
+        lm_chat_template: str = _DEFAULT_CHAT_TEMPLATE,
+        mp_pixel_shuffle_factor: int = 4,
+        mp_image_token_length: int = 64,
+        max_img_size: int = 2048,
+        resize_to_max_side_len: Optional[bool] = None,
+        resize_min_side_len: Optional[int] = None,
+        inference_max_img_size: Optional[int] = None,
+        vlm_extra_tokens: Optional[dict] = None,
+        vlm_load_backbone_weights: bool = True,
+        vlm_checkpoint_path: str = "checkpoints",
+        hf_repo_name: str = "qvac/VisionPsy-Nano-460M",
+        compile_inference: bool = True,
+        compile_inference_mode: str = "reduce-overhead",
+        cuda_graphs_cache_quantum: int = 128,
+        eos_check_interval: int = 16,
+        **kwargs: Any,
+    ):
+        is_flash = resolve_is_flash(
+            is_flash=is_flash,
+            variant=variant,
+            resize_to_max_side_len=resize_to_max_side_len,
+        )
+        resize_to_max_side_len, resize_min_side_len = apply_flash_preprocess(
+            is_flash=is_flash,
+            resize_to_max_side_len=resize_to_max_side_len,
+            resize_min_side_len=resize_min_side_len,
+        )
+
+        if lm_vocab_size is None:
+            lm_vocab_size = lm_base_vocab_size + extra_token_amount
+
+        self.is_flash = bool(is_flash)
+        self.vit_hidden_dim = vit_hidden_dim
+        self.vit_inter_dim = vit_inter_dim
+        self.vit_patch_size = vit_patch_size
+        self.vit_img_size = vit_img_size
+        self.vit_n_heads = vit_n_heads
+        self.vit_dropout = vit_dropout
+        self.vit_n_blocks = vit_n_blocks
+        self.vit_ln_eps = vit_ln_eps
+        self.vit_cls_flag = vit_cls_flag
+        self.vit_model_type = vit_model_type
+
+        self.lm_hidden_dim = lm_hidden_dim
+        self.lm_inter_dim = lm_inter_dim
+        self.lm_rms_eps = lm_rms_eps
+        self.lm_re_base = lm_re_base
+        self.lm_max_position_embeddings = lm_max_position_embeddings
+        self.lm_base_vocab_size = lm_base_vocab_size
+        self.extra_token_amount = extra_token_amount
+        self.lm_vocab_size = lm_vocab_size
+        self.lm_n_heads = lm_n_heads
+        self.lm_n_kv_heads = lm_n_kv_heads
+        self.lm_dropout = lm_dropout
+        self.lm_n_blocks = lm_n_blocks
+        self.lm_attn_scaling = lm_attn_scaling
+        self.lm_max_length = lm_max_length
+        self.lm_use_tokens = lm_use_tokens
+        self.lm_tie_weights = lm_tie_weights
+        self.lm_model_type = lm_model_type
+        self.lm_tokenizer = lm_tokenizer
+        self.lm_chat_template = lm_chat_template
+
+        self.mp_pixel_shuffle_factor = mp_pixel_shuffle_factor
+        self.mp_image_token_length = mp_image_token_length
+        self.max_img_size = max_img_size
+        self.resize_to_max_side_len = bool(resize_to_max_side_len)
+        self.resize_min_side_len = resize_min_side_len
+        self.inference_max_img_size = inference_max_img_size
+        self.vlm_extra_tokens = dict(vlm_extra_tokens or _DEFAULT_EXTRA_TOKENS)
+        self.vlm_load_backbone_weights = vlm_load_backbone_weights
+        self.vlm_checkpoint_path = vlm_checkpoint_path
+        if self.is_flash and hf_repo_name == "qvac/VisionPsy-Nano-460M":
+            self.hf_repo_name = "qvac/VisionPsy-Nano-460M-Flash"
+        else:
+            self.hf_repo_name = hf_repo_name
+
+        self.compile_inference = compile_inference
+        self.compile_inference_mode = compile_inference_mode
+        self.cuda_graphs_cache_quantum = cuda_graphs_cache_quantum
+        self.eos_check_interval = eos_check_interval
+
+        kwargs.pop("text_config", None)
+        kwargs.pop("vision_config", None)
+
+        super().__init__(**kwargs)
+
+    def get_text_config(self, decoder: bool = False, **kwargs):
+        """VisionPsyNano is a single flat config (not text+vision composite)."""
+        return self
+
+    @property
+    def variant(self) -> str:
+        return "flash" if self.is_flash else "nano"
+
+    def to_vlm_config(self):
+        """Convert to the internal VLMConfig dataclass used by the core modules."""
+        from dataclasses import fields as dc_fields
+
+        try:
+            from .vlm_config import VLMConfig
+        except ImportError:
+            from vlm_config import VLMConfig
+
+        valid = {f.name for f in dc_fields(VLMConfig)}
+        payload = {k: getattr(self, k) for k in valid if hasattr(self, k)}
+        return VLMConfig(**payload)
+
+    @classmethod
+    def from_vlm_config(
+        cls, cfg, *, is_flash: Optional[bool] = None, variant: Optional[str] = None
+    ) -> "VisionPsyNanoConfig":
+        from dataclasses import asdict
+
+        data = asdict(cfg)
+        data["is_flash"] = resolve_is_flash(
+            is_flash=is_flash,
+            variant=variant,
+            resize_to_max_side_len=data.get("resize_to_max_side_len"),
+        )
+        data.pop("variant", None)
+        return cls(**data)
+
+    @classmethod
+    def from_legacy_dict(
+        cls,
+        raw: dict,
+        *,
+        is_flash: Optional[bool] = None,
+        variant: Optional[str] = None,
+    ) -> "VisionPsyNanoConfig":
+        """Load an existing VisionPsyNano / nanoVLM config.json (without model_type)."""
+        raw = dict(raw)
+        for key in (
+            "model_type",
+            "architectures",
+            "auto_map",
+            "transformers_version",
+            "text_config",
+            "vision_config",
+        ):
+            raw.pop(key, None)
+        raw["is_flash"] = resolve_is_flash(
+            is_flash=is_flash if is_flash is not None else raw.get("is_flash"),
+            variant=variant if variant is not None else raw.get("variant"),
+            resize_to_max_side_len=raw.get("resize_to_max_side_len"),
+        )
+        raw.pop("variant", None)
+        return cls(**raw)

custom_transforms.py

@@ -0,0 +1,115 @@
+import math
+import torch
+from torchvision.transforms.functional import resize, InterpolationMode
+from einops import rearrange
+from typing import Tuple, Union
+from PIL import Image
+
+
+class DynamicResize(torch.nn.Module):
+    """Resize H/W to patch-aligned sizes within max_side_len."""
+    def __init__(
+        self,
+        patch_size: int,
+        max_side_len: int,
+        resize_to_max_side_len: bool = False,
+        min_side_len: int | None = None,
+        interpolation: InterpolationMode = InterpolationMode.BICUBIC,
+    ) -> None:
+        super().__init__()
+        self.p = int(patch_size)
+        self.m = int(max_side_len)
+        self.interpolation = interpolation
+        self.resize_to_max_side_len = resize_to_max_side_len
+        self.min_side_len = int(min_side_len) if min_side_len else None
+
+    def _get_new_hw(self, h: int, w: int) -> Tuple[int, int]:
+        """Compute target (h, w) divisible by patch_size."""
+        long, short = (w, h) if w >= h else (h, w)
+
+        if (
+            self.min_side_len
+            and not self.resize_to_max_side_len
+            and short < self.min_side_len
+        ):
+            den = short * self.p
+            target_long = min(self.m, -(-(long * self.min_side_len) // den) * self.p)
+            target_short = max(-(-(short * self.min_side_len) // den) * self.p, self.p)
+            return (target_short, target_long) if w >= h else (target_long, target_short)
+
+        target_long = self.m if self.resize_to_max_side_len else min(self.m, math.ceil(long / self.p) * self.p)
+        scale = target_long / long
+        target_short = math.ceil(short * scale / self.p) * self.p
+        target_short = max(target_short, self.p)
+
+        return (target_short, target_long) if w >= h else (target_long, target_short)
+
+    def forward(self, img: Union[Image.Image, torch.Tensor]):
+        if isinstance(img, Image.Image):
+            w, h = img.size
+            new_h, new_w = self._get_new_hw(h, w)
+            return resize(img, [new_h, new_w], interpolation=self.interpolation)
+
+        if not torch.is_tensor(img):
+            raise TypeError(
+                "DynamicResize expects a PIL Image or a torch.Tensor; "
+                f"got {type(img)}"
+            )
+
+        batched = img.ndim == 4
+        if img.ndim not in (3, 4):
+            raise ValueError(
+                "Tensor input must have shape (C,H,W) or (B,C,H,W); "
+                f"got {img.shape}"
+            )
+
+        imgs = img if batched else img.unsqueeze(0)
+        _, _, h, w = imgs.shape
+        new_h, new_w = self._get_new_hw(h, w)
+        out = resize(imgs, [new_h, new_w], interpolation=self.interpolation)
+
+        return out if batched else out.squeeze(0)
+
+
+class SplitImage(torch.nn.Module):
+    """Split (B, C, H, W) image tensor into square patches.
+
+    Returns:
+        patches: (B·n_h·n_w, C, patch_size, patch_size)
+        grid:    (n_h, n_w)  - number of patches along H and W
+    """
+    def __init__(self, patch_size: int) -> None:
+        super().__init__()
+        self.p = patch_size
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        if x.ndim == 3:           
+            x = x.unsqueeze(0)
+
+        b, c, h, w = x.shape
+        if h % self.p or w % self.p:
+            raise ValueError(f'Image size {(h,w)} not divisible by patch_size {self.p}')
+
+        n_h, n_w = h // self.p, w // self.p
+        patches = rearrange(x, 'b c (nh ph) (nw pw) -> (b nh nw) c ph pw',
+                            ph=self.p, pw=self.p)
+        return patches, (n_h, n_w)
+
+
+class GlobalAndSplitImages(torch.nn.Module):
+    def __init__(self, patch_size: int):
+        super().__init__()
+        self.p = patch_size
+        self.splitter = SplitImage(patch_size)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        if x.ndim == 3:
+            x = x.unsqueeze(0)
+
+        patches, grid = self.splitter(x)
+
+        if grid == (1, 1):
+            return patches, grid
+
+        global_patch = resize(x, [self.p, self.p])
+        return torch.cat([global_patch, patches], dim=0), grid

language_model.py

@@ -0,0 +1,557 @@
+import inspect
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def _sdpa_supports_native_gqa() -> bool:
+
+    try:
+        params = inspect.signature(F.scaled_dot_product_attention).parameters
+        return "enable_gqa" in params
+    except (TypeError, ValueError):
+        try:
+            q = torch.zeros(1, 2, 1, 4)
+            k = torch.zeros(1, 1, 1, 4)
+            v = torch.zeros(1, 1, 1, 4)
+            F.scaled_dot_product_attention(q, k, v, enable_gqa=True)
+            return True
+        except TypeError:
+            return False
+        except Exception:
+            return True
+
+_SDPA_HAS_GQA = _sdpa_supports_native_gqa()
+
+class RMSNorm(nn.Module):
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(cfg.lm_hidden_dim))
+        self.eps = cfg.lm_rms_eps
+        self._normalized_shape = (cfg.lm_hidden_dim,)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        return F.rms_norm(x, self._normalized_shape, self.weight, self.eps)
+
+class RotaryEmbedding(nn.Module):
+
+    def __init__(self, cfg):
+        super().__init__()
+        assert cfg.lm_hidden_dim % cfg.lm_n_heads == 0, "Hidden dimension must be divisible by number of heads"
+
+        self.dim = cfg.lm_hidden_dim // cfg.lm_n_heads
+        self.base = cfg.lm_re_base
+        self.max_seq_len = cfg.lm_max_position_embeddings
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self.original_max_seq_len = cfg.lm_max_position_embeddings
+        self.attention_scaling = cfg.lm_attn_scaling
+
+        positions = torch.arange(self.max_seq_len, dtype=torch.float)
+        freqs = positions.unsqueeze(-1) * inv_freq.unsqueeze(0)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        cos_cached = emb.cos() * self.attention_scaling
+        sin_cached = emb.sin() * self.attention_scaling
+        self.register_buffer("cos_cached", cos_cached, persistent=False)
+        self.register_buffer("sin_cached", sin_cached, persistent=False)
+
+    @torch.no_grad()
+    def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+
+        if position_ids.dtype != torch.long:
+            position_ids = position_ids.long()
+        cos = self.cos_cached[position_ids]
+        sin = self.sin_cached[position_ids]
+        return cos, sin
+
+    def extend_cache(self, new_max_seq_len: int) -> None:
+
+        positions = torch.arange(
+            new_max_seq_len, dtype=torch.float, device=self.cos_cached.device,
+        )
+        freqs = positions.unsqueeze(-1) * self.inv_freq.unsqueeze(0)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        self.cos_cached = emb.cos() * self.attention_scaling
+        self.sin_cached = emb.sin() * self.attention_scaling
+        self.original_max_seq_len = new_max_seq_len
+
+def apply_rotary_pos_embd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+
+    half = q.shape[-1] // 2
+    cos_h = cos[..., :half]
+    sin_h = sin[..., :half]
+
+    q1 = q[..., :half]
+    q2 = q[..., half:]
+    k1 = k[..., :half]
+    k2 = k[..., half:]
+
+    q_embed = torch.cat(
+        (q1 * cos_h - q2 * sin_h, q2 * cos_h + q1 * sin_h),
+        dim=-1,
+    )
+    k_embed = torch.cat(
+        (k1 * cos_h - k2 * sin_h, k2 * cos_h + k1 * sin_h),
+        dim=-1,
+    )
+    return q_embed, k_embed
+
+class LanguageModelGroupedQueryAttention(nn.Module):
+
+    def __init__(self, cfg):
+        super().__init__()
+
+        self.n_heads = cfg.lm_n_heads
+        self.n_kv_heads = cfg.lm_n_kv_heads
+        self.embd_dim = cfg.lm_hidden_dim
+        self.dropout = cfg.lm_dropout
+
+        assert self.n_heads % self.n_kv_heads == 0, "n_heads must be divisible by n_kv_heads"
+        assert self.embd_dim % self.n_heads == 0, "embd_dim must be divisible by num_heads"
+
+        self.n_kv_groups = self.n_heads // self.n_kv_heads
+        self.head_dim = self.embd_dim // self.n_heads
+
+        self.q_proj = nn.Linear(self.embd_dim, self.embd_dim, bias=False)
+        self.k_proj = nn.Linear(self.embd_dim, self.head_dim * self.n_kv_heads, bias=False)
+        self.v_proj = nn.Linear(self.embd_dim, self.head_dim * self.n_kv_heads, bias=False)
+        self.out_proj = nn.Linear(self.embd_dim, self.embd_dim, bias=False)
+
+        self.attn_dropout = nn.Dropout(self.dropout)
+        self.resid_dropout = nn.Dropout(self.dropout)
+
+        self.sdpa = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.sdpa:
+            print("Warning: scaled dot product attention not available, using standard attention in LM.")
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, attention_mask=None, block_kv_cache=None, start_pos: int = 0) -> tuple[torch.Tensor, dict]:
+
+        B, T_curr, C = x.size()
+
+        q_curr = self.q_proj(x).view(B, T_curr, self.n_heads, self.head_dim).transpose(1, 2)
+        k_curr = self.k_proj(x).view(B, T_curr, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v_curr = self.v_proj(x).view(B, T_curr, self.n_kv_heads, self.head_dim).transpose(1, 2)
+
+        q, k_rotated = apply_rotary_pos_embd(q_curr, k_curr, cos, sin)
+
+        use_static_cache = (
+            block_kv_cache is not None
+            and block_kv_cache.get('cache_max_length') is not None
+        )
+
+        static_compile_decode = (
+            use_static_cache
+            and T_curr == 1
+            and isinstance(start_pos, torch.Tensor)
+        )
+        sdpa_attn_mask_static: torch.Tensor | None = None
+
+        if static_compile_decode:
+            cache_max_length = block_kv_cache['cache_max_length']
+            idx = start_pos.unsqueeze(0)
+            block_kv_cache['key_cache'].index_copy_(2, idx, k_rotated)
+            block_kv_cache['value_cache'].index_copy_(2, idx, v_curr)
+
+            k = block_kv_cache['key_cache']
+            v = block_kv_cache['value_cache']
+            T_kv = cache_max_length
+
+            arange = torch.arange(
+                cache_max_length, device=q.device, dtype=torch.long,
+            )
+            sdpa_attn_mask_static = (arange <= start_pos).view(1, 1, 1, -1)
+        elif use_static_cache:
+            cache_max_length = block_kv_cache['cache_max_length']
+            if block_kv_cache.get('key_cache') is None:
+                cache_shape = (B, self.n_kv_heads, cache_max_length, self.head_dim)
+                block_kv_cache['key_cache'] = torch.zeros(
+                    cache_shape, dtype=k_rotated.dtype, device=k_rotated.device,
+                )
+                block_kv_cache['value_cache'] = torch.zeros(
+                    cache_shape, dtype=v_curr.dtype, device=v_curr.device,
+                )
+            new_pos = start_pos + T_curr
+            block_kv_cache['key_cache'][:, :, start_pos:new_pos] = k_rotated
+            block_kv_cache['value_cache'][:, :, start_pos:new_pos] = v_curr
+            k = block_kv_cache['key_cache'][:, :, :new_pos]
+            v = block_kv_cache['value_cache'][:, :, :new_pos]
+            T_kv = new_pos
+        else:
+            is_prefill = block_kv_cache is None
+            if not is_prefill and block_kv_cache.get('key') is not None:
+                k = block_kv_cache['key']
+                v = block_kv_cache['value']
+                k = torch.cat([k, k_rotated], dim=2)
+                v = torch.cat([v, v_curr], dim=2)
+                block_kv_cache['key'] = k
+                block_kv_cache['value'] = v
+            else:
+                k = k_rotated
+                v = v_curr
+                block_kv_cache = {'key': k, 'value': v}
+
+        if _SDPA_HAS_GQA and self.sdpa and x.device.type != 'mps':
+            k_exp = k
+            v_exp = v
+        else:
+            k_exp = k.repeat_interleave(self.n_kv_groups, dim=1)
+            v_exp = v.repeat_interleave(self.n_kv_groups, dim=1)
+
+        T_kv = k_exp.size(2)
+
+        if self.sdpa and x.device.type != 'mps':
+            if sdpa_attn_mask_static is not None:
+                sdpa_attn_mask = sdpa_attn_mask_static
+                sdpa_is_causal = False
+            elif attention_mask is None:
+                sdpa_attn_mask = None
+                sdpa_is_causal = (T_curr == T_kv and T_curr > 1)
+            else:
+                if attention_mask.size(-1) < T_kv:
+                    pad = attention_mask.new_ones(
+                        attention_mask.size(0), T_kv - attention_mask.size(-1)
+                    )
+                    key_padding = torch.cat([attention_mask, pad], dim=-1).bool()
+                else:
+                    key_padding = attention_mask[:, :T_kv].bool()
+                if T_curr == T_kv and T_curr > 1:
+                    causal = torch.ones(
+                        T_curr, T_kv, dtype=torch.bool, device=q.device,
+                    ).tril()
+                else:
+                    causal = torch.ones(
+                        T_curr, T_kv, dtype=torch.bool, device=q.device,
+                    )
+                sdpa_attn_mask = (
+                    causal.unsqueeze(0).unsqueeze(0)
+                    & key_padding.unsqueeze(1).unsqueeze(2)
+                )
+                sdpa_is_causal = False
+
+            sdpa_kwargs = {
+                "attn_mask": sdpa_attn_mask,
+                "dropout_p": self.dropout if self.training else 0.0,
+                "is_causal": sdpa_is_causal,
+            }
+            if _SDPA_HAS_GQA and k_exp.size(1) != q.size(1):
+                sdpa_kwargs["enable_gqa"] = True
+
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k_exp, v_exp,
+                **sdpa_kwargs,
+            )
+        else:
+            attn = torch.matmul(q, k_exp.transpose(2, 3)) / math.sqrt(self.head_dim)
+            if T_curr == T_kv and T_curr > 1:
+                causal_mask_val = torch.tril(
+                    torch.ones(T_curr, T_curr, device=x.device, dtype=torch.bool),
+                ).view(1, 1, T_curr, T_curr)
+                attn = attn.masked_fill(~causal_mask_val, float('-inf'))
+
+            if attention_mask is not None:
+                additive = (
+                    1.0 - attention_mask[:, :T_kv].unsqueeze(1).unsqueeze(2).float()
+                ) * torch.finfo(q.dtype).min
+                attn = attn + additive
+
+            attn = F.softmax(attn, dim=-1)
+            attn = self.attn_dropout(attn)
+            y = attn @ v_exp
+        y = y.transpose(1, 2).contiguous().view(B, T_curr, C)
+        y = self.out_proj(y)
+        y = self.resid_dropout(y)
+
+        return y, block_kv_cache
+
+class LanguageModelMLP(nn.Module):
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.embd_dim = cfg.lm_hidden_dim
+        self.inter_dim = cfg.lm_inter_dim
+
+        self.activation_fn = F.silu
+        self.gate_up_proj = nn.Linear(self.embd_dim, 2 * self.inter_dim, bias=False)
+        self.down_proj = nn.Linear(self.inter_dim, self.embd_dim, bias=False)
+
+        self._register_load_state_dict_pre_hook(self._fuse_legacy_gate_up_hook)
+
+    @staticmethod
+    def _fuse_legacy_gate_up_hook(
+        state_dict, prefix, local_metadata, strict,
+        missing_keys, unexpected_keys, error_msgs,
+    ):
+
+        gate_key = f"{prefix}gate_proj.weight"
+        up_key = f"{prefix}up_proj.weight"
+        fused_key = f"{prefix}gate_up_proj.weight"
+        if (
+            gate_key in state_dict
+            and up_key in state_dict
+            and fused_key not in state_dict
+        ):
+            gate_w = state_dict.pop(gate_key)
+            up_w = state_dict.pop(up_key)
+            state_dict[fused_key] = torch.cat([gate_w, up_w], dim=0)
+
+    def forward(self, x):
+
+        gate, up = self.gate_up_proj(x).chunk(2, dim=-1)
+        return self.down_proj(self.activation_fn(gate) * up)
+
+class LanguageModelBlock(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.mlp = LanguageModelMLP(cfg)
+        self.attn = LanguageModelGroupedQueryAttention(cfg)
+        self.norm1 = RMSNorm(cfg)
+        self.norm2 = RMSNorm(cfg)
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, attention_mask: torch.Tensor=None, block_kv_cache: dict=None, start_pos: int = 0):
+
+        res = x
+        x = self.norm1(x)
+        x, block_kv_cache = self.attn(x, cos, sin, attention_mask, block_kv_cache, start_pos=start_pos)
+        x = res + x
+
+        res = x
+        x = self.norm2(x)
+        x = self.mlp(x)
+        x = res + x
+
+        return x, block_kv_cache
+
+class LanguageModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.lm_use_tokens = cfg.lm_use_tokens
+        self.lm_tie_weights = cfg.lm_tie_weights
+
+        self.token_embedding = nn.Embedding(cfg.lm_vocab_size, cfg.lm_hidden_dim)
+        self.rotary_embd = RotaryEmbedding(cfg)
+        self.blocks = nn.ModuleList([
+            LanguageModelBlock(cfg) for _ in range(cfg.lm_n_blocks)
+        ])
+        self.norm = RMSNorm(cfg)
+        self.head = nn.Linear(cfg.lm_hidden_dim, cfg.lm_vocab_size, bias=False)
+        if self.lm_tie_weights:
+            self.head.weight = self.token_embedding.weight
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, RMSNorm):
+            module.weight.data.fill_(1.0)
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor=None, kv_cache: list[dict]=None, start_pos: int=0):
+
+        if self.lm_use_tokens:
+            x = self.token_embedding(x)
+
+        B, T_curr, _ = x.size()
+
+        if isinstance(start_pos, torch.Tensor):
+            offsets = torch.arange(0, T_curr, device=x.device, dtype=torch.long)
+            current_position_ids = (offsets + start_pos).unsqueeze(0).expand(B, -1)
+        else:
+            current_position_ids = torch.arange(start_pos, start_pos + T_curr, device=x.device).unsqueeze(0).expand(B, -1)
+        cos, sin = self.rotary_embd(current_position_ids)
+        cos = cos.unsqueeze(1)
+        sin = sin.unsqueeze(1)
+
+        if kv_cache is None:
+            kv_cache = [None] * len(self.blocks)
+
+        for i, block in enumerate(self.blocks):
+            x, kv_cache[i] = block(x, cos, sin, attention_mask, kv_cache[i], start_pos=start_pos)
+
+        x = self.norm(x)
+
+        if self.lm_use_tokens:
+            x = self.head(x)
+
+        return x, kv_cache
+
+    @torch.inference_mode()
+    def generate(self, inputs: torch.Tensor, max_new_tokens: int=20):
+
+        if inputs.dim() == 1:
+            inputs = inputs.unsqueeze(0)
+        generated_outputs = inputs.clone()
+
+        prompt_output, kv_cache_list = self.forward(
+            generated_outputs,
+            attention_mask=None,
+            kv_cache=None,
+            start_pos=0
+        )
+        last_output = prompt_output[:, -1, :]
+
+        for i in range(max_new_tokens):
+            if self.lm_use_tokens:
+                next_output = torch.argmax(last_output, dim=-1, keepdim=True)
+            else:
+                next_output = last_output.unsqueeze(1)
+
+            generated_outputs = torch.cat((generated_outputs, next_output), dim=1)
+            current_token_start_pos = generated_outputs.size(1) - 1
+
+            if i == max_new_tokens - 1:
+                break
+
+            decode_step_output, kv_cache_list = self.forward(
+                next_output,
+                attention_mask=None,
+                kv_cache=kv_cache_list,
+                start_pos=current_token_start_pos
+            )
+            last_output = decode_step_output[:, -1, :]
+        return generated_outputs
+
+    @classmethod
+    def from_pretrained(cls, cfg):
+        from transformers import AutoConfig
+        from huggingface_hub import hf_hub_download
+        import safetensors
+        import torch.nn.init as init
+        import json
+        from huggingface_hub.utils import EntryNotFoundError
+        hf_config = AutoConfig.from_pretrained(cfg.lm_model_type)
+        original_vocab_size = hf_config.vocab_size
+        cfg.lm_hidden_dim = hf_config.hidden_size
+        cfg.lm_inter_dim = hf_config.intermediate_size
+        cfg.lm_rms_eps = hf_config.rms_norm_eps
+        cfg.lm_re_base = hf_config.rope_theta
+        cfg.lm_max_position_embeddings = hf_config.max_position_embeddings
+        if hasattr(cfg, 'lm_vocab_size'):
+            if cfg.lm_vocab_size < original_vocab_size:
+                raise ValueError(f"Config vocab size ({cfg.lm_vocab_size}) is smaller than pretrained model vocab size ({original_vocab_size})")
+        else:
+            cfg.lm_vocab_size = original_vocab_size
+        cfg.lm_n_heads = hf_config.num_attention_heads
+        cfg.lm_n_kv_heads = hf_config.num_key_value_heads
+        cfg.lm_dropout = hf_config.attention_dropout
+        cfg.lm_n_blocks = hf_config.num_hidden_layers
+        model = cls(cfg)
+        try:
+            index_path = hf_hub_download(repo_id=cfg.lm_model_type, filename="model.safetensors.index.json")
+            with open(index_path, 'r') as f:
+                index = json.load(f)
+            safetensors_filenames = sorted(list(set(index['weight_map'].values())))
+            safetensors_files = [hf_hub_download(repo_id=cfg.lm_model_type, filename=fn) for fn in safetensors_filenames]
+        except EntryNotFoundError:
+            safetensors_files = [hf_hub_download(repo_id=cfg.lm_model_type, filename="model.safetensors")]
+
+        sd = model.state_dict()
+        mapping = {
+            'model.embed_tokens.weight': 'token_embedding.weight',
+            'model.norm.weight': 'norm.weight'
+        }
+        for i in range(cfg.lm_n_blocks):
+            layer_prefix = f'model.layers.{i}.'
+            block_prefix = f'blocks.{i}.'
+            mapping.update({
+                f"{layer_prefix}self_attn.q_proj.weight": f"{block_prefix}attn.q_proj.weight",
+                f"{layer_prefix}self_attn.k_proj.weight": f"{block_prefix}attn.k_proj.weight",
+                f"{layer_prefix}self_attn.v_proj.weight": f"{block_prefix}attn.v_proj.weight",
+                f"{layer_prefix}self_attn.o_proj.weight": f"{block_prefix}attn.out_proj.weight",
+                f"{layer_prefix}mlp.down_proj.weight": f"{block_prefix}mlp.down_proj.weight",
+                f"{layer_prefix}input_layernorm.weight": f"{block_prefix}norm1.weight",
+                f"{layer_prefix}post_attention_layernorm.weight": f"{block_prefix}norm2.weight"
+            })
+        has_extended_embeddings = False
+        loaded_keys = set()
+        for safetensors_file in safetensors_files:
+            with safetensors.safe_open(filename=safetensors_file, framework="pt", device="cpu") as f:
+                for hf_key, our_key in mapping.items():
+                    if our_key in loaded_keys:
+                        continue
+                    if hf_key in f.keys() and our_key in sd:
+                        tensor = f.get_tensor(hf_key)
+                        if hf_key == 'model.embed_tokens.weight' and tensor.shape[0] != sd[our_key].shape[0]:
+                            has_extended_embeddings = True
+                            print(f"Extending token embeddings from {tensor.shape} to {sd[our_key].shape}")
+                            sd[our_key][:tensor.shape[0]].copy_(tensor)
+                            std = 0.02
+                            init.normal_(sd[our_key][tensor.shape[0]:], mean=0.0, std=std)
+                            print(f"Initialized {sd[our_key].shape[0] - tensor.shape[0]} new token embeddings")
+                            sd['head.weight'].copy_(sd[our_key])
+                        elif tensor.shape == sd[our_key].shape:
+                            sd[our_key].copy_(tensor)
+                        else:
+                            print(f"Shape mismatch for {hf_key} -> {our_key}: {tensor.shape} vs {sd[our_key].shape}")
+                        loaded_keys.add(our_key)
+
+        for hf_key, our_key in mapping.items():
+            if our_key not in loaded_keys:
+                if our_key in sd:
+                    print(f"Warning: Key {our_key} not found in any safetensors file (HF key: {hf_key})")
+
+        for i in range(cfg.lm_n_blocks):
+            layer_prefix = f"model.layers.{i}."
+            block_prefix = f"blocks.{i}."
+            fused_param_key = f"{block_prefix}mlp.gate_up_proj.weight"
+            if fused_param_key not in sd or fused_param_key in loaded_keys:
+                continue
+            gate_hf_key = f"{layer_prefix}mlp.gate_proj.weight"
+            up_hf_key = f"{layer_prefix}mlp.up_proj.weight"
+            gate_w = None
+            up_w = None
+            for safetensors_file in safetensors_files:
+                with safetensors.safe_open(filename=safetensors_file, framework="pt", device="cpu") as f:
+                    keys = f.keys()
+                    if gate_w is None and gate_hf_key in keys:
+                        gate_w = f.get_tensor(gate_hf_key)
+                    if up_w is None and up_hf_key in keys:
+                        up_w = f.get_tensor(up_hf_key)
+                if gate_w is not None and up_w is not None:
+                    break
+            if gate_w is None or up_w is None:
+                print(
+                    f"Warning: gate_proj or up_proj missing for layer {i} "
+                    f"(gate={gate_w is not None}, up={up_w is not None}); "
+                    f"{fused_param_key} left at init values."
+                )
+                continue
+            fused = torch.cat([gate_w, up_w], dim=0)
+            if fused.shape != sd[fused_param_key].shape:
+                print(
+                    f"Shape mismatch for fused gate_up at layer {i}: "
+                    f"{fused.shape} vs {sd[fused_param_key].shape}"
+                )
+                continue
+            sd[fused_param_key].copy_(fused)
+            loaded_keys.add(fused_param_key)
+
+        model.load_state_dict(sd)
+        if has_extended_embeddings and hasattr(model, 'head') and 'head.weight' in sd:
+            lm_head_loaded = False
+            for safetensors_file in safetensors_files:
+                with safetensors.safe_open(filename=safetensors_file, framework="pt", device="cpu") as f:
+                    if 'lm_head.weight' in f.keys():
+                        lm_head = f.get_tensor('lm_head.weight')
+                        if lm_head.shape[0] != sd['head.weight'].shape[0]:
+                            print(f"Extending LM head from {lm_head.shape} to {sd['head.weight'].shape}")
+                            sd['head.weight'][:lm_head.shape[0]].copy_(lm_head)
+                            std = 0.02
+                            init.normal_(sd['head.weight'][lm_head.shape[0]:], mean=0.0, std=std)
+                            model.load_state_dict(sd)
+                        lm_head_loaded = True
+                        break
+        if cfg.lm_tie_weights and hasattr(model, 'head') and hasattr(model, 'token_embedding'):
+            model.head.weight = model.token_embedding.weight
+        print(f"Successfully loaded {cfg.lm_model_type} weights from safetensors. Model has {sum(p.numel() for p in model.parameters()):,} parameters.")
+        return model

modality_projector.py

@@ -0,0 +1,39 @@
+import torch.nn as nn
+
+class ModalityProjector(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.input_dim = cfg.vit_hidden_dim * (cfg.mp_pixel_shuffle_factor**2)
+        self.output_dim = cfg.lm_hidden_dim
+        self.scale_factor = cfg.mp_pixel_shuffle_factor
+
+        self.proj = nn.Linear(self.input_dim, self.output_dim, bias=False)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(self.proj.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+    def pixel_shuffle(self, x):
+        bsz, seq, embed_dim = x.size()
+        seq_root = int(seq**0.5)
+        assert seq_root**2 == seq
+        assert seq_root % self.scale_factor == 0
+
+        height = width = seq_root
+        x = x.view(bsz, height, width, embed_dim)
+        h_out = height // self.scale_factor
+        w_out = width // self.scale_factor
+        x = x.reshape(bsz, h_out, self.scale_factor, w_out, self.scale_factor, embed_dim)
+        x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
+        x = x.reshape(bsz, h_out * w_out, embed_dim * self.scale_factor**2)
+        return x
+
+    def forward(self, x):
+        x = self.pixel_shuffle(x)
+        x = self.proj(x)
+
+        return x

model.safetensors

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

model_utils.py

@@ -0,0 +1,21 @@
+import torch
+
+
+def top_k_top_p_filtering(logits, top_k=0, top_p=1.0, filter_value=-float('Inf')):
+    top_k = min(top_k, logits.size(-1))
+
+    if top_k > 0:
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits = logits.masked_fill(indices_to_remove, filter_value)
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
+
+        sorted_indices_to_remove = cumulative_probs > top_p
+
+        sorted_indices_to_remove[..., 0] = False
+        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+        logits = logits.masked_fill(indices_to_remove, filter_value)
+
+    return logits

modeling_visionpsynano.py

@@ -0,0 +1,301 @@
+"""HuggingFace PreTrainedModel wrapper for VisionPsyNano."""
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import torch
+from transformers import PreTrainedModel
+
+
+try:
+    from .configuration_visionpsynano import VisionPsyNanoConfig
+except ImportError:
+    from configuration_visionpsynano import VisionPsyNanoConfig
+
+try:
+    from .processors import get_tokenizer
+    from .language_model import LanguageModel
+    from .modality_projector import ModalityProjector
+    from .vision_language_model import VisionLanguageModel
+    from .vision_transformer import ViT
+except ImportError as exc:
+    raise ImportError(
+        "Failed to import VisionPsyNano core modules (models/, data/). "
+        "For Hub packages use scripts/package_hub_repo.py so imports are flattened."
+    ) from exc
+
+
+class VisionPsyNanoForConditionalGeneration(PreTrainedModel):
+    config_class = VisionPsyNanoConfig
+    base_model_prefix = ""
+    _no_split_modules = ["ViTBlock", "LanguageModelBlock", "ModalityProjector"]
+    main_input_name = "input_ids"
+    supports_gradient_checkpointing = False
+
+    def __init__(self, config: VisionPsyNanoConfig):
+        super().__init__(config)
+        self.cfg = config.to_vlm_config()
+        self.vision_encoder = ViT(self.cfg)
+        self.decoder = LanguageModel(self.cfg)
+        self.MP = ModalityProjector(self.cfg)
+        self.tokenizer = get_tokenizer(
+            self.cfg.lm_tokenizer,
+            self.cfg.vlm_extra_tokens,
+            self.cfg.lm_chat_template,
+        )
+        self._sync_image_token_id()
+        self._compiled_decoder = None
+        self._compiled_vision_encoder = None
+        self.load_backbone = False
+        self.post_init()
+        self.tie_weights()
+
+    def _sync_image_token_id(self) -> None:
+        tok = self.tokenizer
+        if hasattr(tok, "image_token"):
+            tok.image_token_id = tok.convert_tokens_to_ids(tok.image_token)
+        elif "<|image|>" in getattr(tok, "get_vocab", lambda: {})():
+            tok.image_token = "<|image|>"
+            tok.image_token_id = tok.convert_tokens_to_ids("<|image|>")
+
+    def set_tokenizer(self, tokenizer) -> None:
+        self.tokenizer = tokenizer
+        self._sync_image_token_id()
+
+    def _reinit_rotary_cache(self) -> None:
+        re = getattr(self.decoder, "rotary_embd", None)
+        if re is None or not hasattr(re, "extend_cache"):
+            return
+        cos = getattr(re, "cos_cached", None)
+        if cos is None or getattr(cos, "is_meta", False) or cos.device.type == "meta":
+            return
+        length = int(getattr(re, "original_max_seq_len", None) or re.max_seq_len)
+        re.extend_cache(max(length, 1))
+        if bool(torch.isnan(re.cos_cached).any()) or bool(torch.isnan(re.sin_cached).any()):
+            device = re.inv_freq.device
+            positions = torch.arange(length, dtype=torch.float, device=device)
+            freqs = positions.unsqueeze(-1) * re.inv_freq.unsqueeze(0)
+            emb = torch.cat([freqs, freqs], dim=-1)
+            re.cos_cached = emb.cos() * re.attention_scaling
+            re.sin_cached = emb.sin() * re.attention_scaling
+            re.original_max_seq_len = length
+
+    def tie_weights(self, recompute_mapping: bool = True, **kwargs):
+        _ = recompute_mapping, kwargs
+        if not getattr(self.cfg, "lm_tie_weights", True):
+            return
+        head_w = self.decoder.head.weight
+        emb_w = self.decoder.token_embedding.weight
+        if head_w.data_ptr() != emb_w.data_ptr():
+            with torch.no_grad():
+                emb_w.copy_(head_w)
+            self.decoder.head.weight = self.decoder.token_embedding.weight
+
+    def _tie_weights(self):
+        self.tie_weights()
+
+    def _repair_legacy_checkpoint_weights(
+        self, pretrained_model_name_or_path, *, revision: Optional[str] = None
+    ) -> None:
+        weights_path = None
+        if os.path.isdir(pretrained_model_name_or_path):
+            candidate = os.path.join(pretrained_model_name_or_path, "model.safetensors")
+            if os.path.isfile(candidate):
+                weights_path = candidate
+        else:
+            try:
+                from huggingface_hub import hf_hub_download
+
+                weights_path = hf_hub_download(
+                    pretrained_model_name_or_path,
+                    "model.safetensors",
+                    revision=revision,
+                )
+            except Exception:
+                return
+        if not weights_path:
+            return
+
+        from safetensors import safe_open
+
+        with safe_open(weights_path, framework="pt", device="cpu") as f:
+            keys = set(f.keys())
+            for i, block in enumerate(self.decoder.blocks):
+                gate_key = f"decoder.blocks.{i}.mlp.gate_proj.weight"
+                up_key = f"decoder.blocks.{i}.mlp.up_proj.weight"
+                fused_key = f"decoder.blocks.{i}.mlp.gate_up_proj.weight"
+                if gate_key in keys and up_key in keys and fused_key not in keys:
+                    fused = torch.cat([f.get_tensor(gate_key), f.get_tensor(up_key)], dim=0)
+                    param = block.mlp.gate_up_proj.weight
+                    param.data.copy_(fused.to(device=param.device, dtype=param.dtype))
+            if (
+                "decoder.token_embedding.weight" not in keys
+                and "decoder.head.weight" in keys
+            ):
+                head = f.get_tensor("decoder.head.weight")
+                emb = self.decoder.token_embedding.weight
+                emb.data.copy_(head.to(device=emb.device, dtype=emb.dtype))
+
+    def _vision_encoder_for_prefill(self):
+        return VisionLanguageModel._vision_encoder_for_prefill(self)
+
+    def _decoder_for_decode(self):
+        return VisionLanguageModel._decoder_for_decode(self)
+
+    def _replace_img_tokens_with_embd(self, input_ids, token_embd, image_embd):
+        return VisionLanguageModel._replace_img_tokens_with_embd(
+            self, input_ids, token_embd, image_embd
+        )
+
+    def _process_images(self, images, device):
+        return VisionLanguageModel._process_images(self, images, device)
+
+    def apply_deploy_profile(self, device: Optional[torch.device] = None) -> None:
+        try:
+            from .runtime_profile import apply_deploy_profile
+        except ImportError:
+            from runtime_profile import apply_deploy_profile
+
+        apply_deploy_profile(self, device or next(self.parameters()).device)
+
+    def apply_eager_profile(self) -> None:
+        try:
+            from .runtime_profile import apply_eager_profile
+        except ImportError:
+            from runtime_profile import apply_eager_profile
+
+        apply_eager_profile(self)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        images=None,
+        pixel_values=None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        targets: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        images = images if images is not None else pixel_values
+        targets = targets if targets is not None else labels
+        return VisionLanguageModel.forward(
+            self, input_ids, images, attention_mask=attention_mask, targets=targets
+        )
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        images=None,
+        pixel_values=None,
+        attention_mask: Optional[torch.Tensor] = None,
+        max_new_tokens: int = 5,
+        top_k: int = 50,
+        top_p: float = 0.9,
+        temperature: float = 0.5,
+        greedy: bool = False,
+        **kwargs,
+    ):
+        images = images if images is not None else pixel_values
+        _ = kwargs
+        return VisionLanguageModel.generate(
+            self,
+            input_ids,
+            images,
+            attention_mask=attention_mask,
+            max_new_tokens=max_new_tokens,
+            top_k=top_k,
+            top_p=top_p,
+            temperature=temperature,
+            greedy=greedy,
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        model._repair_legacy_checkpoint_weights(
+            pretrained_model_name_or_path, revision=kwargs.get("revision")
+        )
+        model.tie_weights()
+        model._reinit_rotary_cache()
+        try:
+            from transformers import AutoTokenizer
+
+            try:
+                from .processing_visionpsynano import VisionPsyNanoProcessor
+            except ImportError:
+                from processing_visionpsynano import VisionPsyNanoProcessor
+
+            tok = AutoTokenizer.from_pretrained(
+                pretrained_model_name_or_path,
+                trust_remote_code=kwargs.get("trust_remote_code", False),
+            )
+            VisionPsyNanoProcessor._attach_extra_token_attrs(tok, model.config.vlm_extra_tokens)
+            model.set_tokenizer(tok)
+        except Exception:
+            model._sync_image_token_id()
+        return model
+
+    @classmethod
+    def from_legacy_pretrained(
+        cls,
+        repo_id_or_path: str,
+        *,
+        is_flash: Optional[bool] = None,
+        variant: Optional[str] = None,
+        revision: Optional[str] = None,
+        **kwargs,
+    ) -> "VisionPsyNanoForConditionalGeneration":
+        import json
+
+        from huggingface_hub import hf_hub_download
+        from safetensors.torch import load_model
+
+        if os.path.isdir(repo_id_or_path):
+            config_path = os.path.join(repo_id_or_path, "config.json")
+            weights_path = os.path.join(repo_id_or_path, "model.safetensors")
+        else:
+            config_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="config.json", revision=revision
+            )
+            weights_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="model.safetensors", revision=revision
+            )
+
+        with open(config_path, "r") as f:
+            raw = json.load(f)
+
+        config = VisionPsyNanoConfig.from_legacy_dict(
+            raw, is_flash=is_flash, variant=variant
+        )
+        model = cls(config)
+        load_model(model, weights_path)
+        model.tie_weights()
+        model._reinit_rotary_cache()
+        tok_dir = repo_id_or_path if os.path.isdir(repo_id_or_path) else os.path.dirname(config_path)
+        if any(
+            os.path.exists(os.path.join(tok_dir, n))
+            for n in ("tokenizer.json", "tokenizer_config.json")
+        ):
+            from transformers import AutoTokenizer
+
+            try:
+                from .processing_visionpsynano import VisionPsyNanoProcessor
+            except ImportError:
+                from processing_visionpsynano import VisionPsyNanoProcessor
+
+            tok = AutoTokenizer.from_pretrained(tok_dir)
+            VisionPsyNanoProcessor._attach_extra_token_attrs(tok, config.vlm_extra_tokens)
+            model.set_tokenizer(tok)
+        return model
+
+
+try:
+    from transformers import AutoConfig, AutoModel, AutoModelForImageTextToText
+
+    AutoConfig.register("visionpsynano", VisionPsyNanoConfig)
+    AutoModel.register(VisionPsyNanoConfig, VisionPsyNanoForConditionalGeneration)
+    AutoModelForImageTextToText.register(VisionPsyNanoConfig, VisionPsyNanoForConditionalGeneration)
+except Exception:
+    pass

preprocessor_config.json

@@ -0,0 +1,84 @@
+{
+  "processor_class": "VisionPsyNanoProcessor",
+  "auto_map": {
+    "AutoProcessor": "processing_visionpsynano.VisionPsyNanoProcessor"
+  },
+  "is_flash": false,
+  "vit_img_size": 512,
+  "max_img_size": 2048,
+  "inference_max_img_size": null,
+  "resize_to_max_side_len": true,
+  "resize_min_side_len": null,
+  "mp_image_token_length": 64,
+  "lm_max_position_embeddings": 8192,
+  "lm_tokenizer": "HuggingFaceTB/SmolLM2-360M-Instruct",
+  "lm_chat_template": "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
+  "vlm_extra_tokens": {
+    "image_token": "<|image|>",
+    "global_image_token": "<|global_image|>",
+    "r1c1": "<row_1_col_1>",
+    "r1c2": "<row_1_col_2>",
+    "r1c3": "<row_1_col_3>",
+    "r1c4": "<row_1_col_4>",
+    "r1c5": "<row_1_col_5>",
+    "r1c6": "<row_1_col_6>",
+    "r1c7": "<row_1_col_7>",
+    "r1c8": "<row_1_col_8>",
+    "r2c1": "<row_2_col_1>",
+    "r2c2": "<row_2_col_2>",
+    "r2c3": "<row_2_col_3>",
+    "r2c4": "<row_2_col_4>",
+    "r2c5": "<row_2_col_5>",
+    "r2c6": "<row_2_col_6>",
+    "r2c7": "<row_2_col_7>",
+    "r2c8": "<row_2_col_8>",
+    "r3c1": "<row_3_col_1>",
+    "r3c2": "<row_3_col_2>",
+    "r3c3": "<row_3_col_3>",
+    "r3c4": "<row_3_col_4>",
+    "r3c5": "<row_3_col_5>",
+    "r3c6": "<row_3_col_6>",
+    "r3c7": "<row_3_col_7>",
+    "r3c8": "<row_3_col_8>",
+    "r4c1": "<row_4_col_1>",
+    "r4c2": "<row_4_col_2>",
+    "r4c3": "<row_4_col_3>",
+    "r4c4": "<row_4_col_4>",
+    "r4c5": "<row_4_col_5>",
+    "r4c6": "<row_4_col_6>",
+    "r4c7": "<row_4_col_7>",
+    "r4c8": "<row_4_col_8>",
+    "r5c1": "<row_5_col_1>",
+    "r5c2": "<row_5_col_2>",
+    "r5c3": "<row_5_col_3>",
+    "r5c4": "<row_5_col_4>",
+    "r5c5": "<row_5_col_5>",
+    "r5c6": "<row_5_col_6>",
+    "r5c7": "<row_5_col_7>",
+    "r5c8": "<row_5_col_8>",
+    "r6c1": "<row_6_col_1>",
+    "r6c2": "<row_6_col_2>",
+    "r6c3": "<row_6_col_3>",
+    "r6c4": "<row_6_col_4>",
+    "r6c5": "<row_6_col_5>",
+    "r6c6": "<row_6_col_6>",
+    "r6c7": "<row_6_col_7>",
+    "r6c8": "<row_6_col_8>",
+    "r7c1": "<row_7_col_1>",
+    "r7c2": "<row_7_col_2>",
+    "r7c3": "<row_7_col_3>",
+    "r7c4": "<row_7_col_4>",
+    "r7c5": "<row_7_col_5>",
+    "r7c6": "<row_7_col_6>",
+    "r7c7": "<row_7_col_7>",
+    "r7c8": "<row_7_col_8>",
+    "r8c1": "<row_8_col_1>",
+    "r8c2": "<row_8_col_2>",
+    "r8c3": "<row_8_col_3>",
+    "r8c4": "<row_8_col_4>",
+    "r8c5": "<row_8_col_5>",
+    "r8c6": "<row_8_col_6>",
+    "r8c7": "<row_8_col_7>",
+    "r8c8": "<row_8_col_8>"
+  }
+}

processing_visionpsynano.py

@@ -0,0 +1,411 @@
+"""HuggingFace processor for VisionPsyNano."""
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, List, Optional, Union
+
+import torch
+from PIL import Image
+from transformers import AutoTokenizer, BatchFeature, ProcessorMixin
+from transformers.utils import logging
+
+
+try:
+    from .configuration_visionpsynano import (
+        VisionPsyNanoConfig,
+        apply_flash_preprocess,
+        resolve_is_flash,
+    )
+except ImportError:
+    from configuration_visionpsynano import (
+        VisionPsyNanoConfig,
+        apply_flash_preprocess,
+        resolve_is_flash,
+    )
+
+try:
+    from .processors import get_image_processor, get_image_string, get_tokenizer
+except ImportError as exc:
+    raise ImportError(
+        "Failed to import VisionPsyNano data processors. "
+        "For Hub packages use scripts/package_hub_repo.py so imports are flattened."
+    ) from exc
+
+logger = logging.get_logger(__name__)
+
+
+def _as_pil_list(images) -> List[Image.Image]:
+    if images is None:
+        return []
+    if isinstance(images, Image.Image):
+        return [images.convert("RGB")]
+    if isinstance(images, (list, tuple)):
+        out = []
+        for im in images:
+            if isinstance(im, Image.Image):
+                out.append(im.convert("RGB"))
+            else:
+                raise TypeError(f"Expected PIL.Image, got {type(im)}")
+        return out
+    raise TypeError(f"Unsupported images type: {type(images)}")
+
+
+def _extract_user_text(text: Any) -> str:
+    """Normalize text / chat messages into a single user prompt string."""
+    if text is None:
+        return ""
+    if isinstance(text, str):
+        return text
+    if isinstance(text, list) and text and isinstance(text[0], dict) and "role" in text[0]:
+        parts = []
+        for msg in text:
+            content = msg.get("content", "")
+            if isinstance(content, str):
+                parts.append(content)
+            elif isinstance(content, list):
+                for block in content:
+                    if isinstance(block, str):
+                        parts.append(block)
+                    elif isinstance(block, dict) and block.get("type") == "text":
+                        parts.append(block.get("text", ""))
+        return " ".join(p for p in parts if p).strip()
+    if isinstance(text, list) and all(isinstance(t, str) for t in text):
+        if len(text) == 1:
+            return text[0]
+        raise ValueError("Batch text list with >1 items is not supported; call per sample.")
+    raise TypeError(f"Unsupported text type: {type(text)}")
+
+
+class VisionPsyNanoProcessor(ProcessorMixin):
+    """Tokenizer + VisionPsyNano image preprocessing (default or Flash)."""
+
+    attributes = ["tokenizer"]
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(self, tokenizer, chat_template=None, **kwargs):
+        image_keys = {
+            "is_flash",
+            "variant",
+            "vit_img_size",
+            "max_img_size",
+            "inference_max_img_size",
+            "resize_to_max_side_len",
+            "resize_min_side_len",
+            "mp_image_token_length",
+            "lm_max_position_embeddings",
+            "vlm_extra_tokens",
+            "lm_tokenizer",
+            "lm_chat_template",
+        }
+        image_cfg = {k: kwargs.pop(k) for k in list(kwargs) if k in image_keys}
+        super().__init__(tokenizer, chat_template=chat_template, **kwargs)
+
+        flash = resolve_is_flash(
+            is_flash=image_cfg.get("is_flash"),
+            variant=image_cfg.get("variant"),
+            resize_to_max_side_len=image_cfg.get("resize_to_max_side_len"),
+        )
+        resize_to_max, resize_min = apply_flash_preprocess(
+            is_flash=flash,
+            resize_to_max_side_len=image_cfg.get("resize_to_max_side_len"),
+            resize_min_side_len=image_cfg.get("resize_min_side_len"),
+        )
+        self.is_flash = flash
+        self.vit_img_size = int(image_cfg.get("vit_img_size", 512))
+        self.max_img_size = int(image_cfg.get("max_img_size", 2048))
+        self.inference_max_img_size = image_cfg.get("inference_max_img_size")
+        self.resize_to_max_side_len = resize_to_max
+        self.resize_min_side_len = resize_min
+        self.mp_image_token_length = int(image_cfg.get("mp_image_token_length", 64))
+        self.lm_max_position_embeddings = int(
+            image_cfg.get("lm_max_position_embeddings", 8192)
+        )
+        self.vlm_extra_tokens = image_cfg.get("vlm_extra_tokens")
+        self.lm_tokenizer = image_cfg.get("lm_tokenizer")
+        self.lm_chat_template = image_cfg.get("lm_chat_template")
+
+    def _effective_max_img_size(self) -> int:
+        return int(self.inference_max_img_size or self.max_img_size)
+
+    def _build_image_processor(self):
+        return get_image_processor(
+            self._effective_max_img_size(),
+            self.vit_img_size,
+            self.resize_to_max_side_len,
+            self.resize_min_side_len,
+        )
+
+    def set_flash(self, enabled: bool = True) -> None:
+        """Enable/disable Flash optimized image preprocessing in-place."""
+        self.is_flash = bool(enabled)
+        self.resize_to_max_side_len, self.resize_min_side_len = apply_flash_preprocess(
+            is_flash=self.is_flash,
+            resize_to_max_side_len=None,
+            resize_min_side_len=self.resize_min_side_len if self.is_flash else None,
+        )
+
+    def __call__(
+        self,
+        images: Optional[Union[Image.Image, List[Image.Image]]] = None,
+        text: Optional[Any] = None,
+        return_tensors: Optional[str] = "pt",
+        is_flash: Optional[bool] = None,
+        **kwargs,
+    ) -> BatchFeature:
+        if is_flash is not None:
+            self.set_flash(bool(is_flash))
+
+        prompt = _extract_user_text(text)
+        pil_images = _as_pil_list(images)
+
+        image_processor = self._build_image_processor()
+        processed_tensors = []
+        ratios = []
+        for img in pil_images:
+            processed_image, splitted_image_ratio = image_processor(img)
+            if (
+                not hasattr(self.tokenizer, "global_image_token")
+                and splitted_image_ratio[0] * splitted_image_ratio[1]
+                == len(processed_image) - 1
+            ):
+                processed_image = processed_image[1:]
+            processed_tensors.append(processed_image)
+            ratios.append(splitted_image_ratio)
+
+        if processed_tensors:
+            image_string = get_image_string(
+                self.tokenizer, ratios, self.mp_image_token_length
+            )
+            images_tensor = torch.cat(processed_tensors, dim=0)
+        else:
+            image_string = ""
+            images_tensor = None
+
+        messages = [{"role": "user", "content": image_string + prompt}]
+        full_prompt = self.tokenizer.apply_chat_template(
+            [messages], tokenize=False, add_generation_prompt=True
+        )
+        if isinstance(full_prompt, list):
+            full_prompt = full_prompt[0]
+
+        encoded = self.tokenizer(
+            [full_prompt],
+            return_tensors=return_tensors,
+            padding=False,
+            truncation=True,
+            max_length=self.lm_max_position_embeddings,
+            **{k: v for k, v in kwargs.items() if k in ("padding", "truncation", "max_length")},
+        )
+
+        data = {
+            "input_ids": encoded["input_ids"],
+            "attention_mask": encoded.get("attention_mask"),
+        }
+        if images_tensor is not None:
+            data["images"] = images_tensor
+            data["pixel_values"] = images_tensor  # HF-friendly alias
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    def batch_decode(self, *args, **kwargs):
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self.tokenizer.decode(*args, **kwargs)
+
+    def image_config_dict(self) -> dict:
+        return {
+            "processor_class": "VisionPsyNanoProcessor",
+            "auto_map": {
+                "AutoProcessor": "processing_visionpsynano.VisionPsyNanoProcessor",
+            },
+            "is_flash": self.is_flash,
+            "vit_img_size": self.vit_img_size,
+            "max_img_size": self.max_img_size,
+            "inference_max_img_size": self.inference_max_img_size,
+            "resize_to_max_side_len": self.resize_to_max_side_len,
+            "resize_min_side_len": self.resize_min_side_len,
+            "mp_image_token_length": self.mp_image_token_length,
+            "lm_max_position_embeddings": self.lm_max_position_embeddings,
+            "lm_tokenizer": self.lm_tokenizer,
+            "lm_chat_template": self.lm_chat_template,
+            "vlm_extra_tokens": self.vlm_extra_tokens,
+        }
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+        os.makedirs(save_directory, exist_ok=True)
+        super().save_pretrained(save_directory, **kwargs)
+
+        proc_cfg_path = os.path.join(save_directory, "processor_config.json")
+        proc_cfg = {}
+        if os.path.exists(proc_cfg_path):
+            with open(proc_cfg_path) as f:
+                proc_cfg = json.load(f)
+        proc_cfg["processor_class"] = "VisionPsyNanoProcessor"
+        proc_cfg["auto_map"] = {
+            "AutoProcessor": "processing_visionpsynano.VisionPsyNanoProcessor",
+        }
+        with open(proc_cfg_path, "w") as f:
+            json.dump(proc_cfg, f, indent=2)
+            f.write("\n")
+
+        with open(os.path.join(save_directory, "preprocessor_config.json"), "w") as f:
+            json.dump(self.image_config_dict(), f, indent=2)
+            f.write("\n")
+
+    @staticmethod
+    def _attach_extra_token_attrs(tokenizer, vlm_extra_tokens: Optional[dict]):
+        """Restore named attrs (``image_token``, ``r1c1``, ...) used by image strings."""
+        if not vlm_extra_tokens:
+            return tokenizer
+        for name, token in vlm_extra_tokens.items():
+            if not hasattr(tokenizer, name):
+                setattr(tokenizer, name, token)
+        if hasattr(tokenizer, "image_token"):
+            try:
+                tokenizer.image_token_id = tokenizer.convert_tokens_to_ids(tokenizer.image_token)
+            except Exception:
+                pass
+        return tokenizer
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        trust_remote_code = kwargs.pop("trust_remote_code", False)
+        flash_override = kwargs.pop("is_flash", None)
+
+        tokenizer = None
+        prep = {}
+
+        def _prep_from_cfg(cfg: VisionPsyNanoConfig) -> dict:
+            return {
+                "is_flash": cfg.is_flash,
+                "vit_img_size": cfg.vit_img_size,
+                "max_img_size": cfg.max_img_size,
+                "inference_max_img_size": cfg.inference_max_img_size,
+                "resize_to_max_side_len": cfg.resize_to_max_side_len,
+                "resize_min_side_len": cfg.resize_min_side_len,
+                "mp_image_token_length": cfg.mp_image_token_length,
+                "lm_max_position_embeddings": cfg.lm_max_position_embeddings,
+                "lm_tokenizer": cfg.lm_tokenizer,
+                "lm_chat_template": cfg.lm_chat_template,
+                "vlm_extra_tokens": cfg.vlm_extra_tokens,
+            }
+
+        resolved = pretrained_model_name_or_path
+        prep_path = None
+        if os.path.isdir(resolved):
+            prep_path = os.path.join(resolved, "preprocessor_config.json")
+            tok_files_present = any(
+                os.path.exists(os.path.join(resolved, name))
+                for name in ("tokenizer.json", "tokenizer_config.json", "vocab.json")
+            )
+            if tok_files_present:
+                tokenizer = AutoTokenizer.from_pretrained(
+                    resolved, trust_remote_code=trust_remote_code, **kwargs
+                )
+            if prep_path and os.path.exists(prep_path):
+                with open(prep_path) as f:
+                    prep = json.load(f)
+        else:
+            from huggingface_hub import hf_hub_download, list_repo_files
+
+            try:
+                files = set(list_repo_files(resolved))
+            except Exception:
+                files = set()
+            if "preprocessor_config.json" in files:
+                prep_path = hf_hub_download(resolved, "preprocessor_config.json")
+                with open(prep_path) as f:
+                    prep = json.load(f)
+            if any(n in files for n in ("tokenizer.json", "tokenizer_config.json", "vocab.json")):
+                tokenizer = AutoTokenizer.from_pretrained(
+                    resolved, trust_remote_code=trust_remote_code, **kwargs
+                )
+
+        if tokenizer is None:
+            lm_tok = prep.get("lm_tokenizer")
+            extra = prep.get("vlm_extra_tokens")
+            chat = prep.get("lm_chat_template")
+            if lm_tok is None:
+                cfg = VisionPsyNanoConfig.from_pretrained(
+                    pretrained_model_name_or_path, trust_remote_code=trust_remote_code
+                )
+                lm_tok = cfg.lm_tokenizer
+                extra = cfg.vlm_extra_tokens
+                chat = cfg.lm_chat_template
+                if not prep:
+                    prep = _prep_from_cfg(cfg)
+            tokenizer = get_tokenizer(lm_tok, extra, chat)
+
+        if not prep:
+            try:
+                cfg = VisionPsyNanoConfig.from_pretrained(
+                    pretrained_model_name_or_path, trust_remote_code=trust_remote_code
+                )
+                prep = _prep_from_cfg(cfg)
+            except Exception as e:
+                logger.warning("Could not load VisionPsyNanoConfig for processor defaults: %s", e)
+
+        if flash_override is not None:
+            prep["is_flash"] = bool(flash_override)
+
+        extra = prep.get("vlm_extra_tokens")
+        if extra is None:
+            try:
+                cfg = VisionPsyNanoConfig.from_pretrained(
+                    pretrained_model_name_or_path, trust_remote_code=trust_remote_code
+                )
+                extra = cfg.vlm_extra_tokens
+                prep.setdefault("vlm_extra_tokens", extra)
+            except Exception:
+                pass
+        cls._attach_extra_token_attrs(tokenizer, extra)
+
+        chat_template = getattr(tokenizer, "chat_template", None) or prep.get("lm_chat_template")
+        return cls(tokenizer, chat_template=chat_template, **{
+            k: prep[k]
+            for k in (
+                "is_flash",
+                "vit_img_size",
+                "max_img_size",
+                "inference_max_img_size",
+                "resize_to_max_side_len",
+                "resize_min_side_len",
+                "mp_image_token_length",
+                "lm_max_position_embeddings",
+                "vlm_extra_tokens",
+                "lm_tokenizer",
+                "lm_chat_template",
+            )
+            if k in prep
+        })
+
+    @classmethod
+    def from_config(cls, config: VisionPsyNanoConfig, tokenizer=None) -> "VisionPsyNanoProcessor":
+        if tokenizer is None:
+            tokenizer = get_tokenizer(
+                config.lm_tokenizer, config.vlm_extra_tokens, config.lm_chat_template
+            )
+        return cls(
+            tokenizer,
+            chat_template=config.lm_chat_template,
+            is_flash=config.is_flash,
+            vit_img_size=config.vit_img_size,
+            max_img_size=config.max_img_size,
+            inference_max_img_size=config.inference_max_img_size,
+            resize_to_max_side_len=config.resize_to_max_side_len,
+            resize_min_side_len=config.resize_min_side_len,
+            mp_image_token_length=config.mp_image_token_length,
+            lm_max_position_embeddings=config.lm_max_position_embeddings,
+            vlm_extra_tokens=config.vlm_extra_tokens,
+            lm_tokenizer=config.lm_tokenizer,
+            lm_chat_template=config.lm_chat_template,
+        )
+
+
+try:
+    from transformers import AutoProcessor
+
+    AutoProcessor.register(VisionPsyNanoConfig, VisionPsyNanoProcessor)
+except Exception:
+    pass

processor_config.json

@@ -0,0 +1,6 @@
+{
+  "processor_class": "VisionPsyNanoProcessor",
+  "auto_map": {
+    "AutoProcessor": "processing_visionpsynano.VisionPsyNanoProcessor"
+  }
+}

processors.py

@@ -0,0 +1,118 @@
+from typing import Optional
+
+from transformers import AutoTokenizer
+import torchvision.transforms as transforms
+
+from .custom_transforms import DynamicResize, GlobalAndSplitImages
+
+TOKENIZERS_CACHE = {}
+FLASH_MIN_SIDE_LEN = 512
+
+
+# Local copies of configuration_visionpsynano.{resolve_is_flash,apply_flash_preprocess}:
+# the data/ tree is vendored standalone (vLLM plugin), so it must not import from
+# the HF package layout. Keep in sync with configuration_visionpsynano.py.
+def resolve_is_flash(
+    *,
+    is_flash: Optional[bool] = None,
+    variant: Optional[str] = None,
+    resize_to_max_side_len: Optional[bool] = None,
+) -> bool:
+    """Resolve Flash preprocess mode from explicit flag, legacy variant, or resize policy."""
+    if is_flash is not None:
+        return bool(is_flash)
+    if variant is not None:
+        v = str(variant).lower().strip()
+        if v in ("flash", "nano-flash", "visionpsy-nano-flash"):
+            return True
+        if v in ("nano", "plain", "nano-plain", "visionpsy-nano"):
+            return False
+        raise ValueError(
+            "legacy variant must be 'nano' or 'flash' "
+            f"(or aliases); got {variant!r}. Prefer is_flash=True/False."
+        )
+    if resize_to_max_side_len is not None:
+        return not bool(resize_to_max_side_len)
+    return False
+
+
+def apply_flash_preprocess(
+    *,
+    is_flash: bool,
+    resize_to_max_side_len: Optional[bool] = None,
+    resize_min_side_len: Optional[int] = None,
+) -> tuple[bool, Optional[int]]:
+    """Return (resize_to_max_side_len, resize_min_side_len) for the given Flash mode."""
+    if resize_to_max_side_len is None:
+        resize_to_max_side_len = not is_flash
+    resize_to_max_side_len = bool(resize_to_max_side_len)
+    if is_flash:
+        resize_min_side_len = max(int(resize_min_side_len or 0), FLASH_MIN_SIDE_LEN)
+    elif resize_to_max_side_len:
+        resize_min_side_len = None
+    return resize_to_max_side_len, resize_min_side_len
+
+
+def apply_model_preprocess(cfg, *, flash: bool = False) -> bool:
+    is_flash = resolve_is_flash(
+        is_flash=True if flash else getattr(cfg, "is_flash", None),
+        variant=getattr(cfg, "variant", None),
+        resize_to_max_side_len=(
+            False if flash else getattr(cfg, "resize_to_max_side_len", None)
+        ),
+    )
+    resize_to_max, resize_min = apply_flash_preprocess(
+        is_flash=is_flash,
+        resize_to_max_side_len=None if flash else getattr(cfg, "resize_to_max_side_len", None),
+        resize_min_side_len=getattr(cfg, "resize_min_side_len", None),
+    )
+    cfg.resize_to_max_side_len = resize_to_max
+    cfg.resize_min_side_len = resize_min
+    if hasattr(cfg, "is_flash"):
+        try:
+            cfg.is_flash = is_flash
+        except AttributeError:
+            pass
+    return is_flash
+
+
+def get_tokenizer(name, extra_special_tokens=None, chat_template=None):
+    cache_key = (
+        name,
+        tuple(sorted((extra_special_tokens or {}).items())),
+        chat_template,
+    )
+    if cache_key not in TOKENIZERS_CACHE:
+        tokenizer_init_kwargs = {"use_fast": True}
+        if extra_special_tokens is not None:
+            tokenizer_init_kwargs["extra_special_tokens"] = extra_special_tokens
+        if chat_template is not None:
+            tokenizer_init_kwargs["chat_template"] = chat_template
+        tokenizer = AutoTokenizer.from_pretrained(name, **tokenizer_init_kwargs,)
+        tokenizer.pad_token = tokenizer.eos_token
+        TOKENIZERS_CACHE[cache_key] = tokenizer
+    return TOKENIZERS_CACHE[cache_key]
+
+def get_image_processor(max_img_size, splitted_image_size, resize_to_max_side_len=False, min_side_len=None):
+    return transforms.Compose([
+        DynamicResize(splitted_image_size, max_img_size, resize_to_max_side_len, min_side_len),
+        transforms.ToTensor(),
+        GlobalAndSplitImages(splitted_image_size),
+    ])
+
+def get_image_string(tokenizer, splitted_image_counts, mp_image_token_length):
+    image_string = ""
+
+    for idx, (n_h, n_w) in enumerate(splitted_image_counts):
+        if len(splitted_image_counts) > 1:
+            image_string += f"<image: {idx}>"
+        if hasattr(tokenizer, "global_image_token"):
+            image_string += tokenizer.global_image_token
+            image_string += tokenizer.image_token * mp_image_token_length
+            if n_h == 1 and n_w == 1:
+                continue
+        for i in range(n_h):
+            for j in range(n_w):
+                image_string += getattr(tokenizer, f'r{i+1}c{j+1}')
+                image_string += tokenizer.image_token * mp_image_token_length
+    return image_string

runtime_profile.py

@@ -0,0 +1,37 @@
+"""Inference runtime profiles (eager vs torch.compile + CUDA graphs)."""
+from __future__ import annotations
+
+import torch
+
+
+def apply_eager_profile(model) -> None:
+    for cfg in (getattr(model, "cfg", None), getattr(model, "config", None)):
+        if cfg is not None and hasattr(cfg, "compile_inference"):
+            cfg.compile_inference = False
+
+
+def apply_deploy_profile(model, device: torch.device) -> None:
+    if device.type != "cuda":
+        apply_eager_profile(model)
+        return
+    for cfg in (getattr(model, "cfg", None), getattr(model, "config", None)):
+        if cfg is None:
+            continue
+        if hasattr(cfg, "compile_inference"):
+            cfg.compile_inference = True
+        if hasattr(cfg, "compile_inference_mode"):
+            cfg.compile_inference_mode = "reduce-overhead"
+        if hasattr(cfg, "cuda_graphs_cache_quantum"):
+            cfg.cuda_graphs_cache_quantum = 128
+        if hasattr(cfg, "eos_check_interval"):
+            cfg.eos_check_interval = 16
+
+
+def profile_label(model, device: torch.device) -> str:
+    cfg = getattr(model, "cfg", None) or getattr(model, "config", None)
+    if device.type != "cuda":
+        return "cpu-eager"
+    if cfg is not None and getattr(cfg, "compile_inference", False):
+        mode = getattr(cfg, "compile_inference_mode", "default")
+        return f"cuda-compile-{mode}"
+    return "cuda-eager"

tokenizer_config.json

@@ -0,0 +1,150 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|im_start|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "global_image_token": "<|global_image|>",
+  "image_token": "<|image|>",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 8192,
+  "model_specific_special_tokens": {
+    "global_image_token": "<|global_image|>",
+    "image_token": "<|image|>",
+    "r1c1": "<row_1_col_1>",
+    "r1c2": "<row_1_col_2>",
+    "r1c3": "<row_1_col_3>",
+    "r1c4": "<row_1_col_4>",
+    "r1c5": "<row_1_col_5>",
+    "r1c6": "<row_1_col_6>",
+    "r1c7": "<row_1_col_7>",
+    "r1c8": "<row_1_col_8>",
+    "r2c1": "<row_2_col_1>",
+    "r2c2": "<row_2_col_2>",
+    "r2c3": "<row_2_col_3>",
+    "r2c4": "<row_2_col_4>",
+    "r2c5": "<row_2_col_5>",
+    "r2c6": "<row_2_col_6>",
+    "r2c7": "<row_2_col_7>",
+    "r2c8": "<row_2_col_8>",
+    "r3c1": "<row_3_col_1>",
+    "r3c2": "<row_3_col_2>",
+    "r3c3": "<row_3_col_3>",
+    "r3c4": "<row_3_col_4>",
+    "r3c5": "<row_3_col_5>",
+    "r3c6": "<row_3_col_6>",
+    "r3c7": "<row_3_col_7>",
+    "r3c8": "<row_3_col_8>",
+    "r4c1": "<row_4_col_1>",
+    "r4c2": "<row_4_col_2>",
+    "r4c3": "<row_4_col_3>",
+    "r4c4": "<row_4_col_4>",
+    "r4c5": "<row_4_col_5>",
+    "r4c6": "<row_4_col_6>",
+    "r4c7": "<row_4_col_7>",
+    "r4c8": "<row_4_col_8>",
+    "r5c1": "<row_5_col_1>",
+    "r5c2": "<row_5_col_2>",
+    "r5c3": "<row_5_col_3>",
+    "r5c4": "<row_5_col_4>",
+    "r5c5": "<row_5_col_5>",
+    "r5c6": "<row_5_col_6>",
+    "r5c7": "<row_5_col_7>",
+    "r5c8": "<row_5_col_8>",
+    "r6c1": "<row_6_col_1>",
+    "r6c2": "<row_6_col_2>",
+    "r6c3": "<row_6_col_3>",
+    "r6c4": "<row_6_col_4>",
+    "r6c5": "<row_6_col_5>",
+    "r6c6": "<row_6_col_6>",
+    "r6c7": "<row_6_col_7>",
+    "r6c8": "<row_6_col_8>",
+    "r7c1": "<row_7_col_1>",
+    "r7c2": "<row_7_col_2>",
+    "r7c3": "<row_7_col_3>",
+    "r7c4": "<row_7_col_4>",
+    "r7c5": "<row_7_col_5>",
+    "r7c6": "<row_7_col_6>",
+    "r7c7": "<row_7_col_7>",
+    "r7c8": "<row_7_col_8>",
+    "r8c1": "<row_8_col_1>",
+    "r8c2": "<row_8_col_2>",
+    "r8c3": "<row_8_col_3>",
+    "r8c4": "<row_8_col_4>",
+    "r8c5": "<row_8_col_5>",
+    "r8c6": "<row_8_col_6>",
+    "r8c7": "<row_8_col_7>",
+    "r8c8": "<row_8_col_8>"
+  },
+  "pad_token": "<|im_end|>",
+  "processor_class": "VisionPsyNanoProcessor",
+  "r1c1": "<row_1_col_1>",
+  "r1c2": "<row_1_col_2>",
+  "r1c3": "<row_1_col_3>",
+  "r1c4": "<row_1_col_4>",
+  "r1c5": "<row_1_col_5>",
+  "r1c6": "<row_1_col_6>",
+  "r1c7": "<row_1_col_7>",
+  "r1c8": "<row_1_col_8>",
+  "r2c1": "<row_2_col_1>",
+  "r2c2": "<row_2_col_2>",
+  "r2c3": "<row_2_col_3>",
+  "r2c4": "<row_2_col_4>",
+  "r2c5": "<row_2_col_5>",
+  "r2c6": "<row_2_col_6>",
+  "r2c7": "<row_2_col_7>",
+  "r2c8": "<row_2_col_8>",
+  "r3c1": "<row_3_col_1>",
+  "r3c2": "<row_3_col_2>",
+  "r3c3": "<row_3_col_3>",
+  "r3c4": "<row_3_col_4>",
+  "r3c5": "<row_3_col_5>",
+  "r3c6": "<row_3_col_6>",
+  "r3c7": "<row_3_col_7>",
+  "r3c8": "<row_3_col_8>",
+  "r4c1": "<row_4_col_1>",
+  "r4c2": "<row_4_col_2>",
+  "r4c3": "<row_4_col_3>",
+  "r4c4": "<row_4_col_4>",
+  "r4c5": "<row_4_col_5>",
+  "r4c6": "<row_4_col_6>",
+  "r4c7": "<row_4_col_7>",
+  "r4c8": "<row_4_col_8>",
+  "r5c1": "<row_5_col_1>",
+  "r5c2": "<row_5_col_2>",
+  "r5c3": "<row_5_col_3>",
+  "r5c4": "<row_5_col_4>",
+  "r5c5": "<row_5_col_5>",
+  "r5c6": "<row_5_col_6>",
+  "r5c7": "<row_5_col_7>",
+  "r5c8": "<row_5_col_8>",
+  "r6c1": "<row_6_col_1>",
+  "r6c2": "<row_6_col_2>",
+  "r6c3": "<row_6_col_3>",
+  "r6c4": "<row_6_col_4>",
+  "r6c5": "<row_6_col_5>",
+  "r6c6": "<row_6_col_6>",
+  "r6c7": "<row_6_col_7>",
+  "r6c8": "<row_6_col_8>",
+  "r7c1": "<row_7_col_1>",
+  "r7c2": "<row_7_col_2>",
+  "r7c3": "<row_7_col_3>",
+  "r7c4": "<row_7_col_4>",
+  "r7c5": "<row_7_col_5>",
+  "r7c6": "<row_7_col_6>",
+  "r7c7": "<row_7_col_7>",
+  "r7c8": "<row_7_col_8>",
+  "r8c1": "<row_8_col_1>",
+  "r8c2": "<row_8_col_2>",
+  "r8c3": "<row_8_col_3>",
+  "r8c4": "<row_8_col_4>",
+  "r8c5": "<row_8_col_5>",
+  "r8c6": "<row_8_col_6>",
+  "r8c7": "<row_8_col_7>",
+  "r8c8": "<row_8_col_8>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>",
+  "vocab_size": 49152
+}

vision_language_model.py

@@ -0,0 +1,346 @@
+import json
+import os
+import tempfile
+from dataclasses import asdict
+from typing import Optional
+
+from .model_utils import top_k_top_p_filtering
+from .vision_transformer import ViT
+from .language_model import LanguageModel
+from .modality_projector import ModalityProjector
+from .vlm_config import VLMConfig
+
+from .processors import get_tokenizer
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import load_model, save_model
+
+def _drop_if_all_attend(attention_mask: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+    """Drop all-ones masks (avoid decode length mismatch with a fixed prefill mask)."""
+    if attention_mask is None:
+        return None
+    if bool(attention_mask.all()):
+        return None
+    return attention_mask
+
+class VisionLanguageModel(nn.Module):
+    def __init__(self, cfg: VLMConfig, load_backbone=True):
+        super().__init__()
+        self.cfg = cfg
+        if load_backbone:
+            print("Loading from backbone weights")
+            self.vision_encoder = ViT.from_pretrained(cfg)
+            self.decoder = LanguageModel.from_pretrained(cfg)
+        else:
+            self.vision_encoder = ViT(cfg)
+            self.decoder = LanguageModel(cfg)
+        self.MP = ModalityProjector(cfg)
+        self.load_backbone = load_backbone
+        self.tokenizer = get_tokenizer(cfg.lm_tokenizer, cfg.vlm_extra_tokens, cfg.lm_chat_template)
+
+        self._compiled_decoder = None
+        self._compiled_vision_encoder = None
+
+    def _vision_encoder_for_prefill(self):
+        if not getattr(self.cfg, "compile_inference", False):
+            return self.vision_encoder
+        if self._compiled_vision_encoder is None:
+            self._compiled_vision_encoder = torch.compile(
+                self.vision_encoder,
+                mode=getattr(self.cfg, "compile_inference_mode", "default"),
+                dynamic=True,
+                fullgraph=False,
+            )
+        return self._compiled_vision_encoder
+
+    def _decoder_for_decode(self):
+
+        if not getattr(self.cfg, "compile_inference", False):
+            return self.decoder
+        if self._compiled_decoder is None:
+            mode = getattr(self.cfg, "compile_inference_mode", "default")
+            self._compiled_decoder = torch.compile(
+                self.decoder,
+                mode=mode,
+                dynamic=(False if mode == "reduce-overhead" else True),
+                fullgraph=False,
+            )
+        return self._compiled_decoder
+
+    def _replace_img_tokens_with_embd(self, input_ids, token_embd, image_embd):
+        image_token_id = getattr(self.tokenizer, "image_token_id", None)
+        if image_token_id is None and hasattr(self.tokenizer, "image_token"):
+            image_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.image_token)
+        if image_token_id is None:
+            image_token_id = self.tokenizer.convert_tokens_to_ids("<|image|>")
+        mask = (input_ids == image_token_id).unsqueeze(-1)
+        flat = image_embd.view(-1, image_embd.size(-1)).to(token_embd.dtype)
+        token_embd.masked_scatter_(mask, flat)
+        return token_embd
+
+    def _process_images(self, images, device):
+        if isinstance(images, list):
+            if images and isinstance(images[0], list):
+                images = [img for sublist in images for img in sublist]
+
+            if not images:
+                return None
+            else:
+                return torch.cat(images, dim=0).to(device)
+        return images
+
+    def forward(self, input_ids, images, attention_mask=None, targets=None):
+        images_tensor = self._process_images(images, input_ids.device)
+        token_embd = self.decoder.token_embedding(input_ids)
+
+        if images_tensor is not None:
+            image_embd = self.vision_encoder(images_tensor)
+            image_embd = self.MP(image_embd)
+            token_embd = self._replace_img_tokens_with_embd(input_ids, token_embd, image_embd)
+
+        attention_mask = _drop_if_all_attend(attention_mask)
+
+        logits, _ = self.decoder(token_embd, attention_mask=attention_mask)
+
+        loss = None
+        if targets is not None:
+            logits = self.decoder.head(logits)
+            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), ignore_index=-100)
+
+        return logits, loss
+
+    @torch.inference_mode()
+    def generate(self, input_ids, images, attention_mask=None, max_new_tokens=5, top_k=50, top_p=0.9, temperature=0.5, greedy=False):
+        images_tensor = self._process_images(images, input_ids.device)
+        token_embd = self.decoder.token_embedding(input_ids)
+
+        if images_tensor is not None:
+            vision_encoder_fn = self._vision_encoder_for_prefill()
+            image_embd = vision_encoder_fn(images_tensor)
+            image_embd = self.MP(image_embd)
+            token_embd = self._replace_img_tokens_with_embd(input_ids, token_embd, image_embd)
+
+        current_total_seq_len = token_embd.size(1)
+        batch_size = input_ids.size(0)
+        device = input_ids.device
+
+        cache_max_length = current_total_seq_len + max(int(max_new_tokens), 1)
+        head_dim = self.cfg.lm_hidden_dim // self.cfg.lm_n_heads
+        use_cuda_graphs_pre = (
+            getattr(self.cfg, "compile_inference", False)
+            and getattr(self.cfg, "compile_inference_mode", "default") == "reduce-overhead"
+            and device.type == "cuda"
+        )
+        if use_cuda_graphs_pre:
+            quantum = int(getattr(self.cfg, "cuda_graphs_cache_quantum", 128))
+            if quantum > 1:
+                cache_max_length = ((cache_max_length + quantum - 1) // quantum) * quantum
+        cache_shape = (batch_size, self.cfg.lm_n_kv_heads, cache_max_length, head_dim)
+        cache_dtype = next(self.decoder.parameters()).dtype
+        use_cuda_graphs = use_cuda_graphs_pre
+        kv_cache_list = []
+        for _ in range(self.cfg.lm_n_blocks):
+            key_cache = torch.zeros(cache_shape, dtype=cache_dtype, device=device)
+            value_cache = torch.zeros(cache_shape, dtype=cache_dtype, device=device)
+            if use_cuda_graphs:
+                torch._dynamo.mark_static_address(key_cache)
+                torch._dynamo.mark_static_address(value_cache)
+            kv_cache_list.append({
+                "cache_max_length": cache_max_length,
+                "key_cache": key_cache,
+                "value_cache": value_cache,
+            })
+
+        attention_mask = _drop_if_all_attend(attention_mask)
+
+        prefill_output, kv_cache_list = self.decoder(
+            token_embd,
+            attention_mask=attention_mask,
+            kv_cache=kv_cache_list,
+            start_pos=0,
+        )
+
+        last_token_output_from_prefill = prefill_output[:, -1, :]
+
+        if not self.decoder.lm_use_tokens:
+            current_logits = self.decoder.head(last_token_output_from_prefill)
+        else:
+            current_logits = last_token_output_from_prefill
+
+        newly_generated_ids_list = []
+
+        decode_decoder = self._decoder_for_decode()
+
+        if use_cuda_graphs:
+            start_pos_buf = torch.zeros((), dtype=torch.long, device=device)
+            torch._dynamo.mark_static_address(start_pos_buf)
+        else:
+            start_pos_buf = None
+
+        eos_id = self.tokenizer.eos_token_id
+        eos_check_interval = int(getattr(self.cfg, "eos_check_interval", 16))
+        if eos_id is None or eos_check_interval <= 0:
+            eos_check_interval = 0
+
+        for step_idx in range(max_new_tokens):
+            if greedy:
+                next_token_id = torch.argmax(current_logits, dim=-1, keepdim=True)
+            else:
+                filtered_logits = top_k_top_p_filtering(current_logits, top_k=top_k, top_p=top_p)
+                probs = torch.softmax(filtered_logits / temperature, dim=-1)
+                next_token_id = torch.multinomial(probs, num_samples=1)
+
+            newly_generated_ids_list.append(next_token_id)
+
+            if (
+                eos_check_interval > 0
+                and (step_idx + 1) % eos_check_interval == 0
+                and bool((next_token_id == eos_id).all().item())
+            ):
+                break
+
+            next_token_embed = self.decoder.token_embedding(next_token_id)
+
+            current_token_start_pos = current_total_seq_len
+            current_total_seq_len += 1
+
+            if start_pos_buf is not None:
+                start_pos_buf.fill_(current_token_start_pos)
+                start_pos_to_pass = start_pos_buf
+            else:
+                start_pos_to_pass = current_token_start_pos
+
+            decode_step_output, kv_cache_list = decode_decoder(
+                next_token_embed,
+                attention_mask=attention_mask,
+                kv_cache=kv_cache_list,
+                start_pos=start_pos_to_pass,
+            )
+
+            last_token_output = decode_step_output[:, -1, :]
+
+            if not self.decoder.lm_use_tokens:
+                current_logits = self.decoder.head(last_token_output)
+            else:
+                current_logits = last_token_output
+        if not newly_generated_ids_list:
+            return torch.empty((batch_size,0), dtype=torch.long, device=input_ids.device)
+
+        generated_ids = torch.cat(newly_generated_ids_list, dim=1)
+
+        if self.tokenizer.eos_token_id is not None and generated_ids.numel() > 0:
+            seq_len = generated_ids.size(1)
+            device = generated_ids.device
+
+            eos_mask = (generated_ids == self.tokenizer.eos_token_id)
+
+            col_indices_for_min = torch.arange(seq_len, device=device)
+            masked_col_indices = torch.where(eos_mask, col_indices_for_min.unsqueeze(0).expand_as(generated_ids), seq_len + 1)
+
+            first_eos_indices_values = torch.min(masked_col_indices, dim=1).values
+            actual_first_eos_indices = torch.clamp(first_eos_indices_values, max=seq_len)
+
+            col_indices_for_comparison = torch.arange(seq_len, device=device).unsqueeze(0).expand_as(generated_ids)
+            replace_mask = col_indices_for_comparison > actual_first_eos_indices.unsqueeze(1)
+            generated_ids[replace_mask] = self.tokenizer.eos_token_id
+        return generated_ids
+
+    @classmethod
+    def from_pretrained(
+        cls, repo_id_or_path: str, *, revision: Optional[str] = None
+    ) -> "VisionLanguageModel":
+
+        if os.path.exists(repo_id_or_path):
+            config_path = os.path.join(repo_id_or_path, "config.json")
+            weights_path = os.path.join(repo_id_or_path, "model.safetensors")
+
+            if not os.path.exists(config_path):
+                raise ValueError(
+                    f"Config file not found at {config_path}. Please provide a valid path."
+                )
+            if not os.path.exists(weights_path):
+                raise ValueError(
+                    f"Weights file not found at {weights_path}. Please provide a valid path."
+                )
+        else:
+            from huggingface_hub import hf_hub_download
+
+            config_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="config.json", revision=revision
+            )
+            weights_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="model.safetensors", revision=revision
+            )
+
+        with open(config_path, "r") as f:
+            raw_cfg = json.load(f)
+        from dataclasses import fields as dc_fields
+        valid_keys = {fld.name for fld in dc_fields(VLMConfig)}
+        cfg = VLMConfig(**{k: v for k, v in raw_cfg.items() if k in valid_keys})
+
+        model = cls(cfg, load_backbone=False)
+        model._saved_config_keys = set(raw_cfg.keys())
+
+        load_model(model, weights_path)
+
+        return model
+
+    def save_pretrained(self, save_directory: str) -> None:
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        with open(os.path.join(save_directory, "config.json"), "w") as f:
+            f.write(json.dumps(asdict(self.cfg), indent=4))
+
+        save_model(self, os.path.join(save_directory, "model.safetensors"))
+
+    def push_to_hub(self, repo_id: str, private: bool = False) -> None:
+
+        from huggingface_hub import create_repo, upload_folder
+
+        repo_url = create_repo(repo_id=repo_id, private=private, exist_ok=True)
+        repo_id = repo_url.repo_id
+        print("Created repo: ", repo_url)
+
+        with tempfile.TemporaryDirectory() as save_path:
+            self.save_pretrained(save_path)
+
+            with open(os.path.join(save_path, "README.md"), "w") as f:
+                f.write(MODEL_CARD_TEMPLATE.format(repo_id=repo_id))
+
+            return upload_folder(
+                repo_id=repo_id,
+                repo_type="model",
+                folder_path=save_path,
+                commit_message="Upload VisionPsyNano using push_to_hub",
+            )
+
+MODEL_CARD_TEMPLATE = """
+---
+library_name: transformers
+pipeline_tag: image-text-to-text
+tags:
+  - vision-language
+  - multimodal
+---
+
+```python
+from transformers import AutoModelForImageTextToText, AutoProcessor
+from PIL import Image
+
+repo = "{repo_id}"
+model = AutoModelForImageTextToText.from_pretrained(
+    repo, trust_remote_code=True, dtype="auto"
+).cuda().eval()
+model.apply_deploy_profile(model.device)
+processor = AutoProcessor.from_pretrained(repo, trust_remote_code=True)
+
+inputs = processor(images=Image.open("image.jpg"), text="Describe the image.", return_tensors="pt")
+inputs = {{k: v.cuda() if hasattr(v, "cuda") else v for k, v in inputs.items() if v is not None}}
+inputs.pop("pixel_values", None)
+out = model.generate(**inputs, max_new_tokens=64, greedy=True)
+print(processor.batch_decode(out, skip_special_tokens=True)[0])
+```
+"""

vision_transformer.py

@@ -0,0 +1,217 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ViTPatchEmbeddings(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+
+        self.img_size = cfg.vit_img_size
+        self.patch_size = cfg.vit_patch_size
+        self.num_patches = (self.img_size // self.patch_size) ** 2
+        self.cls_flag = cfg.vit_cls_flag
+        self.embd_dim = cfg.vit_hidden_dim
+
+        self.conv = nn.Conv2d(
+            in_channels=3,
+            out_channels=self.embd_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            padding="valid",
+        )
+
+        if self.cls_flag:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embd_dim))
+            self.position_embedding = nn.Parameter(torch.rand(1, self.num_patches + 1, self.embd_dim))
+        else:
+            self.position_embedding = nn.Parameter(torch.rand(1, self.num_patches, self.embd_dim))
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = x.flatten(2)
+        x = x.transpose(1, 2)
+
+        if self.cls_flag:
+            cls_token = self.cls_token.expand(x.shape[0], -1, -1)
+            x = torch.cat((cls_token, x), dim=1)
+        x = x + self.position_embedding
+        return x
+
+class ViTMultiHeadAttention(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+
+        self.n_heads = cfg.vit_n_heads
+        self.embd_dim = cfg.vit_hidden_dim
+        assert self.embd_dim % self.n_heads == 0, "embd_dim must be divisible by num_heads"
+        self.head_dim = self.embd_dim // self.n_heads
+        self.dropout = cfg.vit_dropout
+
+        self.qkv_proj = nn.Linear(self.embd_dim, 3 * self.embd_dim, bias=True)
+        self.out_proj = nn.Linear(self.embd_dim, self.embd_dim, bias=True)
+
+        self.attn_dropout = nn.Dropout(self.dropout)
+        self.resid_dropout = nn.Dropout(self.dropout)
+
+        self.sdpa = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.sdpa:
+            print("Warning: scaled dot product attention not available. Using standard attention in ViT.")
+
+    def forward(self, x):
+        B, T, C = x.size()
+
+        qkv = self.qkv_proj(x)
+        q, k, v = qkv.split(C, dim=2)
+        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+
+        if self.sdpa:
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=None,
+                dropout_p=self.dropout if self.training else 0.0,
+                is_causal=False
+            )
+        else:
+            attn = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            attn = F.softmax(attn, dim=-1)
+            attn = self.attn_dropout(attn)
+            y = attn @ v
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        y = self.out_proj(y)
+        y = self.resid_dropout(y)
+
+        return y
+
+class ViTMLP(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.activation_fn = nn.GELU(approximate='tanh')
+        self.fc1 = nn.Linear(cfg.vit_hidden_dim, cfg.vit_inter_dim)
+        self.fc2 = nn.Linear(cfg.vit_inter_dim, cfg.vit_hidden_dim)
+        self.dropout = nn.Dropout(cfg.vit_dropout)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.activation_fn(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+class ViTBlock(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(cfg.vit_hidden_dim, eps=cfg.vit_ln_eps)
+        self.attn = ViTMultiHeadAttention(cfg)
+        self.ln2 = nn.LayerNorm(cfg.vit_hidden_dim, eps=cfg.vit_ln_eps)
+        self.mlp = ViTMLP(cfg)
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+class ViT(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.patch_embedding = ViTPatchEmbeddings(cfg)
+        self.cls_flag = cfg.vit_cls_flag
+        self.dropout = nn.Dropout(cfg.vit_dropout)
+        self.blocks = nn.ModuleList([ViTBlock(cfg) for _ in range(cfg.vit_n_blocks)])
+        self.layer_norm = nn.LayerNorm(cfg.vit_hidden_dim, eps=cfg.vit_ln_eps)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, nn.Conv2d):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+
+    def forward(self, x):
+        x = self.patch_embedding(x)
+        x = self.dropout(x)
+        for block in self.blocks:
+            x = block(x)
+
+        if self.cls_flag:
+            x = self.layer_norm(x[:, 0])
+        else:
+            x = self.layer_norm(x)
+        return x
+    @classmethod
+    def from_pretrained(cls, cfg):
+        from transformers import SiglipVisionConfig
+        from huggingface_hub import hf_hub_download
+        import safetensors
+
+        hf_config = SiglipVisionConfig.from_pretrained(cfg.vit_model_type)
+        cfg.vit_dropout=hf_config.attention_dropout
+        cfg.vit_hidden_dim=hf_config.hidden_size
+        cfg.vit_img_size=hf_config.image_size
+        cfg.vit_inter_dim=hf_config.intermediate_size
+        cfg.vit_ln_eps=hf_config.layer_norm_eps
+        cfg.vit_n_heads=hf_config.num_attention_heads
+        cfg.vit_n_blocks=hf_config.num_hidden_layers
+        cfg.vit_patch_size=hf_config.patch_size
+        model = cls(cfg)
+        safetensors_file = hf_hub_download(repo_id=cfg.vit_model_type, filename="model.safetensors")
+
+        sd = model.state_dict()
+        mapping = {
+            'vision_model.embeddings.patch_embedding.weight': 'patch_embedding.conv.weight',
+            'vision_model.embeddings.patch_embedding.bias': 'patch_embedding.conv.bias',
+            'vision_model.embeddings.position_embedding.weight': 'patch_embedding.position_embedding',
+            'vision_model.post_layernorm.weight': 'layer_norm.weight',
+            'vision_model.post_layernorm.bias': 'layer_norm.bias',
+        }
+        for i in range(cfg.vit_n_blocks):
+            mapping[f'vision_model.encoder.layers.{i}.layer_norm1.weight'] = f'blocks.{i}.ln1.weight'
+            mapping[f'vision_model.encoder.layers.{i}.layer_norm1.bias'] = f'blocks.{i}.ln1.bias'
+            mapping[f'vision_model.encoder.layers.{i}.layer_norm2.weight'] = f'blocks.{i}.ln2.weight'
+            mapping[f'vision_model.encoder.layers.{i}.layer_norm2.bias'] = f'blocks.{i}.ln2.bias'
+            mapping[f'vision_model.encoder.layers.{i}.mlp.fc1.weight'] = f'blocks.{i}.mlp.fc1.weight'
+            mapping[f'vision_model.encoder.layers.{i}.mlp.fc1.bias'] = f'blocks.{i}.mlp.fc1.bias'
+            mapping[f'vision_model.encoder.layers.{i}.mlp.fc2.weight'] = f'blocks.{i}.mlp.fc2.weight'
+            mapping[f'vision_model.encoder.layers.{i}.mlp.fc2.bias'] = f'blocks.{i}.mlp.fc2.bias'
+            mapping[f'vision_model.encoder.layers.{i}.self_attn.out_proj.weight'] = f'blocks.{i}.attn.out_proj.weight'
+            mapping[f'vision_model.encoder.layers.{i}.self_attn.out_proj.bias'] = f'blocks.{i}.attn.out_proj.bias'
+        with safetensors.safe_open(filename=safetensors_file, framework="pt", device="cpu") as f:
+            for hf_key, our_key in mapping.items():
+                if hf_key in f.keys() and our_key in sd:
+                    tensor = f.get_tensor(hf_key)
+                    if tensor.shape == sd[our_key].shape:
+                        sd[our_key].copy_(tensor)
+                    else:
+                        if 'position_embedding' in hf_key:
+                            sd[our_key].copy_(tensor.unsqueeze(0))
+                        else:
+                            print(f"Shape mismatch for {hf_key} -> {our_key}: {tensor.shape} vs {sd[our_key].shape}")
+                else:
+                    if hf_key not in f.keys():
+                        print(f"Warning: Key {hf_key} not found in safetensors file")
+                    if our_key not in sd:
+                        print(f"Warning: Key {our_key} not found in model state dict")
+            for i in range(model.cfg.vit_n_blocks):
+                q_weight = f.get_tensor(f'vision_model.encoder.layers.{i}.self_attn.q_proj.weight')
+                k_weight = f.get_tensor(f'vision_model.encoder.layers.{i}.self_attn.k_proj.weight')
+                v_weight = f.get_tensor(f'vision_model.encoder.layers.{i}.self_attn.v_proj.weight')
+                qkv_weight = torch.cat((q_weight, k_weight, v_weight), dim=0)
+                sd[f'blocks.{i}.attn.qkv_proj.weight'].copy_(qkv_weight)
+                q_bias = f.get_tensor(f'vision_model.encoder.layers.{i}.self_attn.q_proj.bias')
+                k_bias = f.get_tensor(f'vision_model.encoder.layers.{i}.self_attn.k_proj.bias')
+                v_bias = f.get_tensor(f'vision_model.encoder.layers.{i}.self_attn.v_proj.bias')
+                qkv_bias = torch.cat((q_bias, k_bias, v_bias), dim=0)
+                sd[f'blocks.{i}.attn.qkv_proj.bias'].copy_(qkv_bias)
+        model.load_state_dict(sd)
+        print(f"Successfully loaded {cfg.vit_model_type} weights from safetensors. Model has {sum(p.numel() for p in model.parameters()):,} parameters.")
+        return model

vlm_config.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass, field
+
+@dataclass
+class VLMConfig:
+    vit_hidden_dim: int = 768
+    vit_inter_dim: int = 4 * vit_hidden_dim
+    vit_patch_size: int = 16
+    vit_img_size: int = 512
+    vit_n_heads: int = 12
+    vit_dropout: float = 0.0
+    vit_n_blocks: int = 12
+    vit_ln_eps: float = 1e-6
+    vit_cls_flag: bool = False
+    vit_model_type: str = 'google/siglip2-base-patch16-512'
+
+    lm_hidden_dim: int = 960
+    lm_inter_dim: int = 2560
+    lm_rms_eps: float = 1e-5
+    lm_re_base: int = 100000
+    lm_max_position_embeddings: int = 8192
+    lm_base_vocab_size: int = 49152
+    extra_token_amount: int = 66
+    lm_vocab_size: int = lm_base_vocab_size + extra_token_amount
+    lm_n_heads: int = 15
+    lm_n_kv_heads: int = 5
+    lm_dropout: float = 0.0
+    lm_n_blocks: int = 32
+    lm_attn_scaling: float = 1.0
+    lm_max_length: int = 4096
+    lm_use_tokens: bool = False
+    lm_tie_weights: bool = True
+    lm_model_type: str = 'HuggingFaceTB/SmolLM2-360M-Instruct'
+    lm_tokenizer: str = 'HuggingFaceTB/SmolLM2-360M-Instruct'
+    lm_chat_template: str = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+
+    mp_pixel_shuffle_factor: int = 4
+    mp_image_token_length: int = 64
+
+    max_img_size: int = 2048
+    resize_to_max_side_len: bool = True
+    resize_min_side_len: int | None = None
+
+    inference_max_img_size: int | None = None
+
+    vlm_extra_tokens: dict[str, str] = field(default_factory=lambda: {"image_token": "<|image|>", "global_image_token": "<|global_image|>",
+      "r1c1": "<row_1_col_1>", "r1c2": "<row_1_col_2>", "r1c3": "<row_1_col_3>", "r1c4": "<row_1_col_4>", "r1c5": "<row_1_col_5>", "r1c6": "<row_1_col_6>", "r1c7": "<row_1_col_7>", "r1c8": "<row_1_col_8>",
+      "r2c1": "<row_2_col_1>", "r2c2": "<row_2_col_2>", "r2c3": "<row_2_col_3>", "r2c4": "<row_2_col_4>", "r2c5": "<row_2_col_5>", "r2c6": "<row_2_col_6>", "r2c7": "<row_2_col_7>", "r2c8": "<row_2_col_8>",
+      "r3c1": "<row_3_col_1>", "r3c2": "<row_3_col_2>", "r3c3": "<row_3_col_3>", "r3c4": "<row_3_col_4>", "r3c5": "<row_3_col_5>", "r3c6": "<row_3_col_6>", "r3c7": "<row_3_col_7>", "r3c8": "<row_3_col_8>",
+      "r4c1": "<row_4_col_1>", "r4c2": "<row_4_col_2>", "r4c3": "<row_4_col_3>", "r4c4": "<row_4_col_4>", "r4c5": "<row_4_col_5>", "r4c6": "<row_4_col_6>", "r4c7": "<row_4_col_7>", "r4c8": "<row_4_col_8>",
+      "r5c1": "<row_5_col_1>", "r5c2": "<row_5_col_2>", "r5c3": "<row_5_col_3>", "r5c4": "<row_5_col_4>", "r5c5": "<row_5_col_5>", "r5c6": "<row_5_col_6>", "r5c7": "<row_5_col_7>", "r5c8": "<row_5_col_8>",
+      "r6c1": "<row_6_col_1>", "r6c2": "<row_6_col_2>", "r6c3": "<row_6_col_3>", "r6c4": "<row_6_col_4>", "r6c5": "<row_6_col_5>", "r6c6": "<row_6_col_6>", "r6c7": "<row_6_col_7>", "r6c8": "<row_6_col_8>",
+      "r7c1": "<row_7_col_1>", "r7c2": "<row_7_col_2>", "r7c3": "<row_7_col_3>", "r7c4": "<row_7_col_4>", "r7c5": "<row_7_col_5>", "r7c6": "<row_7_col_6>", "r7c7": "<row_7_col_7>", "r7c8": "<row_7_col_8>",
+      "r8c1": "<row_8_col_1>", "r8c2": "<row_8_col_2>", "r8c3": "<row_8_col_3>", "r8c4": "<row_8_col_4>", "r8c5": "<row_8_col_5>", "r8c6": "<row_8_col_6>", "r8c7": "<row_8_col_7>", "r8c8": "<row_8_col_8>"})
+    vlm_load_backbone_weights: bool = True
+    vlm_checkpoint_path: str = 'checkpoints'
+    hf_repo_name: str = 'qvac/VisionPsy-Nano-460M'
+
+    compile_inference: bool = True
+    compile_inference_mode: str = 'reduce-overhead'
+    cuda_graphs_cache_quantum: int = 128
+
+    eos_check_interval: int = 16