Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,12 @@ 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
+teaser.jpg filter=lfs diff=lfs merge=lfs -text
+assets/decoding_demo.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/demo.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/coco_lvis.png filter=lfs diff=lfs merge=lfs -text
+assets/dense_object_detection.png filter=lfs diff=lfs merge=lfs -text
+assets/layout_ocr.png filter=lfs diff=lfs merge=lfs -text
+assets/referring.png filter=lfs diff=lfs merge=lfs -text
+assets/sspro.png filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,74 @@
+                                   NVIDIA License
+
+1. Definitions
+
+“Licensor” means any person or entity that distributes its Work.
+
+“Work” means (a) the original work of authorship made available under this license, 
+which may include software, documentation, or other files, and (b) any additions to or 
+derivative works  thereof  that are made available under this license.
+
+The terms “reproduce,” “reproduction,” “derivative works,” and “distribution” have the 
+meaning as provided under U.S. copyright law; provided, however, that 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.
+
+Works are “made available” under this license by including in or with the Work either 
+(a) a copyright notice referencing the applicability of this license to the Work, or 
+(b) a copy of this license.
+
+2. License Grant
+
+2.1 Copyright Grant. Subject to the terms and conditions of this license, each Licensor 
+grants to you a perpetual, worldwide, non-exclusive, royalty-free, copyright license to 
+use, reproduce, prepare derivative works of, publicly display, publicly perform, 
+sublicense and distribute its Work and any resulting derivative works in any form.
+
+3. Limitations
+
+3.1 Redistribution. You may reproduce or distribute the Work only if (a) you do so under 
+this license, (b) you include a complete copy of this license with your distribution, 
+and (c) you retain without modification any copyright, patent, trademark, or attribution 
+notices that are present in the Work.
+
+3.2 Derivative Works. You may specify that additional or different terms apply to the 
+use, reproduction, and distribution of your derivative works of the Work (“Your Terms”) 
+only if (a) Your Terms provide that the use limitation in Section 3.3 applies to your 
+derivative works, and (b) you identify the specific derivative works that are subject 
+to Your Terms. Notwithstanding Your Terms, this license (including the redistribution 
+requirements in Section 3.1) will continue to apply to the Work itself.
+
+3.3 Use Limitation. The Work and any derivative works thereof only may be used or 
+intended for use non-commercially. Notwithstanding the foregoing, NVIDIA Corporation 
+and its affiliates may use the Work and any derivative works commercially. As used 
+herein, “non-commercially” means for research or evaluation purposes only.
+
+3.4 Patent Claims. If you bring or threaten to bring a patent claim against any 
+Licensor (including any claim, cross-claim or counterclaim in a lawsuit) to enforce any 
+patents that you allege are infringed by any Work, then your rights under this license 
+from such Licensor (including the grant in Section 2.1) will terminate immediately.
+
+3.5 Trademarks. This license does not grant any rights to use any Licensor’s or its 
+affiliates’ names, logos, or trademarks, except as necessary to reproduce the notices 
+described in this license.
+
+3.6 Termination. If you violate any term of this license, then your rights under this 
+license (including the grant in Section 2.1) will terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED “AS IS” WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER 
+EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF MERCHANTABILITY, FITNESS FOR 
+A PARTICULAR PURPOSE, TITLE OR NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY 
+ACTIVITIES UNDER THIS LICENSE. 
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER 
+IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE SHALL ANY LICENSOR BE LIABLE TO 
+YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL 
+DAMAGES ARISING OUT OF OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK 
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, LOST PROFITS OR 
+DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER DAMAGES OR LOSSES), EVEN IF THE 
+LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+

README.md

@@ -0,0 +1,115 @@
+---
+license: other
+base_model: nvidia/LocateAnything-3B
+base_model_relation: quantized
+quantized_by: blockblockblock
+library_name: exllamav3
+pipeline_tag: text-generation
+tags:
+  - exl3
+  - exllamav3
+  - quantized
+
+quantization_format: exl3
+bits_per_weight: 6.0
+---
+
+<div align="center">
+
+# LocateAnything · 3B
+
+<sub><code>EXL3</code> &nbsp;·&nbsp; <b>6.0&nbsp;bpw</b> &nbsp;·&nbsp; 3.9&nbsp;GB &nbsp;·&nbsp; Dense</sub>
+
+<br/>
+
+[![format](https://img.shields.io/badge/format-EXL3-c63010?style=for-the-badge&labelColor=14120e)](https://github.com/turboderp-org/exllamav3)
+[![bpw](https://img.shields.io/badge/bpw-6.0-6b8a76?style=for-the-badge&labelColor=14120e)](#quants)
+[![size](https://img.shields.io/badge/size-3.9_GB-6b8a76?style=for-the-badge&labelColor=14120e)](#quants)
+[![arch](https://img.shields.io/badge/arch-Dense_3B-c63010?style=for-the-badge&labelColor=14120e)](https://huggingface.co/nvidia/LocateAnything-3B)
+
+[![base model](https://img.shields.io/badge/Base-nvidia%2FLocateAnything--3B-2a2620?style=flat-square&logo=huggingface&logoColor=white)](https://huggingface.co/nvidia/LocateAnything-3B)
+[![quantized by](https://img.shields.io/badge/Quantized_by-blockblockblock-2a2620?style=flat-square&logo=huggingface&logoColor=white)](https://huggingface.co/blockblockblock)
+[![collection](https://img.shields.io/badge/All_bpws-Collection-c63010?style=flat-square&logo=huggingface&logoColor=white)](https://huggingface.co/collections/blockblockblock/locateanything-3b-exl3-6a205fbdd9aee1d0eba42b33)
+
+</div>
+
+---
+
+> [!NOTE]
+> An [ExLlamaV3](https://github.com/turboderp-org/exllamav3) build of [`nvidia/LocateAnything-3B`](https://huggingface.co/nvidia/LocateAnything-3B) at **6.0 bits per weight**. See [Quants](#quants) for sibling repos at other bit&#8209;widths or browse the [collection](https://huggingface.co/collections/blockblockblock/locateanything-3b-exl3-6a205fbdd9aee1d0eba42b33).
+
+## Quants
+
+<div align="center">
+
+| BPW &nbsp; | &nbsp; Head bits &nbsp; | &nbsp; Calibration rows &nbsp; | &nbsp; Size &nbsp; | &nbsp; Status |
+| :---: | :---: | :---: | ---: | :--- |
+| **6.0** | 8 | 250 | **3.9&nbsp;GB** | <kbd>this repo</kbd> |
+
+</div>
+
+## Inference
+
+<table>
+  <thead>
+    <tr>
+      <th align="left" width="32%">Loader</th>
+      <th align="left">Use it for</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td><a href="https://github.com/theroyallab/tabbyAPI"><b>TabbyAPI</b></a></td>
+      <td>OpenAI&#8209;compatible HTTP server. Drop&#8209;in for OpenAI clients.</td>
+    </tr>
+    <tr>
+      <td><a href="https://github.com/oobabooga/text-generation-webui"><b>text&#8209;generation&#8209;webui</b></a></td>
+      <td>Local chat UI. Pick the <i>ExLlamaV3</i> loader from the model dropdown.</td>
+    </tr>
+    <tr>
+      <td><a href="https://github.com/turboderp-org/exllamav3"><b>ExLlamaV3</b></a></td>
+      <td>Direct Python API for embedding the model in your own code or pipeline.</td>
+    </tr>
+  </tbody>
+</table>
+
+## Download
+
+```bash
+pip install -U huggingface_hub
+
+hf download \
+  blockblockblock/LocateAnything-3B-exl3-6.0bpw \
+  --local-dir ./LocateAnything-3B-exl3-6.0bpw
+```
+
+<details>
+<summary><b>Quantization recipe</b> &nbsp;<sub>(advanced, embedded in <code>quantization_config.json</code>)</sub></summary>
+
+<br/>
+
+| Setting | Value |
+| :--- | :--- |
+| Format | `EXL3` |
+| Bits per weight | `6.0` |
+| Head bits | `8` |
+| Calibration rows | `250` |
+| Codebook | `MCG` |
+| Out&#8209;scales | `always` |
+| Parallel mode | `enabled` |
+
+Loaded automatically by every ExLlamaV3 loader; reproduced here for searchability.
+
+</details>
+
+## License &amp; use
+
+> [!IMPORTANT]
+> Use and license **follow the [base model](https://huggingface.co/nvidia/LocateAnything-3B)**.
+> Quantization adds no additional restrictions. Refer to the upstream repository for terms, citation, and safety documentation.
+
+---
+
+<div align="center">
+  <sub><i>Quantized with <a href="https://github.com/Honkware/blockquant"><b>BlockQuant</b></a> &nbsp;·&nbsp; convention&nbsp;<code>{org}/{model}-exl3-{bpw}bpw</code></i></sub>
+</div>

added_tokens.json

@@ -0,0 +1,1040 @@
+{
+  "</box>": 151669,
+  "</c>": 152680,
+  "</img>": 151667,
+  "</interval>": 151675,
+  "</quad>": 151671,
+  "</ref>": 151673,
+  "</tool_call>": 151658,
+  "<0>": 151677,
+  "<1000>": 152677,
+  "<100>": 151777,
+  "<101>": 151778,
+  "<102>": 151779,
+  "<103>": 151780,
+  "<104>": 151781,
+  "<105>": 151782,
+  "<106>": 151783,
+  "<107>": 151784,
+  "<108>": 151785,
+  "<109>": 151786,
+  "<10>": 151687,
+  "<110>": 151787,
+  "<111>": 151788,
+  "<112>": 151789,
+  "<113>": 151790,
+  "<114>": 151791,
+  "<115>": 151792,
+  "<116>": 151793,
+  "<117>": 151794,
+  "<118>": 151795,
+  "<119>": 151796,
+  "<11>": 151688,
+  "<120>": 151797,
+  "<121>": 151798,
+  "<122>": 151799,
+  "<123>": 151800,
+  "<124>": 151801,
+  "<125>": 151802,
+  "<126>": 151803,
+  "<127>": 151804,
+  "<128>": 151805,
+  "<129>": 151806,
+  "<12>": 151689,
+  "<130>": 151807,
+  "<131>": 151808,
+  "<132>": 151809,
+  "<133>": 151810,
+  "<134>": 151811,
+  "<135>": 151812,
+  "<136>": 151813,
+  "<137>": 151814,
+  "<138>": 151815,
+  "<139>": 151816,
+  "<13>": 151690,
+  "<140>": 151817,
+  "<141>": 151818,
+  "<142>": 151819,
+  "<143>": 151820,
+  "<144>": 151821,
+  "<145>": 151822,
+  "<146>": 151823,
+  "<147>": 151824,
+  "<148>": 151825,
+  "<149>": 151826,
+  "<14>": 151691,
+  "<150>": 151827,
+  "<151>": 151828,
+  "<152>": 151829,
+  "<153>": 151830,
+  "<154>": 151831,
+  "<155>": 151832,
+  "<156>": 151833,
+  "<157>": 151834,
+  "<158>": 151835,
+  "<159>": 151836,
+  "<15>": 151692,
+  "<160>": 151837,
+  "<161>": 151838,
+  "<162>": 151839,
+  "<163>": 151840,
+  "<164>": 151841,
+  "<165>": 151842,
+  "<166>": 151843,
+  "<167>": 151844,
+  "<168>": 151845,
+  "<169>": 151846,
+  "<16>": 151693,
+  "<170>": 151847,
+  "<171>": 151848,
+  "<172>": 151849,
+  "<173>": 151850,
+  "<174>": 151851,
+  "<175>": 151852,
+  "<176>": 151853,
+  "<177>": 151854,
+  "<178>": 151855,
+  "<179>": 151856,
+  "<17>": 151694,
+  "<180>": 151857,
+  "<181>": 151858,
+  "<182>": 151859,
+  "<183>": 151860,
+  "<184>": 151861,
+  "<185>": 151862,
+  "<186>": 151863,
+  "<187>": 151864,
+  "<188>": 151865,
+  "<189>": 151866,
+  "<18>": 151695,
+  "<190>": 151867,
+  "<191>": 151868,
+  "<192>": 151869,
+  "<193>": 151870,
+  "<194>": 151871,
+  "<195>": 151872,
+  "<196>": 151873,
+  "<197>": 151874,
+  "<198>": 151875,
+  "<199>": 151876,
+  "<19>": 151696,
+  "<1>": 151678,
+  "<200>": 151877,
+  "<201>": 151878,
+  "<202>": 151879,
+  "<203>": 151880,
+  "<204>": 151881,
+  "<205>": 151882,
+  "<206>": 151883,
+  "<207>": 151884,
+  "<208>": 151885,
+  "<209>": 151886,
+  "<20>": 151697,
+  "<210>": 151887,
+  "<211>": 151888,
+  "<212>": 151889,
+  "<213>": 151890,
+  "<214>": 151891,
+  "<215>": 151892,
+  "<216>": 151893,
+  "<217>": 151894,
+  "<218>": 151895,
+  "<219>": 151896,
+  "<21>": 151698,
+  "<220>": 151897,
+  "<221>": 151898,
+  "<222>": 151899,
+  "<223>": 151900,
+  "<224>": 151901,
+  "<225>": 151902,
+  "<226>": 151903,
+  "<227>": 151904,
+  "<228>": 151905,
+  "<229>": 151906,
+  "<22>": 151699,
+  "<230>": 151907,
+  "<231>": 151908,
+  "<232>": 151909,
+  "<233>": 151910,
+  "<234>": 151911,
+  "<235>": 151912,
+  "<236>": 151913,
+  "<237>": 151914,
+  "<238>": 151915,
+  "<239>": 151916,
+  "<23>": 151700,
+  "<240>": 151917,
+  "<241>": 151918,
+  "<242>": 151919,
+  "<243>": 151920,
+  "<244>": 151921,
+  "<245>": 151922,
+  "<246>": 151923,
+  "<247>": 151924,
+  "<248>": 151925,
+  "<249>": 151926,
+  "<24>": 151701,
+  "<250>": 151927,
+  "<251>": 151928,
+  "<252>": 151929,
+  "<253>": 151930,
+  "<254>": 151931,
+  "<255>": 151932,
+  "<256>": 151933,
+  "<257>": 151934,
+  "<258>": 151935,
+  "<259>": 151936,
+  "<25>": 151702,
+  "<260>": 151937,
+  "<261>": 151938,
+  "<262>": 151939,
+  "<263>": 151940,
+  "<264>": 151941,
+  "<265>": 151942,
+  "<266>": 151943,
+  "<267>": 151944,
+  "<268>": 151945,
+  "<269>": 151946,
+  "<26>": 151703,
+  "<270>": 151947,
+  "<271>": 151948,
+  "<272>": 151949,
+  "<273>": 151950,
+  "<274>": 151951,
+  "<275>": 151952,
+  "<276>": 151953,
+  "<277>": 151954,
+  "<278>": 151955,
+  "<279>": 151956,
+  "<27>": 151704,
+  "<280>": 151957,
+  "<281>": 151958,
+  "<282>": 151959,
+  "<283>": 151960,
+  "<284>": 151961,
+  "<285>": 151962,
+  "<286>": 151963,
+  "<287>": 151964,
+  "<288>": 151965,
+  "<289>": 151966,
+  "<28>": 151705,
+  "<290>": 151967,
+  "<291>": 151968,
+  "<292>": 151969,
+  "<293>": 151970,
+  "<294>": 151971,
+  "<295>": 151972,
+  "<296>": 151973,
+  "<297>": 151974,
+  "<298>": 151975,
+  "<299>": 151976,
+  "<29>": 151706,
+  "<2>": 151679,
+  "<300>": 151977,
+  "<301>": 151978,
+  "<302>": 151979,
+  "<303>": 151980,
+  "<304>": 151981,
+  "<305>": 151982,
+  "<306>": 151983,
+  "<307>": 151984,
+  "<308>": 151985,
+  "<309>": 151986,
+  "<30>": 151707,
+  "<310>": 151987,
+  "<311>": 151988,
+  "<312>": 151989,
+  "<313>": 151990,
+  "<314>": 151991,
+  "<315>": 151992,
+  "<316>": 151993,
+  "<317>": 151994,
+  "<318>": 151995,
+  "<319>": 151996,
+  "<31>": 151708,
+  "<320>": 151997,
+  "<321>": 151998,
+  "<322>": 151999,
+  "<323>": 152000,
+  "<324>": 152001,
+  "<325>": 152002,
+  "<326>": 152003,
+  "<327>": 152004,
+  "<328>": 152005,
+  "<329>": 152006,
+  "<32>": 151709,
+  "<330>": 152007,
+  "<331>": 152008,
+  "<332>": 152009,
+  "<333>": 152010,
+  "<334>": 152011,
+  "<335>": 152012,
+  "<336>": 152013,
+  "<337>": 152014,
+  "<338>": 152015,
+  "<339>": 152016,
+  "<33>": 151710,
+  "<340>": 152017,
+  "<341>": 152018,
+  "<342>": 152019,
+  "<343>": 152020,
+  "<344>": 152021,
+  "<345>": 152022,
+  "<346>": 152023,
+  "<347>": 152024,
+  "<348>": 152025,
+  "<349>": 152026,
+  "<34>": 151711,
+  "<350>": 152027,
+  "<351>": 152028,
+  "<352>": 152029,
+  "<353>": 152030,
+  "<354>": 152031,
+  "<355>": 152032,
+  "<356>": 152033,
+  "<357>": 152034,
+  "<358>": 152035,
+  "<359>": 152036,
+  "<35>": 151712,
+  "<360>": 152037,
+  "<361>": 152038,
+  "<362>": 152039,
+  "<363>": 152040,
+  "<364>": 152041,
+  "<365>": 152042,
+  "<366>": 152043,
+  "<367>": 152044,
+  "<368>": 152045,
+  "<369>": 152046,
+  "<36>": 151713,
+  "<370>": 152047,
+  "<371>": 152048,
+  "<372>": 152049,
+  "<373>": 152050,
+  "<374>": 152051,
+  "<375>": 152052,
+  "<376>": 152053,
+  "<377>": 152054,
+  "<378>": 152055,
+  "<379>": 152056,
+  "<37>": 151714,
+  "<380>": 152057,
+  "<381>": 152058,
+  "<382>": 152059,
+  "<383>": 152060,
+  "<384>": 152061,
+  "<385>": 152062,
+  "<386>": 152063,
+  "<387>": 152064,
+  "<388>": 152065,
+  "<389>": 152066,
+  "<38>": 151715,
+  "<390>": 152067,
+  "<391>": 152068,
+  "<392>": 152069,
+  "<393>": 152070,
+  "<394>": 152071,
+  "<395>": 152072,
+  "<396>": 152073,
+  "<397>": 152074,
+  "<398>": 152075,
+  "<399>": 152076,
+  "<39>": 151716,
+  "<3>": 151680,
+  "<400>": 152077,
+  "<401>": 152078,
+  "<402>": 152079,
+  "<403>": 152080,
+  "<404>": 152081,
+  "<405>": 152082,
+  "<406>": 152083,
+  "<407>": 152084,
+  "<408>": 152085,
+  "<409>": 152086,
+  "<40>": 151717,
+  "<410>": 152087,
+  "<411>": 152088,
+  "<412>": 152089,
+  "<413>": 152090,
+  "<414>": 152091,
+  "<415>": 152092,
+  "<416>": 152093,
+  "<417>": 152094,
+  "<418>": 152095,
+  "<419>": 152096,
+  "<41>": 151718,
+  "<420>": 152097,
+  "<421>": 152098,
+  "<422>": 152099,
+  "<423>": 152100,
+  "<424>": 152101,
+  "<425>": 152102,
+  "<426>": 152103,
+  "<427>": 152104,
+  "<428>": 152105,
+  "<429>": 152106,
+  "<42>": 151719,
+  "<430>": 152107,
+  "<431>": 152108,
+  "<432>": 152109,
+  "<433>": 152110,
+  "<434>": 152111,
+  "<435>": 152112,
+  "<436>": 152113,
+  "<437>": 152114,
+  "<438>": 152115,
+  "<439>": 152116,
+  "<43>": 151720,
+  "<440>": 152117,
+  "<441>": 152118,
+  "<442>": 152119,
+  "<443>": 152120,
+  "<444>": 152121,
+  "<445>": 152122,
+  "<446>": 152123,
+  "<447>": 152124,
+  "<448>": 152125,
+  "<449>": 152126,
+  "<44>": 151721,
+  "<450>": 152127,
+  "<451>": 152128,
+  "<452>": 152129,
+  "<453>": 152130,
+  "<454>": 152131,
+  "<455>": 152132,
+  "<456>": 152133,
+  "<457>": 152134,
+  "<458>": 152135,
+  "<459>": 152136,
+  "<45>": 151722,
+  "<460>": 152137,
+  "<461>": 152138,
+  "<462>": 152139,
+  "<463>": 152140,
+  "<464>": 152141,
+  "<465>": 152142,
+  "<466>": 152143,
+  "<467>": 152144,
+  "<468>": 152145,
+  "<469>": 152146,
+  "<46>": 151723,
+  "<470>": 152147,
+  "<471>": 152148,
+  "<472>": 152149,
+  "<473>": 152150,
+  "<474>": 152151,
+  "<475>": 152152,
+  "<476>": 152153,
+  "<477>": 152154,
+  "<478>": 152155,
+  "<479>": 152156,
+  "<47>": 151724,
+  "<480>": 152157,
+  "<481>": 152158,
+  "<482>": 152159,
+  "<483>": 152160,
+  "<484>": 152161,
+  "<485>": 152162,
+  "<486>": 152163,
+  "<487>": 152164,
+  "<488>": 152165,
+  "<489>": 152166,
+  "<48>": 151725,
+  "<490>": 152167,
+  "<491>": 152168,
+  "<492>": 152169,
+  "<493>": 152170,
+  "<494>": 152171,
+  "<495>": 152172,
+  "<496>": 152173,
+  "<497>": 152174,
+  "<498>": 152175,
+  "<499>": 152176,
+  "<49>": 151726,
+  "<4>": 151681,
+  "<500>": 152177,
+  "<501>": 152178,
+  "<502>": 152179,
+  "<503>": 152180,
+  "<504>": 152181,
+  "<505>": 152182,
+  "<506>": 152183,
+  "<507>": 152184,
+  "<508>": 152185,
+  "<509>": 152186,
+  "<50>": 151727,
+  "<510>": 152187,
+  "<511>": 152188,
+  "<512>": 152189,
+  "<513>": 152190,
+  "<514>": 152191,
+  "<515>": 152192,
+  "<516>": 152193,
+  "<517>": 152194,
+  "<518>": 152195,
+  "<519>": 152196,
+  "<51>": 151728,
+  "<520>": 152197,
+  "<521>": 152198,
+  "<522>": 152199,
+  "<523>": 152200,
+  "<524>": 152201,
+  "<525>": 152202,
+  "<526>": 152203,
+  "<527>": 152204,
+  "<528>": 152205,
+  "<529>": 152206,
+  "<52>": 151729,
+  "<530>": 152207,
+  "<531>": 152208,
+  "<532>": 152209,
+  "<533>": 152210,
+  "<534>": 152211,
+  "<535>": 152212,
+  "<536>": 152213,
+  "<537>": 152214,
+  "<538>": 152215,
+  "<539>": 152216,
+  "<53>": 151730,
+  "<540>": 152217,
+  "<541>": 152218,
+  "<542>": 152219,
+  "<543>": 152220,
+  "<544>": 152221,
+  "<545>": 152222,
+  "<546>": 152223,
+  "<547>": 152224,
+  "<548>": 152225,
+  "<549>": 152226,
+  "<54>": 151731,
+  "<550>": 152227,
+  "<551>": 152228,
+  "<552>": 152229,
+  "<553>": 152230,
+  "<554>": 152231,
+  "<555>": 152232,
+  "<556>": 152233,
+  "<557>": 152234,
+  "<558>": 152235,
+  "<559>": 152236,
+  "<55>": 151732,
+  "<560>": 152237,
+  "<561>": 152238,
+  "<562>": 152239,
+  "<563>": 152240,
+  "<564>": 152241,
+  "<565>": 152242,
+  "<566>": 152243,
+  "<567>": 152244,
+  "<568>": 152245,
+  "<569>": 152246,
+  "<56>": 151733,
+  "<570>": 152247,
+  "<571>": 152248,
+  "<572>": 152249,
+  "<573>": 152250,
+  "<574>": 152251,
+  "<575>": 152252,
+  "<576>": 152253,
+  "<577>": 152254,
+  "<578>": 152255,
+  "<579>": 152256,
+  "<57>": 151734,
+  "<580>": 152257,
+  "<581>": 152258,
+  "<582>": 152259,
+  "<583>": 152260,
+  "<584>": 152261,
+  "<585>": 152262,
+  "<586>": 152263,
+  "<587>": 152264,
+  "<588>": 152265,
+  "<589>": 152266,
+  "<58>": 151735,
+  "<590>": 152267,
+  "<591>": 152268,
+  "<592>": 152269,
+  "<593>": 152270,
+  "<594>": 152271,
+  "<595>": 152272,
+  "<596>": 152273,
+  "<597>": 152274,
+  "<598>": 152275,
+  "<599>": 152276,
+  "<59>": 151736,
+  "<5>": 151682,
+  "<600>": 152277,
+  "<601>": 152278,
+  "<602>": 152279,
+  "<603>": 152280,
+  "<604>": 152281,
+  "<605>": 152282,
+  "<606>": 152283,
+  "<607>": 152284,
+  "<608>": 152285,
+  "<609>": 152286,
+  "<60>": 151737,
+  "<610>": 152287,
+  "<611>": 152288,
+  "<612>": 152289,
+  "<613>": 152290,
+  "<614>": 152291,
+  "<615>": 152292,
+  "<616>": 152293,
+  "<617>": 152294,
+  "<618>": 152295,
+  "<619>": 152296,
+  "<61>": 151738,
+  "<620>": 152297,
+  "<621>": 152298,
+  "<622>": 152299,
+  "<623>": 152300,
+  "<624>": 152301,
+  "<625>": 152302,
+  "<626>": 152303,
+  "<627>": 152304,
+  "<628>": 152305,
+  "<629>": 152306,
+  "<62>": 151739,
+  "<630>": 152307,
+  "<631>": 152308,
+  "<632>": 152309,
+  "<633>": 152310,
+  "<634>": 152311,
+  "<635>": 152312,
+  "<636>": 152313,
+  "<637>": 152314,
+  "<638>": 152315,
+  "<639>": 152316,
+  "<63>": 151740,
+  "<640>": 152317,
+  "<641>": 152318,
+  "<642>": 152319,
+  "<643>": 152320,
+  "<644>": 152321,
+  "<645>": 152322,
+  "<646>": 152323,
+  "<647>": 152324,
+  "<648>": 152325,
+  "<649>": 152326,
+  "<64>": 151741,
+  "<650>": 152327,
+  "<651>": 152328,
+  "<652>": 152329,
+  "<653>": 152330,
+  "<654>": 152331,
+  "<655>": 152332,
+  "<656>": 152333,
+  "<657>": 152334,
+  "<658>": 152335,
+  "<659>": 152336,
+  "<65>": 151742,
+  "<660>": 152337,
+  "<661>": 152338,
+  "<662>": 152339,
+  "<663>": 152340,
+  "<664>": 152341,
+  "<665>": 152342,
+  "<666>": 152343,
+  "<667>": 152344,
+  "<668>": 152345,
+  "<669>": 152346,
+  "<66>": 151743,
+  "<670>": 152347,
+  "<671>": 152348,
+  "<672>": 152349,
+  "<673>": 152350,
+  "<674>": 152351,
+  "<675>": 152352,
+  "<676>": 152353,
+  "<677>": 152354,
+  "<678>": 152355,
+  "<679>": 152356,
+  "<67>": 151744,
+  "<680>": 152357,
+  "<681>": 152358,
+  "<682>": 152359,
+  "<683>": 152360,
+  "<684>": 152361,
+  "<685>": 152362,
+  "<686>": 152363,
+  "<687>": 152364,
+  "<688>": 152365,
+  "<689>": 152366,
+  "<68>": 151745,
+  "<690>": 152367,
+  "<691>": 152368,
+  "<692>": 152369,
+  "<693>": 152370,
+  "<694>": 152371,
+  "<695>": 152372,
+  "<696>": 152373,
+  "<697>": 152374,
+  "<698>": 152375,
+  "<699>": 152376,
+  "<69>": 151746,
+  "<6>": 151683,
+  "<700>": 152377,
+  "<701>": 152378,
+  "<702>": 152379,
+  "<703>": 152380,
+  "<704>": 152381,
+  "<705>": 152382,
+  "<706>": 152383,
+  "<707>": 152384,
+  "<708>": 152385,
+  "<709>": 152386,
+  "<70>": 151747,
+  "<710>": 152387,
+  "<711>": 152388,
+  "<712>": 152389,
+  "<713>": 152390,
+  "<714>": 152391,
+  "<715>": 152392,
+  "<716>": 152393,
+  "<717>": 152394,
+  "<718>": 152395,
+  "<719>": 152396,
+  "<71>": 151748,
+  "<720>": 152397,
+  "<721>": 152398,
+  "<722>": 152399,
+  "<723>": 152400,
+  "<724>": 152401,
+  "<725>": 152402,
+  "<726>": 152403,
+  "<727>": 152404,
+  "<728>": 152405,
+  "<729>": 152406,
+  "<72>": 151749,
+  "<730>": 152407,
+  "<731>": 152408,
+  "<732>": 152409,
+  "<733>": 152410,
+  "<734>": 152411,
+  "<735>": 152412,
+  "<736>": 152413,
+  "<737>": 152414,
+  "<738>": 152415,
+  "<739>": 152416,
+  "<73>": 151750,
+  "<740>": 152417,
+  "<741>": 152418,
+  "<742>": 152419,
+  "<743>": 152420,
+  "<744>": 152421,
+  "<745>": 152422,
+  "<746>": 152423,
+  "<747>": 152424,
+  "<748>": 152425,
+  "<749>": 152426,
+  "<74>": 151751,
+  "<750>": 152427,
+  "<751>": 152428,
+  "<752>": 152429,
+  "<753>": 152430,
+  "<754>": 152431,
+  "<755>": 152432,
+  "<756>": 152433,
+  "<757>": 152434,
+  "<758>": 152435,
+  "<759>": 152436,
+  "<75>": 151752,
+  "<760>": 152437,
+  "<761>": 152438,
+  "<762>": 152439,
+  "<763>": 152440,
+  "<764>": 152441,
+  "<765>": 152442,
+  "<766>": 152443,
+  "<767>": 152444,
+  "<768>": 152445,
+  "<769>": 152446,
+  "<76>": 151753,
+  "<770>": 152447,
+  "<771>": 152448,
+  "<772>": 152449,
+  "<773>": 152450,
+  "<774>": 152451,
+  "<775>": 152452,
+  "<776>": 152453,
+  "<777>": 152454,
+  "<778>": 152455,
+  "<779>": 152456,
+  "<77>": 151754,
+  "<780>": 152457,
+  "<781>": 152458,
+  "<782>": 152459,
+  "<783>": 152460,
+  "<784>": 152461,
+  "<785>": 152462,
+  "<786>": 152463,
+  "<787>": 152464,
+  "<788>": 152465,
+  "<789>": 152466,
+  "<78>": 151755,
+  "<790>": 152467,
+  "<791>": 152468,
+  "<792>": 152469,
+  "<793>": 152470,
+  "<794>": 152471,
+  "<795>": 152472,
+  "<796>": 152473,
+  "<797>": 152474,
+  "<798>": 152475,
+  "<799>": 152476,
+  "<79>": 151756,
+  "<7>": 151684,
+  "<800>": 152477,
+  "<801>": 152478,
+  "<802>": 152479,
+  "<803>": 152480,
+  "<804>": 152481,
+  "<805>": 152482,
+  "<806>": 152483,
+  "<807>": 152484,
+  "<808>": 152485,
+  "<809>": 152486,
+  "<80>": 151757,
+  "<810>": 152487,
+  "<811>": 152488,
+  "<812>": 152489,
+  "<813>": 152490,
+  "<814>": 152491,
+  "<815>": 152492,
+  "<816>": 152493,
+  "<817>": 152494,
+  "<818>": 152495,
+  "<819>": 152496,
+  "<81>": 151758,
+  "<820>": 152497,
+  "<821>": 152498,
+  "<822>": 152499,
+  "<823>": 152500,
+  "<824>": 152501,
+  "<825>": 152502,
+  "<826>": 152503,
+  "<827>": 152504,
+  "<828>": 152505,
+  "<829>": 152506,
+  "<82>": 151759,
+  "<830>": 152507,
+  "<831>": 152508,
+  "<832>": 152509,
+  "<833>": 152510,
+  "<834>": 152511,
+  "<835>": 152512,
+  "<836>": 152513,
+  "<837>": 152514,
+  "<838>": 152515,
+  "<839>": 152516,
+  "<83>": 151760,
+  "<840>": 152517,
+  "<841>": 152518,
+  "<842>": 152519,
+  "<843>": 152520,
+  "<844>": 152521,
+  "<845>": 152522,
+  "<846>": 152523,
+  "<847>": 152524,
+  "<848>": 152525,
+  "<849>": 152526,
+  "<84>": 151761,
+  "<850>": 152527,
+  "<851>": 152528,
+  "<852>": 152529,
+  "<853>": 152530,
+  "<854>": 152531,
+  "<855>": 152532,
+  "<856>": 152533,
+  "<857>": 152534,
+  "<858>": 152535,
+  "<859>": 152536,
+  "<85>": 151762,
+  "<860>": 152537,
+  "<861>": 152538,
+  "<862>": 152539,
+  "<863>": 152540,
+  "<864>": 152541,
+  "<865>": 152542,
+  "<866>": 152543,
+  "<867>": 152544,
+  "<868>": 152545,
+  "<869>": 152546,
+  "<86>": 151763,
+  "<870>": 152547,
+  "<871>": 152548,
+  "<872>": 152549,
+  "<873>": 152550,
+  "<874>": 152551,
+  "<875>": 152552,
+  "<876>": 152553,
+  "<877>": 152554,
+  "<878>": 152555,
+  "<879>": 152556,
+  "<87>": 151764,
+  "<880>": 152557,
+  "<881>": 152558,
+  "<882>": 152559,
+  "<883>": 152560,
+  "<884>": 152561,
+  "<885>": 152562,
+  "<886>": 152563,
+  "<887>": 152564,
+  "<888>": 152565,
+  "<889>": 152566,
+  "<88>": 151765,
+  "<890>": 152567,
+  "<891>": 152568,
+  "<892>": 152569,
+  "<893>": 152570,
+  "<894>": 152571,
+  "<895>": 152572,
+  "<896>": 152573,
+  "<897>": 152574,
+  "<898>": 152575,
+  "<899>": 152576,
+  "<89>": 151766,
+  "<8>": 151685,
+  "<900>": 152577,
+  "<901>": 152578,
+  "<902>": 152579,
+  "<903>": 152580,
+  "<904>": 152581,
+  "<905>": 152582,
+  "<906>": 152583,
+  "<907>": 152584,
+  "<908>": 152585,
+  "<909>": 152586,
+  "<90>": 151767,
+  "<910>": 152587,
+  "<911>": 152588,
+  "<912>": 152589,
+  "<913>": 152590,
+  "<914>": 152591,
+  "<915>": 152592,
+  "<916>": 152593,
+  "<917>": 152594,
+  "<918>": 152595,
+  "<919>": 152596,
+  "<91>": 151768,
+  "<920>": 152597,
+  "<921>": 152598,
+  "<922>": 152599,
+  "<923>": 152600,
+  "<924>": 152601,
+  "<925>": 152602,
+  "<926>": 152603,
+  "<927>": 152604,
+  "<928>": 152605,
+  "<929>": 152606,
+  "<92>": 151769,
+  "<930>": 152607,
+  "<931>": 152608,
+  "<932>": 152609,
+  "<933>": 152610,
+  "<934>": 152611,
+  "<935>": 152612,
+  "<936>": 152613,
+  "<937>": 152614,
+  "<938>": 152615,
+  "<939>": 152616,
+  "<93>": 151770,
+  "<940>": 152617,
+  "<941>": 152618,
+  "<942>": 152619,
+  "<943>": 152620,
+  "<944>": 152621,
+  "<945>": 152622,
+  "<946>": 152623,
+  "<947>": 152624,
+  "<948>": 152625,
+  "<949>": 152626,
+  "<94>": 151771,
+  "<950>": 152627,
+  "<951>": 152628,
+  "<952>": 152629,
+  "<953>": 152630,
+  "<954>": 152631,
+  "<955>": 152632,
+  "<956>": 152633,
+  "<957>": 152634,
+  "<958>": 152635,
+  "<959>": 152636,
+  "<95>": 151772,
+  "<960>": 152637,
+  "<961>": 152638,
+  "<962>": 152639,
+  "<963>": 152640,
+  "<964>": 152641,
+  "<965>": 152642,
+  "<966>": 152643,
+  "<967>": 152644,
+  "<968>": 152645,
+  "<969>": 152646,
+  "<96>": 151773,
+  "<970>": 152647,
+  "<971>": 152648,
+  "<972>": 152649,
+  "<973>": 152650,
+  "<974>": 152651,
+  "<975>": 152652,
+  "<976>": 152653,
+  "<977>": 152654,
+  "<978>": 152655,
+  "<979>": 152656,
+  "<97>": 151774,
+  "<980>": 152657,
+  "<981>": 152658,
+  "<982>": 152659,
+  "<983>": 152660,
+  "<984>": 152661,
+  "<985>": 152662,
+  "<986>": 152663,
+  "<987>": 152664,
+  "<988>": 152665,
+  "<989>": 152666,
+  "<98>": 151775,
+  "<990>": 152667,
+  "<991>": 152668,
+  "<992>": 152669,
+  "<993>": 152670,
+  "<994>": 152671,
+  "<995>": 152672,
+  "<996>": 152673,
+  "<997>": 152674,
+  "<998>": 152675,
+  "<999>": 152676,
+  "<99>": 151776,
+  "<9>": 151686,
+  "<IMG_CONTEXT>": 151665,
+  "<box>": 151668,
+  "<img>": 151666,
+  "<interval>": 151674,
+  "<null>": 152678,
+  "<quad>": 151670,
+  "<ref>": 151672,
+  "<switch>": 152679,
+  "<text_mask>": 151676,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

all_results.json

@@ -0,0 +1,9 @@
+{
+    "epoch": 0.069,
+    "total_flos": 5.6478980657499236e+20,
+    "train_loss": 0.02277738807797432,
+    "train_runtime": 3535.7374,
+    "train_samples": "streaming",
+    "train_samples_per_second": 362.018,
+    "train_steps_per_second": 1.414
+}

chat_template.jinja

@@ -0,0 +1,54 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\n<tool_call>\n{"name": "' }}
+            {{- tool_call.name }}
+            {{- '", "arguments": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\n</tool_call>' }}
+        {%- endfor %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}

chat_template.json

@@ -0,0 +1,4 @@
+{
+    "chat_template": "{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}<image {{ image_count.value }}>{% endif %}<image-{{ image_count.value }}>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}<video {{ video_count.value }}>{% endif %}<video-{{ video_count.value }}>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
+}
+  

config.json

@@ -0,0 +1,94 @@
+{
+    "_attn_implementation": "magi",
+    "_commit_hash": null,
+    "architectures": [
+        "LocateAnythingForConditionalGeneration"
+    ],
+    "auto_map": {
+        "AutoConfig": "configuration_locateanything.LocateAnythingConfig",
+        "AutoModel": "modeling_locateanything.LocateAnythingForConditionalGeneration"
+    },
+    "box_end_token_id": 151669,
+    "box_start_token_id": 151668,
+    "coord_end_token_id": 152677,
+    "coord_start_token_id": 151677,
+    "image_token_index": 151665,
+    "mlp_checkpoint": false,
+    "mlp_connector_layers": 2,
+    "model_type": "locateanything",
+    "none_token_id": 4064,
+    "ref_end_token_id": 151673,
+    "ref_start_token_id": 151672,
+    "template": null,
+    "text_config": {
+        "_attn_implementation_autoset": true,
+        "_name_or_path": "Qwen/Qwen2.5-3B-Instruct",
+        "architectures": [
+            "Qwen2ForCausalLM"
+        ],
+        "attention_dropout": 0.0,
+        "block_size": 6,
+        "bos_token_id": 151643,
+        "causal_attn": false,
+        "eos_token_id": 151645,
+        "hidden_act": "silu",
+        "hidden_size": 2048,
+        "initializer_range": 0.02,
+        "intermediate_size": 11008,
+        "max_position_embeddings": 32768,
+        "max_window_layers": 70,
+        "model_type": "qwen2",
+        "null_token_id": 152678,
+        "num_attention_heads": 16,
+        "num_hidden_layers": 36,
+        "num_key_value_heads": 2,
+        "rms_norm_eps": 1e-06,
+        "rope_scaling": null,
+        "rope_theta": 1000000.0,
+        "sliding_window": 32768,
+        "switch_token_id": 152679,
+        "text_mask_token_id": 151676,
+        "tie_word_embeddings": true,
+        "torch_dtype": "bfloat16",
+        "use_cache": false,
+        "use_sliding_window": false,
+        "vocab_size": 152681
+    },
+    "torch_dtype": "bfloat16",
+    "transformers_version": null,
+    "use_backbone_lora": 0,
+    "use_llm_lora": 0,
+    "vision_config": {
+        "_attn_implementation_autoset": true,
+        "_name_or_path": "moonshotai/MoonViT-SO-400M",
+        "auto_map": {
+            "AutoConfig": "moonshotai/MoonViT-SO-400M--configuration_moonvit.MoonViTConfig",
+            "AutoModel": "moonshotai/MoonViT-SO-400M--modeling_moonvit.MoonVitPretrainedModel"
+        },
+        "hidden_size": 1152,
+        "init_pos_emb_height": 64,
+        "init_pos_emb_width": 64,
+        "intermediate_size": 4304,
+        "merge_kernel_size": [
+            2,
+            2
+        ],
+        "model_type": "moonvit",
+        "num_attention_heads": 16,
+        "num_hidden_layers": 27,
+        "patch_size": 14,
+        "torch_dtype": "bfloat16"
+    },
+    "quantization_config": {
+        "quant_method": "exl3",
+        "version": "0.0.39",
+        "bits": 6.0,
+        "head_bits": 8,
+        "calibration": {
+            "rows": 250,
+            "cols": 2048
+        },
+        "out_scales": "always",
+        "codebook": "mcg"
+    }
+}

configuration_locateanything.py

@@ -0,0 +1,131 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import copy
+
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+logger = logging.get_logger(__name__)
+
+class MoonViTConfig(PretrainedConfig):
+    model_type = "moonvit"
+
+    def __init__(
+        self,
+        patch_size: int = 14,
+        init_pos_emb_height: int = 64,
+        init_pos_emb_width: int = 64,
+        num_attention_heads: int = 16,
+        num_hidden_layers: int = 27,
+        hidden_size: int = 1152,
+        intermediate_size: int = 4304,
+        merge_kernel_size: tuple[int, int] = (2, 2),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.patch_size = patch_size
+        # Positional embedding config
+        self.init_pos_emb_height = init_pos_emb_height
+        self.init_pos_emb_width = init_pos_emb_width
+        # Transformer config
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        # Patch merger config
+        self.merge_kernel_size = merge_kernel_size
+
+
+class LocateAnythingConfig(PretrainedConfig):
+    model_type = 'locateanything'
+    is_composition = True
+    sub_configs = {"vision_config": MoonViTConfig, "text_config": Qwen2Config}
+    def __init__(
+            self,
+            vision_config=None,
+            text_config=None,
+            use_backbone_lora=0,
+            use_llm_lora=0,
+            downsample_ratio=0.5,
+            template=None,
+            loss_version='v1',
+            mlp_checkpoint=False,
+            image_token_index=151667,
+            box_start_token_id=151668,
+            box_end_token_id=151669,
+            coord_start_token_id=151677,
+            coord_end_token_id=152677,
+            ref_start_token_id=151672,
+            ref_end_token_id=151673,
+            none_token_id=4064,
+            **kwargs):
+        super().__init__(**kwargs)
+
+        if vision_config is None:
+            vision_config = {'model_type': 'moonvit'}
+            logger.info('vision_config is None. Initializing the MoonViTConfig with default values.')
+
+        if text_config is None:
+            text_config = {'architectures': ['Qwen2ForCausalLM']}
+            logger.info('text_config is None. Initializing the Qwen2Config config with default values.')
+
+        if vision_config['model_type'] == 'moonvit':
+            self.vision_config = MoonViTConfig(**vision_config)
+        else:
+            raise ValueError('Unsupported model_type: {}. Only moonvit is supported.'.format(vision_config['model_type']))
+
+
+        if text_config['architectures'][0] == 'Qwen2ForCausalLM':
+            self.text_config = Qwen2Config(**text_config)
+        elif text_config['architectures'][0] == 'Qwen3ForCausalLM':
+            self.text_config = Qwen3Config(**text_config)
+        else:
+            raise ValueError('Unsupported architecture: {}. Only Qwen2ForCausalLM and Qwen3ForCausalLM are supported.'.format(text_config['architectures'][0]))
+        self.use_backbone_lora = use_backbone_lora
+        self.use_llm_lora = use_llm_lora
+        self.mlp_checkpoint = mlp_checkpoint
+        self.downsample_ratio = downsample_ratio
+        self.template = template
+        self.loss_version = loss_version
+        self.tie_word_embeddings = self.text_config.tie_word_embeddings
+        self.image_token_index = image_token_index
+        self.box_start_token_id = box_start_token_id
+        self.box_end_token_id = box_end_token_id
+        self.coord_start_token_id = coord_start_token_id
+        self.coord_end_token_id = coord_end_token_id
+        self.ref_start_token_id = ref_start_token_id
+        self.ref_end_token_id = ref_end_token_id
+        self.none_token_id = none_token_id
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
+
+        Returns:
+            `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+        output['vision_config'] = self.vision_config.to_dict()
+        output['text_config'] = self.text_config.to_dict()
+        output['model_type'] = self.__class__.model_type
+        output['use_backbone_lora'] = self.use_backbone_lora
+        output['use_llm_lora'] = self.use_llm_lora
+        output['downsample_ratio'] = self.downsample_ratio
+        output['template'] = self.template
+        output['image_token_index'] = self.image_token_index
+        output['box_start_token_id'] = self.box_start_token_id
+        output['box_end_token_id'] = self.box_end_token_id
+        output['coord_start_token_id'] = self.coord_start_token_id
+        output['coord_end_token_id'] = self.coord_end_token_id
+        output['ref_start_token_id'] = self.ref_start_token_id
+        output['ref_end_token_id'] = self.ref_end_token_id
+        output['none_token_id'] = self.none_token_id
+        output['_attn_implementation'] = self._attn_implementation
+        if hasattr(self, '_attn_implementation_autoset'):
+            output['_attn_implementation_autoset'] = self._attn_implementation_autoset
+        return output

configuration_qwen2.py

@@ -0,0 +1,148 @@
+# coding=utf-8
+# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+""" Qwen2 model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+QWEN2_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "Qwen/Qwen2-7B-beta": "https://huggingface.co/Qwen/Qwen2-7B-beta/resolve/main/config.json",
+}
+
+
+class Qwen2Config(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Qwen2Model`]. It is used to instantiate a
+    Qwen2 model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of
+    Qwen2-7B-beta [Qwen/Qwen2-7B-beta](https://huggingface.co/Qwen/Qwen2-7B-beta).
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 151936):
+            Vocabulary size of the Qwen2 model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`Qwen2Model`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 22016):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 32):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 32768):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+        max_window_layers (`int`, *optional*, defaults to 28):
+            The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+    ```python
+    >>> from transformers import Qwen2Model, Qwen2Config
+
+    >>> # Initializing a Qwen2 style configuration
+    >>> configuration = Qwen2Config()
+
+    >>> # Initializing a model from the Qwen2-7B style configuration
+    >>> model = Qwen2Model(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "qwen2"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=4096,
+        intermediate_size=22016,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=32,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        attention_dropout=0.0,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window
+        self.max_window_layers = max_window_layers
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout
+        if kwargs.get('attn_implementation', None) is None:
+            self.attn_implementation = kwargs['attn_implementation'] = 'flash_attention_2'
+        else:
+            self.attn_implementation = kwargs['attn_implementation']
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

generate_utils.py

@@ -0,0 +1,504 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import torch
+import torch.nn.functional as F
+import torch.distributions as dists
+from typing import Dict, Optional
+
+
+def get_token_ids_from_config(config) -> Dict[str, int]:
+    """Extract all token IDs from the configuration object.
+    
+    Args:
+        config: Configuration object (LocateAnythingConfig or similar)
+        
+    Returns:
+        Dictionary containing all token IDs
+    """
+    token_ids = {}
+    
+    # Get from main config
+    token_ids['box_start_token_id'] = getattr(config, 'box_start_token_id', 151668)
+    token_ids['box_end_token_id'] = getattr(config, 'box_end_token_id', 151669)
+    token_ids['coord_start_token_id'] = getattr(config, 'coord_start_token_id', 151677)
+    token_ids['coord_end_token_id'] = getattr(config, 'coord_end_token_id', 152677)
+    token_ids['ref_start_token_id'] = getattr(config, 'ref_start_token_id', 151672)
+    token_ids['ref_end_token_id'] = getattr(config, 'ref_end_token_id', 151673)
+    token_ids['none_token_id'] = getattr(config, 'none_token_id', 4064)
+    
+    # Get from text_config
+    text_config = getattr(config, 'text_config', None)
+    if text_config is not None:
+        token_ids['null_token_id'] = getattr(text_config, 'null_token_id', 152678)
+        token_ids['im_end_token_id'] = getattr(text_config, 'eos_token_id', 151645)
+        token_ids['switch_token_id'] = getattr(text_config, 'switch_token_id', 152679)
+        token_ids['default_mask_token_id'] = getattr(text_config, 'text_mask_token_id', 151676)
+    else:
+        token_ids['null_token_id'] = 152678
+        token_ids['im_end_token_id'] = 151645
+        token_ids['switch_token_id'] = 152679
+        token_ids['default_mask_token_id'] = 151676
+    
+    return token_ids
+
+
+def top_p_logits(
+    logits: torch.Tensor,
+    top_p: float = None
+) -> torch.Tensor:
+    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+    sorted_indices_to_remove = cumulative_probs > top_p
+    # Shift the indices to the right to keep the first token above the threshold
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = 0
+
+    mask = torch.zeros_like(logits, dtype=torch.bool, device=logits.device)
+    mask = mask.scatter_(-1, sorted_indices, sorted_indices_to_remove)
+    logits = logits.masked_fill(mask, torch.finfo(logits.dtype).min)
+    return logits
+
+
+def top_k_logits(
+    logits: torch.Tensor,
+    top_k: int = None
+) -> torch.Tensor:
+    top_k = min(top_k, logits.size(-1))  # Safety check
+    # Remove all tokens with a probability less than the last token of the top-k
+    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+    logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min)
+    return logits
+
+
+def apply_repetition_penalty(
+    logits: torch.Tensor,
+    input_ids: torch.Tensor,
+    repetition_penalty: float = 1.0
+) -> torch.Tensor:
+    """
+    Apply repetition penalty to logits.
+    
+    Args:
+        logits: Shape [batch_size, seq_len, vocab_size] or [batch_size, vocab_size]
+        input_ids: Previously generated token ids, shape [batch_size, seq_len]
+        repetition_penalty: Penalty factor. > 1.0 penalizes repetition, < 1.0 encourages it.
+    
+    Returns:
+        Modified logits with repetition penalty applied.
+    """
+    if repetition_penalty == 1.0:
+        return logits
+    
+    # Convert to 3D for vectorized computation
+    if logits.dim() == 2:
+        logits = logits.unsqueeze(1)  # [B, 1, V]
+        squeeze_back = True
+    else:
+        squeeze_back = False
+    
+    batch_size, seq_len, vocab_size = logits.shape
+
+    # Construct [B, V] bool mask marking tokens that have appeared in each batch
+    device = logits.device
+    token_mask = torch.zeros(batch_size, vocab_size, dtype=torch.bool, device=device)
+    for b in range(batch_size):
+        # Apply penalty only based on tokens already generated in this batch
+        unique_tokens = input_ids[b].unique()
+        # Prevent out-of-bounds: only keep IDs within vocab range
+        valid_tokens = unique_tokens[(unique_tokens >= 0) & (unique_tokens < vocab_size)]
+        if valid_tokens.numel() > 0:
+            token_mask[b, valid_tokens] = True
+
+    # Expand to [B, L, V] to align with logits
+    token_mask = token_mask.unsqueeze(1).expand(-1, seq_len, -1)
+
+    # Divide positive values by penalty, multiply negative values by penalty
+    positive = logits > 0
+    negative = ~positive
+
+    # Apply penalty only at mask positions
+    logits = torch.where(token_mask & positive, logits / repetition_penalty, logits)
+    logits = torch.where(token_mask & negative, logits * repetition_penalty, logits)
+    
+    if squeeze_back:
+        logits = logits.squeeze(1)
+    
+    return logits
+
+
+def sample_tokens(
+    logits: torch.Tensor,
+    generated: torch.Tensor,
+    token_ids: Dict[str, int],
+    **generate_kwargs,
+):  
+    batch_size, seq_len, vocab_size = logits.shape
+
+    repetition_penalty = generate_kwargs.get('repetition_penalty', 1.0)
+    temperature = generate_kwargs.get('temperature', 0)
+    top_p = generate_kwargs.get('top_p', None)
+    top_k = generate_kwargs.get('top_k', None)
+
+    # Apply repetition penalty based on all previously generated tokens
+    if repetition_penalty != 1.0:
+        logits = apply_repetition_penalty(logits, generated, repetition_penalty)
+
+    if temperature > 0:
+            logits = logits / temperature
+    if top_p is not None and top_p < 1:
+        logits = top_p_logits(logits, top_p)
+    if top_k is not None:
+        logits = top_k_logits(logits, top_k)
+    
+    probs = torch.softmax(logits, dim=-1)
+
+    if temperature > 0:
+        try:
+            x0 = dists.Categorical(probs=probs).sample()
+            confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)
+        except Exception:
+            confidence, x0 = probs.max(dim=-1)
+    else:
+        confidence, x0 = probs.max(dim=-1)
+
+    if seq_len == 1:
+        return probs, confidence, x0, None
+
+    box_avg = []
+    fallback_box = torch.zeros(1, dtype=x0.dtype, device=x0.device)
+
+    for b in range(batch_size):
+        decoded_box = decode_bbox_avg(
+            logits[b], probs[b], token_ids, keep_k=generate_kwargs.get('keep_k_avg', 4),
+            generation_mode=generate_kwargs.get('generation_mode', 'hybrid'),
+        )
+        if decoded_box is not None:
+            box_avg.append(decoded_box)
+        else:
+            out_ref = decode_ref(logits[b], probs[b], token_ids)
+            if out_ref is not None:
+                box_avg.append(torch.tensor(out_ref, dtype=x0.dtype, device=x0.device))
+            else:
+                box_avg.append(fallback_box)
+
+    box_avg = torch.stack(box_avg)
+
+    return probs, confidence, x0, box_avg
+
+
+def sample_tokens_ar(
+    logits: torch.Tensor,
+    generated: torch.Tensor,
+    token_ids: Dict[str, int],
+    **generate_kwargs,
+):
+    """
+    Lightweight sampling function for AR single-step sampling only.
+    
+    Args:
+        logits: [batch_size, vocab_size] or [batch_size, 1, vocab_size]
+        generated: [batch_size, seq_len]
+    """
+    # Convert to 3D for reusing repetition penalty and clipping logic
+    if logits.dim() == 2:
+        logits = logits.unsqueeze(1)  # [B, 1, V]
+    batch_size, seq_len, vocab_size = logits.shape
+    assert seq_len == 1, "sample_tokens_ar only supports single-step AR sampling (seq_len == 1)"
+
+    repetition_penalty = generate_kwargs.get('repetition_penalty', 1.0)
+    temperature = generate_kwargs.get('temperature', 0)
+    top_p = generate_kwargs.get('top_p', None)
+    top_k = generate_kwargs.get('top_k', None)
+
+    # Apply repetition penalty only based on historically generated tokens
+    if repetition_penalty != 1.0:
+        logits = apply_repetition_penalty(logits, generated, repetition_penalty)
+
+    if temperature > 0:
+        logits = logits / temperature
+    if top_p is not None and top_p < 1:
+        logits = top_p_logits(logits, top_p)
+    if top_k is not None:
+        logits = top_k_logits(logits, top_k)
+
+    probs = torch.softmax(logits, dim=-1)
+
+    if temperature > 0:
+        try:
+            x0 = dists.Categorical(probs=probs).sample()
+            confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)
+        except Exception:
+            confidence, x0 = probs.max(dim=-1)
+    else:
+        # For greedy: directly take the token with maximum probability
+        confidence, x0 = probs.max(dim=-1)
+
+    # Keep interface consistent with sample_tokens: return [B, 1, V] / [B, 1] shape
+    return probs, confidence, x0, None, None
+
+
+def is_valid_box_frame(
+    probs,
+    token_ids: Dict[str, int],
+    start_thresh=0.6,
+    end_thresh=0.2,
+    topk=5,
+):
+    box_start_token_id = token_ids['box_start_token_id']
+    box_end_token_id = token_ids['box_end_token_id']
+    null_token_id = token_ids['null_token_id']
+    im_end_token_id = token_ids['im_end_token_id']
+    none_token_id = token_ids['none_token_id'] # none
+
+    p_start = probs[0, box_start_token_id]
+    if p_start >= start_thresh:
+        if (probs[1, none_token_id] > 0.2 and 
+            probs[2, box_end_token_id] > 0.2 and 
+            probs[3, null_token_id] > 0.1 and 
+            probs[4, null_token_id] > 0.1):
+            return 'empty_box'
+
+    end_target_ids = torch.tensor([box_end_token_id, null_token_id, im_end_token_id], device=probs.device)
+    end_score = probs[5, end_target_ids].sum()
+
+    if end_score >= end_thresh:
+            return 'legal_box'
+
+    return 'illegal_box'
+
+
+def decode_bbox_avg(
+    logits,
+    probs,
+    token_ids: Dict[str, int],
+    keep_k=5,
+    start_thresh=0.7,
+    end_thresh=0.2,
+    generation_mode: str = 'hybrid',
+):
+    """
+    Decode bounding box coordinates using top-k weighted average.
+    
+    Args:
+        logits: Logits of shape (6, vocab_size)
+        probs: Probability distribution of shape (6, vocab_size)
+        token_ids: Dictionary containing all token IDs
+        keep_k: Number of top-k candidate tokens to keep at each position
+        start_thresh: Confidence threshold for box start token
+        end_thresh: Confidence threshold for box end token
+        
+    Returns:
+        Decoded bounding box coordinate list in format [box_start, x1, x2, y1, y2, box_end],
+        or None if decoding fails
+    """
+    coord_start_token_id = token_ids['coord_start_token_id']
+    coord_end_token_id = token_ids['coord_end_token_id']
+    box_start_token_id = token_ids['box_start_token_id']
+    box_end_token_id = token_ids['box_end_token_id']
+    none_token_id = token_ids['none_token_id']
+
+    device = logits.device
+
+    box_type = is_valid_box_frame(
+        probs,
+        token_ids,
+        start_thresh=start_thresh,
+        end_thresh=end_thresh,
+        topk=keep_k
+    )
+    if box_type == 'empty_box':
+        # Handle the <box>none</box> case first
+        return torch.tensor([
+            box_start_token_id, 
+            none_token_id, 
+            box_end_token_id, 
+            token_ids['null_token_id'], 
+            token_ids['null_token_id'], 
+            token_ids['null_token_id']
+        ], dtype=torch.long, device=probs.device)
+    elif box_type == 'illegal_box':
+        return None
+
+    # Extract probabilities at positions 1-4 and compute Top-K for all 4 positions at once
+    pos_probs, pos_ids = torch.topk(probs[1:5], k=keep_k, dim=-1)
+    mask = (pos_ids >= coord_start_token_id) & (pos_ids <= coord_end_token_id)
+    has_valid = mask.any(dim=-1) # shape: [4]
+    if not has_valid.all():
+        return None # not a box, exit...
+
+    first_valid_idx = mask.long().argmax(dim=-1, keepdim=True) # [4, 1]
+    # Extract highest-probability valid_probs[0] and corresponding valid_ids[0]
+    first_valid_probs = pos_probs.gather(-1, first_valid_idx).squeeze(-1) # [4]
+    first_valid_ids = pos_ids.gather(-1, first_valid_idx).squeeze(-1) # [4]
+    if generation_mode == 'hybrid':
+        valid_counts = mask.sum(dim=-1) # [4]
+        # Compute max/min of valid ids: fill invalid positions with extreme values to avoid interfering with max/min
+        LARGE_NUM, SMALL_NUM = 999999, -999999
+        valid_ids_for_max = torch.where(mask, pos_ids, torch.tensor(SMALL_NUM, device=device))
+        valid_ids_for_min = torch.where(mask, pos_ids, torch.tensor(LARGE_NUM, device=device))
+
+        valid_max = valid_ids_for_max.max(dim=-1)[0]
+        valid_min = valid_ids_for_min.min(dim=-1)[0]
+
+        is_abnormal = (first_valid_probs < 0.9) & (valid_counts > 1) & ((valid_max - valid_min) > 60)
+        # is_abnormal = (first_valid_probs < 0.7) & (valid_counts > 1) & ((valid_max - valid_min) > 80)
+
+        # Normal positions take top-1 (first_valid_ids); abnormal positions are replaced with 0
+        final_coords = torch.where(is_abnormal, torch.tensor(0, device=pos_ids.device), first_valid_ids)
+    elif generation_mode == 'fast':
+        final_coords = first_valid_ids
+
+
+    start_t = torch.tensor([box_start_token_id], dtype=final_coords.dtype, device=device)
+    end_t = torch.tensor([box_end_token_id], dtype=final_coords.dtype, device=device)
+
+    return torch.cat([start_t, final_coords, end_t])
+    
+
+def decode_ref(
+    logits,
+    probs,
+    token_ids: Dict[str, int],
+    keep_k=5,
+    start_thresh=0.6,
+):
+    ref_start_token_id = token_ids.get('ref_start_token_id')
+    coord_start_token_id = token_ids['coord_start_token_id']
+    coord_end_token_id = token_ids['coord_end_token_id']
+    device = probs.device
+    L = probs.size(0)
+
+    # 1. Check if the first position is <ref> and its probability meets start_thresh
+    # Note: we directly use the probability of the ref token at position 0 for the check
+    if probs[0, ref_start_token_id] < start_thresh:
+        return None
+    
+    # 2. Extract Top-K probabilities and token IDs for all subsequent positions
+    pos_probs, pos_ids = torch.topk(probs[1:], k=keep_k, dim=-1) # shape: [L-1, keep_k]
+    
+    # 3. Build mask: identify coordinate tokens (<0> ~ <1000>)
+    is_coord = (pos_ids >= coord_start_token_id) & (pos_ids <= coord_end_token_id)
+    # Invert: valid tokens are non-coordinate tokens
+    is_valid = ~is_coord # shape: [L-1, keep_k]
+
+    # Ensure each position has at least one non-coordinate valid token in its Top-K
+    has_valid = is_valid.any(dim=-1) # shape: [L-1]
+    if not has_valid.all():
+        return None
+
+    # 4. Get the highest-probability valid token
+    # Since topk results are sorted in descending order of probability,
+    # argmax returns the first index where is_valid is True, i.e., the index of the most probable valid token
+    first_valid_idx = is_valid.long().argmax(dim=-1, keepdim=True) # shape: [L-1, 1]
+    
+    # Extract the final token IDs
+    final_text_ids = pos_ids.gather(-1, first_valid_idx).squeeze(-1) # shape: [L-1]
+
+    start_t = torch.tensor([ref_start_token_id], dtype=final_text_ids.dtype, device=device)
+    
+    return torch.cat([start_t, final_text_ids])
+
+
+def handle_pattern(x0, token_ids: Dict[str, int], generation_mode: str = 'hybrid'):
+    """
+    Args:
+        x0: Token ID list of length 6
+        token_ids: Dictionary containing all token IDs
+    """
+    null_token_id = token_ids['null_token_id']
+    im_end_token_id = token_ids['im_end_token_id']
+    box_start_token_id = token_ids['box_start_token_id']
+    box_end_token_id = token_ids['box_end_token_id']
+    none_token_id = token_ids['none_token_id']
+    coord_start_token_id = token_ids['coord_start_token_id']
+    coord_end_token_id = token_ids['coord_end_token_id']
+    ref_end_token_id = token_ids['ref_end_token_id']
+    
+    x0 = x0.tolist()
+
+    if x0[0] == null_token_id:
+        return {
+            "type": "im_end",
+            "tokens": [im_end_token_id],
+            "need_switch_to_ar": False,
+            "is_terminal": True,
+        }
+    elif x0[0] == im_end_token_id:
+        return {
+            "type": "im_end",
+            "tokens": [im_end_token_id],
+            "need_switch_to_ar": False,
+            "is_terminal": True,
+        }
+    elif x0[:2] == [box_start_token_id, none_token_id]:
+        return {
+            "type": "empty_box",
+            "tokens": [box_start_token_id, none_token_id, box_end_token_id],
+            "need_switch_to_ar": False,
+            "is_terminal": False,
+        }
+    elif x0[0] == box_start_token_id:
+        coord_ix = 1
+        for coord in x0[1:5]:
+            if coord_start_token_id <= coord <= coord_end_token_id:
+                coord_ix += 1
+            else:
+                break
+
+        # Standard 4-coordinate bbox: <box><x1><x2><y1><y2></box>
+        if coord_ix == 5 and x0[5] == box_end_token_id:
+            return {
+                "type": "coord_box",
+                "tokens": x0,
+                "need_switch_to_ar": False,
+                "is_terminal": False,
+            }
+        # Two-coordinate pointing: <box><x><y></box>
+        # Convention: the first two coordinates are valid coord tokens, the third token is box_end.
+        # Remaining positions (if any) are not part of the pattern; truncate at box_end.
+        elif coord_ix == 3 and x0[3] == box_end_token_id:
+            return {
+                "type": "point_box",
+                "tokens": x0[:4],
+                "need_switch_to_ar": False,
+                "is_terminal": False,
+            }
+        else:
+            if generation_mode == 'fast':
+                # fast mode: treat as coord_box, stay in MTP
+                return {
+                    "type": "coord_box",
+                    "tokens": x0,
+                    "need_switch_to_ar": False,
+                    "is_terminal": False,
+                }
+            else:
+                # hybrid mode: error_box, switch to AR
+                return {
+                    "type": "error_box",
+                    "tokens": x0[:coord_ix],
+                    "need_switch_to_ar": True,
+                    "is_terminal": False,
+                }
+
+    else:
+        for i, token in enumerate(x0):
+            if token == null_token_id:
+                x0 = x0[:i]
+                break
+
+        if len(x0) >= 2 and x0[-1] == x0[-2] == ref_end_token_id:
+            x0 = x0[:-1]
+
+        return {
+            "type": "ref_object",
+            "tokens": x0,
+            "need_switch_to_ar": False,
+            "is_terminal": False,
+        }

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "transformers_version": "4.51.0"
+}

image_processing_locateanything.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Image processor class for KimiVL."""
+
+import math
+import numpy as np
+from PIL import Image
+from typing import Optional, Union
+
+import torch
+from torchvision.transforms import functional as TF
+from transformers.image_utils import ImageInput, make_list_of_images, valid_images
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+from transformers.utils import TensorType
+from transformers import AutoImageProcessor
+
+MEAN = (0.5, 0.5, 0.5)
+STD = (0.5, 0.5, 0.5)
+
+
+class LocateAnythingImageProcessor(BaseImageProcessor):
+    model_type = "locateanything"
+
+    def __init__(
+        self,
+        patch_size: int = 14,
+        image_mean: tuple[float, float, float] = MEAN,
+        image_std: tuple[float, float, float] = STD,
+        in_token_limit: int = 4096,
+        merge_kernel_size: list[int, int] = [2, 2],
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_token_limit = in_token_limit
+        self.patch_size = patch_size
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.merge_kernel_size = merge_kernel_size
+
+    def rescale(
+        self, image: Image.Image, merge_kernel_size: list[int, int] = [2, 2]
+    ) -> Image.Image:
+        w, h = image.size
+        patch_size = self.patch_size
+
+        if (w // patch_size) * (h // patch_size) > self.in_token_limit:
+            scale = math.sqrt(self.in_token_limit / ((w // patch_size) * (h // patch_size)))
+            new_w, new_h = int(w * scale), int(h * scale)
+            image = image.resize((new_w, new_h), Image.Resampling.BICUBIC)
+
+        new_w, new_h = image.size
+        pad_size_h = merge_kernel_size[0] * patch_size
+        pad_size_w = merge_kernel_size[1] * patch_size
+
+        target_w = math.ceil(new_w / pad_size_w) * pad_size_w
+        target_h = math.ceil(new_h / pad_size_h) * pad_size_h
+
+        if target_w != new_w or target_h != new_h:
+            image = image.resize((target_w, target_h), Image.Resampling.BICUBIC)
+
+        w, h = image.size
+        if w // patch_size >= 512 or h // patch_size >= 512:
+            raise ValueError("Exceed pos emb")
+
+        return image
+
+    def to_tensor(self, image: Image.Image) -> torch.Tensor:
+        return TF.to_tensor(image.convert("RGB"))
+
+    def normalize(self, image: torch.Tensor) -> torch.Tensor:
+        return TF.normalize(image, self.image_mean, self.image_std)
+
+    def patchify(self, image: torch.Tensor) -> tuple[torch.Tensor, list[int, int]]:
+        patch_size = self.patch_size
+        C, H, W = image.shape
+        patches = image.reshape(C, H // patch_size, patch_size, W // patch_size, patch_size)
+        patches = patches.permute(1, 3, 0, 2, 4)
+        patches = patches.contiguous().view(-1, C, patch_size, patch_size)
+        grid_hw = (H // patch_size, W // patch_size)
+        return patches, grid_hw
+
+    def _preprocess(self, image: ImageInput) -> tuple[torch.Tensor, list[int, int]]:
+        """
+        Preprocess image and patchify it.
+        Args:
+            image (`ImageInput`):
+                Image to preprocess. Expects pixel values ranging from 0 to 255. If pixel values range from 0 to 1, set `do_rescale=False`.
+        Returns:
+            patches: torch.Tensor
+            grid_hw: list[int, int]
+        """
+        image = self.rescale(image, self.merge_kernel_size)
+        image = self.to_tensor(image)
+        image = self.normalize(image)
+        patches, grid_hw = self.patchify(image)
+        return patches, grid_hw
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+    ) -> BatchFeature:
+        images = make_list_of_images(images)
+
+        if not valid_images(images):
+            raise ValueError(
+                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
+                "torch.Tensor, tf.Tensor or jax.ndarray."
+            )
+
+        pixel_values, image_grid_hws = [], []
+        for image in images:
+            patches, image_grid_hw = self._preprocess(image)
+            pixel_values.append(patches)
+            image_grid_hws.append(image_grid_hw)
+        pixel_values = torch.concat(pixel_values, dim=0)
+        image_grid_hws = np.array(image_grid_hws)
+        data = {"pixel_values": pixel_values, "image_grid_hws": image_grid_hws}
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+AutoImageProcessor.register("LocateAnythingImageProcessor", LocateAnythingImageProcessor)

mask_magi_utils.py

@@ -0,0 +1,101 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import torch
+
+# MagiAttention attn_type_map convention
+FULL, CAUSAL = 0, 1
+
+def build_magi_ranges(kv_len: int, q_len: int, block_size: int, ar_decode: bool=False, device: str = "cpu"):
+    """
+    Fixed strategy:
+      - use_cache=True: Mask blocked_k = (kv_len - block_size - 1) column
+      - causal_attn=False: Window interior is FULL (bidirectional)
+      - If q_len==kv_len: Use coarse prefix version (fewer ranges)
+      - Otherwise: General decode version (recompute rows expanding visible region row by row)
+
+    Conventions:
+      - K/V global length kv_len: [0, kv_len)
+      - Current Q is "last q_len tokens"
+      - First r=q_len-block_size rows are recomputed; last block_size rows are window
+    """
+    assert 0 < q_len <= kv_len
+
+    if ar_decode:
+        return {
+            "q_ranges": torch.tensor([[0, q_len]], dtype=torch.int32, device=device).contiguous(),
+            "k_ranges": torch.tensor([[0, kv_len]], dtype=torch.int32, device=device).contiguous(),
+            "attn_type_map": torch.tensor([CAUSAL],    dtype=torch.int32, device=device).contiguous(),
+        }
+
+
+    assert 0 < block_size <= q_len <= kv_len
+    B = block_size
+    r = q_len - B
+    q_global_start = kv_len - q_len
+
+    window_start_k = kv_len - B
+    blocked_k = window_start_k - 1  # The column that is blocked
+
+    q_ranges, k_ranges, types = [], [], []
+
+    # -------- prefix (q_len == kv_len) coarse-grained --------
+    if q_len == kv_len:
+        prefix_len = window_start_k  # kv_len - B
+
+        # prefix->prefix: causal
+        if prefix_len > 0:
+            q_ranges += [[0, prefix_len]]
+            k_ranges += [[0, prefix_len]]
+            types   += [CAUSAL]
+
+        # window->prefix: full, but exclude blocked_k => keys [0, blocked_k)
+        if prefix_len > 0 and blocked_k > 0:
+            q_ranges += [[prefix_len, kv_len]]
+            k_ranges += [[0, blocked_k]]
+            types   += [FULL]
+
+        # window->window: full
+        q_ranges += [[prefix_len, kv_len]]
+        k_ranges += [[prefix_len, kv_len]]
+        types += [FULL]
+
+        return {
+            "q_ranges": torch.tensor(q_ranges, dtype=torch.int32, device=device).contiguous(),
+            "k_ranges": torch.tensor(k_ranges, dtype=torch.int32, device=device).contiguous(),
+            "attn_type_map": torch.tensor(types,    dtype=torch.int32, device=device).contiguous(),
+        }
+
+    # -------- decode / general (q_len < kv_len) --------
+
+    # A) Recomputed rows: expand visible key cutoff row by row (use FULL + single-row q_range for precise shape)
+    for i in range(r):
+        g = q_global_start + i
+        q_ranges.append([i, i + 1])
+        k_ranges.append([0, g + 1])   # Allow keys [0, g]
+        types.append(FULL)
+
+    # B) Window rows: allow prefix but block blocked_k; window interior is full
+    q_win = [r, q_len]
+
+    # prefix keys [0, blocked_k)
+    if blocked_k > 0:
+        q_ranges.append(q_win)
+        k_ranges.append([0, blocked_k])
+        types.append(FULL)
+
+    # window keys [window_start_k, kv_len)
+    q_ranges.append(q_win)
+    k_ranges.append([window_start_k, kv_len])
+    types.append(FULL)
+
+    return {
+        "q_ranges": torch.tensor(q_ranges, dtype=torch.int32, device=device).contiguous(),
+        "k_ranges": torch.tensor(k_ranges, dtype=torch.int32, device=device).contiguous(),
+        "attn_type_map": torch.tensor(types, dtype=torch.int32, device=device).contiguous(),
+    }

mask_sdpa_utils.py

@@ -0,0 +1,232 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import torch
+
+
+def find_prefix_seq_length_by_pe(
+    pe: torch.Tensor
+) -> torch.Tensor:
+    """
+    Find the sequence length where position encoding drops (indicating prefix boundary).
+    Args:
+        pe: Position encoding tensor of shape [Batch size, Sequence length ]
+            Contains position indices for each token in the sequence.
+    Returns:
+        torch.Tensor: A tensor of shape [B] containing:
+            - The index where position encoding drops for each sequence
+            - -1 if no drop occurs in the sequence
+    """
+    batch_size, seq_len = pe.shape
+    prev = pe[:, :-1]
+    curr = pe[:, 1:]
+    drop_mask = curr < prev  #  [batch_size, seq_len-1]
+
+    seq_len = torch.full((batch_size,), -1, dtype=torch.long)
+    
+    for b in range(batch_size):
+        drop_pos = torch.nonzero(drop_mask[b], as_tuple=False)
+        if drop_pos.numel() > 0:
+            i = drop_pos[0].item() + 1  # Take first drop position (+1 because we compared shifted sequences)
+            seq_len[b] = i
+
+    return seq_len
+
+
+
+def update_causal_mask_with_pad_non_visible_2d(
+    input_ids: torch.Tensor,
+    attn_mask_2d: torch.Tensor,
+    text_mask_token_id: int,
+    block_size: int = 4,
+    causal_attn: bool = False
+) -> torch.Tensor:
+    """
+    Updates a 2D attention mask for hole sequence through input_ids and text_mask_token_id
+
+    Args:
+        input_ids: Input token IDs (unused in current implementation)
+        attn_mask_2d: 2D attention mask matrix of shape [seq_len, seq_len] where:
+            - 0.0 indicates allowed attention
+            - -inf indicates masked attention
+        text_mask_token_id: ID representing masked tokens
+        block_size: Size of the diffusion window
+        causal_attn: If True, maintains strict causal masking throughout
+    
+    Returns:
+        Modified attention mask with updated visibility patterns
+    """
+    seq_len = input_ids.shape[0]
+    device = input_ids.device
+
+    # Identify masked tokens and their preceding positions
+    input_mask = input_ids.eq(text_mask_token_id)
+    input_before_mask = torch.zeros_like(input_mask)
+    input_before_mask[:-1] = input_mask[1:]
+    mask_cols = (input_mask | input_before_mask)
+    non_mask = ~mask_cols
+
+    rows = torch.arange(seq_len, device=device)[:, None]
+    cols = torch.arange(seq_len, device=device)
+
+    indices = torch.arange(seq_len, device=device)
+    prev_non_mask = (indices * non_mask).cummax(dim=0).values
+
+    max_value = torch.iinfo(indices.dtype).max
+    mask_indices = torch.where(non_mask, indices, torch.full_like(indices, max_value))
+    reversed_mask_indices = torch.flip(mask_indices, dims=[0])
+    reversed_cummin = reversed_mask_indices.cummin(dim=0).values
+    next_non_mask = torch.flip(reversed_cummin, dims=[0])
+
+    infra_mask = (
+        (cols > prev_non_mask) &
+        (rows >= next_non_mask[None, :]) &
+        mask_cols[None, :]
+    )
+    attn_mask_2d.masked_fill_(infra_mask, -float('inf'))
+
+    if not causal_attn:
+        visible_mask = (
+            (rows > prev_non_mask[None, :]) &
+            (rows < cols) &
+            mask_cols[None, :]
+        )
+        attn_mask_2d.masked_fill_(visible_mask, 0.0)
+
+    return attn_mask_2d
+
+
+def update_causal_mask_for_one_gen_window_2d(
+    input_ids: torch.Tensor,
+    attn_mask_2d: torch.Tensor,
+    block_size: int = 4,
+    use_cache: bool = True,
+    causal_attn: bool = False
+) -> torch.Tensor:
+    """
+    Updates a 2D attention mask for a diffusion window in transformer inference.
+
+    Args:
+        input_ids: Input token IDs (unused in current implementation)
+        attn_mask_2d: 2D attention mask matrix of shape [seq_len, seq_len] where:
+            - 0.0 indicates allowed attention
+            - -inf indicates masked attention
+        block_size: Size of the diffusion window
+        use_cache: Whether key-value cache is being used
+        causal_attn: If True, maintains strict causal masking throughout
+    
+    Returns:
+        Modified attention mask with updated visibility patterns
+    """
+
+    if not causal_attn:
+        # Make the diffusion window (last block_size tokens) fully visible to itself
+        # This allows bidirectional attention within the diffusion window
+        attn_mask_2d[-block_size:, -block_size:] = 0.0
+    if use_cache:
+        # Mask the last token from previous round to prevent recomputation and maintain generation consistency.
+        attn_mask_2d[-block_size:, -block_size-1] = -float('inf')
+
+    return attn_mask_2d
+
+
+def create_block_diff_mask_by_pe_4d(
+    block_size: int, 
+    x0_len_list: torch.Tensor, 
+    position_ids: torch.Tensor, 
+    causal_attn: bool = False
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Generates a 4D attention mask for block-difference attention patterns.
+
+    The mask consists of three regions:
+    1. Causal block (top-left): Standard causal attention for `x0` tokens.
+    2. Mutual block (bottom-right): Non-causal attention within the same block for non-`x0` tokens.
+    3. Prefix block (bottom-left): Non-`x0` tokens can attend to a prefix of `x0` tokens.
+
+    Args:
+        block_size (int): Size of processing blocks for non-`x0` tokens.
+        x0_len_list (torch.Tensor): Tensor of shape [B] containing lengths of `x0` segments per batch.
+        position_ids (torch.Tensor): Tensor of shape [B, seq_len] containing position IDs.
+        causal_attn (bool, optional): If True, enforces causal masking in mutual blocks. Defaults to False.
+
+    Returns:
+        tuple[torch.Tensor, torch.Tensor]:
+            - A float mask of shape [batch_size, 1, seq_len, seq_len] with `-inf` for masked positions (non visiable).
+            - A boolean mask of shape [batch_size, 1, seq_len, seq_len] indicating allowed attention positions.
+    """
+    batch_size, seq_len = position_ids.shape
+    device = position_ids.device
+    
+    # Create position indices [batch_size, seq_len, seq_len]
+    q_idx = torch.arange(seq_len, device=device).view(1, seq_len, 1)  # [1, seq_len, 1]
+    kv_idx = torch.arange(seq_len, device=device).view(1, 1, seq_len)  # [1, 1, seq_len]
+    
+    # Broadcast to [B, seq_len, seq_len]
+    x0_len = x0_len_list.view(batch_size, 1, 1)  # [batch_size, 1, 1]
+    x0_flag_q = q_idx < x0_len           # [batch_size, seq_len, seq_len]
+    x0_flag_kv = kv_idx < x0_len         
+    
+    # Block indices calculation [batch_size, seq_len, seq_len]
+    q_block_idx = (q_idx - x0_len) // block_size
+    kv_block_idx = (kv_idx - x0_len) // block_size
+    
+    # causal block (top-left)
+    block_causal = x0_flag_q & x0_flag_kv & (q_idx >= kv_idx)
+
+    mutual_condition = (q_idx >= kv_idx) if causal_attn else torch.ones_like(q_idx, dtype=torch.bool)
+    block_mutual = (
+        ~x0_flag_q & ~x0_flag_kv &
+        (q_block_idx == kv_block_idx) &
+        mutual_condition
+    )
+
+    q_blk = torch.div(q_idx - x0_len, block_size, rounding_mode='floor')
+    q_blk_start = (x0_len_list.view(batch_size, 1) + q_blk[:, :, 0] * block_size).clamp(min=0, max=seq_len - 1)
+    prefix_len = position_ids.gather(1, q_blk_start)
+    prefix_len = prefix_len.unsqueeze(2)
+    block_prefix = (~x0_flag_q & x0_flag_kv) & (kv_idx < prefix_len)
+
+    final_mask = (block_causal | block_mutual | block_prefix)
+    customized_mask = torch.full_like(final_mask, float('-inf'), dtype=torch.bfloat16)
+    customized_mask.masked_fill_(final_mask, 0.0)
+    
+    return customized_mask.unsqueeze(1).to(device=device), final_mask.unsqueeze(1).to(device=device)
+
+
+def find_pred_pos_from_input_ids(
+    input_ids: torch.LongTensor = None,
+    text_mask_token_id: int = None,
+) -> torch.Tensor:
+    """Compute the relative prediction positions for masked tokens in a sequence.
+
+    For non-masked positions, the output is 0. For masked positions, the value increments
+    by 1 for each consecutive mask token, indicating how many steps ahead the prediction is.
+
+    Args:
+        input_ids (torch.LongTensor): Input token IDs of shape [batch_size, seq_len].
+        text_mask_token_id (int, optional): Token ID representing masked positions. Defaults to 151666.
+
+    Returns:
+        torch.Tensor: A tensor of shape [batch_size, seq_len] where:
+            - 0 indicates a non-masked token.
+            - n > 0 indicates the nth consecutive masked token (e.g., 1 = first mask, 2 = second mask, etc.).
+    """
+    batch_size, seq_len = input_ids.shape
+    device = input_ids.device
+
+    is_mask = (input_ids == text_mask_token_id)
+
+    base_mask = torch.zeros((batch_size, seq_len), dtype=torch.int8, device=device)
+
+    for b in range(batch_size):
+        for ix in range(1, seq_len):
+            if is_mask[b][ix] == True:
+                # Increment counter if current token is masked
+                base_mask[b][ix] = base_mask[b][ix-1] + 1
+
+    return base_mask

model.safetensors

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

modeling_locateanything.py

@@ -0,0 +1,537 @@
+# --------------------------------------------------------
+# NVIDIA
+# Copyright (c) 2025 NVIDIA
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import time
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import add_start_docstrings, is_flash_attn_2_available, logging
+from peft import LoraConfig, get_peft_model
+
+from .configuration_locateanything import LocateAnythingConfig
+from .modeling_qwen2 import Qwen2ForCausalLM
+from .modeling_vit import MoonVitPretrainedModel
+from transformers.models.qwen3.modeling_qwen3 import Qwen3ForCausalLM
+from .mask_sdpa_utils import *
+from .mask_magi_utils import *
+from .configuration_qwen2 import Qwen2Config
+
+from .generate_utils import (
+    sample_tokens,
+    handle_pattern,
+    get_token_ids_from_config,
+)
+
+logger = logging.get_logger(__name__)
+
+
+LOCATEANYTHING_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+    Parameters:
+        config ([`LocateAnythingConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+@add_start_docstrings(
+    "The bare LocateAnything Model outputting raw hidden-states without any specific head on top.",
+    LOCATEANYTHING_START_DOCSTRING,
+)
+class LocateAnythingPreTrainedModel(PreTrainedModel):
+    config_class = LocateAnythingConfig
+    base_model_prefix = "model"
+    main_input_name = 'input_ids'
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen2DecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_cache_class = True
+    _supports_static_cache = True
+    _supports_quantized_cache = True
+    _supports_sdpa = True
+
+    @classmethod
+    def _autoset_attn_implementation(cls, config, *args, **kwargs):
+        if getattr(config, '_attn_implementation', None) == 'magi':
+            return config
+        return super()._autoset_attn_implementation(config, *args, **kwargs)
+
+    def _check_and_adjust_attn_implementation(self, attn_implementation, is_init_check=False):
+        if attn_implementation == "magi":
+            return "magi"
+        return super()._check_and_adjust_attn_implementation(attn_implementation, is_init_check)
+    
+    def _init_weights(self, module):
+        std = getattr(self.config, 'initializer_range', None) or self.config.text_config.initializer_range
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+class LocateAnythingForConditionalGeneration(LocateAnythingPreTrainedModel, GenerationMixin):
+    config_class = LocateAnythingConfig
+    def __init__(self, config: LocateAnythingConfig, vision_model=None, language_model=None):
+        super().__init__(config)
+
+        self.template = config.template
+        self.mlp_checkpoint = config.mlp_checkpoint
+
+        logger.info(f'mlp_checkpoint: {self.mlp_checkpoint}')
+        if vision_model is not None:
+            self.vision_model = vision_model
+        else:
+            if config.vision_config.model_type == 'moonvit':
+                vision_attn_impl = getattr(config.vision_config, '_attn_implementation', None) or 'flash_attention_2'
+                if vision_attn_impl == 'flash_attention_2' and not is_flash_attn_2_available():
+                    logger.warning_once(
+                        "flash_attn is not available for MoonViT inference; falling back to sdpa."
+                    )
+                    vision_attn_impl = 'sdpa'
+                config.vision_config._attn_implementation = vision_attn_impl
+                self.vision_model = MoonVitPretrainedModel(config.vision_config)
+            else:
+                raise ValueError(f'Unsupported vision model type: {config.vision_config.model_type}. Only moonvit is supported.')
+
+        text_attn_impl = (
+            getattr(config.text_config, '_attn_implementation', None)
+            or getattr(config, '_attn_implementation', None)
+            or 'magi'
+        )
+        config.text_config._attn_implementation = text_attn_impl
+
+        if language_model is not None:
+            self.language_model = language_model
+        else:
+            if config.text_config.architectures[0] == 'Qwen2ForCausalLM':
+                self.language_model = Qwen2ForCausalLM(config.text_config)
+            elif config.text_config.architectures[0] == 'Qwen3ForCausalLM':
+                self.language_model = Qwen3ForCausalLM(config.text_config)
+            else:
+                raise ValueError(f'Unsupported language model architecture: {config.text_config.architectures[0]}. Only Qwen2ForCausalLM and Qwen3ForCausalLM are supported.')
+
+        vit_hidden_size = config.vision_config.hidden_size
+        llm_hidden_size = config.text_config.hidden_size
+
+        # MLP for moonvit (without pixel_shuffle_back, direct mapping)
+        self.mlp1 = nn.Sequential(
+                nn.LayerNorm(vit_hidden_size*4),
+                nn.Linear(vit_hidden_size*4, llm_hidden_size),
+                nn.GELU(),
+                nn.Linear(llm_hidden_size, llm_hidden_size)
+            )
+        self.image_token_index = config.image_token_index
+        self.neftune_alpha = None
+
+        if config.use_backbone_lora:
+            self.wrap_backbone_lora(r=config.use_backbone_lora, lora_alpha=2 * config.use_backbone_lora)
+
+        self.use_llm_lora = config.use_llm_lora
+        if config.use_llm_lora:
+            self.wrap_llm_lora(r=config.use_llm_lora, lora_alpha=2 * config.use_llm_lora)
+
+        self.token_ids = get_token_ids_from_config(config)
+
+        # Set _no_split_modules dynamically based on the actual LLM architecture
+        arch = config.text_config.architectures[0] if hasattr(config.text_config, 'architectures') and config.text_config.architectures else 'Qwen2ForCausalLM'
+        if 'Qwen3' in arch:
+            self._no_split_modules = ["Qwen3DecoderLayer"]
+        else:
+            self._no_split_modules = ["Qwen2DecoderLayer"]
+
+        
+    def wrap_backbone_lora(self, r=128, lora_alpha=256, lora_dropout=0.05):
+        lora_config = LoraConfig(
+            r=r,
+            target_modules=['self_attn.q_proj', 'self_attn.k_proj', 'self_attn.v_proj', 'self_attn.out_proj',
+                            'mlp.fc1', 'mlp.fc2'],
+            lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
+        )
+        self.vision_model = get_peft_model(self.vision_model, lora_config)
+        self.vision_model.print_trainable_parameters()
+
+    def wrap_llm_lora(self, r=128, lora_alpha=256, lora_dropout=0.05):
+        lora_config = LoraConfig(
+            r=r,
+            target_modules=['self_attn.q_proj', 'self_attn.k_proj', 'self_attn.v_proj', 'self_attn.o_proj',
+                            'mlp.gate_proj', 'mlp.down_proj', 'mlp.up_proj'],
+            lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
+            task_type='CAUSAL_LM'
+        )
+        self.language_model = get_peft_model(self.language_model, lora_config)
+        self.language_model.enable_input_require_grads()
+        self.language_model.print_trainable_parameters()
+        self.use_llm_lora = True
+        
+    
+    def forward(
+            self,
+            pixel_values: List[torch.FloatTensor],
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            image_grid_hws: Optional[torch.Tensor] = None,
+            image_flags: Optional[torch.Tensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        input_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        has_images = image_flags is not None and image_flags.sum() > 0
+        
+        vit_embeds = self.extract_feature(pixel_values, image_grid_hws)
+            
+        B, N, C = input_embeds.shape
+        input_embeds = input_embeds.reshape(B * N, C)
+
+        if has_images:
+            filtered_vit_embeds = []
+            idx = 0
+            for flag in image_flags:
+                flag_val = flag.item()
+                if flag_val != 0:
+                    filtered_vit_embeds.extend(vit_embeds[idx:idx + flag_val])
+                    idx += flag_val
+                else:
+                    idx += 1
+
+            vit_embeds = filtered_vit_embeds
+            vit_embeds = torch.cat(vit_embeds, dim=0)
+
+            vit_embeds = self.mlp1(vit_embeds)
+            input_ids = input_ids.reshape(B * N)
+            selected = (input_ids == self.image_token_index)
+
+            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:selected.sum()]
+        else:
+            if vit_embeds:
+                vit_embeds = torch.cat(vit_embeds, dim=0)
+                vit_embeds = self.mlp1(vit_embeds)
+                input_ids = input_ids.reshape(B * N)
+                selected = (input_ids == self.image_token_index)
+                if selected.sum() > 0:
+                    input_embeds[selected] = vit_embeds[:selected.sum()]
+
+        input_embeds = input_embeds.reshape(B, N, C)
+
+        outputs = self.language_model(
+            inputs_embeds=input_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        logits = outputs.logits
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    
+    def extract_feature(self, pixel_values, image_grid_hws):
+        vit_embeds = self.vision_model(pixel_values=pixel_values, grid_hws=image_grid_hws)
+
+        return vit_embeds
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def set_decoder(self, decoder):
+        self.language_model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.language_model.get_decoder()
+
+    @torch.no_grad()
+    def generate(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        input_ids: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        visual_features: Optional[torch.FloatTensor] = None,
+        image_grid_hws: Optional[torch.Tensor] = None,
+        tokenizer = None, 
+        n_future_tokens: int = 6,
+        **generate_kwargs,
+    ) -> torch.LongTensor:
+
+        verbose = generate_kwargs.pop('verbose', False)
+        start_time = time.time()
+        prefill_time = None
+
+        pixel_values = pixel_values.to(self.language_model.dtype)
+        # Convert numpy array to tensor if needed
+        if isinstance(image_grid_hws, np.ndarray):
+            image_grid_hws = torch.from_numpy(image_grid_hws).to(pixel_values.device, dtype=torch.int32)
+
+        batch_size, seq_len = input_ids.shape
+        assert batch_size == 1, 'only batch size = 1 is supported now'
+        assert generate_kwargs.get('use_cache', False), "Only use_cache=True is supported."
+
+        generated = input_ids.clone()
+        total_gen_length = min(tokenizer.model_max_length, seq_len + generate_kwargs.get('max_new_tokens', 2048))
+        iter_round = 0
+        past_key_values = None
+
+        # Extract visual features once before the loop
+        if visual_features is not None:
+            vit_embeds = visual_features
+        elif pixel_values is not None:
+            vit_embeds = self.extract_feature(pixel_values, image_grid_hws)
+        else:
+            vit_embeds = None
+        
+        if image_grid_hws is not None:
+            vit_embeds = torch.cat(vit_embeds, dim=0)
+            vit_embeds = self.mlp1(vit_embeds)
+
+        # ==================== Generation Mode ====================
+        # 'fast'   : MTP only, never fall back to AR
+        # 'slow'   : AR only, pure auto-regressive decoding
+        # 'hybrid' : MTP first, fall back to AR on error, switch back on box_end
+        generation_mode = generate_kwargs.get('generation_mode', 'hybrid')
+        assert generation_mode in ('fast', 'slow', 'hybrid'), \
+            f"Unsupported generation_mode='{generation_mode}'. Use 'fast', 'slow', or 'hybrid'."
+
+        sampling_history = []
+
+
+        use_mtp = generation_mode in ('fast', 'hybrid')
+        switch_to_ar_count = 0
+
+        # Pre-allocate mask tokens and position ids
+        default_mask_token_id = self.token_ids['default_mask_token_id']
+        pre_mask_tokens = torch.full(
+            (batch_size, n_future_tokens - 1),
+            default_mask_token_id,
+            dtype=generated.dtype,
+            device=generated.device
+        )
+        max_possible_len = total_gen_length + n_future_tokens
+        full_position_ids = torch.arange(0, max_possible_len, device=generated.device).unsqueeze(0)
+
+
+        def _prepare_inputs_in_mtp(generated):
+            generated_with_mask = torch.cat(
+                (
+                    generated, 
+                    generated[:, -1].unsqueeze(1),
+                    pre_mask_tokens
+                ), 
+                dim=1
+            ) # [batch_size, seq_len + 1 +  n_future_tokens - 1]
+
+            # Update pe for kvcache
+            start_idx = past_key_values[0][0].size(2) if past_key_values is not None else 0
+            position_ids = full_position_ids[:, start_idx : generated_with_mask.size(1)].clone()
+            position_ids[0, -n_future_tokens:] -= 1
+
+            prepare_inputs = self.language_model.prepare_inputs_for_generation(
+                generated_with_mask,
+                past_key_values,
+                None,
+                inputs_embeds=None,
+                use_cache=True,
+                position_ids=position_ids
+            )
+            return prepare_inputs
+
+
+        def _prepare_input_in_ar(generated):
+            start_idx = past_key_values[0][0].size(2) if past_key_values is not None else 0
+            position_ids = full_position_ids[:, start_idx : generated.size(1)]
+            prepare_inputs = self.language_model.prepare_inputs_for_generation(
+                generated,
+                past_key_values,
+                None,
+                inputs_embeds=None,
+                use_cache=True,
+                position_ids=position_ids
+            )
+            return prepare_inputs
+
+
+        def _sample_token_in_mtp(generated, outputs):
+            """Sample tokens using MTP (Multi-Token Prediction) mode."""
+            next_token_logits = outputs.logits[:, -n_future_tokens:, :]
+            probs, confidence, x0, box_avg = sample_tokens(
+                next_token_logits, generated, self.token_ids, keep_k=5, **generate_kwargs
+            )
+
+            is_box_empty = (box_avg[0] == 0).all()
+            new_tokens = x0[0] if is_box_empty else box_avg[0]
+
+            out_pattern = handle_pattern(new_tokens, self.token_ids, generation_mode)
+            out_type = out_pattern['type']
+            out_token = torch.tensor(out_pattern['tokens'], dtype=x0.dtype, device=x0.device)
+
+            return out_type, out_token
+
+
+        def _sample_token_in_ar(generated, outputs):
+            """Sample a single token using AR (Auto-Regressive) mode."""
+            next_token_logits = outputs.logits[:, -1:, :]
+            probs, confidence, x0, _ = sample_tokens(
+                next_token_logits, generated, self.token_ids, **generate_kwargs
+            )
+
+            out_token = x0[0]
+            out_type = 'continue_ar'
+            token_val = out_token[0].item()
+
+            box_end_token_id = self.token_ids['box_end_token_id']
+            coord_start_token_id = self.token_ids['coord_start_token_id']
+            coord_end_token_id = self.token_ids['coord_end_token_id']
+            none_token_id = self.token_ids['none_token_id']
+            im_end_token_id = self.token_ids['im_end_token_id']
+
+            if generation_mode == 'hybrid':
+                # Hybrid AR phase: detect box boundaries to switch back to MTP
+                if token_val == box_end_token_id:
+                    out_type = 'box_end_ar'
+                elif coord_start_token_id <= token_val <= coord_end_token_id or token_val == none_token_id:
+                    out_type = 'coord_ar'
+                else:
+                    out_type = 'im_end'
+            else:
+                # Slow mode: pure AR, only stop on im_end
+                if token_val == im_end_token_id:
+                    out_type = 'im_end'
+
+            return out_type, out_token
+
+
+        # Generate loop
+        while generated.size(1) < total_gen_length:
+            iter_round += 1
+
+            # Step 1: Prepare inputs
+            if use_mtp:
+                prepare_inputs = _prepare_inputs_in_mtp(generated)
+            else:
+                prepare_inputs = _prepare_input_in_ar(generated)
+
+            if iter_round == 1:
+                prepare_inputs.update({
+                    'visual_features': vit_embeds,
+                    'image_token_index': self.config.image_token_index,
+                })
+
+            # Step 2: Model forward & update KV cache
+            with torch.no_grad():
+                outputs = self.language_model(**prepare_inputs)
+
+            past_key_values = tuple(
+                (kv[0][:, :, :generated.shape[1], :], kv[1][:, :, :generated.shape[1], :])
+                for kv in outputs.past_key_values
+            )
+
+            # Step 3: Sample tokens
+            if use_mtp:
+                out_type, out_token = _sample_token_in_mtp(generated, outputs)
+            else:
+                out_type, out_token = _sample_token_in_ar(generated, outputs)
+
+            if verbose:
+                sampling_history.append(('ar' if 'ar' in out_type else 'mtp', tokenizer.decode(out_token, skip_special_tokens=False)))
+
+            generated = torch.cat([generated, out_token.unsqueeze(0)], dim=1)
+
+            # Step 4: Mode switching & termination
+            if out_type == 'im_end':
+                break
+
+            if generation_mode == 'hybrid':
+                if out_type == 'error_box':
+                    use_mtp = False
+                    switch_to_ar_count += 1
+                elif out_type == 'box_end_ar':
+                    use_mtp = True
+            # fast mode: use_mtp stays True always
+            # slow mode: use_mtp stays False always
+
+            if prefill_time is None:
+                prefill_time = time.time() - start_time
+
+        # Decode and return
+        generated_ids = generated[:, seq_len:]
+        response = tokenizer.batch_decode(generated_ids, skip_special_tokens=False)
+
+        if verbose:
+            end_time = time.time()
+            num_tokens = generated_ids.size(1)
+            num_boxes = response[0].count("<box>")
+            total_time = end_time - start_time
+
+            out_info =  f"\nStatistic Info, num_tokens={num_tokens}; " + \
+                    f"generate_time(s)={total_time:.4f}; " + \
+                    f"tps={(num_tokens / total_time):.4f}; " + \
+                    f"forward_step={iter_round}; " + \
+                    f"num_boxes={num_boxes}; " + \
+                    f"bps={(num_boxes / total_time):.4f}; " + \
+                    f"prefill_time={(prefill_time):.4f}; " + \
+                    f"switch_to_ar={switch_to_ar_count}\n"
+            print(out_info)
+
+            return response[0], sampling_history, out_info
+
+        return response[0]

modeling_qwen2.py

@@ -0,0 +1,1738 @@
+# coding=utf-8
+# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+""" PyTorch Qwen2 model."""
+import inspect
+import math
+import copy
+import warnings
+from functools import partial
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask, _prepare_4d_causal_attention_mask_for_sdpa
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+    replace_return_docstrings,
+)
+from .configuration_qwen2 import Qwen2Config
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+
+    _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
+
+
+logger = logging.get_logger(__name__)
+
+# Magi Attention Supported
+_MAGI_AVAILABLE = False
+try:
+    from magi_attention.functional.flex_flash_attn import flex_flash_attn_func
+    _MAGI_AVAILABLE = True
+except ImportError:
+    flex_flash_attn_func = None
+
+
+_CHECKPOINT_FOR_DOC = "Qwen/Qwen2-7B-beta"
+_CONFIG_FOR_DOC = "Qwen2Config"
+
+QWEN2_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "Qwen/Qwen2-7B-beta",
+    # See all Qwen2 models at https://huggingface.co/models?filter=qwen2
+]
+
+from .mask_sdpa_utils import (
+    find_prefix_seq_length_by_pe,
+    update_causal_mask_with_pad_non_visible_2d,
+    update_causal_mask_for_one_gen_window_2d,
+    create_block_diff_mask_by_pe_4d,
+    find_pred_pos_from_input_ids
+)
+
+from .mask_magi_utils import build_magi_ranges
+
+# Copied from transformers.models.llama.modeling_llama._get_unpad_data
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Qwen2
+class Qwen2RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        Qwen2RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Qwen2
+class Qwen2RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Build here to make `torch.jit.trace` work.
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
+        )
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if seq_len > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
+
+        return (
+            self.cos_cached[:seq_len].to(dtype=x.dtype),
+            self.sin_cached[:seq_len].to(dtype=x.dtype),
+        )
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`):
+            The position indices of the tokens corresponding to the query and key tensors. For example, this can be
+            used to pass offsetted position ids when working with a KV-cache.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+# Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Qwen2
+class Qwen2MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class Qwen2Attention(nn.Module):
+    """
+    Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
+    and "Generating Long Sequences with Sparse Transformers".
+    """
+
+    def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.is_causal = True
+        self.attention_dropout = config.attention_dropout
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+        self.rotary_emb = Qwen2RotaryEmbedding(
+            self.head_dim,
+            max_position_embeddings=self.max_position_embeddings,
+            base=self.rope_theta,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+            )
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            if self.layer_idx is None:
+                raise ValueError(
+                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+                    "with a layer index."
+                )
+            kv_seq_len += past_key_value.get_seq_length(self.layer_idx)
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        # repeat k/v heads if n_kv_heads < n_heads
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class Qwen2FlashAttention2(Qwen2Attention):
+    """
+    Qwen2 flash attention module, following Qwen2 attention module. This module inherits from `Qwen2Attention`
+    as the weights of the module stays untouched. The only required change would be on the forward pass
+    where it needs to correctly call the public API of flash attention and deal with padding tokens
+    in case the input contains any of them. Additionally, for sliding window attention, we apply SWA only to the bottom
+    config.max_window_layers layers.
+    """
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ):
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+            )
+
+            # overwrite attention_mask with padding_mask
+            attention_mask = kwargs.pop("padding_mask")
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            if self.layer_idx is None:
+                raise ValueError(
+                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+                    "with a layer index."
+                )
+            kv_seq_len += past_key_value.get_seq_length(self.layer_idx)
+
+        # Because the input can be padded, the absolute sequence length depends on the max position id.
+        rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1
+        cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        use_sliding_windows = (
+            _flash_supports_window_size
+            and getattr(self.config, "sliding_window", None) is not None
+            and kv_seq_len > self.config.sliding_window
+            and self.config.use_sliding_window
+        )
+
+        if not _flash_supports_window_size:
+            logger.warning_once(
+                "The current flash attention version does not support sliding window attention, for a more memory efficient implementation"
+                " make sure to upgrade flash-attn library."
+            )
+
+        if past_key_value is not None:
+            # Activate slicing cache only if the config has a value `sliding_windows` attribute
+            cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0
+            if (
+                getattr(self.config, "sliding_window", None) is not None
+                and kv_seq_len > self.config.sliding_window
+                and cache_has_contents
+            ):
+                slicing_tokens = 1 - self.config.sliding_window
+
+                past_key = past_key_value[self.layer_idx][0]
+                past_value = past_key_value[self.layer_idx][1]
+
+                past_key = past_key[:, :, slicing_tokens:, :].contiguous()
+                past_value = past_value[:, :, slicing_tokens:, :].contiguous()
+
+                if past_key.shape[-2] != self.config.sliding_window - 1:
+                    raise ValueError(
+                        f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
+                        f" {past_key.shape}"
+                    )
+
+                if attention_mask is not None:
+                    attention_mask = attention_mask[:, slicing_tokens:]
+                    attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1)
+
+            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        # repeat k/v heads if n_kv_heads < n_heads
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        dropout_rate = 0.0 if not self.training else self.attention_dropout
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in float16 just to be sure everything works as expected.
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        # Reashape to the expected shape for Flash Attention
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        attn_output = self._flash_attention_forward(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            q_len,
+            dropout=dropout_rate,
+            use_sliding_windows=use_sliding_windows,
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+    def _flash_attention_forward(
+        self,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        query_length,
+        dropout=0.0,
+        softmax_scale=None,
+        use_sliding_windows=False,
+    ):
+        """
+        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+        first unpad the input, then computes the attention scores and pad the final attention scores.
+
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`int`, *optional*):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+            use_sliding_windows (`bool`, *optional*):
+                Whether to activate sliding window attention.
+        """
+        if not self._flash_attn_uses_top_left_mask:
+            causal = self.is_causal
+        else:
+            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
+            causal = self.is_causal and query_length != 1
+
+        # Decide whether to use SWA or not by layer index.
+        if use_sliding_windows and self.layer_idx >= self.config.max_window_layers:
+            use_sliding_windows = False
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query_states, key_states, value_states, attention_mask, query_length
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            if not use_sliding_windows:
+                attn_output_unpad = flash_attn_varlen_func(
+                    query_states,
+                    key_states,
+                    value_states,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=max_seqlen_in_batch_q,
+                    max_seqlen_k=max_seqlen_in_batch_k,
+                    dropout_p=dropout,
+                    softmax_scale=softmax_scale,
+                    causal=causal,
+                )
+            else:
+                attn_output_unpad = flash_attn_varlen_func(
+                    query_states,
+                    key_states,
+                    value_states,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=max_seqlen_in_batch_q,
+                    max_seqlen_k=max_seqlen_in_batch_k,
+                    dropout_p=dropout,
+                    softmax_scale=softmax_scale,
+                    causal=causal,
+                    window_size=(self.config.sliding_window, self.config.sliding_window),
+                )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            if not use_sliding_windows:
+                attn_output = flash_attn_func(
+                    query_states,
+                    key_states,
+                    value_states,
+                    dropout,
+                    softmax_scale=softmax_scale,
+                    causal=causal,
+                )
+            else:
+                attn_output = flash_attn_func(
+                    query_states,
+                    key_states,
+                    value_states,
+                    dropout,
+                    softmax_scale=softmax_scale,
+                    causal=causal,
+                    window_size=(self.config.sliding_window, self.config.sliding_window),
+                )
+
+        return attn_output
+
+    # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._upad_input
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
+
+        # On the first iteration we need to properly re-create the padding mask
+        # by slicing it on the proper place
+        if kv_seq_len != attention_mask.shape[-1]:
+            attention_mask_num_tokens = attention_mask.shape[-1]
+            attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]
+
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+
+        key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
+        value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
+
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Qwen2
+class Qwen2SdpaAttention(Qwen2Attention):
+    """
+    Qwen2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+    `Qwen2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
+    SDPA API.
+    """
+
+    # Adapted from Qwen2Attention.forward
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if output_attentions:
+            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
+            logger.warning_once(
+                "Qwen2Model is using Qwen2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
+                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+            )
+
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value.get_seq_length(self.layer_idx)
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+
+        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
+        # Reference: https://github.com/pytorch/pytorch/issues/112577.
+        if query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=False,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+class Qwen2SdpaAttentionGqa(Qwen2Attention):
+    """
+    Qwen2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+    `Qwen2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
+    SDPA API.
+    """
+
+    # Adapted from Qwen2Attention.forward
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if output_attentions:
+            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
+            logger.warning_once(
+                "Qwen2Model is using Qwen2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
+                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+            )
+
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value.get_seq_length(self.layer_idx)
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        # key_states = repeat_kv(key_states, self.num_key_value_groups)
+        # value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+
+        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
+        # Reference: https://github.com/pytorch/pytorch/issues/112577.
+        if query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        with torch.backends.cuda.sdp_kernel(enable_flash=True,   
+                                    enable_math=True, 
+                                    enable_mem_efficient=False):
+
+            attn_output = torch.nn.functional.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=attention_mask,
+                enable_gqa=True,
+                dropout_p=self.attention_dropout if self.training else 0.0,
+                is_causal=False,
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+class Qwen2MagiAttention(Qwen2Attention):
+    """
+    Qwen2 attention using MagiAttention for efficient training with MTP packing support.
+    
+    MagiAttention uses range-based sparse attention patterns:
+    - q_ranges/k_ranges define which query/key ranges attend to each other
+    - attn_type_map specifies causal(1) or full(0) attention for each range pair
+    """
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if not _MAGI_AVAILABLE:
+            raise ImportError(
+                "magi_attention is not installed. Install with: pip install magi-attention"
+            )
+        self.softmax_scale = self.head_dim ** -0.5
+    
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[dict] = None,  # magi_plan dict
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if output_attentions:
+            raise NotImplementedError('MagiAttention does not support output_attentions=True')
+
+        bsz, q_len, _ = hidden_states.size()
+        assert bsz == 1, "MagiAttention only supports batch_size=1 (use packing instead)"
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # Magi expects [T, H, D] format (no batch dimension)
+        query_states = query_states.view(q_len, self.num_heads, self.head_dim)
+        key_states = key_states.view(q_len, self.num_key_value_heads, self.head_dim)
+        value_states = value_states.view(q_len, self.num_key_value_heads, self.head_dim)
+
+        kv_seq_len = q_len
+        if past_key_value is not None:
+            if self.layer_idx is None:
+                raise ValueError(
+                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+                    "with a layer index."
+                )
+            kv_seq_len += past_key_value.get_seq_length(self.layer_idx)
+        
+        cos, sin = self.rotary_emb(value_states.unsqueeze(0).transpose(1, 2), seq_len=kv_seq_len)
+        
+        # Apply RoPE: need [B, H, L, D] format for apply_rotary_pos_emb
+        q_for_rope = query_states.unsqueeze(0).transpose(1, 2)  # [1, H, L, D]
+        k_for_rope = key_states.unsqueeze(0).transpose(1, 2)    # [1, Hkv, L, D]
+        q_for_rope, k_for_rope = apply_rotary_pos_emb(q_for_rope, k_for_rope, cos, sin, position_ids)
+        
+        # Back to [T, H, D]
+        query_states = q_for_rope.squeeze(0).transpose(0, 1).contiguous()  # [L, H, D]
+        key_states = k_for_rope.squeeze(0).transpose(0, 1).contiguous()    # [L, Hkv, D]
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos}
+            # Note: Magi doesn't support KV cache in training, this is for potential future use
+            key_states_4d = key_states.unsqueeze(0).transpose(1, 2)
+            value_states_4d = value_states.unsqueeze(0).transpose(1, 2)
+            key_states_4d, value_states_4d = past_key_value.update(
+                key_states_4d, value_states_4d, self.layer_idx, cache_kwargs
+            )
+            key_states = key_states_4d.squeeze(0).transpose(0, 1).contiguous()
+            value_states = value_states_4d.squeeze(0).transpose(0, 1).contiguous()
+
+        # Run Magi Attention
+        # attention_mask is a magi_plan dict with q_ranges, k_ranges, attn_type_map, etc.
+
+        attn_output, _ = flex_flash_attn_func(
+            query_states.contiguous(),
+            key_states.contiguous(),
+            value_states.contiguous(),
+            q_ranges=attention_mask["q_ranges"],
+            k_ranges=attention_mask["k_ranges"],
+            attn_type_map=attention_mask["attn_type_map"],
+            softmax_scale=self.softmax_scale,
+            softcap=0.0,
+            deterministic=False,
+        )  # [T, H, D]
+        
+        # Reshape to [B, L, H*D]
+        attn_output = attn_output.view(1, q_len, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+QWEN2_ATTENTION_CLASSES = {
+    "eager": Qwen2Attention,
+    "flash_attention_2": Qwen2FlashAttention2,
+    "sdpa": Qwen2SdpaAttention,
+    "magi": Qwen2MagiAttention,
+}
+
+
+class Qwen2DecoderLayer(nn.Module):
+    def __init__(self, config: Qwen2Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        if config._attn_implementation == 'magi' and not _MAGI_AVAILABLE:
+            if is_flash_attn_2_available():
+                logger.warning_once(
+                    'magi_attention not available, falling back to flash_attention_2'
+                )
+                config._attn_implementation = 'flash_attention_2'
+            else:
+                logger.warning_once(
+                    'magi_attention not available, falling back to sdpa'
+                )
+                config._attn_implementation = 'sdpa'
+        if config._attn_implementation == 'flash_attention_2' and not is_flash_attn_2_available():
+            logger.warning_once(
+                'flash_attn is not available, falling back to sdpa'
+            )
+            config._attn_implementation = 'sdpa'
+
+        self.self_attn = QWEN2_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
+
+        self.mlp = Qwen2MLP(config)
+        self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. "
+                "Please make sure use `attention_mask` instead.`"
+            )
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, sequence_length)` where padding elements are indicated by 0.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+        """
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+QWEN2_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`Qwen2Config`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The bare Qwen2 Model outputting raw hidden-states without any specific head on top.",
+    QWEN2_START_DOCSTRING,
+)
+class Qwen2PreTrainedModel(PreTrainedModel):
+    config_class = Qwen2Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen2DecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    @classmethod
+    def _autoset_attn_implementation(cls, config, *args, **kwargs):
+        if getattr(config, '_attn_implementation', None) == 'magi':
+            return config
+        return super()._autoset_attn_implementation(config, *args, **kwargs)
+
+    def _check_and_adjust_attn_implementation(self, attn_implementation, is_init_check=False):
+        if attn_implementation == "magi":
+            return "magi"
+        return super()._check_and_adjust_attn_implementation(attn_implementation, is_init_check)
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+QWEN2_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
+            returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
+
+            Two formats are allowed:
+            - a [`~cache_utils.Cache`] instance;
+            - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+            shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
+            cache format.
+
+            The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
+            legacy cache format will be returned.
+
+            If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+            have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+            of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare Qwen2 Model outputting raw hidden-states without any specific head on top.",
+    QWEN2_START_DOCSTRING,
+)
+class Qwen2Model(Qwen2PreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2DecoderLayer`]
+
+    Args:
+        config: Qwen2Config
+    """
+
+    def __init__(self, config: Qwen2Config):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self._attn_implementation = config._attn_implementation
+        self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+        
+        self.block_size = getattr(config, 'block_size', 6)
+        self.causal_attn = getattr(config, 'causal_attn', False)
+        self.text_mask_token_id = getattr(config, 'text_mask_token_id', 151676)
+
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def image_processing(self, input_ids, visual_features, image_token_index):
+        if visual_features is not None:
+            input_embeds = self.get_input_embeddings()(input_ids)
+            B, N, C = input_embeds.shape
+            input_embeds = input_embeds.reshape(B * N, C)
+
+            input_ids = input_ids.reshape(B * N)
+            selected = (input_ids == image_token_index)
+            assert selected.sum() != 0
+            input_embeds[selected] = visual_features.reshape(-1, C).to(input_embeds.device)
+            input_embeds = input_embeds.reshape(B, N, C)
+        else:
+            input_embeds = self.get_input_embeddings()(input_ids)
+        return input_embeds
+        
+    @add_start_docstrings_to_model_forward(QWEN2_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        visual_features: Optional[torch.FloatTensor] = None,
+        image_token_index: int = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        past_key_values_length = 0
+
+        if use_cache:
+            use_legacy_cache = not isinstance(past_key_values, Cache)
+            if use_legacy_cache:
+                if past_key_values is None:
+                    past_key_values = DynamicCache()
+                else:
+                    past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+            past_key_values_length = past_key_values.get_seq_length()
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+        else:
+            position_ids = position_ids.view(-1, seq_length).long()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.image_processing(input_ids, visual_features, image_token_index)
+
+        if attention_mask is not None and self._attn_implementation == "magi" and use_cache:
+            is_padding_right = attention_mask[:, -1].sum().item() != batch_size
+            if is_padding_right:
+                raise ValueError(
+                    "You are attempting to perform batched generation with padding_side='right'"
+                    " this may lead to unexpected behaviour for Flash Attention version of Qwen2. Make sure to "
+                    " call `tokenizer.padding_side  = 'left'` before tokenizing the input. "
+                )
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        x0_len = find_prefix_seq_length_by_pe(position_ids).to(device=device)
+
+        def _prepare_block_mask_for_inference(attention_mask):
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask,
+                (batch_size, seq_length),
+                inputs_embeds,
+                past_key_values_length,
+                sliding_window=self.config.sliding_window,
+            )
+            # switch to ar mode
+            if seq_length == 1 or (input_ids is not None and input_ids[0][-1].item() != self.text_mask_token_id):
+                return attention_mask
+
+            
+            if attention_mask is None or len(attention_mask.shape) != 4:
+                return attention_mask
+
+            # For SDLM, the generation window should set to bidirectional attention
+            if use_cache:
+                update_mask_func = partial(
+                    update_causal_mask_for_one_gen_window_2d,
+                    block_size=self.block_size,
+                    use_cache=use_cache,
+                    causal_attn=self.causal_attn,
+                )
+            else:
+                update_mask_func = partial(
+                    update_causal_mask_with_pad_non_visible_2d,
+                    block_size=self.block_size,
+                    text_mask_token_id=self.text_mask_token_id,
+                    causal_attn=self.causal_attn,
+                )
+
+            new_attention_mask = []
+            for b in range(attention_mask.shape[0]):
+                new_attention_mask.append(
+                    update_mask_func(
+                        input_ids[b],
+                        attention_mask[b][0],
+                    ).unsqueeze(0)
+                )
+            return torch.stack(new_attention_mask, dim=0)
+
+        def _prepare_block_mask_for_training():
+            block_mask, _ = create_block_diff_mask_by_pe_4d(
+                block_size=self.block_size,
+                x0_len_list=x0_len,
+                position_ids=position_ids,
+                causal_attn=self.causal_attn,
+            )
+            return block_mask
+
+        if self._attn_implementation == "magi":
+            ar_decode = seq_length == 1 or (input_ids is not None and input_ids[0][-1].item() != self.text_mask_token_id)
+            attention_mask = build_magi_ranges(
+                kv_len=seq_length + past_key_values_length,
+                q_len=seq_length,
+                block_size=self.block_size,
+                ar_decode=ar_decode,
+                device=device
+            )
+
+        elif self._attn_implementation == "sdpa":
+            attention_mask = _prepare_block_mask_for_training() if self.training else _prepare_block_mask_for_inference(attention_mask)
+
+        else:
+            raise NotImplementedError(f'{self._attn_implementation=}')
+
+
+        hidden_states = inputs_embeds
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class Qwen2ForCausalLM(Qwen2PreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Qwen2Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.text_mask_token_id = getattr(config, 'text_mask_token_id', 151676)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @add_start_docstrings_to_model_forward(QWEN2_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        visual_features: Optional[torch.FloatTensor] = None,
+        image_token_index: int = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, Qwen2ForCausalLM
+
+        >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            visual_features=visual_features,
+            image_token_index=image_token_index,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+        logits = logits.float()
+
+        loss = None
+        if labels is not None:
+
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+
+            shift_labels = shift_labels.view(-1)
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+            pos_masks = find_pred_pos_from_input_ids(input_ids, text_mask_token_id=self.text_mask_token_id)
+            shift_input_ids = input_ids[..., :-1].contiguous()
+            shift_pos_masks = pos_masks[:, :-1]
+            shift_input_ids = shift_input_ids.view(-1)
+            max_n_future_tokens = min(4, self.model.block_size)
+            pos_loss_list = torch.zeros(max_n_future_tokens, device=shift_logits.device)
+            shift_pos_masks = shift_pos_masks.reshape(-1)
+
+            for ix in range(max_n_future_tokens):
+                seg_loss = F.cross_entropy(
+                    shift_logits[shift_pos_masks == ix],
+                    shift_labels[shift_pos_masks == ix],
+                    reduction='mean'
+                )
+                pos_loss_list[ix] = seg_loss
+
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+        
+        if self.training:
+            return CausalLMOutputWithPast(
+                loss=loss,
+                logits=logits,
+                past_key_values=outputs.past_key_values,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+            ), pos_loss_list
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        # Omit tokens covered by past_key_values
+        if past_key_values is not None:
+            if isinstance(past_key_values, Cache):
+                cache_length = past_key_values.get_seq_length()
+                past_length = past_key_values.seen_tokens
+                max_cache_length = past_key_values.get_max_length()
+            else:
+                cache_length = past_length = past_key_values[0][0].shape[2]
+                max_cache_length = None
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
+            # input)
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past
+
+
+@add_start_docstrings(
+    """
+    The Qwen2 Model transformer with a sequence classification head on top (linear layer).
+
+    [`Qwen2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+    (e.g. GPT-2) do.
+
+    Since it does classification on the last token, it requires to know the position of the last token. If a
+    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+    each row of the batch).
+    """,
+    QWEN2_START_DOCSTRING,
+)
+class Qwen2ForSequenceClassification(Qwen2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = Qwen2Model(config)
+        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(QWEN2_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        transformer_outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+        if self.config.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
+                sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
+                sequence_lengths = sequence_lengths % input_ids.shape[-1]
+                sequence_lengths = sequence_lengths.to(logits.device)
+            else:
+                sequence_lengths = -1
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(pooled_logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(pooled_logits, labels)
+        if not return_dict:
+            output = (pooled_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )

modeling_vit.py

@@ -0,0 +1,615 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import math
+from copy import deepcopy
+from typing import Union, Tuple, Sequence, Optional, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+try:
+    from transformers.activations import PytorchGELUTanh
+except ImportError:
+    PytorchGELUTanh = lambda: nn.GELU(approximate='tanh')
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import is_flash_attn_2_available, logging
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_varlen_func
+else:
+    flash_attn_varlen_func = None
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class MoonViTConfig(PretrainedConfig):
+    model_type = "moonvit"
+
+    def __init__(
+        self,
+        patch_size: int = 14,
+        init_pos_emb_height: int = 64,
+        init_pos_emb_width: int = 64,
+        num_attention_heads: int = 16,
+        num_hidden_layers: int = 27,
+        hidden_size: int = 1152,
+        intermediate_size: int = 4304,
+        merge_kernel_size: tuple[int, int] = (2, 2),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.patch_size = patch_size
+        # Positional embedding config
+        self.init_pos_emb_height = init_pos_emb_height
+        self.init_pos_emb_width = init_pos_emb_width
+        # Transformer config
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        # Patch merger config
+        self.merge_kernel_size = merge_kernel_size
+
+
+def multihead_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_cu_seqlens: Optional[torch.Tensor] = None,
+    k_cu_seqlens: Optional[torch.Tensor] = None,
+):
+    """Multi-head attention using flash attention 2.
+
+    Args:
+        q, k, v: tensor of shape (batch_size, seqlen, num_heads, head_dim),
+            or (tot_seqlens, num_heads, head_dim) if packing.
+        q_cu_seqlens (torch.Tensor): cumulative sequence lengths of q.
+            The first element should be 0 and the last element should be q.shape[0].
+        k_cu_seqlens (torch.Tensor): cumulative sequence lengths of k.
+            The first element should be 0 and the last element should be k.shape[0].
+
+    Returns:
+        output: shape (batch_size, seqlen, dim) or (tot_seqlens, dim) if packing,
+            where dim = num_heads * head_dim
+    """
+    if flash_attn_varlen_func is None:
+        logger.warning_once(
+            "flash_attn is not available for MoonViT; falling back to sdpa attention."
+        )
+        return sdpa_attention(
+            q,
+            k,
+            v,
+            q_cu_seqlens=q_cu_seqlens,
+            k_cu_seqlens=k_cu_seqlens,
+        )
+
+    # Unified format legal check
+    assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must have 3 dims"
+    assert q_cu_seqlens[-1] == q.shape[0], "q_cu_seqlens must sum to q.shape[0]"
+    assert (
+        k_cu_seqlens[-1] == k.shape[0] == v.shape[0]
+    ), "k_cu_seqlens must sum to k.shape[0]"
+    assert q.dtype in [
+        torch.bfloat16,
+        torch.float16,
+    ], f"unsupported dtype {q.dtype} for multihead attn"
+
+    max_seqlen_q = (q_cu_seqlens[1:] - q_cu_seqlens[:-1]).max().item()
+    max_seqlen_k = (k_cu_seqlens[1:] - k_cu_seqlens[:-1]).max().item()
+    attn_out = flash_attn_varlen_func(
+        q,
+        k,
+        v,
+        q_cu_seqlens,
+        k_cu_seqlens,
+        max_seqlen_q,
+        max_seqlen_k,
+        causal=False,
+    )
+    attn_out = attn_out.flatten(start_dim=-2)
+
+    return attn_out
+
+
+def sdpa_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_cu_seqlens: Optional[torch.Tensor] = None,
+    k_cu_seqlens: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """SDPA attention.
+
+    Args:
+        q, k, v: tensor of shape (batch_size, seqlen, num_heads, head_dim),
+            or (tot_seqlens, num_heads, head_dim) if packing.
+    """
+    seq_length = q.shape[0]
+    attention_mask = torch.zeros(
+        [1, seq_length, seq_length], device=q.device, dtype=torch.bool
+    )
+    for i in range(1, len(q_cu_seqlens)):
+        attention_mask[
+            ...,
+            q_cu_seqlens[i - 1] : q_cu_seqlens[i],
+            q_cu_seqlens[i - 1] : q_cu_seqlens[i],
+        ] = True
+    q = q.transpose(0, 1)
+    k = k.transpose(0, 1)
+    v = v.transpose(0, 1)
+    attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0)
+    attn_output = attn_output.transpose(0, 1)
+    attn_output = attn_output.reshape(seq_length, -1)
+    return attn_output
+
+
+def eager_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_cu_seqlens: Optional[torch.Tensor] = None,
+    k_cu_seqlens: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    seq_length = q.shape[0]
+    attention_mask = torch.zeros(
+        [1, seq_length, seq_length], device=q.device, dtype=torch.bool
+    )
+    for i in range(1, len(q_cu_seqlens)):
+        attention_mask[
+            ...,
+            q_cu_seqlens[i - 1] : q_cu_seqlens[i],
+            q_cu_seqlens[i - 1] : q_cu_seqlens[i],
+        ] = True
+    q = q.transpose(0, 1)
+    k = k.transpose(0, 1)
+    v = v.transpose(0, 1)
+
+    attn_weight = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
+    attn_weight += attention_mask
+    attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32).to(q.dtype)
+
+    attn_output = attn_weight @ v
+    attn_output = attn_output.transpose(0, 1)
+    attn_output = attn_output.reshape(seq_length, -1)
+    return attn_output
+
+
+VL_VISION_ATTENTION_FUNCTIONS = {
+    "flash_attention_2": multihead_attention,
+    "sdpa": sdpa_attention,
+    "eager": eager_attention,
+}
+
+
+def _apply_rope_input_validation(x, freqs_cis):
+    assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape)
+    assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape)
+    assert x.shape[-1] == 2 * freqs_cis.shape[-1], (x.shape, freqs_cis.shape)
+    assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype
+
+
+def apply_rope(
+    xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Args: (The leading dimensions of all inputs should be the same)
+        xq: query, tensor of shape (..., num_heads, head_dim)
+        xk: key, tensor of shape (..., num_heads, head_dim)
+        freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid.
+    Returns:
+        xq_out, xk_out: tensors of shape (..., num_heads, head_dim)
+    """
+    _apply_rope_input_validation(xq, freqs_cis)
+    _apply_rope_input_validation(xk, freqs_cis)
+
+    freqs_cis = freqs_cis.unsqueeze(-2)  # ..., 1, head_dim/2
+    # ..., num_heads, head_dim/2
+    xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2))
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2)  # ..., num_heads, head_dim
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2)  # ..., num_heads, head_dim
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class Learnable2DInterpPosEmb(nn.Module):
+    def __init__(
+        self, height: int, width: int, dim: int, interpolation_mode: str = "bicubic"
+    ) -> None:
+        super().__init__()
+        self.height = height
+        self.width = width
+        self.interpolation_mode = interpolation_mode
+        self.weight = nn.Parameter(torch.empty(height, width, dim))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.normal_(self.weight)
+
+    def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor:
+        pos_embs = []
+        for shape in grid_hws.tolist():
+            if shape == self.weight.shape[:-1]:
+                pos_embs.append(self.weight.flatten(end_dim=1))
+            else:
+                pos_embs.append(
+                    F.interpolate(
+                        self.weight.permute((2, 0, 1)).unsqueeze(0),
+                        size=shape,
+                        mode=self.interpolation_mode,
+                    )
+                    .squeeze(0)
+                    .permute((1, 2, 0))
+                    .flatten(end_dim=1)
+                )
+        out = x + torch.cat(pos_embs)
+        return out
+
+
+class MoonVisionPatchEmbed(nn.Module):
+
+    def __init__(
+        self,
+        out_dim: int,
+        in_dim: int = 3,
+        patch_size: Union[int, Tuple[int, int]] = (14, 14),
+        pos_emb_height: int = 14,
+        pos_emb_width: int = 14,
+    ):
+        super().__init__()
+        assert isinstance(
+            patch_size, (int, Sequence)
+        ), f"Invalid patch_size type: {type(patch_size)}"
+        if isinstance(patch_size, int):
+            patch_size = (patch_size, patch_size)
+        assert (
+            len(patch_size) == 2
+        ), f"Expected patch_size to be a tuple of 2, got {patch_size}"
+        self.patch_size = patch_size
+
+        self.proj = nn.Conv2d(
+            in_dim, out_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+        self.pos_emb = Learnable2DInterpPosEmb(
+            height=pos_emb_height, width=pos_emb_width, dim=out_dim
+        )
+
+    def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x (L, Channels): input tensor
+            grid_hws (N, 2): grid height and width
+
+        Returns:
+            (L, Cout) tensor
+        """
+        x = self.proj(x).view(x.size(0), -1)
+        # apply positional embedding
+        x = self.pos_emb(x, grid_hws)
+        return x
+
+
+class Rope2DPosEmb(nn.Module):
+    """2D rotary position embedding with multi-resolution support.
+
+    This class is intended to be used in the following way:
+    1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis.
+    2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration.
+    3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation.
+        The rope is shared across all attention layers and all heads.
+
+    Refs:
+    - RoFormer: https://arxiv.org/abs/2104.09864
+    - VisionLLaMA: https://arxiv.org/abs/2403.00522
+    - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py
+
+    Args:
+        dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed)
+        max_height (int): the maximum height of the 2D grid
+        max_width (int): the maximum width of the 2D grid
+        theta_base (float): the base of the theta
+        device (str): the device to store the precomputed cis
+    """
+
+    def __init__(self, dim: int, max_height: int, max_width: int, theta_base=10000):
+        super().__init__()
+        self.dim = dim
+        assert self.dim % 4 == 0, "dim must be divisible by 4"
+        self.max_height = max_height
+        self.max_width = max_width
+        self.theta_base = theta_base
+
+        self.freqs_cis = None
+
+    def extra_repr(self):
+        return f"dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}"
+
+    def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor:
+        """Calculate the cis(freqs) for each position in the 2D grid.
+
+        Return: complex tensor of shape (max_height, max_width, dim//2) and value:
+            height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim))
+            weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim))   with (i in [0, dim//4))
+            note: `cis` is a mathematical notation defined by cis x = cos x + i sin x,
+        """
+        N = self.max_height * self.max_width
+        flat_pos = torch.arange(0, N).float().to(device)
+        x_pos = flat_pos % self.max_width
+        y_pos = flat_pos // self.max_width
+        dim_range = (
+            torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device)
+        )  # C/4
+        freqs = 1.0 / (self.theta_base ** (dim_range / self.dim))
+        x_freqs = torch.outer(x_pos, freqs).float()  # N, C/4
+        y_freqs = torch.outer(y_pos, freqs).float()  # N, C/4
+        x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs)  # N, C/4
+        y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs)  # N, C/4
+        # N, C/4, 2
+        freqs_cis = torch.cat(
+            [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1
+        )
+        # max_height, max_width, C/2
+        freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1)
+        return freqs_cis
+
+    def get_freqs_cis(self, grid_hws: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            grid_hws (torch.Tensor): grid height and width
+
+        Returns:
+            freqs_cis: tensor of shape (sum(t * height * width), dim//2)
+        """
+        if self.freqs_cis is None:
+            self.freqs_cis = self._precompute_freqs_cis(grid_hws.device)
+
+        shapes = grid_hws.tolist()
+        assert all(
+            1 <= h <= self.max_height and 1 <= w <= self.max_width for h, w in shapes
+        ), (
+            shapes,
+            self.max_height,
+            self.max_width,
+        )
+        freqs_cis = torch.cat(
+            [self.freqs_cis[:h, :w].reshape(-1, self.dim // 2) for h, w in shapes],
+            dim=0,
+        )
+        return freqs_cis
+
+
+class MLP2(nn.Module):
+    """
+    Args:
+        dims: [in_dim, hidden_dim, out_dim]
+        bias: whether to use bias in linear layer.
+    """
+
+    def __init__(self, dims: list[int], activation, bias=True):
+        super().__init__()
+        assert len(dims) == 3
+        self.fc0 = nn.Linear(dims[0], dims[1], bias=bias)
+        self.fc1 = nn.Linear(dims[1], dims[2], bias=bias)
+        self.activation = activation
+        for m in [self.fc0, self.fc1]:
+            nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features))
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc0(x)
+        x = self.activation(x)
+        return self.fc1(x)
+
+
+class MoonVitEncoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        *,
+        attn_implementation: str = "eager",
+        activation=F.gelu,
+        attn_bias: bool = False,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads
+        self.attn_implementation = attn_implementation
+
+        self.norm0 = nn.LayerNorm(hidden_dim)
+        self.norm1 = nn.LayerNorm(hidden_dim)
+        self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation)
+        self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias)
+        self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias)
+
+    def attention_qkvpacked(
+        self,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rope_freqs_cis: Optional[torch.Tensor] = None,
+    ):
+        """
+        Args:
+            x (torch.Tensor): (batch_size, seqlen, hidden_dim)
+            cu_seqlens (torch.Tensor):
+        """
+        xqkv = self.wqkv(x)
+
+        qkv_shape = xqkv.size()[:-1] + (
+            3,
+            self.num_heads,
+            self.hidden_size_per_attention_head,
+        )
+        # xqkv: (batch_size, seqlen, 3, nheads, headdim)
+        xqkv = xqkv.view(*qkv_shape)
+        xq, xk, xv = torch.unbind(xqkv, dim=-3)
+
+        xq, xk = apply_rope(xq, xk, rope_freqs_cis)
+
+        attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation]
+        attn_out = attn_func(
+            xq, xk, xv, q_cu_seqlens=cu_seqlens, k_cu_seqlens=cu_seqlens
+        )
+
+        attn_out = self.wo(attn_out)
+        return attn_out
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rope_freqs_cis: Union[torch.Tensor, None] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: non-packed (B, N, D) or packed (L, D). if non-packed, seqlens should be None, if packed, seqlens should be set
+
+        Returns:
+            output: same shape of input, non-packed (B, N, D) for non-packed input, (L, D) for packed input
+        """
+        residual = hidden_states
+        hidden_states = self.norm0(hidden_states)
+        attn_out = self.attention_qkvpacked(
+            hidden_states, cu_seqlens, rope_freqs_cis=rope_freqs_cis
+        )
+        hidden_states = residual + attn_out
+
+        residual = hidden_states
+        hidden_states = self.mlp(self.norm1(hidden_states))
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class MoonVitEncoder(nn.Module):
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_layers: int,
+        block_cfg: dict,
+    ) -> None:
+        super().__init__()
+
+        self.rope_2d = Rope2DPosEmb(
+            block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512
+        )
+        self.blocks = nn.ModuleList(
+            [MoonVitEncoderLayer(**block_cfg) for _ in range(num_layers)]
+        )
+        self.final_layernorm = nn.LayerNorm(hidden_dim)
+
+    def forward(
+        self, hidden_states: torch.Tensor, grid_hws: torch.Tensor
+    ) -> torch.Tensor:
+        rope_freqs_cis = self.rope_2d.get_freqs_cis(grid_hws=grid_hws)
+
+        lengths = torch.cat(
+            (
+                torch.zeros(1, device=hidden_states.device, dtype=grid_hws.dtype),
+                grid_hws[:, 0] * grid_hws[:, 1],
+            )
+        )
+        cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32)
+
+        for _, block in enumerate(self.blocks):
+            hidden_states = block(
+                hidden_states, cu_seqlens, rope_freqs_cis=rope_freqs_cis
+            )
+
+        hidden_states = self.final_layernorm(hidden_states)
+
+        return hidden_states
+
+
+def patch_merger(
+    x: torch.Tensor,
+    grid_hws: torch.Tensor,
+    merge_kernel_size: list[int, int] = (2, 2),
+) -> List[torch.Tensor]:
+    d_model = x.size(-1)
+
+    outputs = []
+    pre_sum = 0
+    for x_shape in grid_hws.tolist():
+        height, width = x_shape[0], x_shape[1]
+        # Get the current sequence
+        seq = x[pre_sum : pre_sum + height * width]
+        # Reshape along self.merge_kernel_size and concat to the last dimension
+        kernel_height, kernel_width = merge_kernel_size
+        new_height, new_width = height // kernel_height, width // kernel_width
+        reshaped_seq = seq.view(
+            new_height, kernel_height, new_width, kernel_width, d_model
+        )
+        reshaped_seq = reshaped_seq.permute(0, 2, 1, 3, 4).contiguous()
+        padded_seq = reshaped_seq.view(
+            new_height * new_width, -1
+        )
+        outputs.append(padded_seq)
+        pre_sum += height * width
+
+    return outputs
+
+
+class MoonVitPretrainedModel(PreTrainedModel):
+    config_class = MoonViTConfig
+    model_type = "moonvit"
+    _no_split_modules = ["PackingTransformer"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+
+    def __init__(self, config: MoonViTConfig, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+        config = deepcopy(config)
+        self.merge_kernel_size = config.merge_kernel_size
+        self.patch_size = config.patch_size
+        self.patch_embed = MoonVisionPatchEmbed(
+            out_dim=config.hidden_size,
+            patch_size=config.patch_size,
+            pos_emb_height=config.init_pos_emb_height,
+            pos_emb_width=config.init_pos_emb_width,
+        )
+
+        self.encoder = MoonVitEncoder(
+            hidden_dim=config.hidden_size,
+            num_layers=config.num_hidden_layers,
+            block_cfg={
+                "num_heads": config.num_attention_heads,
+                "hidden_dim": config.hidden_size,
+                "mlp_dim": config.intermediate_size,
+                "activation": PytorchGELUTanh(),
+                "attn_bias": True,
+                "attn_implementation": config._attn_implementation,
+            },
+        )
+
+    def forward(
+        self, pixel_values: torch.Tensor, grid_hws: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Args:
+            pixel_values (torch.Tensor): The input pixel values.
+            grid_hws (torch.Tensor): The grid height and width.
+
+        Returns:
+            torch.Tensor: The output tokens.
+        """
+        hidden_states = self.patch_embed(pixel_values, grid_hws)
+        hidden_states = self.encoder(hidden_states, grid_hws)
+        hidden_states = patch_merger(
+            hidden_states, grid_hws, merge_kernel_size=self.merge_kernel_size
+        )
+        return hidden_states

preprocessor_config.json

@@ -0,0 +1,24 @@
+{
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_locateanything.LocateAnythingImageProcessor",
+    "AutoProcessor": "processing_locateanything.LocateAnythingProcessor"
+  },
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "LocateAnythingImageProcessor",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "in_token_limit": 25600,
+  "merge_kernel_size": [
+    2,
+    2
+  ],
+  "patch_size": 14,
+  "processor_class": "LocateAnythingProcessor"
+}

processing_locateanything.py

@@ -0,0 +1,678 @@
+# coding=utf-8
+# Copyright 2024 The HuggingFace Inc. team.
+#
+# 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.
+"""
+Processor class for LocateAnything.
+"""
+
+import math
+import os
+from typing import Iterable, List, Union, Literal
+import base64
+import sys
+import time
+import warnings
+from functools import lru_cache
+from io import BytesIO
+import re
+import requests
+import torch
+import torchvision
+from packaging import version
+from PIL import Image
+from torchvision import io
+from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+from typing import Optional, Any
+import numpy as np
+
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.image_utils import ImageInput
+try:
+    from transformers.image_utils import VideoInput
+except ImportError:
+    VideoInput = None
+from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+from transformers.utils import logging
+import lmdb
+import cv2
+import pickle
+import decord
+
+logger = logging.get_logger(__name__)
+
+FPS = 2.0
+MAX_FRAMES = 64 
+VIDEO_TOTAL_PIXELS = int(float(os.environ.get('VIDEO_MAX_PIXELS', 32000 * 28 * 28 * 0.9)))
+logger.info(f"set VIDEO_TOTAL_PIXELS: {VIDEO_TOTAL_PIXELS}")
+
+
+def to_rgb(pil_image: Image.Image) -> Image.Image:
+      if pil_image.mode == 'RGBA':
+          white_background = Image.new("RGB", pil_image.size, (255, 255, 255))
+          white_background.paste(pil_image, mask=pil_image.split()[3])  # Use alpha channel as mask
+          return white_background
+      else:
+          return pil_image.convert("RGB")
+
+def read_img_from_lmdb_v2(image_data):
+    # special case for AgiBotWorld
+    lmdb_file, lmdb_key = image_data['lmdb_file'], image_data['lmdb_key']
+    key = lmdb_key.encode('ascii')
+    env = lmdb.open(lmdb_file, max_readers=10240, readonly=True, lock=False, readahead=False, meminit=False)
+    txn = env.begin()
+    value = txn.get(key)
+    if value is None:
+        print(f"Warning: Key {key} not found.")
+        return None
+    record = pickle.loads(value)
+    image_bgr = cv2.imdecode(np.frombuffer(record['image'], dtype=np.uint8), cv2.IMREAD_COLOR)
+    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+    image = Image.fromarray(image_rgb)
+        
+    return image
+
+def parse_lmdb_image_data(image_data):
+    lmdb_file = image_data['lmdb_file']
+    if not os.path.exists(lmdb_file):
+        if "/home/zhidingy/workspace/libs/eagle/Eagle2/" in lmdb_file:
+            image_data['lmdb_file'] = lmdb_file.replace("/home/zhidingy/workspace/libs/eagle/Eagle2/", "")
+        else:
+            raise ValueError(f"LMDB file {lmdb_file} does not exist")
+    # special case for AgiBotWorld
+    if 'AgiBotWorld' in image_data['lmdb_file']:
+        return read_img_from_lmdb_v2(image_data)
+    
+    try:
+        env = lmdb.open(image_data['lmdb_file'], readonly=True, lock=False, max_readers=10240)
+    except Exception as e:
+        print(f"Failed to open lmdb file {image_data['lmdb_file']}. Error message: {e}", flush=True)
+        raise e
+
+    with env.begin(write=False) as txn:
+        try:
+            image_bin = txn.get(image_data['lmdb_key'].encode('ascii'))
+            buf = BytesIO(image_bin)
+        except Exception as e:
+            print(f"Failed to get image from lmdb file {image_data['lmdb_file']}. Error message: {e}", flush=True)
+            raise e
+    try:
+        image = Image.open(buf)
+    except Exception as e:
+        image_np = np.frombuffer(image_bin, dtype=np.uint8)
+        image_bgr = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
+        image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+        image = Image.fromarray(image_rgb)
+    return image
+
+def fetch_image(ele: dict[str, str | Image.Image]) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif isinstance(image, dict) and 'lmdb_file' in image:
+        image_obj = parse_lmdb_image_data(image)
+    elif image.startswith("http://") or image.startswith("https://"):
+        response = requests.get(image, stream=True)
+        image_obj = Image.open(BytesIO(response.content))
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
+    image = to_rgb(image_obj)
+
+    return image
+
+
+def get_video_frame_indices(
+    ele: dict,
+    total_frames: int,
+    video_fps: int | float,
+) -> tuple[torch.Tensor, float]:
+    target_fps = ele.get("fps", FPS)
+    max_frames = ele.get("max_frames", MAX_FRAMES)
+    
+    nframes = (total_frames / video_fps) * target_fps
+    nframes = int(round(nframes))
+    nframes = max(1, nframes)
+    
+    if nframes > max_frames:
+        nframes = max_frames
+    
+    nframes = min(nframes, total_frames)
+    
+    if nframes == total_frames:
+        idx = torch.arange(total_frames).long()
+    else:
+        idx = torch.linspace(0, total_frames - 1, nframes).round().long()
+        
+    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
+    
+    return idx, sample_fps
+
+def _read_video_torchvision(
+    ele: dict,
+) -> (torch.Tensor, float, list):
+    """read video using torchvision.io.read_video and return also per-frame timestamps"""
+    video_path = ele["video"]
+    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
+        if "http://" in video_path or "https://" in video_path:
+            warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.")
+        if "file://" in video_path:
+            video_path = video_path[7:]
+    st = time.time()
+    
+    video, audio, info = io.read_video(
+        video_path,
+        start_pts=ele.get("video_start", 0.0),
+        end_pts=ele.get("video_end", None),
+        pts_unit="sec",
+        output_format="TCHW",
+    )
+    total_frames, video_fps = video.size(0), info["video_fps"]
+    logger.info(f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+    
+    idx, sample_fps = get_video_frame_indices(ele, total_frames, video_fps)
+    
+    start_time = ele.get("video_start", 0.0)
+    timestamps = (start_time + idx.to(torch.float32) / video_fps).tolist()
+    
+    video = video[idx]
+    return video, sample_fps, timestamps
+
+
+def is_decord_available() -> bool:
+    import importlib.util
+    return importlib.util.find_spec("decord") is not None
+
+def _read_video_decord(
+    ele: dict,
+) -> (torch.Tensor, float, list):
+    """read video using decord.VideoReader and return also per-frame timestamps"""
+    video_path = ele["video"]
+    st = time.time()
+    vr = decord.VideoReader(video_path)
+    
+    total_frames, video_fps = len(vr), vr.get_avg_fps()
+    logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+    
+    idx_tensor, sample_fps = get_video_frame_indices(ele, total_frames, video_fps)
+    idx = idx_tensor.tolist()
+    
+    start_time = ele.get("video_start", 0.0)
+    timestamps = [start_time + i / video_fps for i in idx]
+    
+    video = vr.get_batch(idx).asnumpy()
+    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
+    
+    return video, sample_fps, timestamps
+
+
+VIDEO_READER_BACKENDS = {
+    "decord": _read_video_decord,
+    "torchvision": _read_video_torchvision,
+}
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+    if is_decord_available():
+        video_reader_backend = "decord"
+    else:
+        video_reader_backend = "torchvision"
+    return video_reader_backend
+
+
+def fetch_video(ele: dict, return_video_sample_fps: bool = False, video_reader_backend: str = "torchvision") -> torch.Tensor | list[Image.Image]:
+    """
+    Fetches video, samples frames, resizes based on video_total_pixels, and returns as Tensor (TCHW).
+    """
+    if isinstance(ele["video"], str):
+        video_reader_backend = video_reader_backend if video_reader_backend is not None else get_video_reader_backend()
+        try:
+            video, sample_fps, timestamps = VIDEO_READER_BACKENDS[video_reader_backend](ele)
+        except Exception as e:
+            logger.warning(f"video_reader_backend {video_reader_backend} error, use torchvision as default, msg: {e}")
+            video, sample_fps, timestamps = VIDEO_READER_BACKENDS["torchvision"](ele)
+
+        nframes, _, height, width = video.shape
+        
+        video_total_pixels = ele.get("video_total_pixels", VIDEO_TOTAL_PIXELS)
+        current_pixels = nframes * height * width
+        
+        if current_pixels > video_total_pixels:
+            scale_factor = math.sqrt(video_total_pixels / current_pixels)
+            new_height = int(height * scale_factor)
+            new_width = int(width * scale_factor)
+            
+            video = transforms.functional.resize(
+                video,
+                [new_height, new_width],
+                interpolation=InterpolationMode.BICUBIC,
+                antialias=True,
+            ).float()
+        else:
+            video = video.float()
+
+        if return_video_sample_fps:
+            return video, sample_fps, timestamps
+        return video
+
+    else:
+        assert isinstance(ele["video"], (list, tuple))
+        process_info = ele.copy()
+        process_info.pop("type", None)
+        process_info.pop("video", None)
+        
+        images = [
+            fetch_image({"image": video_element, **process_info})
+            for video_element in ele["video"]
+        ]
+        
+        nframes = len(images)
+        timestamps = [-1 for i in range(nframes)] 
+        
+        # For list of images, we return list of PIL images directly, 
+        # the processor will handle conversion to tensor later.
+        if return_video_sample_fps:
+            return images, process_info.get("fps", 2.0), timestamps
+        return images
+
+class LocateAnythingProcessorKwargs(ProcessingKwargs, total=False):
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+        },
+        "images_kwargs": {},
+        "videos_kwargs": {},
+    }
+
+
+class LocateAnythingProcessor(ProcessorMixin):
+    attributes = ["image_processor", "tokenizer"]
+    valid_kwargs = [
+        "chat_template",
+        "num_image_tokens",
+        "image_token",
+        "video_token",
+        "images_kwargs",
+        "videos_kwargs",
+        "text_kwargs",
+    ]
+    image_processor_class = "AutoImageProcessor" 
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor=None,
+        tokenizer=None,
+        chat_template=None,
+        image_token='<IMG_CONTEXT>',
+        video_token='<IMG_CONTEXT>',
+        merge_kernel_size=[2, 2], # Note: This might need adjustment based on your patch_size (14*14)
+        image_placeholder='image',
+        video_placeholder='video',
+        image_start_token='<img>',
+        image_end_token='</img>',
+        **kwargs,
+    ):  
+        self.image_token = tokenizer.image_token if hasattr(tokenizer, "image_token") else image_token
+        self.video_token = tokenizer.video_token if hasattr(tokenizer, "video_token") else video_token
+        self.image_token_id = (
+            tokenizer.image_token_id
+            if getattr(tokenizer, "image_token_id", None)
+            else tokenizer.convert_tokens_to_ids(self.image_token)
+        )
+        self.video_token_id = (
+            tokenizer.video_token_id
+            if getattr(tokenizer, "video_token_id", None)
+            else tokenizer.convert_tokens_to_ids(self.video_token)
+        )
+        self.image_placeholder = image_placeholder
+        self.video_placeholder = video_placeholder
+        self.merge_kernel_size = merge_kernel_size
+        self.image_start_token = image_start_token
+        self.image_end_token = image_end_token
+        if 'auto_map' in kwargs:
+            self.auto_map = kwargs['auto_map']
+        super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+    
+    def replace_media_placeholder(self, text, image_list, video_list, timestamps_list, fps_list, **output_kwargs):
+
+        num_of_images_in_this_sample = 0
+        num_of_videos_in_this_sample = 0
+        pattern = re.compile(rf"<({self.image_placeholder}|{self.video_placeholder})-(\d+)>")
+        unified_frame_list = []
+        
+        def replace_in_text(text):
+            def repl(match):
+                nonlocal unified_frame_list
+                nonlocal num_of_images_in_this_sample
+                nonlocal num_of_videos_in_this_sample
+                media_type = match.group(1)
+                idx_in_list = int(match.group(2)) - 1
+                idx_mapper = {0: "first", 1: "second", 2: "third", 3: "fourth", 4: "fifth", 5: "sixth", 6: "seventh", 7: "eighth", 8: "ninth", 9: "tenth"}  
+                
+                if media_type == 'image':
+                    # Call LocateAnythingImageProcessor with a single image in a list
+                    image_inputs = self.image_processor(images=[image_list[idx_in_list]], **output_kwargs["images_kwargs"])
+                    
+                    num_of_tokens_list = [int(h * w) // (self.image_processor.merge_kernel_size[0] * self.image_processor.merge_kernel_size[1]) for h, w in image_inputs['image_grid_hws']]
+                    
+                    special_placeholder = f"<image {idx_in_list+1}>{self.image_start_token}{self.image_token * num_of_tokens_list[0]}{self.image_end_token}"
+                    unified_frame_list.append(image_inputs)
+                    num_of_images_in_this_sample += 1
+                    
+                elif media_type == 'video':
+                    video_obj = video_list[idx_in_list]
+                    
+                    # Convert Tensor TCHW to list of PIL Images for the ImageProcessor
+                    if isinstance(video_obj, torch.Tensor):
+                        # video_obj is [T, C, H, W], float, likely 0-255 or standardized
+                        # LocateAnythingImageProcessor expects PIL or 0-255 inputs usually.
+                        # We need to convert back to PIL or List[Tensor] compatible with make_list_of_images
+                        video_frames = []
+                        for i in range(video_obj.shape[0]):
+                            frame = video_obj[i] # [C, H, W]
+                            # Assuming fetch_video returns float tensors.
+                            # If they are 0-255, convert to uint8.
+                            if frame.dtype.is_floating_point and frame.max() > 1.0:
+                                frame = frame.byte()
+                            elif frame.dtype.is_floating_point:
+                                frame = (frame * 255).byte()
+                            
+                            img = transforms.ToPILImage()(frame)
+                            video_frames.append(img)
+                    elif isinstance(video_obj, list):
+                        # Already list of PIL images
+                        video_frames = video_obj
+                    else:
+                        raise ValueError("Unsupported video format")
+
+                    # Call ImageProcessor with list of frames
+                    video_inputs = self.image_processor(images=video_frames, **output_kwargs["videos_kwargs"])
+                    
+                    # Calculate tokens per frame
+                    num_of_tokens_list = [int(h * w) // (self.image_processor.merge_kernel_size[0] * self.image_processor.merge_kernel_size[1]) for h, w in video_inputs['image_grid_hws']]
+                    
+                    if timestamps_list is not None and -1 not in timestamps_list:
+                        frame_timestamps = timestamps_list[idx_in_list]
+                    else:
+                        frame_timestamps = None
+                    sampled_fps = fps_list[idx_in_list] if fps_list is not None else None
+                    
+                    if frame_timestamps is not None:
+                        # Ensure lengths match (sometimes rounding might cause off-by-one if not careful, but usually safe here)
+                        if len(frame_timestamps) != len(num_of_tokens_list):
+                             logger.warning(f"Timestamp mismatch: {len(frame_timestamps)} vs {len(num_of_tokens_list)}")
+                             min_len = min(len(frame_timestamps), len(num_of_tokens_list))
+                             frame_timestamps = frame_timestamps[:min_len]
+                             num_of_tokens_list = num_of_tokens_list[:min_len]
+
+                        special_placeholder = [f"Frame-{i+1}-{frame_timestamps[i]:.2f}s: {self.image_start_token}{self.image_token * num_of_tokens}{self.image_end_token}" for i, num_of_tokens in enumerate(num_of_tokens_list)]
+                    else:
+                        special_placeholder = [f"Frame-{i+1}: {self.image_start_token}{self.image_token * num_of_tokens}{self.image_end_token}" for i, num_of_tokens in enumerate(num_of_tokens_list)]
+                    
+                    if sampled_fps is not None:
+                        special_placeholder = f"The {idx_mapper[idx_in_list]} video sampled with {sampled_fps:.2f} fps: " + "".join(special_placeholder)
+                    else:
+                        special_placeholder = f"The {idx_mapper[idx_in_list]} video: " + "".join(special_placeholder)
+                    
+                    unified_frame_list.append(video_inputs)
+                    num_of_videos_in_this_sample += 1
+                else:
+                    raise ValueError(f'Unknown media type: {media_type}')
+                return special_placeholder
+            return pattern.sub(repl, text)
+        
+        text = replace_in_text(text)
+        
+        if len(unified_frame_list) > 0:
+            # Concatenate all pixel values from all images/videos in this sample
+            pixel_values = torch.cat([frame['pixel_values'] for frame in unified_frame_list], dim=0)
+            # Concatenate grid hws
+            image_grid_hws = np.concatenate([frame['image_grid_hws'] for frame in unified_frame_list], axis=0)
+        else:
+            pixel_values = torch.empty(0)
+            image_grid_hws = np.empty(0)
+            
+        return text, pixel_values, image_grid_hws, num_of_images_in_this_sample, num_of_videos_in_this_sample
+    
+    def __call__(
+        self,
+        images: ImageInput = None,
+        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
+        audio=None,
+        videos: VideoInput = None,
+        **kwargs: Unpack[LocateAnythingProcessorKwargs],
+    ) -> BatchFeature:
+        output_kwargs = self._merge_kwargs(
+            LocateAnythingProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+        
+        if isinstance(text, str):
+            text_list = [text]
+        elif not isinstance(text, list) and not isinstance(text[0], str):
+            raise ValueError("Invalid input text. Please provide a string, or a list of strings")
+        elif isinstance(text, list) and isinstance(text[0], str):
+            text_list = text
+        
+        if images is None: images = []
+        if videos is None: videos = []
+        
+        pixel_values_list = []
+        image_grid_hws_list = []
+        new_sample_list = []
+        image_start_idx = 0
+        video_start_idx = 0
+        timestamps_batch = output_kwargs['videos_kwargs'].pop("timestamps", None)
+        fps_batch = output_kwargs['videos_kwargs'].pop("fps", None)
+        
+        for sample in text_list:
+            timestamps_list = timestamps_batch[video_start_idx:] if timestamps_batch is not None else None
+            fps_list = fps_batch[video_start_idx:] if fps_batch is not None else None
+            
+            sample, pixel_values, image_grid_hws, num_of_images_in_this_sample, num_of_videos_in_this_sample = self.replace_media_placeholder(
+                sample, images[image_start_idx:], videos[video_start_idx:], timestamps_list, fps_list, **output_kwargs
+            )
+            new_sample_list.append(sample)
+            
+            if pixel_values.numel() > 0:
+                pixel_values_list.append(pixel_values)
+                image_grid_hws_list.append(image_grid_hws)
+
+            image_start_idx += num_of_images_in_this_sample
+            video_start_idx += num_of_videos_in_this_sample
+
+        image_inputs = {}
+        if len(pixel_values_list) > 0:
+            # Concatenate across the batch
+            image_inputs['pixel_values'] = torch.cat(pixel_values_list, dim=0)
+            image_inputs['image_grid_hws'] = np.concatenate(image_grid_hws_list, axis=0)
+        
+        video_inputs = {} # Video data is merged into image_inputs now
+        text_inputs = self.tokenizer(new_sample_list, **output_kwargs["text_kwargs"])
+        
+        return BatchFeature(data={**text_inputs, **image_inputs, **video_inputs})
+
+    def batch_decode(self, *args, **kwargs):
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
+
+    def save_pretrained(self, save_directory, **kwargs):
+        if os.path.isfile(save_directory):
+            raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file")
+        os.makedirs(save_directory, exist_ok=True)
+        outputs = super().save_pretrained(save_directory, **kwargs)
+        return outputs
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        processor = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
+        if isinstance(processor, tuple):
+            processor = processor[0]
+        return processor
+
+    def process_vision_info(
+        self,
+        conversations: list[dict] | list[list[dict]],
+        return_video_kwargs: bool = False,
+        video_reader_backend: str = "torchvision",
+    ) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None, Optional[dict]]:
+
+        vision_infos = self.extract_vision_info(conversations)
+        image_inputs = []
+        video_inputs = []
+        video_sample_fps_list = []
+        video_timestamps_list = []
+        
+        for vision_info in vision_infos:
+            if "image" in vision_info or "image_url" in vision_info:
+                image_inputs.append(fetch_image(vision_info))
+            elif "video" in vision_info:
+                video_input, video_sample_fps, video_timestamps = fetch_video(vision_info, return_video_sample_fps=True, video_reader_backend=video_reader_backend)
+                video_sample_fps_list.append(video_sample_fps)
+                video_inputs.append(video_input)
+                video_timestamps_list.append(video_timestamps)
+            else:
+                raise ValueError("image, image_url or video should in content.")
+        
+        if len(image_inputs) == 0:
+            image_inputs = None
+        if len(video_inputs) == 0:
+            video_inputs = None
+            
+        if return_video_kwargs:
+            return image_inputs, video_inputs, {'fps': video_sample_fps_list, 'timestamps': video_timestamps_list}
+        return image_inputs, video_inputs
+
+    def extract_vision_info(self, conversations: list[dict] | list[list[dict]]) -> list[dict]:
+        vision_infos = []
+        if isinstance(conversations[0], dict):
+            conversations = [conversations]
+        for conversation in conversations:
+            for message in conversation:
+                if isinstance(message["content"], list):
+                    for ele in message["content"]:
+                        if (
+                            "image" in ele
+                            or "image_url" in ele
+                            or "video" in ele
+                            or ele["type"] in ("image", "image_url", "video")
+                        ):
+                            vision_infos.append(ele)
+        return vision_infos
+    
+    def py_apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False):
+        assert tokenize == False, "tokenize is not supported yet"
+        result = ""
+        image_count = 0
+        video_count = 0
+        
+        message_text = ""
+        for idx, message in enumerate(messages):
+            if message.get('role') != 'user': continue
+            content = message.get('content')
+            if isinstance(content, str):
+                message_text += content
+            elif isinstance(content, list):
+                for item in content:
+                    if isinstance(item, dict) and "text" in item:
+                        message_text += item["text"]
+                    elif isinstance(item, str):
+                        message_text += item
+                        
+        for idx, message in enumerate(messages):
+            if idx == 0 and message.get('role') != 'system':
+                result += "<|im_start|>system\n"
+                result += "You are a helpful assistant.\n"
+                result += "<|im_end|>\n"
+
+            result += f"<|im_start|>{message.get('role', '')}\n"
+            content = message.get('content')
+
+            if isinstance(content, str):
+                result += content
+                result += "<|im_end|>\n"
+            else:
+                for item in content:
+                    if (isinstance(item, dict) and (item.get('type') == 'image' or 'image' in item or 'image_url' in item)):
+                        image_count += 1
+                        candidate_token = f"<image-{image_count}>"
+                        if candidate_token not in message_text:
+                            result += candidate_token
+                    elif (isinstance(item, dict) and (item.get('type') == 'video' or 'video' in item)):
+                        video_count += 1
+                        candidate_token = f"<video-{video_count}>"
+                        if candidate_token not in message_text:
+                            result += candidate_token
+                    elif isinstance(item, dict) and 'text' in item:
+                        result += item['text']
+                    elif isinstance(item, str):
+                        result += item
+                result += "<|im_end|>\n"
+
+        if add_generation_prompt:
+            result += "<|im_start|>assistant\n"
+
+        return result
+
+
+    @classmethod
+    def from_args_and_dict(cls, args, processor_dict: dict[str, Any], **kwargs):
+        processor_dict = processor_dict.copy()
+        return_unused_kwargs = kwargs.pop("return_unused_kwargs", False)
+
+        if "processor_class" in processor_dict:
+            del processor_dict["processor_class"]
+
+        unused_kwargs = cls.validate_init_kwargs(processor_config=processor_dict, valid_kwargs=cls.valid_kwargs)
+        processor = cls(*args, **processor_dict)
+
+        for key in set(kwargs.keys()):
+            if hasattr(processor, key):
+                setattr(processor, key, kwargs.pop(key))
+
+        if isinstance(unused_kwargs, dict):
+            kwargs.update(unused_kwargs)
+        logger.info(f"Processor {processor}")
+        if return_unused_kwargs:
+            return processor, kwargs
+        else:
+            return processor
+        
+        
+__all__ = ["LocateAnythingProcessor"]

processor_config.json

@@ -0,0 +1,18 @@
+{
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_locateanything.LocateAnythingImageProcessor",
+    "AutoProcessor": "processing_locateanything.LocateAnythingProcessor"
+  },
+  "image_end_token": "</img>",
+  "image_placeholder": "image",
+  "image_start_token": "<img>",
+  "image_token": "<IMG_CONTEXT>",
+  "merge_kernel_size": [
+    2,
+    2
+  ],
+  "processor_class": "LocateAnythingProcessor",
+  "video_placeholder": "video",
+  "video_token": "<IMG_CONTEXT>",
+  "patch_size": 14
+}

special_tokens_map.json

@@ -0,0 +1,1053 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<IMG_CONTEXT>",
+    "<img>",
+    "</img>",
+    "<box>",
+    "</box>",
+    "<quad>",
+    "</quad>",
+    "<ref>",
+    "</ref>",
+    "<interval>",
+    "</interval>",
+    "<text_mask>",
+    "<0>",
+    "<1>",
+    "<2>",
+    "<3>",
+    "<4>",
+    "<5>",
+    "<6>",
+    "<7>",
+    "<8>",
+    "<9>",
+    "<10>",
+    "<11>",
+    "<12>",
+    "<13>",
+    "<14>",
+    "<15>",
+    "<16>",
+    "<17>",
+    "<18>",
+    "<19>",
+    "<20>",
+    "<21>",
+    "<22>",
+    "<23>",
+    "<24>",
+    "<25>",
+    "<26>",
+    "<27>",
+    "<28>",
+    "<29>",
+    "<30>",
+    "<31>",
+    "<32>",
+    "<33>",
+    "<34>",
+    "<35>",
+    "<36>",
+    "<37>",
+    "<38>",
+    "<39>",
+    "<40>",
+    "<41>",
+    "<42>",
+    "<43>",
+    "<44>",
+    "<45>",
+    "<46>",
+    "<47>",
+    "<48>",
+    "<49>",
+    "<50>",
+    "<51>",
+    "<52>",
+    "<53>",
+    "<54>",
+    "<55>",
+    "<56>",
+    "<57>",
+    "<58>",
+    "<59>",
+    "<60>",
+    "<61>",
+    "<62>",
+    "<63>",
+    "<64>",
+    "<65>",
+    "<66>",
+    "<67>",
+    "<68>",
+    "<69>",
+    "<70>",
+    "<71>",
+    "<72>",
+    "<73>",
+    "<74>",
+    "<75>",
+    "<76>",
+    "<77>",
+    "<78>",
+    "<79>",
+    "<80>",
+    "<81>",
+    "<82>",
+    "<83>",
+    "<84>",
+    "<85>",
+    "<86>",
+    "<87>",
+    "<88>",
+    "<89>",
+    "<90>",
+    "<91>",
+    "<92>",
+    "<93>",
+    "<94>",
+    "<95>",
+    "<96>",
+    "<97>",
+    "<98>",
+    "<99>",
+    "<100>",
+    "<101>",
+    "<102>",
+    "<103>",
+    "<104>",
+    "<105>",
+    "<106>",
+    "<107>",
+    "<108>",
+    "<109>",
+    "<110>",
+    "<111>",
+    "<112>",
+    "<113>",
+    "<114>",
+    "<115>",
+    "<116>",
+    "<117>",
+    "<118>",
+    "<119>",
+    "<120>",
+    "<121>",
+    "<122>",
+    "<123>",
+    "<124>",
+    "<125>",
+    "<126>",
+    "<127>",
+    "<128>",
+    "<129>",
+    "<130>",
+    "<131>",
+    "<132>",
+    "<133>",
+    "<134>",
+    "<135>",
+    "<136>",
+    "<137>",
+    "<138>",
+    "<139>",
+    "<140>",
+    "<141>",
+    "<142>",
+    "<143>",
+    "<144>",
+    "<145>",
+    "<146>",
+    "<147>",
+    "<148>",
+    "<149>",
+    "<150>",
+    "<151>",
+    "<152>",
+    "<153>",
+    "<154>",
+    "<155>",
+    "<156>",
+    "<157>",
+    "<158>",
+    "<159>",
+    "<160>",
+    "<161>",
+    "<162>",
+    "<163>",
+    "<164>",
+    "<165>",
+    "<166>",
+    "<167>",
+    "<168>",
+    "<169>",
+    "<170>",
+    "<171>",
+    "<172>",
+    "<173>",
+    "<174>",
+    "<175>",
+    "<176>",
+    "<177>",
+    "<178>",
+    "<179>",
+    "<180>",
+    "<181>",
+    "<182>",
+    "<183>",
+    "<184>",
+    "<185>",
+    "<186>",
+    "<187>",
+    "<188>",
+    "<189>",
+    "<190>",
+    "<191>",
+    "<192>",
+    "<193>",
+    "<194>",
+    "<195>",
+    "<196>",
+    "<197>",
+    "<198>",
+    "<199>",
+    "<200>",
+    "<201>",
+    "<202>",
+    "<203>",
+    "<204>",
+    "<205>",
+    "<206>",
+    "<207>",
+    "<208>",
+    "<209>",
+    "<210>",
+    "<211>",
+    "<212>",
+    "<213>",
+    "<214>",
+    "<215>",
+    "<216>",
+    "<217>",
+    "<218>",
+    "<219>",
+    "<220>",
+    "<221>",
+    "<222>",
+    "<223>",
+    "<224>",
+    "<225>",
+    "<226>",
+    "<227>",
+    "<228>",
+    "<229>",
+    "<230>",
+    "<231>",
+    "<232>",
+    "<233>",
+    "<234>",
+    "<235>",
+    "<236>",
+    "<237>",
+    "<238>",
+    "<239>",
+    "<240>",
+    "<241>",
+    "<242>",
+    "<243>",
+    "<244>",
+    "<245>",
+    "<246>",
+    "<247>",
+    "<248>",
+    "<249>",
+    "<250>",
+    "<251>",
+    "<252>",
+    "<253>",
+    "<254>",
+    "<255>",
+    "<256>",
+    "<257>",
+    "<258>",
+    "<259>",
+    "<260>",
+    "<261>",
+    "<262>",
+    "<263>",
+    "<264>",
+    "<265>",
+    "<266>",
+    "<267>",
+    "<268>",
+    "<269>",
+    "<270>",
+    "<271>",
+    "<272>",
+    "<273>",
+    "<274>",
+    "<275>",
+    "<276>",
+    "<277>",
+    "<278>",
+    "<279>",
+    "<280>",
+    "<281>",
+    "<282>",
+    "<283>",
+    "<284>",
+    "<285>",
+    "<286>",
+    "<287>",
+    "<288>",
+    "<289>",
+    "<290>",
+    "<291>",
+    "<292>",
+    "<293>",
+    "<294>",
+    "<295>",
+    "<296>",
+    "<297>",
+    "<298>",
+    "<299>",
+    "<300>",
+    "<301>",
+    "<302>",
+    "<303>",
+    "<304>",
+    "<305>",
+    "<306>",
+    "<307>",
+    "<308>",
+    "<309>",
+    "<310>",
+    "<311>",
+    "<312>",
+    "<313>",
+    "<314>",
+    "<315>",
+    "<316>",
+    "<317>",
+    "<318>",
+    "<319>",
+    "<320>",
+    "<321>",
+    "<322>",
+    "<323>",
+    "<324>",
+    "<325>",
+    "<326>",
+    "<327>",
+    "<328>",
+    "<329>",
+    "<330>",
+    "<331>",
+    "<332>",
+    "<333>",
+    "<334>",
+    "<335>",
+    "<336>",
+    "<337>",
+    "<338>",
+    "<339>",
+    "<340>",
+    "<341>",
+    "<342>",
+    "<343>",
+    "<344>",
+    "<345>",
+    "<346>",
+    "<347>",
+    "<348>",
+    "<349>",
+    "<350>",
+    "<351>",
+    "<352>",
+    "<353>",
+    "<354>",
+    "<355>",
+    "<356>",
+    "<357>",
+    "<358>",
+    "<359>",
+    "<360>",
+    "<361>",
+    "<362>",
+    "<363>",
+    "<364>",
+    "<365>",
+    "<366>",
+    "<367>",
+    "<368>",
+    "<369>",
+    "<370>",
+    "<371>",
+    "<372>",
+    "<373>",
+    "<374>",
+    "<375>",
+    "<376>",
+    "<377>",
+    "<378>",
+    "<379>",
+    "<380>",
+    "<381>",
+    "<382>",
+    "<383>",
+    "<384>",
+    "<385>",
+    "<386>",
+    "<387>",
+    "<388>",
+    "<389>",
+    "<390>",
+    "<391>",
+    "<392>",
+    "<393>",
+    "<394>",
+    "<395>",
+    "<396>",
+    "<397>",
+    "<398>",
+    "<399>",
+    "<400>",
+    "<401>",
+    "<402>",
+    "<403>",
+    "<404>",
+    "<405>",
+    "<406>",
+    "<407>",
+    "<408>",
+    "<409>",
+    "<410>",
+    "<411>",
+    "<412>",
+    "<413>",
+    "<414>",
+    "<415>",
+    "<416>",
+    "<417>",
+    "<418>",
+    "<419>",
+    "<420>",
+    "<421>",
+    "<422>",
+    "<423>",
+    "<424>",
+    "<425>",
+    "<426>",
+    "<427>",
+    "<428>",
+    "<429>",
+    "<430>",
+    "<431>",
+    "<432>",
+    "<433>",
+    "<434>",
+    "<435>",
+    "<436>",
+    "<437>",
+    "<438>",
+    "<439>",
+    "<440>",
+    "<441>",
+    "<442>",
+    "<443>",
+    "<444>",
+    "<445>",
+    "<446>",
+    "<447>",
+    "<448>",
+    "<449>",
+    "<450>",
+    "<451>",
+    "<452>",
+    "<453>",
+    "<454>",
+    "<455>",
+    "<456>",
+    "<457>",
+    "<458>",
+    "<459>",
+    "<460>",
+    "<461>",
+    "<462>",
+    "<463>",
+    "<464>",
+    "<465>",
+    "<466>",
+    "<467>",
+    "<468>",
+    "<469>",
+    "<470>",
+    "<471>",
+    "<472>",
+    "<473>",
+    "<474>",
+    "<475>",
+    "<476>",
+    "<477>",
+    "<478>",
+    "<479>",
+    "<480>",
+    "<481>",
+    "<482>",
+    "<483>",
+    "<484>",
+    "<485>",
+    "<486>",
+    "<487>",
+    "<488>",
+    "<489>",
+    "<490>",
+    "<491>",
+    "<492>",
+    "<493>",
+    "<494>",
+    "<495>",
+    "<496>",
+    "<497>",
+    "<498>",
+    "<499>",
+    "<500>",
+    "<501>",
+    "<502>",
+    "<503>",
+    "<504>",
+    "<505>",
+    "<506>",
+    "<507>",
+    "<508>",
+    "<509>",
+    "<510>",
+    "<511>",
+    "<512>",
+    "<513>",
+    "<514>",
+    "<515>",
+    "<516>",
+    "<517>",
+    "<518>",
+    "<519>",
+    "<520>",
+    "<521>",
+    "<522>",
+    "<523>",
+    "<524>",
+    "<525>",
+    "<526>",
+    "<527>",
+    "<528>",
+    "<529>",
+    "<530>",
+    "<531>",
+    "<532>",
+    "<533>",
+    "<534>",
+    "<535>",
+    "<536>",
+    "<537>",
+    "<538>",
+    "<539>",
+    "<540>",
+    "<541>",
+    "<542>",
+    "<543>",
+    "<544>",
+    "<545>",
+    "<546>",
+    "<547>",
+    "<548>",
+    "<549>",
+    "<550>",
+    "<551>",
+    "<552>",
+    "<553>",
+    "<554>",
+    "<555>",
+    "<556>",
+    "<557>",
+    "<558>",
+    "<559>",
+    "<560>",
+    "<561>",
+    "<562>",
+    "<563>",
+    "<564>",
+    "<565>",
+    "<566>",
+    "<567>",
+    "<568>",
+    "<569>",
+    "<570>",
+    "<571>",
+    "<572>",
+    "<573>",
+    "<574>",
+    "<575>",
+    "<576>",
+    "<577>",
+    "<578>",
+    "<579>",
+    "<580>",
+    "<581>",
+    "<582>",
+    "<583>",
+    "<584>",
+    "<585>",
+    "<586>",
+    "<587>",
+    "<588>",
+    "<589>",
+    "<590>",
+    "<591>",
+    "<592>",
+    "<593>",
+    "<594>",
+    "<595>",
+    "<596>",
+    "<597>",
+    "<598>",
+    "<599>",
+    "<600>",
+    "<601>",
+    "<602>",
+    "<603>",
+    "<604>",
+    "<605>",
+    "<606>",
+    "<607>",
+    "<608>",
+    "<609>",
+    "<610>",
+    "<611>",
+    "<612>",
+    "<613>",
+    "<614>",
+    "<615>",
+    "<616>",
+    "<617>",
+    "<618>",
+    "<619>",
+    "<620>",
+    "<621>",
+    "<622>",
+    "<623>",
+    "<624>",
+    "<625>",
+    "<626>",
+    "<627>",
+    "<628>",
+    "<629>",
+    "<630>",
+    "<631>",
+    "<632>",
+    "<633>",
+    "<634>",
+    "<635>",
+    "<636>",
+    "<637>",
+    "<638>",
+    "<639>",
+    "<640>",
+    "<641>",
+    "<642>",
+    "<643>",
+    "<644>",
+    "<645>",
+    "<646>",
+    "<647>",
+    "<648>",
+    "<649>",
+    "<650>",
+    "<651>",
+    "<652>",
+    "<653>",
+    "<654>",
+    "<655>",
+    "<656>",
+    "<657>",
+    "<658>",
+    "<659>",
+    "<660>",
+    "<661>",
+    "<662>",
+    "<663>",
+    "<664>",
+    "<665>",
+    "<666>",
+    "<667>",
+    "<668>",
+    "<669>",
+    "<670>",
+    "<671>",
+    "<672>",
+    "<673>",
+    "<674>",
+    "<675>",
+    "<676>",
+    "<677>",
+    "<678>",
+    "<679>",
+    "<680>",
+    "<681>",
+    "<682>",
+    "<683>",
+    "<684>",
+    "<685>",
+    "<686>",
+    "<687>",
+    "<688>",
+    "<689>",
+    "<690>",
+    "<691>",
+    "<692>",
+    "<693>",
+    "<694>",
+    "<695>",
+    "<696>",
+    "<697>",
+    "<698>",
+    "<699>",
+    "<700>",
+    "<701>",
+    "<702>",
+    "<703>",
+    "<704>",
+    "<705>",
+    "<706>",
+    "<707>",
+    "<708>",
+    "<709>",
+    "<710>",
+    "<711>",
+    "<712>",
+    "<713>",
+    "<714>",
+    "<715>",
+    "<716>",
+    "<717>",
+    "<718>",
+    "<719>",
+    "<720>",
+    "<721>",
+    "<722>",
+    "<723>",
+    "<724>",
+    "<725>",
+    "<726>",
+    "<727>",
+    "<728>",
+    "<729>",
+    "<730>",
+    "<731>",
+    "<732>",
+    "<733>",
+    "<734>",
+    "<735>",
+    "<736>",
+    "<737>",
+    "<738>",
+    "<739>",
+    "<740>",
+    "<741>",
+    "<742>",
+    "<743>",
+    "<744>",
+    "<745>",
+    "<746>",
+    "<747>",
+    "<748>",
+    "<749>",
+    "<750>",
+    "<751>",
+    "<752>",
+    "<753>",
+    "<754>",
+    "<755>",
+    "<756>",
+    "<757>",
+    "<758>",
+    "<759>",
+    "<760>",
+    "<761>",
+    "<762>",
+    "<763>",
+    "<764>",
+    "<765>",
+    "<766>",
+    "<767>",
+    "<768>",
+    "<769>",
+    "<770>",
+    "<771>",
+    "<772>",
+    "<773>",
+    "<774>",
+    "<775>",
+    "<776>",
+    "<777>",
+    "<778>",
+    "<779>",
+    "<780>",
+    "<781>",
+    "<782>",
+    "<783>",
+    "<784>",
+    "<785>",
+    "<786>",
+    "<787>",
+    "<788>",
+    "<789>",
+    "<790>",
+    "<791>",
+    "<792>",
+    "<793>",
+    "<794>",
+    "<795>",
+    "<796>",
+    "<797>",
+    "<798>",
+    "<799>",
+    "<800>",
+    "<801>",
+    "<802>",
+    "<803>",
+    "<804>",
+    "<805>",
+    "<806>",
+    "<807>",
+    "<808>",
+    "<809>",
+    "<810>",
+    "<811>",
+    "<812>",
+    "<813>",
+    "<814>",
+    "<815>",
+    "<816>",
+    "<817>",
+    "<818>",
+    "<819>",
+    "<820>",
+    "<821>",
+    "<822>",
+    "<823>",
+    "<824>",
+    "<825>",
+    "<826>",
+    "<827>",
+    "<828>",
+    "<829>",
+    "<830>",
+    "<831>",
+    "<832>",
+    "<833>",
+    "<834>",
+    "<835>",
+    "<836>",
+    "<837>",
+    "<838>",
+    "<839>",
+    "<840>",
+    "<841>",
+    "<842>",
+    "<843>",
+    "<844>",
+    "<845>",
+    "<846>",
+    "<847>",
+    "<848>",
+    "<849>",
+    "<850>",
+    "<851>",
+    "<852>",
+    "<853>",
+    "<854>",
+    "<855>",
+    "<856>",
+    "<857>",
+    "<858>",
+    "<859>",
+    "<860>",
+    "<861>",
+    "<862>",
+    "<863>",
+    "<864>",
+    "<865>",
+    "<866>",
+    "<867>",
+    "<868>",
+    "<869>",
+    "<870>",
+    "<871>",
+    "<872>",
+    "<873>",
+    "<874>",
+    "<875>",
+    "<876>",
+    "<877>",
+    "<878>",
+    "<879>",
+    "<880>",
+    "<881>",
+    "<882>",
+    "<883>",
+    "<884>",
+    "<885>",
+    "<886>",
+    "<887>",
+    "<888>",
+    "<889>",
+    "<890>",
+    "<891>",
+    "<892>",
+    "<893>",
+    "<894>",
+    "<895>",
+    "<896>",
+    "<897>",
+    "<898>",
+    "<899>",
+    "<900>",
+    "<901>",
+    "<902>",
+    "<903>",
+    "<904>",
+    "<905>",
+    "<906>",
+    "<907>",
+    "<908>",
+    "<909>",
+    "<910>",
+    "<911>",
+    "<912>",
+    "<913>",
+    "<914>",
+    "<915>",
+    "<916>",
+    "<917>",
+    "<918>",
+    "<919>",
+    "<920>",
+    "<921>",
+    "<922>",
+    "<923>",
+    "<924>",
+    "<925>",
+    "<926>",
+    "<927>",
+    "<928>",
+    "<929>",
+    "<930>",
+    "<931>",
+    "<932>",
+    "<933>",
+    "<934>",
+    "<935>",
+    "<936>",
+    "<937>",
+    "<938>",
+    "<939>",
+    "<940>",
+    "<941>",
+    "<942>",
+    "<943>",
+    "<944>",
+    "<945>",
+    "<946>",
+    "<947>",
+    "<948>",
+    "<949>",
+    "<950>",
+    "<951>",
+    "<952>",
+    "<953>",
+    "<954>",
+    "<955>",
+    "<956>",
+    "<957>",
+    "<958>",
+    "<959>",
+    "<960>",
+    "<961>",
+    "<962>",
+    "<963>",
+    "<964>",
+    "<965>",
+    "<966>",
+    "<967>",
+    "<968>",
+    "<969>",
+    "<970>",
+    "<971>",
+    "<972>",
+    "<973>",
+    "<974>",
+    "<975>",
+    "<976>",
+    "<977>",
+    "<978>",
+    "<979>",
+    "<980>",
+    "<981>",
+    "<982>",
+    "<983>",
+    "<984>",
+    "<985>",
+    "<986>",
+    "<987>",
+    "<988>",
+    "<989>",
+    "<990>",
+    "<991>",
+    "<992>",
+    "<993>",
+    "<994>",
+    "<995>",
+    "<996>",
+    "<997>",
+    "<998>",
+    "<999>",
+    "<1000>",
+    "<null>",
+    "<switch>",
+    {
+      "content": "</c>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,19 @@
+{
+  "add_prefix_space": false,
+  "auto_map": {
+    "AutoProcessor": "processing_locateanything.LocateAnythingProcessor"
+  },
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 16384,
+  "pad_token": "<|endoftext|>",
+  "processor_class": "LocateAnythingProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}