update model

.gitattributes

@@ -33,3 +33,5 @@ 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
+model-00001-of-00002.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00002-of-00002.safetensors filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,222 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+---
+# Megrez-3B-Omni: 首个端侧全模态理解开源模型
+<p align="center">
+    <img src="assets/megrez_logo.png" width="400"/>
+<p>
+<p align="center">
+        🔗 <a href="https://github.com/infinigence/Infini-Megrez-Omni">GitHub</a>&nbsp&nbsp | &nbsp&nbsp🏠 <a href="https://huggingface.co/spaces/Infinigence/Megrez-3B-Omni">Demo</a>&nbsp&nbsp | &nbsp&nbsp📖 <a href="assets/png/wechat-official.jpg">WeChat Official</a>&nbsp&nbsp | &nbsp&nbsp💬 <a href="assets/wechat-group.jpg">WeChat Groups</a>&nbsp&nbsp   
+</p>
+<h4 align="center">
+    <p>
+        <b>中文</b> | <a href="https://huggingface.co/Infinigence/Megrez-3B-Omni/blob/main/README_EN.md">English</a>
+    <p>
+</h4>
+
+
+## 模型简介
+Megrez-3B-Omni是由无问芯穹（[Infinigence AI](https://cloud.infini-ai.com/platform/ai)）研发的**端侧全模态**理解模型，基于无问大语言模型Megrez-3B-Instruct扩展，同时具备图片、文本、音频三种模态数据的理解分析能力，在三个方面均取得最优精度
+- 在图像理解方面，基于SigLip-400M构建图像Token，在OpenCompass榜单上（综合8个主流多模态评测基准）平均得分66.2，超越LLaVA-NeXT-Yi-34B等更大参数规模的模型。Megrez-3B-Omni也是在MME、MMMU、OCRBench等测试集上目前精度最高的图像理解模型之一，在场景理解、OCR等方面具有良好表现。
+- 在语言理解方面，Megrez-3B-Omni并未牺牲模型的文本处理能力，综合能力较单模态版本（Megrez-3B-Instruct）精度变化小于2%，保持在C-EVAL、MMLU/MMLU Pro、AlignBench等多个测试集上的最优精度优势，依然取得超越上一代14B模型的能力表现
+- 在语音理解方面，采用Qwen2-Audio/whisper-large-v3的Encoder作为语音输入，支持中英文语音输入及多轮对话，支持对输入图片的语音提问，根据语音指令直接响应文本，在多项基准任务上取得了领先的结果
+
+
+## 基础信息
+<table>
+  <thead>
+    <tr>
+      <th></th>
+      <th>Language Module</th>
+      <th>Vision Module</th>
+      <th>Audio Module</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>Architecture</td>
+      <td>Llama-2 with GQA</td>
+      <td>SigLip-SO400M</td>
+      <td>Whisper-large-v3
+(encoder-only)</td>
+    </tr>
+    <tr>
+      <td># Params (Backbone)</td>
+      <td>2.29B</td>
+      <td>0.42B</td>
+      <td>0.64B</td>
+    </tr>
+    <tr>
+      <td>Connector</td>
+      <td>-</td>
+      <td>Cross Attention</td>
+      <td>Linear</td>
+    </tr>
+    <tr>
+      <td># Params (Others)</td>
+      <td>Emb: 0.31B<br>Softmax: 0.31B</td>
+      <td>Connector: 0.036B</td>
+      <td>Connector: 0.003B</td>
+    </tr>
+    <tr>
+      <td># Params (Total)</td>
+      <td colspan="3">4B</td>
+    </tr>
+    <tr>
+      <td># Vocab Size</td>
+      <td>122880</td>
+      <td>64 tokens/slice</td>
+      <td>-</td>
+    </tr>
+    <tr>
+      <td>Context length</td>
+      <td colspan="3">4K tokens</td>
+    </tr>
+    <tr>
+      <td>Supported languages</td>
+      <td colspan="3">Chinese & English</td>
+    </tr>
+  </tbody>
+</table>
+
+### 图片理解能力
+- 上图为Megrez-3B-Omni与其他开源模型在主流图片多模态任务上的性能比较
+- 下图为Megrez-3B-Omni在OpenCompass测试集上表现，图片引用自： [InternVL 2.5 Blog Post](https://internvl.github.io/blog/2024-12-05-InternVL-2.5/)
+ <!-- <div style="display: flex; justify-content: space-between;">
+  <img src="assets/multitask.jpg" alt="Image 1" style="width: 45%;">
+  <img src="assets/opencompass.jpg" alt="Image 2" style="width: 45%;">
+</div> -->
+
+![Multitask](assets/multitask.jpg)
+![OpencompassBmk](assets/opencompass.jpg)
+
+| model                 | basemodel             | 发布时间       | OpenCompass | MME      | MMMU val  | OCRBench | MathVista | RealWorldQA | MMVet  | hallusionBench | MMB TEST (en) | MMB TEST (zh) | TextVQA val | AI2D_TEST | MMstar    | DocVQA_TEST |
+|-----------------------|-----------------------|----------------|--------------------|----------|-----------|----------|-----------------|-------------|--------|----------------|--------------|--------------|-------------|-----------|-----------|-------------|
+| **Megrez-3B-Omni**    | **Megrez-3B**         | **2024.12.16** | **66.2**           | **2315** | **51.89** | **82.8** | **62**          | **71.89**   | **60** | **50.12**      | **80.8**     | **82.3**     | **80.3**    | **82.05** | **60.46** | **91.62**   |
+| Qwen2-VL-2B-Instruct  | Qwen2-1.5B            | 2024.08.28     | 57.2               | 1872     | 41.1      | 79.4     | 43              | 62.9        | 49.5   | 41.7           | 74.9         | 73.5         | 79.7        | 74.7      | 48        | 90.1        |
+| InternVL2.5-2B        | Internlm2.5-1.8B-chat | 2024.12.06     | 59.9               | 2138     | 43.6      | 80.4     | 51.3            | 60.1        | 60.8   | 42.6           | 74.7         | 71.9         | 74.3        | 74.9      | 53.7      | 88.7        |
+| BlueLM-V-3B           | -                     | 2024.11.29     | 66.1               | -        | 45.1      | 82.9     | 60.8            | 66.7        | 61.8   | 48             | 83           | 80.5         | 78.4        | 85.3      | 62.3      | 87.8        |
+| InternVL2.5-4B        | Qwen2.5-3B-Instruct   | 2024.12.06     | 65.1               | 2337     | 52.3      | 82.8     | 60.5            | 64.3        | 60.6   | 46.3           | 81.1         | 79.3         | 76.8        | 81.4      | 58.3      | 91.6        |
+| Baichuan-Omni         | Unknown-7B            | 2024.10.11     | -                  | 2186     | 47.3      | 70.0     | 51.9            | 62.6        | 65.4   | 47.8           | 76.2         | 74.9         | 74.3        | -         | -         | -           |
+| MiniCPM-V-2.6         | Qwen2-7B              | 2024.08.06     | 65.2               | 2348     | 49.8      | 85.2     | 60.6            | 69.7        | 60     | 48.1           | 81.2         | 79           | 80.1        | 82.1      | 57.26     | 90.8        |
+| Qwen2-VL-7B-Instruct  | Qwen2-7B              | 2024.08.28     | 67                 | 2326     | 54.1      | 84.5     | 58.2            | 70.1        | 62     | 50.6           | 83           | 80.5         | 84.3        | 83        | 60.7      | 94.5        |
+| MiniCPM-Llama3-V-2.5  | Llama3-Instruct 8B    | 2024.05.20     | 58.8               | 2024     | 45.8      | 72.5     | 54.3            | 63.5        | 52.8   | 42.4           | 77.2         | 74.2         | 76.6        | 78.4      | -         | 84.8        |
+| VITA                  | Mixtral 8x7B          | 2024.08.12     | -                  | 2097     | 47.3      | 67.8     | 44.9            | 59          | 41.6   | 39.7           | 74.7         | 71.4         | 71.8        | -         | -         | -           |
+| GLM-4V-9B             | GLM-4-9B              | 2024.06.04     | 59.1               | 2018     | 46.9      | 77.6     | 51.1            | -           | 58     | 46.6           | 81.1         | 79.4         | -           | 81.1      | 58.7      | -           |
+| LLaVA-NeXT-Yi-34B     | Yi-34B                | 2024.01.18     | 55                 | 2006     | 48.8      | 57.4     | 40.4            | 66          | 50.7   | 34.8           | 81.1         | 79           | 69.3        | 78.9      | 51.6      | -           |
+| Qwen2-VL-72B-Instruct | Qwen2-72B             | 2024.08.28     | 74.8               | 2482     | 64.5      | 87.7     | 70.5            | 77.8        | 74     | 58.1           | 86.5         | 86.6         | 85.5        | 88.1      | 68.3      | 96.5        |
+
+### 文本处理能力
+|                       |          |             |                                       | 对话&指令 |                 |        | 中文&英文任务 |            |       |          |  代码任务 |       | 数学任务 |       |
+|:---------------------:|:--------:|:-----------:|:-------------------------------------:|:---------:|:---------------:|:------:|:-------------:|:----------:|:-----:|:--------:|:---------:|:-----:|:--------:|:-----:|
+|         models        | 指令模型 |   发布时间  | # Non-Emb Params |  MT-Bench | AlignBench (ZH) | IFEval |  C-EVAL (ZH)  | CMMLU (ZH) | MMLU  | MMLU-Pro | HumanEval |  MBPP |   GSM8K  |  MATH |
+| Megrez-3B-Omni        |     Y    |  2024.12.16 |                  2.3                  |    8.4    |       6.94       |  66.5  |     84.0      |    75.3    | 73.3  |   45.2   |   72.6    | 60.6  |   63.8   | 27.3  |
+| Megrez-3B-Instruct    |     Y    |  2024.12.16 |                  2.3                  |   8.64    |      7.06       |  68.6  |     84.8      |    74.7    | 72.8  |   46.1   |   78.7    | 71.0  |   65.5   | 28.3  |
+| Baichuan-Omni         |     Y    |  2024.10.11 |                  7.0                  |     -     |        -        |    -   |     68.9      |    72.2    |  65.3 |     -    |     -     |   -   |     -    |   -   |
+| VITA                  |     Y    |  2024.08.12 |                 12.9                  |     -     |        -        |    -   |     56.7      |    46.6    | 71.0  |     -    |     -     |   -   |   75.7   |   -   |
+| Qwen1.5-7B            |          |  2024.02.04 |                  6.5                  |     -     |        -        |    -   |     74.1      |    73.1    | 61.0  |   29.9   |   36.0    | 51.6  |   62.5   | 20.3  |
+| Qwen1.5-7B-Chat       |     Y    |  2024.02.04 |                  6.5                  |   7.60    |      6.20       |    -   |     67.3      |      -     | 59.5  |   29.1   |   46.3    | 48.9  |   60.3   | 23.2  |
+| Qwen1.5-14B           |          |  2024.02.04 |                 12.6                  |     -     |        -        |    -   |     78.7      |    77.6    | 67.6  |     -    |   37.8    | 44.0  |   70.1   | 29.2  |
+| Qwen1.5-14B-Chat      |     Y    |  2024.02.04 |                 12.6                  |    7.9    |        -        |    -   |       -       |      -     |   -   |     -    |     -     |   -   |     -    |   -   |
+| Qwen2-7B              |          |  2024.06.07 |                  6.5                  |     -     |        -        |    -   |     83.2      |    83.9    | 70.3  |   40.0   |   51.2    | 65.9  |   79.9   | 44.2  |
+| Qwen2-7b-Instruct     |     Y    |  2024.06.07 |                  6.5                  |   8.41    |      7.21       |  51.4  |     80.9      |    77.2    | 70.5  |   44.1   |   79.9    | 67.2  |   85.7   | 52.9  |
+| Qwen2.5-3B-Instruct   |     Y    |  2024.9.19  |                  2.8                  |     -     |        -        |    -   |       -       |      -     |   -   |   43.7   |   74.4    | 72.7  |   86.7   | 65.9  |
+| Qwen2.5-7B            |          |  2024.9.19  |                  6.5                  |     -     |        -        |    -   |       -       |      -     | 74.2  |   45.0   |   57.9    | 74.9  |   85.4   | 49.8  |
+| Qwen2.5-7B-Instruct   |     Y    |  2024.09.19 |                  6.5                  |   8.75    |        -        |  74.9  |       -       |      -     |   -   |   56.3   |   84.8    | 79.2  |   91.6   | 75.5  |
+| Llama-3.1-8B          |          |  2024.07.23 |                  7.0                  |    8.3    |       5.7       |  71.5  |     55.2      |    55.8    | 66.7  |   37.1   |     -     |   -   |   84.5   | 51.9  |
+| Llama-3.2-3B          |          |  2024.09.25 |                  2.8                  |     -     |        -        |  77.4  |       -       |      -     | 63.4  |     -    |     -     |   -   |   77.7   | 48.0  |
+| Phi-3.5-mini-instruct |     Y    |  2024.08.23 |                  3.6                  |    8.6    |       5.7       |  49.4  |     46.1      |    46.9    | 69.0  |   47.4   |   62.8    | 69.6  |   86.2   | 48.5  |
+| MiniCPM3-4B           |     Y    |  2024.09.05 |                  3.9                  |   8.41    |      6.74       |  68.4  |     73.6      |    73.3    | 67.2  |     -    |   74.4    | 72.5  |   81.1   | 46.6  |
+| Yi-1.5-6B-Chat        |     Y    |  2024.05.11 |                  5.5                  |   7.50    |      6.20       |    -   |     74.2      |    74.7    | 61.0  |     -    |   64.0    | 70.9  |   78.9   | 40.5  |
+| GLM-4-9B-chat         |     Y    |  2024.06.04 |                  8.2                  |   8.35    |      7.01       |  64.5  |     75.6      |    71.5    | 72.4  |     -    |   71.8    |   -   |   79.6   | 50.6  |
+| Baichuan2-13B-Base    |          |  2023.09.06 |                 12.6                  |     -     |      5.25       |    -   |     58.1      |    62.0    | 59.2  |     -    |   17.1    | 30.2  |   52.8   | 10.1  |
+
+注：Qwen2-1.5B模型的指标在论文和Qwen2.5报告中点数不一致，当前采用原始论文中的精度
+
+### 语音理解能力
+|       Model      |     Base model     | Release Time | Fleurs test-zh | WenetSpeech test_net | WenetSpeech test_meeting |
+|:----------------:|:------------------:|:-------------:|:--------------:|:--------------------:|:------------------------:|
+|  Megrez-3B-Omni  | Megrez-3B-Instruct |   2024.12.16  |      10.8      |           -          |           16.4           |
+| Whisper-large-v3 |          -         |   2023.11.06  |      12.4      |         17.5         |           30.8           |
+|  Qwen2-Audio-7B  |      Qwen2-7B      |   2024.08.09  |        9       |          11          |           10.7           |
+|  Baichuan2-omni  |     Unknown-7B     |   2024.10.11  |        7       |          6.9         |            8.4           |
+|       VITA       |    Mixtral 8x7B    |   2024.08.12  |        -       |      -/12.2(CER)     |        -/16.5(CER)       |
+
+### 速度
+|                | image_tokens | prefill (tokens/s) | decode (tokens/s) |
+|----------------|:------------:|:------------------:|:-----------------:|
+| Megrez-3B-Omni |      448     |       6312.66      |       1294.9      |
+| Qwen2-VL-2B    |     1378     |       7349.39      |       685.66      |
+| MiniCPM-V-2_6  |      448     |       2167.09      |       452.51      |
+
+实验设置： 
+- 测试环境：NVIDIA H100，vLLM下输入128个Text token和一张1480x720大小图片，输出128个token，num_seqs固定为8
+- Qwen2-VL-2B虽然其具备更小尺寸的基座模型，但编码上述大小图片后的image_token相较Megrez-3B-Omni多很多，导致此实验下的decode速度小于Megrez-3B-Omni
+
+## 快速上手
+
+### 在线体验
+[HF Chat Demo](https://huggingface.co/spaces/Infinigence/Megrez-3B-Omni)
+
+### 本地部署
+环境安装和vLLM推理代码等部署问题可以参考 [Infini-Megrez-Omni](https://github.com/infinigence/Infini-Megrez-Omni)
+
+如下是一个使用transformers进行推理的例子，通过在content字段中分别传入text、image和audio，可以图文/图音等多种模态和模型进行交互。
+```python
+import torch
+from transformers import AutoModelForCausalLM
+
+path = "{{PATH_TO_PRETRAINED_MODEL}}"  # Change this to the path of the model.
+
+model = (
+    AutoModelForCausalLM.from_pretrained(
+        path,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+    )
+    .eval()
+    .cuda()
+)
+
+# Chat with text and image
+messages = [
+    {
+        "role": "user",
+        "content": {
+            "text": "Please describe the content of the image.",
+            "image": "./data/sample_image.jpg",
+        },
+    },
+]
+
+# Chat with audio and image
+messages = [
+    {
+        "role": "user",
+        "content": {
+            "image": "./data/sample_image.jpg",
+            "audio": "./data/sample_audio.m4a",
+        },
+    },
+]
+
+MAX_NEW_TOKENS = 100
+response = model.chat(
+    messages,
+    sampling=False,
+    max_new_tokens=MAX_NEW_TOKENS,
+    temperature=0,
+)
+print(response)
+```
+
+## 注意事项
+1. 请将图片尽量在首轮输入以保证推理效果，语音和文本无此限制，可以自由切换
+2. 语音识别（ASR）场景下，只需要将content['text']修改为“将语音转化为文字。”
+3. OCR场景下开启采样可能会引入语言模型幻觉导致的文字变化，可考虑关闭采样进行推理（sampling=False），但关闭采样可能引入模型复读
+
+## 开源协议及使用声明
+- 协议：本仓库中代码依照 [Apache-2.0](https://www.apache.org/licenses/LICENSE-2.0) 协议开源。
+- 幻觉：大模型天然存在幻觉问题，用户使用过程中请勿完全相信模型生成的内容。
+- 价值观及安全性：本模型已尽全力确保训练过程中使用的数据的合规性，但由于数据的大体量及复杂性，仍有可能存在一些无法预见的问题。如果出现使用本开源模型而导致的任何问题，包括但不限于数据安全问题、公共舆论风险，或模型被误导、滥用、传播或不当利用所带来的任何风险和问题，我们将不承担任何责任。

README_EN.md

@@ -0,0 +1,217 @@
+---
+license: apache-2.0
+---
+# Megrez-3B-Omni: The First Open-Source On-device LLM with Full Modality Understanding
+<p align="center">
+    <img src="assets/megrez_logo.png" width="400"/>
+<p>
+<p align="center">
+        🔗 <a href="https://github.com/infinigence/Infini-Megrez-Omni">GitHub</a>&nbsp&nbsp | &nbsp&nbsp🏠 <a href="https://huggingface.co/spaces/Infinigence/Megrez-3B-Omni">Demo</a>&nbsp&nbsp | &nbsp&nbsp📖 <a href="assets/wechat-official.jpg">WeChat Official</a>&nbsp&nbsp | &nbsp&nbsp💬 <a href="assets/wechat-group.jpg">WeChat Groups</a>&nbsp&nbsp   
+</p>
+<h4 align="center">
+    <p>
+        <a href="https://huggingface.co/Infinigence/Megrez-3B-Omni/blob/main/README.md">中文</a> | <b>English</b> 
+    <p>
+</h4>
+
+## Introduction
+**Megrez-3B-Omni** is an on-device multimodal understanding LLM model developed by **Infinigence AI** ([Infinigence AI](https://cloud.infini-ai.com/platform/ai)). It is an extension of the Megrez-3B-Instruct model and supports analysis of image, text, and audio modalities. The model achieves state-of-the-art accuracy in all three domains:
+- Image Understanding: By utilizing SigLip-400M for constructing image tokens, Megrez-3B-Omni outperforms models with more parameters such as LLaVA-NeXT-Yi-34B. It is one of the best image understanding models among multiple mainstream benchmarks, including MME, MMMU, and OCRBench. It demonstrates excellent performance in tasks such as scene understanding and OCR.
+- Language Understanding: Megrez-3B-Omni retains text understanding capabilities without significant trade-offs. Compared to its single-modal counterpart (Megrez-3B-Instruct), the accuracy variation is less than 2%, maintaining state-of-the-art performance on benchmarks like C-EVAL, MMLU/MMLU Pro, and AlignBench. It also outperforms previous-generation models with 14B parameters.
+- Speech Understanding: Equipped with the encoder head of Qwen2-Audio/whisper-large-v3, the model supports both Chinese and English speech input, multi-turn conversations, and voice-based questions about input images. It can directly respond to voice commands with text and achieved leading results across multiple benchmarks.
+
+## Model Info
+<table>
+  <thead>
+    <tr>
+      <th></th>
+      <th>Language Module</th>
+      <th>Vision Module</th>
+      <th>Audio Module</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>Architecture</td>
+      <td>Llama-2 with GQA</td>
+      <td>SigLip-SO400M</td>
+      <td>Whisper-large-v3
+(encoder-only)</td>
+    </tr>
+    <tr>
+      <td># Params (Backbone)</td>
+      <td>2.29B</td>
+      <td>0.42B</td>
+      <td>0.64B</td>
+    </tr>
+    <tr>
+      <td>Connector</td>
+      <td>-</td>
+      <td>Cross Attention</td>
+      <td>Linear</td>
+    </tr>
+    <tr>
+      <td># Params (Others)</td>
+      <td>Emb: 0.31B<br>Softmax: 0.31B</td>
+      <td>Connector: 0.036B</td>
+      <td>Connector: 0.003B</td>
+    </tr>
+    <tr>
+      <td># Params (Total)</td>
+      <td colspan="3">4B</td>
+    </tr>
+    <tr>
+      <td># Vocab Size</td>
+      <td>122880</td>
+      <td>64 tokens/slice</td>
+      <td>-</td>
+    </tr>
+    <tr>
+      <td>Context length</td>
+      <td colspan="3">4K tokens</td>
+    </tr>
+    <tr>
+      <td>Supported languages</td>
+      <td colspan="3">Chinese & English</td>
+    </tr>
+  </tbody>
+</table>
+
+### Image Understanding
+- The above image compares the performance of Megrez-3B-Omni with other open-source models on mainstream image multimodal tasks.
+- The below image shows the performance of Megrez-3B-Omni on the OpenCompass test set. Image reference: [InternVL 2.5 Blog Post](https://internvl.github.io/blog/2024-12-05-InternVL-2.5/)
+
+![Multitask](assets/multitask.jpg)
+![OpencompassBmk](assets/opencompass.jpg)
+
+| model                 | basemodel             | release time       | OpenCompass | MME      | MMMU val  | OCRBench | MathVista | RealWorldQA | MMVet  | hallusionBench | MMB TEST (en) | MMB TEST (zh) | TextVQA val | AI2D_TEST | MMstar    | DocVQA_TEST |
+|-----------------------|-----------------------|----------------|--------------------|----------|-----------|----------|-----------------|-------------|--------|----------------|--------------|--------------|-------------|-----------|-----------|-------------|
+| **Megrez-3B-Omni**    | **Megrez-3B**         | **2024.12.16** | **66.2**           | **2315** | **51.89** | **82.8** | **62**          | **71.89**   | **60** | **50.12**      | **80.8**     | **82.3**     | **80.3**    | **82.05** | **60.46** | **91.62**   |
+| Qwen2-VL-2B-Instruct  | Qwen2-1.5B            | 2024.08.28     | 57.2               | 1872     | 41.1      | 79.4     | 43              | 62.9        | 49.5   | 41.7           | 74.9         | 73.5         | 79.7        | 74.7      | 48        | 90.1        |
+| InternVL2.5-2B        | Internlm2.5-1.8B-chat | 2024.12.06     | 59.9               | 2138     | 43.6      | 80.4     | 51.3            | 60.1        | 60.8   | 42.6           | 74.7         | 71.9         | 74.3        | 74.9      | 53.7      | 88.7        |
+| BlueLM-V-3B           | -                     | 2024.11.29     | 66.1               | -        | 45.1      | 82.9     | 60.8            | 66.7        | 61.8   | 48             | 83           | 80.5         | 78.4        | 85.3      | 62.3      | 87.8        |
+| InternVL2.5-4B        | Qwen2.5-3B-Instruct   | 2024.12.06     | 65.1               | 2337     | 52.3      | 82.8     | 60.5            | 64.3        | 60.6   | 46.3           | 81.1         | 79.3         | 76.8        | 81.4      | 58.3      | 91.6        |
+| Baichuan-Omni         | Unknown-7B            | 2024.10.11     | -                  | 2186     | 47.3      | 70.0     | 51.9            | 62.6        | 65.4   | 47.8           | 76.2         | 74.9         | 74.3        | -         | -         | -           |
+| MiniCPM-V-2.6         | Qwen2-7B              | 2024.08.06     | 65.2               | 2348     | 49.8      | 85.2     | 60.6            | 69.7        | 60     | 48.1           | 81.2         | 79           | 80.1        | 82.1      | 57.26     | 90.8        |
+| Qwen2-VL-7B-Instruct  | Qwen2-7B              | 2024.08.28     | 67                 | 2326     | 54.1      | 84.5     | 58.2            | 70.1        | 62     | 50.6           | 83           | 80.5         | 84.3        | 83        | 60.7      | 94.5        |
+| MiniCPM-Llama3-V-2.5  | Llama3-Instruct 8B    | 2024.05.20     | 58.8               | 2024     | 45.8      | 72.5     | 54.3            | 63.5        | 52.8   | 42.4           | 77.2         | 74.2         | 76.6        | 78.4      | -         | 84.8        |
+| VITA                  | Mixtral 8x7B          | 2024.08.12     | -                  | 2097     | 47.3      | 67.8     | 44.9            | 59          | 41.6   | 39.7           | 74.7         | 71.4         | 71.8        | -         | -         | -           |
+| GLM-4V-9B             | GLM-4-9B              | 2024.06.04     | 59.1               | 2018     | 46.9      | 77.6     | 51.1            | -           | 58     | 46.6           | 81.1         | 79.4         | -           | 81.1      | 58.7      | -           |
+| LLaVA-NeXT-Yi-34B     | Yi-34B                | 2024.01.18     | 55                 | 2006     | 48.8      | 57.4     | 40.4            | 66          | 50.7   | 34.8           | 81.1         | 79           | 69.3        | 78.9      | 51.6      | -           |
+| Qwen2-VL-72B-Instruct | Qwen2-72B             | 2024.08.28     | 74.8               | 2482     | 64.5      | 87.7     | 70.5            | 77.8        | 74     | 58.1           | 86.5         | 86.6         | 85.5        | 88.1      | 68.3      | 96.5        |
+
+### Text Understanding
+|                       |          |             |                                       | Chat&Instruction |                 |        | Zh&En Tasks |            |       |          |  Code |       | Math |       |
+|:---------------------:|:--------:|:-----------:|:-------------------------------------:|:---------:|:---------------:|:------:|:-------------:|:----------:|:-----:|:--------:|:---------:|:-----:|:--------:|:-----:|
+|         models        | Instruction |   Release Time  | Non-Emb Params |  MT-Bench | AlignBench (ZH) | IFEval |  C-EVAL (ZH)  | CMMLU (ZH) | MMLU  | MMLU-Pro | HumanEval |  MBPP |   GSM8K  |  MATH |
+| Megrez-3B-Omni        |     Y    |  2024.12.16 |                  2.3                  |    8.4    |       6.94       |  66.5  |     84.0      |    75.3    | 73.3  |   45.2   |   72.6    | 60.6  |   63.8   | 27.3  |
+| Megrez-3B-Instruct    |     Y    |  2024.12.16 |                  2.3                  |   8.64    |      7.06       |  68.6  |     84.8      |    74.7    | 72.8  |   46.1   |   78.7    | 71.0  |   65.5   | 28.3  |
+| Baichuan-Omni         |     Y    |  2024.10.11 |                  7.0                  |     -     |        -        |    -   |     68.9      |    72.2    |  65.3 |     -    |     -     |   -   |     -    |   -   |
+| VITA                  |     Y    |  2024.08.12 |                 12.9                  |     -     |        -        |    -   |     56.7      |    46.6    | 71.0  |     -    |     -     |   -   |   75.7   |   -   |
+| Qwen1.5-7B            |          |  2024.02.04 |                  6.5                  |     -     |        -        |    -   |     74.1      |    73.1    | 61.0  |   29.9   |   36.0    | 51.6  |   62.5   | 20.3  |
+| Qwen1.5-7B-Chat       |     Y    |  2024.02.04 |                  6.5                  |   7.60    |      6.20       |    -   |     67.3      |      -     | 59.5  |   29.1   |   46.3    | 48.9  |   60.3   | 23.2  |
+| Qwen1.5-14B           |          |  2024.02.04 |                 12.6                  |     -     |        -        |    -   |     78.7      |    77.6    | 67.6  |     -    |   37.8    | 44.0  |   70.1   | 29.2  |
+| Qwen1.5-14B-Chat      |     Y    |  2024.02.04 |                 12.6                  |    7.9    |        -        |    -   |       -       |      -     |   -   |     -    |     -     |   -   |     -    |   -   |
+| Qwen2-7B              |          |  2024.06.07 |                  6.5                  |     -     |        -        |    -   |     83.2      |    83.9    | 70.3  |   40.0   |   51.2    | 65.9  |   79.9   | 44.2  |
+| Qwen2-7b-Instruct     |     Y    |  2024.06.07 |                  6.5                  |   8.41    |      7.21       |  51.4  |     80.9      |    77.2    | 70.5  |   44.1   |   79.9    | 67.2  |   85.7   | 52.9  |
+| Qwen2.5-3B-Instruct   |     Y    |  2024.9.19  |                  2.8                  |     -     |        -        |    -   |       -       |      -     |   -   |   43.7   |   74.4    | 72.7  |   86.7   | 65.9  |
+| Qwen2.5-7B            |          |  2024.9.19  |                  6.5                  |     -     |        -        |    -   |       -       |      -     | 74.2  |   45.0   |   57.9    | 74.9  |   85.4   | 49.8  |
+| Qwen2.5-7B-Instruct   |     Y    |  2024.09.19 |                  6.5                  |   8.75    |        -        |  74.9  |       -       |      -     |   -   |   56.3   |   84.8    | 79.2  |   91.6   | 75.5  |
+| Llama-3.1-8B          |          |  2024.07.23 |                  7.0                  |    8.3    |       5.7       |  71.5  |     55.2      |    55.8    | 66.7  |   37.1   |     -     |   -   |   84.5   | 51.9  |
+| Llama-3.2-3B          |          |  2024.09.25 |                  2.8                  |     -     |        -        |  77.4  |       -       |      -     | 63.4  |     -    |     -     |   -   |   77.7   | 48.0  |
+| Phi-3.5-mini-instruct |     Y    |  2024.08.23 |                  3.6                  |    8.6    |       5.7       |  49.4  |     46.1      |    46.9    | 69.0  |   47.4   |   62.8    | 69.6  |   86.2   | 48.5  |
+| MiniCPM3-4B           |     Y    |  2024.09.05 |                  3.9                  |   8.41    |      6.74       |  68.4  |     73.6      |    73.3    | 67.2  |     -    |   74.4    | 72.5  |   81.1   | 46.6  |
+| Yi-1.5-6B-Chat        |     Y    |  2024.05.11 |                  5.5                  |   7.50    |      6.20       |    -   |     74.2      |    74.7    | 61.0  |     -    |   64.0    | 70.9  |   78.9   | 40.5  |
+| GLM-4-9B-chat         |     Y    |  2024.06.04 |                  8.2                  |   8.35    |      7.01       |  64.5  |     75.6      |    71.5    | 72.4  |     -    |   71.8    |   -   |   79.6   | 50.6  |
+| Baichuan2-13B-Base    |          |  2023.09.06 |                 12.6                  |     -     |      5.25       |    -   |     58.1      |    62.0    | 59.2  |     -    |   17.1    | 30.2  |   52.8   | 10.1  |
+
+- The metrics for the Qwen2-1.5B model differ between the original paper and the Qwen2.5 report. Currently, the accuracy figures from the original paper are being used.
+
+### Audio Understanding
+|       Model      |     Base model     | Release Time | Fleurs test-zh | WenetSpeech test_net | WenetSpeech test_meeting |
+|:----------------:|:------------------:|:-------------:|:--------------:|:--------------------:|:------------------------:|
+|  Megrez-3B-Omni  | Megrez-3B-Instruct |   2024.12.16  |      10.8      |           -          |           16.4           |
+| Whisper-large-v3 |          -         |   2023.11.06  |      12.4      |         17.5         |           30.8           |
+|  Qwen2-Audio-7B  |      Qwen2-7B      |   2024.08.09  |        9       |          11          |           10.7           |
+|  Baichuan2-omni  |     Unknown-7B     |   2024.10.11  |        7       |          6.9         |            8.4           |
+|       VITA       |    Mixtral 8x7B    |   2024.08.12  |        -       |      -/12.2(CER)     |        -/16.5(CER)       |
+
+### Inference Speed
+|                | image_tokens | prefill (tokens/s) | decode (tokens/s) |
+|----------------|:------------:|:------------------:|:-----------------:|
+| Megrez-3B-Omni |      448     |       6312.66      |       1294.9      |
+| Qwen2-VL-2B    |     1378     |       7349.39      |       685.66      |
+| MiniCPM-V-2_6  |      448     |       2167.09      |       452.51      |
+
+Setup:
+- The testing environment utilizes an NVIDIA H100 GPU with vLLM. Each test includes 128 text tokens and a 720×1480 image as input, producing 128 output tokens, with `num_seqs` fixed at 8.  
+- Under this setup, the decoding speed of Qwen2-VL-2B is slower than Megrez-3B-Omni, despite having a smaller base LLM. This is due to the larger number of image tokens generated when encoding images of the specified size, which impacts actual inference speed.  
+
+## Quickstart
+
+### Online Experience
+[HF Chat Demo](https://huggingface.co/spaces/Infinigence/Megrez-3B-Omni)(recommend) 
+
+### Local Deployment
+For environment installation and vLLM inference code deployment, refer to [Infini-Megrez-Omni](https://github.com/infinigence/Infini-Megrez-Omni)
+
+Below is an example of using transformers for inference. By passing text, image, and audio in the content field, you can interact with various modalities and models.
+```python
+import torch
+from transformers import AutoModelForCausalLM
+
+path = "{{PATH_TO_PRETRAINED_MODEL}}"  # Change this to the path of the model.
+
+model = (
+    AutoModelForCausalLM.from_pretrained(
+        path,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+    )
+    .eval()
+    .cuda()
+)
+
+# Chat with text and image
+messages = [
+    {
+        "role": "user",
+        "content": {
+            "text": "Please describe the content of the image.",
+            "image": "./data/sample_image.jpg",
+        },
+    },
+]
+
+# Chat with audio and image
+messages = [
+    {
+        "role": "user",
+        "content": {
+            "image": "./data/sample_image.jpg",
+            "audio": "./data/sample_audio.m4a",
+        },
+    },
+]
+
+MAX_NEW_TOKENS = 100
+response = model.chat(
+    messages,
+    sampling=False,
+    max_new_tokens=MAX_NEW_TOKENS,
+    temperature=0,
+)
+print(response)
+```
+
+## Notes
+1. We recommend to put the images in the first round of chat for better inference results. There are no such restrictions for audio and text, which can be switched freely.
+2. In the Automatic Speech Recognition (ASR) scenario, simply change content['text'] to "Convert speech to text."
+3. In the OCR scenario, enabling sampling may introduce language model hallucinations which cause text changes. Users may consider disabling sampling in inference (sampling=False). However, disabling sampling may introduce model repetition.
+
+
+## Open Source License and Usage Statement
+- **License**: The code in this repository is open-sourced under the [Apache-2.0](https://www.apache.org/licenses/LICENSE-2.0) license.  
+- **Hallucination**: Large models inherently have hallucination issues. Users should not completely trust the content generated by the model. 
+- **Values and Safety**: While we have made every effort to ensure compliance of the data used during training, the large volume and complexity of the data may still lead to unforeseen issues. We disclaim any liability for problems arising from the use of this open-source model, including but not limited to data security issues, public opinion risks, or risks and problems caused by misleading, misuse, propagation, or improper utilization of the model.  

audio.py

@@ -0,0 +1,228 @@
+# -*- encoding: utf-8 -*-
+# File: audio.py
+# Description: None
+
+
+from typing import Iterable, List, Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class LayerNorm(nn.LayerNorm):
+    def forward(self, x: Tensor) -> Tensor:
+        return super().forward(x).type(x.dtype)
+
+
+class Linear(nn.Linear):
+    def forward(self, x: Tensor) -> Tensor:
+        return F.linear(
+            x,
+            self.weight.to(x.dtype),
+            None if self.bias is None else self.bias.to(x.dtype),
+        )
+
+
+class Conv1d(nn.Conv1d):
+    def _conv_forward(self, x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> Tensor:
+        return super()._conv_forward(x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype))
+
+
+def sinusoids(length, channels, max_timescale=10000):
+    """Returns sinusoids for positional embedding"""
+    assert channels % 2 == 0
+    log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
+    inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
+    scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
+    return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, n_state: int, n_head: int):
+        super().__init__()
+        self.n_head = n_head
+        self.query = Linear(n_state, n_state)
+        self.key = Linear(n_state, n_state, bias=False)
+        self.value = Linear(n_state, n_state)
+        self.out = Linear(n_state, n_state)
+
+    def forward(
+        self,
+        x: Tensor,
+        xa: Optional[Tensor] = None,
+        mask: Optional[Tensor] = None,
+        kv_cache: Optional[dict] = None,
+    ):
+        q = self.query(x)
+
+        if kv_cache is None or xa is None or self.key not in kv_cache:
+            # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;
+            # otherwise, perform key/value projections for self- or cross-attention as usual.
+            k = self.key(x if xa is None else xa)
+            v = self.value(x if xa is None else xa)
+        else:
+            # for cross-attention, calculate keys and values once and reuse in subsequent calls.
+            k = kv_cache[self.key]
+            v = kv_cache[self.value]
+
+        wv, qk = self.qkv_attention(q, k, v, mask)
+        return self.out(wv), qk
+
+    def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None):
+        n_batch, n_ctx, n_state = q.shape
+        scale = (n_state // self.n_head) ** -0.25
+        q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale
+        k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale
+        v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
+
+        qk = q @ k
+        if mask is not None:
+            qk += mask
+
+        w = F.softmax(qk, dim=-1).to(q.dtype)
+        return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach()
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, n_state: int, n_head: int, cross_attention: bool = False):
+        super().__init__()
+
+        self.attn = MultiHeadAttention(n_state, n_head)
+        self.attn_ln = LayerNorm(n_state)
+
+        self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None
+        self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None
+
+        n_mlp = n_state * 4
+        self.mlp = nn.Sequential(Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state))
+        self.mlp_ln = LayerNorm(n_state)
+
+    def forward(
+        self,
+        x: Tensor,
+        xa: Optional[Tensor] = None,
+        mask: Optional[Tensor] = None,
+        kv_cache: Optional[dict] = None,
+    ):
+        x = x + self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache)[0]
+        if self.cross_attn:
+            x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache)[0]
+        x = x + self.mlp(self.mlp_ln(x))
+        return x
+
+
+class AudioEncoder(nn.Module):
+    def __init__(
+        self,
+        n_mels: int,
+        n_ctx: int,
+        n_state: int,
+        n_head: int,
+        n_layer: int,
+        output_dim: int = 512,
+        avg_pool: bool = True,
+        add_audio_bos_eos_token: bool = True,
+        **kwargs,
+    ):
+        super().__init__()
+        self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, padding=1)
+        self.conv2 = Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1)
+        self.register_buffer("positional_embedding", sinusoids(n_ctx, n_state))
+
+        self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(
+            [ResidualAttentionBlock(n_state, n_head) for _ in range(n_layer)]
+        )
+        self.ln_post = LayerNorm(n_state)
+
+        if avg_pool:
+            self.avg_pooler = nn.AvgPool1d(2, stride=2)
+        else:
+            self.avg_pooler = None
+        self.proj = nn.Linear(n_state, output_dim)
+        if add_audio_bos_eos_token:
+            self.audio_bos_eos_token = nn.Embedding(2, output_dim)
+        else:
+            self.audio_bos_eos_token = None
+        self.output_dim = output_dim
+        self.n_head = n_head
+
+    def forward(self, x: Tensor, padding_mask: Tensor = None, audio_lengths: Tensor = None):
+        """
+        x : torch.Tensor, shape = (batch_size, n_mels, n_ctx)
+            the mel spectrogram of the audio
+        """
+        x = x.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device)
+        if audio_lengths is not None:
+            input_mel_len = audio_lengths[:, 0] * 2
+            max_mel_len_in_batch = input_mel_len.max()
+            x = x[:, :, :max_mel_len_in_batch]
+        x = F.gelu(self.conv1(x))
+        x = F.gelu(self.conv2(x))
+        x = x.permute(0, 2, 1)  # B, L, D
+        bsz = x.size(0)
+        src_len = x.size(1)
+
+        self.input_positional_embedding = self.positional_embedding[:src_len]
+        assert (
+            x.shape[1:] == self.input_positional_embedding.shape
+        ), f"incorrect audio shape: {x.shape[1:], self.input_positional_embedding.shape}"
+        x = (x + self.input_positional_embedding).to(x.dtype)
+        if padding_mask is not None:
+            padding_mask = padding_mask.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device)
+            batch_src_len = padding_mask.size(1)
+            x = x[:, :batch_src_len, :]
+            padding_mask = padding_mask.view(bsz, -1, batch_src_len)
+            padding_mask_ = padding_mask.all(1)
+            x[padding_mask_] = 0
+            key_padding_mask = (
+                padding_mask_.view(bsz, 1, 1, batch_src_len)
+                .expand(-1, self.n_head, -1, -1)
+                .reshape(bsz, self.n_head, 1, batch_src_len)
+            )
+            new_padding_mask = torch.zeros_like(key_padding_mask, dtype=x.dtype)
+            padding_mask = new_padding_mask.masked_fill(key_padding_mask, float("-inf"))
+
+        for block in self.blocks:
+            x = block(x, mask=padding_mask)
+
+        if self.avg_pooler:
+            x = x.permute(0, 2, 1)
+            x = self.avg_pooler(x)
+            x = x.permute(0, 2, 1)
+
+        x = self.ln_post(x)
+        x = self.proj(x)
+
+        if self.audio_bos_eos_token is not None:
+            bos = self.audio_bos_eos_token.weight[0][None, :]
+            eos = self.audio_bos_eos_token.weight[1][None, :]
+        else:
+            bos, eos = None, None
+        return x, bos, eos
+
+    def encode(
+        self,
+        input_audios: Tensor,
+        input_audio_lengths: Tensor,
+        audio_span_tokens: List,
+    ):
+        real_input_audio_lens = input_audio_lengths[:, 0].tolist()
+        max_len_in_batch = max(real_input_audio_lens)
+        padding_mask = torch.ones([input_audios.size(0), max_len_in_batch]).to(
+            dtype=self.conv1.weight.dtype, device=self.conv1.weight.device
+        )
+        for index in range(len(input_audios)):
+            padding_mask[index, : input_audio_lengths[index][0].item()] = 0
+        x, bos, eos = self(input_audios, padding_mask, input_audio_lengths)
+        output_audios = []
+        for i in range(len(audio_span_tokens)):
+            audio_span = audio_span_tokens[i]
+            audio = x[i][: audio_span - 2]
+            if bos is not None:
+                audio = torch.concat([bos, audio, eos])
+            assert len(audio) == audio_span
+            output_audios.append(audio)
+        return output_audios

config.json

@@ -0,0 +1,125 @@
+{
+  "_name_or_path": "/mnt/public/algm/lizhiyuan/models/megrez-o-release/pretrain_audio_stage3_vision_stage3_merge",
+  "architectures": [
+    "MegrezO"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "audio_config": {
+    "add_audio_bos_eos_token": true,
+    "avg_pool": true,
+    "n_ctx": 1500,
+    "n_head": 20,
+    "n_layer": 32,
+    "n_mels": 128,
+    "n_state": 1280,
+    "output_dim": 2560
+  },
+  "auto_map": {
+    "AutoModel": "modeling_megrezo.MegrezO",
+    "AutoModelForCausalLM": "modeling_megrezo.MegrezO",
+    "AutoConfig": "configuration_megrezo.MegrezOConfig",
+    "AutoProcessor": "processing_megrezo.MegrezOProcessor",
+    "AutoImageProcessor": "image_processing_megrezo.MegrezOImageProcessor"
+  },
+  "bos_token_id": null,
+  "drop_vision_last_layer": false,
+  "eos_token_id": 120005,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 7168,
+  "max_position_embeddings": 4096,
+  "max_sequence_length": 4096,
+  "mlp_bias": false,
+  "model_type": "megrezo",
+  "num_attention_heads": 40,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 10,
+  "pad_token_id": 0,
+  "patch_size": 14,
+  "pretraining_tp": 1,
+  "query_num": 64,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 5000000.0,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.42.0",
+  "use_cache": false,
+  "vision_batch_size": 8,
+  "vision_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 980,
+    "intermediate_size": 4304,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "siglip_vision_model",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_channels": 3,
+    "num_hidden_layers": 27,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 14,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  },
+  "vocab_size": 122880
+}

configuration_megrezo.py

@@ -0,0 +1,87 @@
+"""MegrezO model configuration"""
+
+from typing import Optional
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.utils import logging
+
+from .modeling_navit_siglip import SiglipVisionConfig
+
+logger = logging.get_logger(__name__)
+
+
+class AudioConfig(PretrainedConfig):
+    model_type = "megrezo"
+
+    def __init__(
+        self,
+        n_mels: int = 128,
+        n_ctx: int = 1500,
+        n_state: int = 1280,
+        n_head: int = 20,
+        n_layer: int = 32,
+        output_dim: int = 2560,
+        avg_pool: bool = True,
+        add_audio_bos_eos_token: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.n_mels = n_mels
+        self.n_ctx = n_ctx
+        self.n_state = n_state
+        self.n_head = n_head
+        self.n_layer = n_layer
+        self.output_dim = output_dim
+        self.avg_pool = avg_pool
+        self.add_audio_bos_eos_token = add_audio_bos_eos_token
+
+
+class MegrezOConfig(LlamaConfig):
+    model_type = "megrezo"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    is_composition = True
+
+    _default_audio_config = {
+        "n_mels": 128,
+        "n_ctx": 1500,
+        "n_state": 1280,
+        "n_head": 20,
+        "n_layer": 32,
+        "output_dim": 2560,
+        "avg_pool": True,
+        "add_audio_bos_eos_token": True,
+    }
+
+    _default_vision_config = {
+        "intermediate_size": 4304,
+        "num_hidden_layers": 27,
+        "num_attention_heads": 16,
+        "image_size": 980,
+        "hidden_size": 1152,
+        "patch_size": 16,
+        "model_type": "siglip_vision_model",
+    }
+
+    def __init__(
+        self,
+        audio_config: Optional[AudioConfig] = None,
+        vision_config: Optional[SiglipVisionConfig] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if audio_config is None:
+            self.audio_config = AudioConfig(**self._default_audio_config)
+        elif isinstance(audio_config, dict):
+            self.audio_config = AudioConfig(**audio_config)
+        elif isinstance(audio_config, AudioConfig):
+            self.audio_config = audio_config
+
+        if vision_config is None:
+            self.vision_config = SiglipVisionConfig(**self._default_vision_config)
+        elif isinstance(vision_config, dict):
+            self.vision_config = SiglipVisionConfig(**vision_config)
+        elif isinstance(vision_config, SiglipVisionConfig):
+            self.vision_config = vision_config

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": [120000, 120005],
+  "pad_token_id": 120002,
+  "transformers_version": "4.42.0",
+  "use_cache": false
+}

image_processing_megrezo.py

@@ -0,0 +1,386 @@
+from typing import Optional, Union, Dict, Any, List
+
+import torch
+import math
+import PIL.Image
+import PIL.ImageSequence
+import numpy as np
+import PIL
+from PIL import Image
+
+from transformers.utils import TensorType, requires_backends, is_torch_dtype, is_torch_device
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+from transformers import AutoImageProcessor
+from transformers.image_transforms import to_channel_dimension_format
+from transformers.image_utils import (
+    valid_images,
+    is_torch_tensor,
+    to_numpy_array,
+    infer_channel_dimension_format,
+    ChannelDimension,
+)
+
+
+def recursive_converter(converter, value):
+    if isinstance(value, list):
+        new_value = []
+        for v in value:
+            new_value += [recursive_converter(converter, v)]
+        return new_value
+    else:
+        return converter(value)
+
+
+class MegrezOBatchFeature(BatchFeature):
+    r"""
+    Extend from BatchFeature for supporting various image size
+    """
+
+    def __init__(self, data: Optional[Dict[str, Any]] = None, tensor_type: Union[None, str, TensorType] = None):
+        super().__init__(data)
+        self.convert_to_tensors(tensor_type=tensor_type)
+
+    def convert_to_tensors(self, tensor_type: Optional[Union[str, TensorType]] = None):
+        if tensor_type is None:
+            return self
+
+        is_tensor, as_tensor = self._get_is_as_tensor_fns(tensor_type)
+
+        def converter(value):
+            try:
+                if not is_tensor(value):
+                    tensor = as_tensor(value)
+                    return tensor
+            except:  # noqa E722
+                if key == "overflowing_values":
+                    raise ValueError("Unable to create tensor returning overflowing values of different lengths. ")
+                raise ValueError(
+                    "Unable to create tensor, you should probably activate padding "
+                    "with 'padding=True' to have batched tensors with the same length."
+                )
+
+        for key, value in self.items():
+            self[key] = recursive_converter(converter, value)
+        return self
+
+    def to(self, *args, **kwargs) -> "MegrezOBatchFeature":
+        requires_backends(self, ["torch"])
+        import torch
+
+        def cast_tensor(v):
+            # check if v is a floating point
+            if torch.is_floating_point(v):
+                # cast and send to device
+                return v.to(*args, **kwargs)
+            elif device is not None:
+                return v.to(device=device)
+            else:
+                return v
+
+        new_data = {}
+        device = kwargs.get("device")
+        # Check if the args are a device or a dtype
+        if device is None and len(args) > 0:
+            # device should be always the first argument
+            arg = args[0]
+            if is_torch_dtype(arg):
+                # The first argument is a dtype
+                pass
+            elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):
+                device = arg
+            else:
+                # it's something else
+                raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.")
+        # We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor`
+        for k, v in self.items():
+            new_data[k] = recursive_converter(cast_tensor, v)
+        self.data = new_data
+        return self
+
+
+class MegrezOImageProcessor(BaseImageProcessor):
+    model_input_names = ["pixel_values"]
+
+    def __init__(self, max_slice_nums=9, scale_resolution=448, patch_size=14, **kwargs):
+        super().__init__(**kwargs)
+        self.max_slice_nums = max_slice_nums
+        self.scale_resolution = scale_resolution
+        self.patch_size = patch_size
+        self.use_image_id = kwargs.pop("use_image_id", False)
+        self.image_feature_size = kwargs.pop("image_feature_size", 64)
+        self.im_start_token = kwargs.pop("im_start", "<|image_start|>")
+        self.im_end_token = kwargs.pop("im_end", "<|image_end|>")
+        self.slice_start_token = kwargs.pop("slice_start", "<|slice_start|>")
+        self.slice_end_token = kwargs.pop("slice_end", "<|slice_end|>")
+        self.unk_token = kwargs.pop("unk", "<|unk|>")
+        self.im_id_start = kwargs.pop("im_id_start", "<|image_id_start|>")
+        self.im_id_end = kwargs.pop("im_id_end", "<|image_id_end|>")
+        self.slice_mode = kwargs.pop("slice_mode", True)
+        self.mean = np.array(kwargs.pop("norm_mean", [0.5, 0.5, 0.5]))
+        self.std = np.array(kwargs.pop("norm_std", [0.5, 0.5, 0.5]))
+        self.version = kwargs.pop("version", 2.0)
+
+    def ensure_divide(self, length, patch_size):
+        return max(round(length / patch_size) * patch_size, patch_size)
+
+    def find_best_resize(self, original_size, scale_resolution, patch_size, allow_upscale=False):
+        width, height = original_size
+        if (width * height > scale_resolution * scale_resolution) or allow_upscale:
+            r = width / height
+            height = int(scale_resolution / math.sqrt(r))
+            width = int(height * r)
+        best_width = self.ensure_divide(width, patch_size)
+        best_height = self.ensure_divide(height, patch_size)
+        return (best_width, best_height)
+
+    def get_refine_size(self, original_size, grid, scale_resolution, patch_size, allow_upscale=False):
+        width, height = original_size
+        grid_x, grid_y = grid
+
+        refine_width = self.ensure_divide(width, grid_x)
+        refine_height = self.ensure_divide(height, grid_y)
+
+        grid_width = refine_width / grid_x
+        grid_height = refine_height / grid_y
+
+        best_grid_size = self.find_best_resize(
+            (grid_width, grid_height), scale_resolution, patch_size, allow_upscale=allow_upscale
+        )
+        refine_size = (best_grid_size[0] * grid_x, best_grid_size[1] * grid_y)
+        return refine_size
+
+    def split_to_patches(self, image, grid):
+        patches = []
+        width, height = image.size
+        grid_x = int(width / grid[0])
+        grid_y = int(height / grid[1])
+        for i in range(0, height, grid_y):
+            images = []
+            for j in range(0, width, grid_x):
+                box = (j, i, j + grid_x, i + grid_y)
+                patch = image.crop(box)
+                images.append(patch)
+            patches.append(images)
+        return patches
+
+    def slice_image(self, image, max_slice_nums=9, scale_resolution=448, patch_size=14, never_split=False):
+        original_size = image.size
+        source_image = None
+        best_grid = self.get_sliced_grid(original_size, max_slice_nums, never_split)
+        patches = []
+
+        if best_grid is None:
+            # dont need to slice, upsample
+            best_size = self.find_best_resize(original_size, scale_resolution, patch_size, allow_upscale=True)
+            source_image = image.resize(best_size, resample=Image.Resampling.BILINEAR)
+        else:
+            # source image, down-sampling and ensure divided by patch_size
+            best_resize = self.find_best_resize(original_size, scale_resolution, patch_size)
+            source_image = image.copy().resize(best_resize, resample=Image.Resampling.BILINEAR)
+            refine_size = self.get_refine_size(
+                original_size, best_grid, scale_resolution, patch_size, allow_upscale=True
+            )
+            refine_image = image.resize(refine_size, resample=Image.Resampling.BILINEAR)
+            patches = self.split_to_patches(refine_image, best_grid)
+
+        return source_image, patches, best_grid
+
+    def get_grid_placeholder(self, grid):
+        if grid is None:
+            return ""
+        slice_image_placeholder = (
+            self.slice_start_token + self.unk_token * self.image_feature_size + self.slice_end_token
+        )
+
+        cols = grid[0]
+        rows = grid[1]
+        slices = []
+        for i in range(rows):
+            lines = []
+            for j in range(cols):
+                lines.append(slice_image_placeholder)
+            slices.append("".join(lines))
+
+        slice_placeholder = "\n".join(slices)
+        return slice_placeholder
+
+    def get_image_id_placeholder(self, idx=0):
+        return f"{self.im_id_start}{idx}{self.im_id_end}"
+
+    def get_sliced_images(self, image, max_slice_nums=None):
+        slice_images = []
+
+        if not self.slice_mode:
+            return [image]
+
+        max_slice_nums = self.max_slice_nums if max_slice_nums is None else int(max_slice_nums)
+        assert max_slice_nums > 0
+        source_image, patches, sliced_grid = self.slice_image(
+            image, max_slice_nums, self.scale_resolution, self.patch_size  # default: 9  # default: 448  # default: 14
+        )
+
+        slice_images.append(source_image)
+        if len(patches) > 0:
+            for i in range(len(patches)):
+                for j in range(len(patches[0])):
+                    slice_images.append(patches[i][j])
+        return slice_images
+
+    def get_sliced_grid(self, image_size, max_slice_nums, nerver_split=False):
+        original_width, original_height = image_size
+        log_ratio = math.log(original_width / original_height)
+        ratio = original_width * original_height / (self.scale_resolution * self.scale_resolution)
+        multiple = min(math.ceil(ratio), max_slice_nums)
+        if multiple <= 1 or nerver_split:
+            return None
+        candidate_split_grids_nums = []
+        for i in [multiple - 1, multiple, multiple + 1]:
+            if i == 1 or i > max_slice_nums:
+                continue
+            candidate_split_grids_nums.append(i)
+
+        candidate_grids = []
+        for split_grids_nums in candidate_split_grids_nums:
+            m = 1
+            while m <= split_grids_nums:
+                if split_grids_nums % m == 0:
+                    candidate_grids.append([m, split_grids_nums // m])
+                m += 1
+
+        best_grid = [1, 1]
+        min_error = float("inf")
+        for grid in candidate_grids:
+            error = abs(log_ratio - math.log(grid[0] / grid[1]))
+            if error < min_error:
+                best_grid = grid
+                min_error = error
+
+        return best_grid
+
+    def get_slice_image_placeholder(self, image_size, image_idx=0, max_slice_nums=None, use_image_id=None):
+        max_slice_nums = self.max_slice_nums if max_slice_nums is None else int(max_slice_nums)
+        assert max_slice_nums > 0
+        grid = self.get_sliced_grid(image_size=image_size, max_slice_nums=max_slice_nums)
+
+        image_placeholder = self.im_start_token + self.unk_token * self.image_feature_size + self.im_end_token
+        use_image_id = self.use_image_id if use_image_id is None else bool(use_image_id)
+        if use_image_id:
+            final_placeholder = self.get_image_id_placeholder(image_idx) + image_placeholder
+        else:
+            final_placeholder = image_placeholder
+
+        if self.slice_mode:
+            final_placeholder = final_placeholder + self.get_grid_placeholder(grid=grid)
+        return final_placeholder
+
+    def to_pil_image(self, image, rescale=None) -> PIL.Image.Image:
+        """
+        Converts `image` to a PIL Image. Optionally rescales it and puts the channel dimension back as the last axis if
+        needed.
+        Args:
+            image (`PIL.Image.Image` or `numpy.ndarray` or `torch.Tensor`):
+                The image to convert to the PIL Image format.
+            rescale (`bool`, *optional*):
+                Whether or not to apply the scaling factor (to make pixel values integers between 0 and 255). Will
+                default to `True` if the image type is a floating type, `False` otherwise.
+        """
+        if isinstance(image, PIL.Image.Image):
+            return image
+        if is_torch_tensor(image):
+            image = image.numpy()
+
+        if isinstance(image, np.ndarray):
+            if rescale is None:
+                # rescale default to the array being of floating type.
+                rescale = isinstance(image.flat[0], np.floating)
+            # If the channel as been moved to first dim, we put it back at the end.
+            if image.ndim == 3 and image.shape[0] in [1, 3]:
+                image = image.transpose(1, 2, 0)
+            if rescale:
+                image = image * 255
+            image = image.astype(np.uint8)
+            return PIL.Image.fromarray(image)
+        return image
+
+    def reshape_by_patch(self, image):
+        """
+        :param image: shape [3, H, W]
+        :param patch_size:
+        :return: [3, patch_size, HW/patch_size]
+        """
+        image = torch.from_numpy(image)
+        patch_size = self.patch_size
+        patches = torch.nn.functional.unfold(image, (patch_size, patch_size), stride=(patch_size, patch_size))
+
+        patches = patches.reshape(image.size(0), patch_size, patch_size, -1)
+        patches = patches.permute(0, 1, 3, 2).reshape(image.size(0), patch_size, -1)
+        return patches.numpy()
+
+    def preprocess(
+        self,
+        images: Union[Image.Image, List[Image.Image], List[List[Image.Image]]],
+        do_pad: Optional[bool] = True,
+        max_slice_nums: int = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        **kwargs,
+    ) -> MegrezOBatchFeature:
+        if isinstance(images, Image.Image):
+            images_list = [[images]]
+        elif isinstance(images[0], Image.Image):
+            images_list = [images]
+        else:
+            images_list = images
+
+        new_images_list = []
+        image_sizes_list = []
+        tgt_sizes_list = []
+
+        for _images in images_list:
+            if _images is None or len(_images) == 0:
+                new_images_list.append([])
+                image_sizes_list.append([])
+                tgt_sizes_list.append([])
+                continue
+            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."
+                )
+
+            _images = [self.to_pil_image(image).convert("RGB") for image in _images]
+            input_data_format = infer_channel_dimension_format(np.array(_images[0]))
+
+            new_images = []
+            image_sizes = np.array([image.size for image in _images])
+            tgt_sizes = []
+            for image in _images:
+                image_patches = self.get_sliced_images(image, max_slice_nums)
+                image_patches = [to_numpy_array(image).astype(np.float32) for image in image_patches]
+                # image_patches = [to_numpy_array(image).astype(np.float32) / 255 for image in image_patches]
+                # image_patches = [
+                #     self.normalize(image=image, mean=self.mean, std=self.std, input_data_format=input_data_format)
+                #         for image in image_patches
+                # ]
+                image_patches = [
+                    to_channel_dimension_format(image, ChannelDimension.FIRST, input_channel_dim=input_data_format)
+                    for image in image_patches
+                ]
+                for slice_image in image_patches:
+                    new_images.append(self.reshape_by_patch(slice_image))
+                    tgt_sizes.append(
+                        np.array((slice_image.shape[1] // self.patch_size, slice_image.shape[2] // self.patch_size))
+                    )
+
+            if tgt_sizes:
+                tgt_sizes = np.vstack(tgt_sizes)
+
+            new_images_list.append(new_images)
+            image_sizes_list.append(image_sizes)
+            tgt_sizes_list.append(tgt_sizes)
+        return MegrezOBatchFeature(
+            data={"pixel_values": new_images_list, "image_sizes": image_sizes_list, "tgt_sizes": tgt_sizes_list},
+            tensor_type=return_tensors,
+        )
+
+
+AutoImageProcessor.register("MegrezOImageProcessor", MegrezOImageProcessor)

model-00001-of-00002.safetensors

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

model-00002-of-00002.safetensors

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

modeling_megrezo.py

@@ -0,0 +1,328 @@
+# -*- encoding: utf-8 -*-
+# File: modeling_megrezo.py
+# Description: This file contains the implementation of the Megrez-Omni model.
+
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from transformers import AutoProcessor
+from transformers import LlamaForCausalLM
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import add_start_docstrings
+from transformers.utils import add_start_docstrings_to_model_forward
+from transformers.utils import is_flash_attn_2_available
+from transformers.utils import is_flash_attn_greater_or_equal_2_10
+from transformers.utils import logging
+from transformers.utils import replace_return_docstrings
+
+from .audio import AudioEncoder
+from .configuration_megrezo import MegrezOConfig
+from .modeling_navit_siglip import SiglipVisionTransformer
+from .resampler import Resampler
+
+
+def insert_audio_embeddings(text_embeddings, inserted_embeddings, inserted_bounds):
+
+    inserted_bounds = inserted_bounds.long()
+
+    for idx in range(len(inserted_embeddings)):
+        bid = inserted_bounds[idx][0]
+        start_id = inserted_bounds[idx][1]
+        end_id = inserted_bounds[idx][2]
+        embedding = inserted_embeddings[idx]
+        text_embeddings[bid, start_id + 1 : end_id] = embedding
+
+    return text_embeddings
+
+
+def insert_image_embeddings(text_embeddings, inserted_embeddings, inserted_bounds):
+
+    inserted_bounds = inserted_bounds.long()
+    for idx in range(len(inserted_embeddings)):
+        bid = inserted_bounds[idx][0]
+        start_id = inserted_bounds[idx][1]
+        end_id = inserted_bounds[idx][2]
+        embedding = inserted_embeddings[idx]
+        text_embeddings[bid, start_id:end_id] = embedding
+
+    return text_embeddings
+
+
+MegrezO_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 ([`MegrezOConfig`]):
+            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 MegrezO Model outputting raw hidden-states without any specific head on top.",
+    MegrezO_START_DOCSTRING,
+)
+class MegrezOPreTrainedModel(PreTrainedModel):
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    config_class = MegrezOConfig
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+
+
+class AudioModel(torch.nn.Module):
+
+    def __init__(self, config: MegrezOConfig):
+        super(AudioModel, self).__init__()
+        self.config = config
+        self.audio = AudioEncoder(**config.audio_config.to_dict())
+
+    def forward(self, audio_info):
+        audios = audio_info["input_audios"]
+        input_audio_lengths = audio_info["input_audio_lengths"]
+        audio_span_tokens = audio_info["audio_span_tokens"]
+        audios_features = self.audio.encode(audios, input_audio_lengths, audio_span_tokens)
+        return audios_features
+
+
+class VisionModel(torch.nn.Module):
+
+    def __init__(self, config: MegrezOConfig):
+        super(VisionModel, self).__init__()
+        self.config = config
+        self.vpm = self.init_vision_module()
+        self.resampler = self.init_resampler(self.config.hidden_size, self.vpm.embed_dim)
+
+    def init_vision_module(self):
+        if self.config._attn_implementation == "flash_attention_2":
+            self.config.vision_config._attn_implementation = "flash_attention_2"
+        else:
+            # not suport sdpa
+            self.config.vision_config._attn_implementation = "eager"
+        model = SiglipVisionTransformer(self.config.vision_config)
+        if self.config.drop_vision_last_layer:
+            model.encoder.layers = model.encoder.layers[:-1]
+
+        setattr(model, "embed_dim", model.embeddings.embed_dim)
+        setattr(model, "patch_size", model.embeddings.patch_size)
+
+        return model
+
+    def init_resampler(self, embed_dim, vision_dim):
+        return Resampler(
+            num_queries=self.config.query_num,
+            embed_dim=embed_dim,
+            num_heads=embed_dim // 128,
+            kv_dim=vision_dim,
+            adaptive=True,
+        )
+
+    def get_vision_embedding(
+        self,
+        all_pixel_values: torch.Tensor,
+        patch_attention_mask: torch.Tensor,
+        tgt_sizes: torch.Tensor,
+    ):
+        B = all_pixel_values.size(0)
+        vision_batch_size = self.config.vision_batch_size
+        if B > vision_batch_size:
+            hs = []
+            for i in range(0, B, vision_batch_size):
+                start_idx = i
+                end_idx = i + vision_batch_size
+                tmp_hs = self.vpm(
+                    all_pixel_values[start_idx:end_idx],
+                    patch_attention_mask=patch_attention_mask[start_idx:end_idx],
+                    tgt_sizes=tgt_sizes[start_idx:end_idx],
+                ).last_hidden_state
+                hs.append(tmp_hs)
+            vision_embedding = torch.cat(hs, dim=0)
+        else:
+            vision_embedding = self.vpm(
+                all_pixel_values,
+                patch_attention_mask=patch_attention_mask,
+                tgt_sizes=tgt_sizes,
+            ).last_hidden_state
+
+        return vision_embedding
+
+    def _prepare_vision_input(self, images, patch_attention_mask, tgt_sizes):
+        # (TODO) Move to processor
+        device = self.vpm.device
+        dtype = self.vpm.dtype
+
+        pixel_values = torch.stack([(image.to(device) - 127.5) / 127.5 for image in images]).type(dtype)
+        patch_attention_mask = patch_attention_mask.to(device)
+        return pixel_values, patch_attention_mask, tgt_sizes
+
+    def forward(self, images, tgt_sizes, patch_attention_mask):
+        pixel_values, patch_attention_mask, tgt_sizes = self._prepare_vision_input(
+            images, patch_attention_mask, tgt_sizes
+        )
+        embedding = self.get_vision_embedding(pixel_values, patch_attention_mask, tgt_sizes)
+        embedding = self.resampler(embedding, tgt_sizes)
+        return embedding
+
+
+class MegrezO(MegrezOPreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.llm = LlamaForCausalLM(config)
+        self.vision = VisionModel(config)
+        self.audio = AudioModel(config)
+        self.post_init()
+        self.processor = None
+
+        # Will be set in the training script
+        self.tune_vision = False
+        self.tune_audio = False
+
+    def _get_or_init_processor(self):
+
+        if self.processor is None:
+            self.processor = AutoProcessor.from_pretrained(
+                self.config._name_or_path,
+                trust_remote_code=True,
+            )
+
+        return self.processor
+
+    def convert_to_device(self, mini_batch):
+        for key in mini_batch:
+            if isinstance(mini_batch[key], torch.Tensor):
+                mini_batch[key] = mini_batch[key].to(self.device)
+            if isinstance(mini_batch[key], list):
+                return_value = []
+                for value in mini_batch[key]:
+                    if isinstance(value, torch.Tensor):
+                        value = value.to(self.device)
+                    return_value.append(value)
+                mini_batch[key] = return_value
+
+        return mini_batch
+
+    def compose_embeddings(self, mini_batch):
+        position_ids = mini_batch["position_ids"]
+        input_ids = mini_batch["input_ids"]
+        image_encoding = mini_batch.get("image_encoding")
+        audio_encoding = mini_batch.get("audio_encoding")
+        if position_ids.dtype != torch.int64:
+            position_ids = position_ids.long()
+
+        embeddings_text = self.llm.model.embed_tokens(input_ids)
+        input_embeds = embeddings_text
+        if image_encoding:
+            pixel_values = image_encoding["pixel_values"]
+            tgt_sizes = image_encoding["tgt_sizes"]
+            patch_attention_mask = image_encoding["patch_attention_mask"]
+            bounds_image = image_encoding["image_bounds"]
+            embeddings_image = self.vision(pixel_values, tgt_sizes, patch_attention_mask=patch_attention_mask)
+            input_embeds = insert_image_embeddings(embeddings_text, embeddings_image, bounds_image)
+        elif self.training and self.tune_vision:
+            pixel_values = torch.zeros((3, 14, 3584), dtype=torch.float32)
+            tgt_sizes = torch.tensor([[16, 16]], dtype=torch.int64)
+            patch_attention_mask = torch.ones((3, 14), dtype=torch.float32)
+            embeddings_image = self.vision(pixel_values, tgt_sizes, patch_attention_mask=patch_attention_mask)
+            input_embeds += embeddings_image[0].sum() * 0.0
+
+        if audio_encoding:
+            embeddings_audio = self.audio(audio_encoding)
+            bounds_audio = audio_encoding["audio_bounds"]
+            input_embeds = insert_audio_embeddings(embeddings_text, embeddings_audio, bounds_audio)
+        elif self.training and self.tune_audio:
+            dummy_audio = torch.zeros((1, 128, 3000), dtype=torch.float32)
+            dummy_audio_lengths = torch.tensor([[125, 62]], dtype=torch.int32)
+            dummy_span_tokens = [64]
+            dummy_audio_encoding = [
+                {
+                    "input_audios": dummy_audio,
+                    "input_audio_lengths": dummy_audio_lengths,
+                    "audio_span_tokens": dummy_span_tokens,
+                }
+            ]
+            embeddings_audio = self.audio(dummy_audio_encoding)
+            input_embeds += embeddings_audio[0].sum() * 0.0
+
+        return input_ids, input_embeds, position_ids
+
+    def forward(self, data, **kwargs):
+        if self.training:
+            _, input_embeds, position_ids = self.compose_embeddings(data)
+            return self.llm.forward(
+                input_ids=None,
+                position_ids=position_ids,
+                inputs_embeds=input_embeds,
+                **kwargs,
+            )
+        return self.llm.forward(**kwargs)
+
+    def generate(
+        self,
+        input_ids,
+        position_ids,
+        attention_mask,
+        image_encoding=None,
+        audio_encoding=None,
+        **kwargs,
+    ):
+        tokenizer = self._get_or_init_processor().tokenizer
+        data = {
+            "input_ids": input_ids,
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "image_encoding": image_encoding,
+            "audio_encoding": audio_encoding,
+        }
+        data = self.convert_to_device(data)
+        input_ids, input_embeds, position_ids = self.compose_embeddings(data)
+
+        output = self.llm.generate(
+            inputs_embeds=input_embeds,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            **kwargs,
+        )
+        return output
+
+    def trim_stop_words(self, response, stop_words):
+        if stop_words:
+            for stop in stop_words:
+                idx = response.find(stop)
+                if idx != -1:
+                    response = response[:idx]
+        return response
+
+    @torch.inference_mode()
+    def chat(self, input_msgs, processor=None, sampling=False, **kwargs):
+        if processor is None:
+            processor = self._get_or_init_processor()
+
+        if sampling:
+            generation_config = {
+                "top_p": 0.8,
+                "top_k": 100,
+                "temperature": 0.7,
+                "do_sample": True,
+                "repetition_penalty": 1.05,
+            }
+        else:
+            generation_config = {
+                "num_beams": 1,
+                "repetition_penalty": 1.2,
+            }
+
+        generation_config.update(kwargs)
+        if generation_config.get("temperature") == 0:
+            generation_config["do_sample"] = False
+
+        data = processor(input_msgs)
+        output_ids = self.generate(**data, **generation_config)
+        tokenizer = processor.tokenizer
+        answer = tokenizer.decode(output_ids[0])
+        return answer

modeling_navit_siglip.py

@@ -0,0 +1,937 @@
+# coding=utf-8
+# Copyright 2024 Google AI and The HuggingFace 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.
+""" PyTorch Siglip model. """
+# Copied from  HuggingFaceM4/siglip-so400m-14-980-flash-attn2-navit and add tgt_sizes
+
+
+import os
+import math
+import warnings
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn.init import _calculate_fan_in_and_fan_out
+
+from transformers.activations import ACT2FN
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from transformers.modeling_utils import PreTrainedModel
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import (
+    ModelOutput,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class SiglipVisionConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`SiglipVisionModel`]. It is used to instantiate a
+    Siglip vision encoder according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the vision encoder of the Siglip
+    [google/siglip-base-patch16-224](https://huggingface.co/google/siglip-base-patch16-224) architecture.
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_channels (`int`, *optional*, defaults to 3):
+            Number of channels in the input images.
+        image_size (`int`, *optional*, defaults to 224):
+            The size (resolution) of each image.
+        patch_size (`int`, *optional*, defaults to 16):
+            The size (resolution) of each patch.
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` ``"quick_gelu"` are supported.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the layer normalization layers.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+    Example:
+    ```python
+    >>> from transformers import SiglipVisionConfig, SiglipVisionModel
+    >>> # Initializing a SiglipVisionConfig with google/siglip-base-patch16-224 style configuration
+    >>> configuration = SiglipVisionConfig()
+    >>> # Initializing a SiglipVisionModel (with random weights) from the google/siglip-base-patch16-224 style configuration
+    >>> model = SiglipVisionModel(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "siglip_vision_model"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        intermediate_size=3072,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        num_channels=3,
+        image_size=224,
+        patch_size=16,
+        hidden_act="gelu_pytorch_tanh",
+        layer_norm_eps=1e-6,
+        attention_dropout=0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        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.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
+        cls._set_token_in_kwargs(kwargs)
+
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        # get the vision config dict if we are loading from SiglipConfig
+        if config_dict.get("model_type") == "siglip":
+            config_dict = config_dict["vision_config"]
+
+        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+_CHECKPOINT_FOR_DOC = "google/siglip-base-patch16-224"
+
+SIGLIP_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "google/siglip-base-patch16-224",
+    # See all SigLIP models at https://huggingface.co/models?filter=siglip
+]
+
+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
+
+
+# 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,
+    )
+
+
+def _trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+            "The distribution of values may be incorrect.",
+            stacklevel=2,
+        )
+
+    # Values are generated by using a truncated uniform distribution and
+    # then using the inverse CDF for the normal distribution.
+    # Get upper and lower cdf values
+    l = norm_cdf((a - mean) / std)
+    u = norm_cdf((b - mean) / std)
+
+    # Uniformly fill tensor with values from [l, u], then translate to
+    # [2l-1, 2u-1].
+    tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+    # Use inverse cdf transform for normal distribution to get truncated
+    # standard normal
+    if tensor.dtype in [torch.float16, torch.bfloat16]:
+        # The `erfinv_` op is not (yet?) defined in float16+cpu, bfloat16+gpu
+        og_dtype = tensor.dtype
+        tensor = tensor.to(torch.float32)
+        tensor.erfinv_()
+        tensor = tensor.to(og_dtype)
+    else:
+        tensor.erfinv_()
+
+    # Transform to proper mean, std
+    tensor.mul_(std * math.sqrt(2.0))
+    tensor.add_(mean)
+
+    # Clamp to ensure it's in the proper range
+    if tensor.dtype == torch.float16:
+        # The `clamp_` op is not (yet?) defined in float16+cpu
+        tensor = tensor.to(torch.float32)
+        tensor.clamp_(min=a, max=b)
+        tensor = tensor.to(torch.float16)
+    else:
+        tensor.clamp_(min=a, max=b)
+
+
+def trunc_normal_tf_(
+    tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, a: float = -2.0, b: float = 2.0
+) -> torch.Tensor:
+    """Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \\leq \text{mean} \\leq b`.
+    NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the
+    bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0
+    and the result is subsquently scaled and shifted by the mean and std args.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    """
+    with torch.no_grad():
+        _trunc_normal_(tensor, 0, 1.0, a, b)
+        tensor.mul_(std).add_(mean)
+
+
+def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"):
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    if mode == "fan_in":
+        denom = fan_in
+    elif mode == "fan_out":
+        denom = fan_out
+    elif mode == "fan_avg":
+        denom = (fan_in + fan_out) / 2
+
+    variance = scale / denom
+
+    if distribution == "truncated_normal":
+        # constant is stddev of standard normal truncated to (-2, 2)
+        trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978)
+    elif distribution == "normal":
+        with torch.no_grad():
+            tensor.normal_(std=math.sqrt(variance))
+    elif distribution == "uniform":
+        bound = math.sqrt(3 * variance)
+        with torch.no_grad():
+            tensor.uniform_(-bound, bound)
+    else:
+        raise ValueError(f"invalid distribution {distribution}")
+
+
+def lecun_normal_(tensor):
+    variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal")
+
+
+def default_flax_embed_init(tensor):
+    variance_scaling_(tensor, mode="fan_in", distribution="normal")
+
+
+@dataclass
+# Copied from transformers.models.clip.modeling_clip.CLIPVisionModelOutput with CLIP->Siglip
+class SiglipVisionModelOutput(ModelOutput):
+    """
+    Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states.
+    Args:
+        image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
+            The image embeddings obtained by applying the projection layer to the pooler_output.
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    image_embeds: Optional[torch.FloatTensor] = None
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+class SiglipVisionEmbeddings(nn.Module):
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            padding="valid",
+        )
+
+        self.num_patches_per_side = self.image_size // self.patch_size
+        self.num_patches = self.num_patches_per_side**2
+        self.num_positions = self.num_patches
+        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        patch_attention_mask: torch.BoolTensor,
+        tgt_sizes: Optional[torch.IntTensor] = None,
+    ) -> torch.Tensor:
+        batch_size = pixel_values.size(0)
+
+        patch_embeds = self.patch_embedding(pixel_values)
+        embeddings = patch_embeds.flatten(2).transpose(1, 2)
+
+        max_im_h, max_im_w = pixel_values.size(2), pixel_values.size(3)
+        max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size
+        boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side)
+        position_ids = torch.full(
+            size=(
+                batch_size,
+                max_nb_patches_h * max_nb_patches_w,
+            ),
+            fill_value=0,
+        )
+
+        for batch_idx, p_attn_mask in enumerate(patch_attention_mask):
+            if tgt_sizes is not None:
+                nb_patches_h = tgt_sizes[batch_idx][0]
+                nb_patches_w = tgt_sizes[batch_idx][1]
+            else:
+                nb_patches_h = p_attn_mask[:, 0].sum()
+                nb_patches_w = p_attn_mask[0].sum()
+
+            fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h)
+            fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w)
+
+            bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True)
+            bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True)
+
+            pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten()
+            position_ids[batch_idx][p_attn_mask.view(-1).cpu()] = pos_ids
+
+        position_ids = position_ids.to(self.position_embedding.weight.device)
+
+        embeddings = embeddings + self.position_embedding(position_ids)
+        return embeddings
+
+
+class SiglipAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    # Copied from transformers.models.clip.modeling_clip.CLIPAttention.__init__
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+        self.scale = self.head_dim**-0.5
+        self.dropout = config.attention_dropout
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """Input shape: Batch x Time x Channel"""
+
+        batch_size, 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(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        k_v_seq_len = key_states.shape[-2]
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale
+
+        if attn_weights.size() != (batch_size, self.num_heads, q_len, k_v_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(batch_size, self.num_heads, q_len, k_v_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (batch_size, 1, q_len, k_v_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(batch_size, 1, q_len, k_v_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.dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (batch_size, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(batch_size, 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(batch_size, q_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+class SiglipFlashAttention2(SiglipAttention):
+    """
+    Llama flash attention module. This module inherits from `LlamaAttention` 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.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.is_causal = False  # Hack to make sure we don't use a causal mask
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        output_attentions = False
+
+        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)
+
+        # Flash attention requires the input to have the shape
+        # batch_size x seq_length x head_dim x hidden_dim
+        # therefore we just need to keep the original shape
+        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_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_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_usable_length(kv_seq_len, 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)
+
+        # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+        # to be able to avoid many of these transpose/reshape/view.
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        dropout_rate = self.dropout if self.training else 0.0
+
+        # 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 the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (LlamaRMSNorm handles it correctly)
+
+        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(
+                "The input hidden states seems to be silently casted in float32, this might be related to 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)
+
+        attn_output = self._flash_attention_forward(
+            query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, self.embed_dim).contiguous()
+        attn_output = self.out_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights
+
+    def _flash_attention_forward(
+        self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+    ):
+        """
+        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)
+        """
+
+        # 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
+
+        # 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
+
+            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,
+            )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
+            )
+
+        return attn_output
+
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.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.clip.modeling_clip.CLIPMLP with CLIP->Siglip
+class SiglipMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->Siglip
+class SiglipEncoderLayer(nn.Module):
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        self.self_attn = SiglipAttention(config) if not self._use_flash_attention_2 else SiglipFlashAttention2(config)
+        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.mlp = SiglipMLP(config)
+        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`):
+                Input to the layer of shape `(batch, seq_len, embed_dim)`.
+            attention_mask (`torch.FloatTensor`):
+                Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*, defaults to `False`):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        residual = hidden_states
+
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+class SiglipPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = SiglipVisionConfig
+    base_model_prefix = "siglip"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+
+        if isinstance(module, SiglipVisionEmbeddings):
+            width = self.config.hidden_size
+            nn.init.normal_(module.position_embedding.weight, std=1 / np.sqrt(width))
+        elif isinstance(module, nn.Embedding):
+            default_flax_embed_init(module.weight)
+        elif isinstance(module, SiglipAttention):
+            nn.init.normal_(module.q_proj.weight)
+            nn.init.normal_(module.k_proj.weight)
+            nn.init.normal_(module.v_proj.weight)
+            nn.init.normal_(module.out_proj.weight)
+            nn.init.zeros_(module.q_proj.bias)
+            nn.init.zeros_(module.k_proj.bias)
+            nn.init.zeros_(module.v_proj.bias)
+            nn.init.zeros_(module.out_proj.bias)
+        elif isinstance(module, SiglipMLP):
+            nn.init.normal_(module.fc1.weight)
+            nn.init.normal_(module.fc2.weight)
+            nn.init.normal_(module.fc1.bias, std=1e-6)
+            nn.init.normal_(module.fc2.bias, std=1e-6)
+        elif isinstance(module, (nn.Linear, nn.Conv2d)):
+            lecun_normal_(module.weight)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+SIGLIP_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 ([`SiglipVisionConfig`]): 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.
+"""
+
+
+SIGLIP_VISION_INPUTS_DOCSTRING = r"""
+    Args:
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
+            [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details.
+        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.
+"""
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->Siglip
+class SiglipEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`SiglipEncoderLayer`].
+    Args:
+        config: SiglipConfig
+    """
+
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    # Ignore copy
+    def forward(
+        self,
+        inputs_embeds,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                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.
+            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)
+            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.
+        """
+        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
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        hidden_states = inputs_embeds
+        for encoder_layer in self.layers:
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    encoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    output_attentions,
+                )
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                    attention_mask,
+                    output_attentions=output_attentions,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
+        return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)
+
+
+@add_start_docstrings("""The vision model from SigLIP without any head or projection on top.""", SIGLIP_START_DOCSTRING)
+class SiglipVisionTransformer(SiglipPreTrainedModel):
+    config_class = SiglipVisionConfig
+    main_input_name = "pixel_values"
+    _supports_flash_attn_2 = True
+
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__(config)
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.embeddings = SiglipVisionEmbeddings(config)
+        self.encoder = SiglipEncoder(config)
+        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.embeddings.patch_embedding
+
+    @add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=SiglipVisionConfig)
+    def forward(
+        self,
+        pixel_values,
+        patch_attention_mask: Optional[torch.BoolTensor] = None,
+        tgt_sizes: Optional[torch.IntTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+        """
+        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
+
+        batch_size = pixel_values.size(0)
+        if patch_attention_mask is None:
+            patch_attention_mask = torch.ones(
+                size=(
+                    batch_size,
+                    pixel_values.size(2) // self.config.patch_size,
+                    pixel_values.size(3) // self.config.patch_size,
+                ),
+                dtype=torch.bool,
+                device=pixel_values.device,
+            )
+
+        hidden_states = self.embeddings(
+            pixel_values=pixel_values, patch_attention_mask=patch_attention_mask, tgt_sizes=tgt_sizes
+        )
+
+        patch_attention_mask = patch_attention_mask.view(batch_size, -1)
+        # The call to `_upad_input` in `_flash_attention_forward` is expensive
+        # So when the `patch_attention_mask` is full of 1s (i.e. attending to the whole sequence),
+        # avoiding passing the attention_mask, which is equivalent to attending to the full sequence
+        if not torch.any(~patch_attention_mask):
+            attention_mask = None
+        else:
+            attention_mask = (
+                _prepare_4d_attention_mask(patch_attention_mask, hidden_states.dtype)
+                if not self._use_flash_attention_2
+                else patch_attention_mask
+            )
+
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[0]
+        last_hidden_state = self.post_layernorm(last_hidden_state)
+
+        if not return_dict:
+            return (last_hidden_state, None) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=None,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )

preprocessor_config.json

@@ -0,0 +1,46 @@
+{
+  "im_end": "<|image_end|>",
+  "im_end_token": "<|image_end|>",
+  "im_id_end": "<|image_id_end|>",
+  "im_id_start": "<|image_id_start|>",
+  "im_start": "<|image_start|>",
+  "im_start_token": "<|image_start|>",
+  "image_feature_size": 64,
+  "image_processor_type": "MegrezOImageProcessor",
+  "audio_feature_extractor_type": "WhisperFeatureExtractor",
+  "feature_size": 128,
+  "max_slice_nums": 9,
+  "mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "norm_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "norm_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "patch_size": 14,
+  "processor_class": "MegrezOProcessor",
+  "sampling_rate": 16000,
+  "scale_resolution": 448,
+  "slice_end": "<|slice_end|>",
+  "slice_end_token": "<|slice_end|>",
+  "slice_mode": true,
+  "slice_start": "<|slice_start|>",
+  "slice_start_token": "<|slice_start|>",
+  "std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "unk": "<|unk|>",
+  "unk_token": "<|unk|>",
+  "pad_token": "<|pad|>",
+  "use_image_id": true
+}

processing_megrezo.py

@@ -0,0 +1,587 @@
+# -*- encoding: utf-8 -*-
+# File: processing_megrezo.py
+# Description: None
+
+import io
+import re
+import subprocess
+from collections import UserDict
+from typing import List, Literal, Optional, Tuple, Union
+
+import numpy as np
+import PIL
+import PIL.Image
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from transformers import TensorType
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.image_utils import ImageInput
+from transformers.processing_utils import ProcessorMixin
+
+from .image_processing_megrezo import MegrezOImageProcessor  # noqa: F401
+
+AudioInput = Union[str, bytes, "np.ndarray", List[str], List[bytes], List["np.ndarray"]]
+ReturnTensorType = Union[str, TensorType]
+
+
+class ImageBatchFeature(BatchFeature):
+    r"""
+    Holds the image features of a batch of images.
+    """
+
+    pixel_values: Union[np.ndarray, torch.Tensor]
+    image_sizes: Union[np.ndarray, torch.Tensor]
+    tgt_sizes: Union[np.ndarray, torch.Tensor]
+    patch_attention_mask: Union[np.ndarray, torch.Tensor]
+    image_bounds: Union[np.ndarray, torch.Tensor]
+
+
+class AudioBatchFeature(BatchFeature):
+    r"""
+    Holds the audio features of a batch of audio.
+    """
+
+    input_audios: List[Union[np.ndarray, torch.Tensor]]
+    input_audio_lengths: List[Union[np.ndarray, torch.Tensor]]
+    audio_span_tokens: List[Union[np.ndarray, torch.Tensor]]
+    audio_bounds: Union[np.ndarray, torch.Tensor]
+
+
+class ConvContent(UserDict):
+    text: Optional[str] = None
+    image: Optional[ImageInput] = None
+    audio: Optional[Union[str, bytes, List[Union[str, bytes]]]] = None
+
+
+class Conversation(UserDict):
+    role: Literal["user", "assistant"]
+    content: Union[str, dict, ConvContent]
+
+
+def load_audio(
+    audio: Union[str, bytes],
+    sample_rate: int = 16000,
+) -> "np.ndarray":
+    """Load audio from a file path or bytes and return as a numpy array.
+
+    Args:
+        audio (Union[str, bytes]): path to a audio file or audio bytes.
+        sample_rate (int, optional): sample rate. Defaults to 16000.
+
+    Raises:
+        ValueError: if the input audio is neither a path nor bytes.
+
+    Returns:
+        np.ndarray: the audio as a numpy array.
+    """
+    if isinstance(audio, str):
+        inp = audio
+        out = "-"
+        cmd_inp = None
+    elif isinstance(audio, bytes):
+        inp = "pipe:"
+        out = "pipe:"
+        cmd_inp = audio
+    else:
+        raise ValueError("input audio must be either a path or bytes")
+
+    cmd = [
+        "ffmpeg",
+        "-nostdin",
+        "-threads",
+        "0",
+        "-i",
+        inp,
+        "-f",
+        "s16le",
+        "-ac",
+        "1",
+        "-acodec",
+        "pcm_s16le",
+        "-ar",
+        str(sample_rate),
+        out,
+    ]
+
+    out = subprocess.check_output(cmd, input=cmd_inp, stderr=subprocess.PIPE)
+    arr = np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+    return arr
+
+
+def load_image(
+    image: Union[str, bytes, PIL.Image.Image],
+) -> PIL.Image.Image:
+    """Load image from a file path or bytes and return as a PIL image.
+
+    Args:
+        image (Union[str, bytes, PIL.Image.Image]): path to an image file, image bytes or a PIL image.
+
+    Raises:
+        ValueError: if the input image is neither a path nor bytes.
+
+    Returns:
+        PIL.Image.Image: the image as a PIL image.
+    """
+    if isinstance(image, PIL.Image.Image):
+        return image
+
+    if isinstance(image, str):
+        img = PIL.Image.open(image)
+    elif isinstance(image, bytes):
+        img = PIL.Image.open(io.BytesIO(image))
+    else:
+        raise ValueError("input image must be either a path or bytes")
+
+    return img
+
+
+class MegrezOProcessor(ProcessorMixin):
+    attributes = ["image_processor", "audio_feature_extractor", "tokenizer"]
+    image_processor_class = "AutoImageProcessor"
+    audio_feature_extractor_class = "WhisperFeatureExtractor"
+    tokenizer_class = "AutoTokenizer"
+
+    _image_placeholder = r"(<image>./</image>)"
+    _audio_placeholder = r"(<audio>./</audio>)"
+
+    def __init__(self, image_processor=None, audio_feature_extractor=None, tokenizer=None):
+        super().__init__(image_processor, audio_feature_extractor, tokenizer)
+        self.chat_template = self.tokenizer.chat_template
+
+    def _parse_and_check_inputs(self, inputs) -> List[Conversation]:
+        if not isinstance(inputs, list):
+            raise ValueError("inputs must be a list of conversations")
+
+        conversations = []
+        images = []
+        audios = []
+
+        for input in inputs:
+            if not isinstance(input, dict) and not isinstance(input, Conversation):
+                raise ValueError("each element of inputs must be a dictionary or a Conversation object")
+
+            role = input.get("role")
+            content = input.get("content")
+            if role is None or content is None:
+                raise ValueError("role and content must be provided in each conversation")
+
+            if isinstance(content, str):
+                content = content
+            elif isinstance(content, dict):
+                content = ConvContent({**content})
+            elif not isinstance(content, ConvContent):
+                raise ValueError("content must be a dictionary or a ConvContent object")
+
+            if not isinstance(content, str):
+                if content.get("image") is not None:
+                    images.extend(content["image"] if isinstance(content["image"], list) else [content["image"]])
+
+                if content.get("audio") is not None:
+                    audios.extend(content["audio"] if isinstance(content["audio"], list) else [content["audio"]])
+
+            conv = Conversation({"role": role, "content": content})
+            conversations.append(conv)
+
+        return conversations, images, audios
+
+    def __call__(
+        self,
+        conversations: List[Conversation],
+        apply_chat_template: bool = True,
+        max_length: Optional[int] = None,
+        return_tensors: ReturnTensorType = TensorType.PYTORCH,
+        apply_data_collator: bool = True,
+        **kwargs,
+    ):
+        assert return_tensors is TensorType.PYTORCH, "Only PyTorch tensors are supported for now."
+        convs, images, audios = self._parse_and_check_inputs(conversations)
+        add_generation_prompt = kwargs.pop("add_generation_prompt", True)
+        if apply_chat_template:
+            prompt = self.tokenizer.apply_chat_template(
+                convs,
+                tokenize=False,
+                add_generation_prompt=add_generation_prompt,
+            )
+        else:  # (TODO) For clarification temporarily. Check whether this needs to be removed.
+            prompt = "\n".join([conv["content"] for conv in convs])
+
+        prompt, multimodal_inputs = self.process_multimodal_inputs(
+            prompt,
+            images=images,
+            audios=audios,
+            return_tensors=return_tensors,
+            **kwargs,
+        )
+        text_encodings = self.tokenizer(
+            prompt,
+            return_tensors=return_tensors,
+            max_length=max_length,
+            padding=True,
+            padding_side="left",
+            truncation=True,
+            **kwargs,
+        )
+
+        merged = self.merge_encodings(text_encodings, multimodal_inputs)
+
+        if apply_data_collator:
+            return self.data_collator([merged])
+
+        return merged
+
+    def merge_encodings(self, text_encodings, multimodal_inputs):
+
+        result = {
+            "image_encoding": None,
+            "audio_encoding": None,
+        }
+
+        result["input_ids"] = text_encodings["input_ids"].reshape(-1).to(torch.int32)
+        result["attention_mask"] = result["input_ids"].ne(0)
+        result["position_ids"] = torch.arange(result["input_ids"].size(0)).long()
+
+        if "image_encoding" in multimodal_inputs and multimodal_inputs["image_encoding"]:
+            result["image_encoding"] = multimodal_inputs["image_encoding"]
+            result["image_encoding"]["image_bounds"] = self.compute_bounds_image(result["input_ids"])
+
+        if "audio_encoding" in multimodal_inputs and multimodal_inputs["audio_encoding"]:
+            result["audio_encoding"] = multimodal_inputs["audio_encoding"]
+            result["audio_encoding"]["audio_bounds"] = self.compute_bounds_audio(result["input_ids"])
+
+        return result
+
+    def compute_bounds_image(self, input_ids: torch.Tensor) -> List[torch.Tensor]:
+        image_start_ids = (
+            torch.where((input_ids == self.tokenizer.im_start_id) | (input_ids == self.tokenizer.slice_start_id))[0] + 1
+        )
+        image_end_ids = torch.where(
+            (input_ids == self.tokenizer.im_end_id) | (input_ids == self.tokenizer.slice_end_id)
+        )[0]
+
+        valid_image_nums = max(len(image_start_ids), len(image_end_ids))
+        bounds_image = torch.hstack(
+            [
+                image_start_ids[:valid_image_nums].unsqueeze(-1),
+                image_end_ids[:valid_image_nums].unsqueeze(-1),
+            ]
+        )
+        return bounds_image
+
+    def compute_bounds_audio(self, input_ids: torch.Tensor) -> torch.Tensor:
+        audio_bos_ids = torch.where(input_ids == self.tokenizer.audio_start_id)[0]
+        audio_eos_ids = torch.where(input_ids == self.tokenizer.audio_end_id)[0]
+        bounds_audio = torch.stack([audio_bos_ids, audio_eos_ids], 1)
+        return bounds_audio
+
+    def process_multimodal_inputs(
+        self,
+        text: str,
+        images: Optional[ImageInput] = None,
+        audios: Optional[Union[str, bytes, List[Union[str, bytes]]]] = None,
+        return_tensors: ReturnTensorType = TensorType.PYTORCH,
+        **kwargs,
+    ):
+        # (NOTE) Only single pair of multimodal input is allowed currently.
+        # (TODO) Check whether single multimodal input is allowed.
+        if text is None and images is None and audios is None:
+            raise ValueError("At least one of text, images or audio must be provided")
+
+        image_processor_kwargs, audio_feature_extractor_kwargs = {}, {}
+        if kwargs:
+            image_processor_kwargs = {
+                k: v for k, v in kwargs.items() if k in self.image_processor._valid_processor_keys
+            }
+            audio_feature_extractor_kwargs = {
+                k: v for k, v in kwargs.items() if k in self.audio_feature_extractor._valid_processor_keys
+            }
+
+        multimodal_encodings = {
+            "image_encoding": None,
+            "audio_encoding": None,
+        }
+
+        if images:
+            image_encoding = self.process_image(
+                images,
+                return_tensors=return_tensors,
+                **image_processor_kwargs,
+            )
+            text = self.insert_image_feature_placeholders(text, image_encoding)
+            multimodal_encodings["image_encoding"] = image_encoding
+
+        if audios:
+            audio_encoding = self.process_audio(
+                audios,
+                **audio_feature_extractor_kwargs,
+            )
+            text = self.insert_audio_feature_placeholders(text, audio_encoding)
+            multimodal_encodings["audio_encoding"] = audio_encoding
+
+        return text, multimodal_encodings
+
+    def insert_image_feature_placeholders(
+        self,
+        prompt: str,
+        image_features: ImageBatchFeature,
+        max_slice_nums: Optional[int] = None,
+        use_image_id: Optional[bool] = None,
+    ) -> List[str]:
+        # Check the number of image tags and the number of images.
+        img_tags = re.findall(self._image_placeholder, prompt)
+        assert len(img_tags) == len(
+            image_features.image_sizes
+        ), f"the number of image tags must match the number of images, got {len(img_tags)} and {len(image_features.image_sizes)}"
+
+        # Replace image tags with image placeholders.
+        text_chunks = prompt.split(self._image_placeholder)
+        final_text = ""
+        for i in range(len(img_tags)):
+            final_text += text_chunks[i] + self.image_processor.get_slice_image_placeholder(
+                image_features.image_sizes[i],
+                i,
+                max_slice_nums,
+                use_image_id,
+            )
+        final_text += text_chunks[-1]
+
+        return final_text
+
+    def insert_audio_feature_placeholders(
+        self,
+        prompt: str,
+        audio_features: AudioBatchFeature,
+    ) -> List[str]:
+        # Check the number of audio tags and the number of audios.
+        audio_tags = re.findall(self._audio_placeholder, prompt)
+        assert len(audio_tags) == len(
+            audio_features.input_audios
+        ), "the number of audio tags must match the number of audios"
+
+        # Replace audio tags with audio placeholders.
+        text_chunks = prompt.split(self._audio_placeholder)
+        final_text = ""
+        for idx in range(len(audio_features.input_audios)):
+            final_text += text_chunks[idx] + (
+                self.tokenizer.audio_start
+                + self.tokenizer.unk_token * audio_features.audio_span_tokens[idx]
+                + self.tokenizer.audio_end
+            )
+        final_text += text_chunks[-1]
+
+        return final_text
+
+    def process_audio(
+        self,
+        audio_input: AudioInput,
+        return_tensors: ReturnTensorType = TensorType.PYTORCH,
+        **kwargs,
+    ) -> AudioBatchFeature:
+        if isinstance(audio_input, list):
+            inputs = [load_audio(x) for x in audio_input]
+        elif isinstance(audio_input, (str, bytes, "np.ndarray")):
+            inputs = [load_audio(audio_input)]
+        else:
+            raise ValueError("audio_input must be a path or bytes or a list of paths/bytes")
+
+        features = self.audio_feature_extractor(
+            inputs,
+            sampling_rate=self.audio_feature_extractor.sampling_rate,
+            return_attention_mask=True,
+            return_token_timestamps=True,
+            padding="max_length",
+            return_tensors=return_tensors,
+            **kwargs,
+        )
+
+        input_lengths = features["num_frames"]
+        input_lengths = (input_lengths - 1) // 2 + 1
+        output_lengths = (input_lengths - 2) // 2 + 1
+        input_audio_lengths = torch.stack([input_lengths, output_lengths], dim=1)
+        audio_span_tokens = (output_lengths + 2).tolist()  # add bos and eos tokens
+
+        data = {
+            "input_audios": features["input_features"],
+            "input_audio_lengths": input_audio_lengths,
+            "audio_span_tokens": audio_span_tokens,
+        }
+
+        # tensor types are already converted in `self.audio_feature_extractor`.
+        return AudioBatchFeature(data={**data})
+
+    def pad_images(
+        self,
+        pixel_values_list: List[torch.Tensor],
+        tgt_sizes: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Pad images to the same size and return the padded pixel values and patch attention mask.
+
+        Sliced pataches may have different sizes. We pad them to the same size and return the padded pixel values and corresponding patch attention mask.
+        """
+
+        all_pixel_values = []
+        for pixel_value in pixel_values_list:
+            all_pixel_values.append(pixel_value.flatten(end_dim=1).permute(1, 0))
+
+        max_patches = torch.max(tgt_sizes[:, 0] * tgt_sizes[:, 1])
+        all_pixel_values = torch.nn.utils.rnn.pad_sequence(all_pixel_values, batch_first=True, padding_value=0.0)
+        B, L, _ = all_pixel_values.shape
+        all_pixel_values = all_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L)
+
+        patch_attention_mask = torch.zeros((B, 1, max_patches), dtype=torch.bool)
+        for i in range(B):
+            patch_attention_mask[i, 0, : tgt_sizes[i][0] * tgt_sizes[i][1]] = True
+
+        return all_pixel_values, patch_attention_mask
+
+    def process_image(
+        self,
+        image_input: ImageInput,
+        do_pad: bool = True,
+        max_slice_nums: Optional[int] = None,
+        return_tensors: ReturnTensorType = TensorType.PYTORCH,
+        **kwargs,
+    ) -> ImageBatchFeature:
+        if isinstance(image_input, list):
+            image_input = [load_image(x) for x in image_input]
+        elif isinstance(image_input, (str, bytes, PIL.Image.Image)):
+            image_input = [load_image(image_input)]
+        else:
+            raise ValueError(f"image_input must be a path or bytes or a list of paths/bytes, not: {type(image_input)}")
+
+        image_features = self.image_processor(
+            image_input,
+            do_pad=do_pad,
+            max_slice_nums=max_slice_nums,
+            return_tensors=return_tensors,
+            **kwargs,
+        )
+
+        # Multiple images are packed into first element of the list. We unpack them here.
+        assert len(image_features.pixel_values) == 1, "images should be packed into one list."
+        pixel_values = image_features.pixel_values[0]
+        tgt_sizes = image_features.tgt_sizes[0]
+        image_sizes = image_features.image_sizes[0]
+
+        pixel_values, patch_attention_mask = self.pad_images(pixel_values, tgt_sizes)
+
+        data = {
+            "pixel_values": pixel_values,
+            "image_sizes": image_sizes,
+            "tgt_sizes": tgt_sizes,
+            "patch_attention_mask": patch_attention_mask,
+        }
+
+        # tensor types are already converted in `self.image_processor`.
+        return ImageBatchFeature(data=data)
+
+    def data_collator(self, examples, padding_value=0, max_length=4096, collate_labels=False):
+        """Collate data for MegrezO model.
+
+        Batch data for MegrezO model. This function trims and pads the input_ids, position_ids, and attention_mask tensors. For bounds tensors, it adds batch index to the bounds.
+        """
+        # (TODO) Remove this function?
+
+        def trim_and_pad(seq, batch_first, padding_value):
+            return pad_sequence(
+                [s[:max_length] for s in seq],
+                batch_first=True,
+                padding_value=padding_value,
+            )
+
+        input_ids = trim_and_pad(
+            [example["input_ids"] for example in examples],
+            batch_first=True,
+            padding_value=padding_value,
+        )
+        position_ids = trim_and_pad(
+            [example["position_ids"] for example in examples],
+            batch_first=True,
+            padding_value=padding_value,
+        )
+
+        attention_mask = trim_and_pad(
+            [example["attention_mask"] for example in examples],
+            batch_first=True,
+            padding_value=padding_value,
+        )
+
+        image_encoding_list = {
+            "pixel_values": [],
+            "image_bounds": [],
+            "tgt_sizes": [],
+            "patch_attention_mask": [],
+        }
+        for bid, example in enumerate(examples):
+            image_encoding = example.get("image_encoding")
+            if not image_encoding:
+                continue
+
+            image_encoding_list["pixel_values"].append(image_encoding["pixel_values"])
+            image_encoding_list["tgt_sizes"].append(image_encoding["tgt_sizes"])
+            image_encoding_list["patch_attention_mask"].append(image_encoding["patch_attention_mask"])
+
+            # (TODO) Remove?
+            # add batch index to bounds (bid, start, end)
+            bounds_with_bid = image_encoding["image_bounds"].clone()
+            bounds_with_bid = torch.hstack(
+                [
+                    torch.full((bounds_with_bid.size(0), 1), bid, dtype=bounds_with_bid.dtype),
+                    bounds_with_bid,
+                ]
+            )
+            image_encoding_list["image_bounds"].append(bounds_with_bid)
+
+        audio_encoding_list = {
+            "input_audios": [],
+            "input_audio_lengths": [],
+            "audio_span_tokens": [],
+            "audio_bounds": [],
+        }
+        for bid, example in enumerate(examples):
+            audio_encoding = example.get("audio_encoding")
+            if not audio_encoding:
+                continue
+
+            audio_encoding_list["input_audios"].append(audio_encoding["input_audios"])
+            audio_encoding_list["input_audio_lengths"].append(audio_encoding["input_audio_lengths"])
+            audio_encoding_list["audio_span_tokens"].extend(audio_encoding["audio_span_tokens"])
+            bounds_with_bid = audio_encoding["audio_bounds"].clone()
+            bounds_with_bid = torch.hstack(
+                [
+                    torch.full((bounds_with_bid.size(0), 1), bid, dtype=bounds_with_bid.dtype),
+                    bounds_with_bid,
+                ]
+            )
+            audio_encoding_list["audio_bounds"].append(bounds_with_bid)
+
+        result = {
+            "input_ids": input_ids,
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "image_encoding": None,
+            "audio_encoding": None,
+        }
+
+        if collate_labels:
+            labels = trim_and_pad(
+                [example["labels"] for example in examples],
+                batch_first=True,
+                padding_value=-100,
+            )
+            result["labels"] = labels
+
+        if any(image_encoding_list.values()):
+            result["image_encoding"] = {
+                "pixel_values": torch.vstack(image_encoding_list["pixel_values"]),
+                "tgt_sizes": torch.vstack(image_encoding_list["tgt_sizes"]),
+                "patch_attention_mask": torch.vstack(image_encoding_list["patch_attention_mask"]),
+                "image_bounds": torch.vstack(image_encoding_list["image_bounds"]),
+            }
+        if any(audio_encoding_list.values()):
+            result["audio_encoding"] = {
+                "input_audios": torch.vstack(audio_encoding_list["input_audios"]),
+                "input_audio_lengths": torch.vstack(audio_encoding_list["input_audio_lengths"]),
+                "audio_span_tokens": audio_encoding_list["audio_span_tokens"],
+                "audio_bounds": torch.vstack(audio_encoding_list["audio_bounds"]),
+            }
+        return result

processor_config.json

@@ -0,0 +1,7 @@
+{
+  "auto_map": {
+    "AutoProcessor": "processing_megrezo.MegrezOProcessor",
+    "AutoImageProcessor": "image_processing_megrezo.MegrezOImageProcessor"
+  },
+  "processor_class": "MegrezOProcessor"
+}

resampler.py

@@ -0,0 +1,783 @@
+from functools import partial
+from typing import Optional, Tuple
+import numpy as np
+import warnings
+
+import torch
+from torch import nn
+from torch import Tensor
+import torch.nn.functional as F
+from torch.nn.functional import *
+from torch.nn.modules.activation import *
+from torch.nn.init import trunc_normal_, constant_, xavier_normal_, xavier_uniform_
+
+from transformers.integrations import is_deepspeed_zero3_enabled
+
+def get_2d_sincos_pos_embed(embed_dim, image_size):
+    """
+    image_size: image_size or (image_height, image_width)
+    return:
+    pos_embed: [image_height, image_width, embed_dim]
+    """
+    if isinstance(image_size, int):
+        grid_h_size, grid_w_size = image_size, image_size
+    else:
+        grid_h_size, grid_w_size = image_size[0], image_size[1]
+
+    grid_h = np.arange(grid_h_size, dtype=np.float32)
+    grid_w = np.arange(grid_w_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid_new(embed_dim // 2, grid[0])  # (H, W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid_new(embed_dim // 2, grid[1])  # (H, W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=-1)  # (H, W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid_new(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (H, W)
+    out: (H, W, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float32)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000 ** omega  # (D/2,)
+
+    out = np.einsum('hw,d->hwd', pos, omega)  # (H, W, D/2), outer product
+
+    emb_sin = np.sin(out)  # (H, W, D/2)
+    emb_cos = np.cos(out)  # (H, W, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=-1)  # (H, W, D)
+    return emb
+
+
+class Resampler(nn.Module):
+    """
+    A 2D perceiver-resampler network with one cross attention layers by
+       given learnable queries and 2d sincos pos_emb
+    Outputs:
+        A tensor with the shape of (batch_size, num_queries, embed_dim)
+    """
+
+    def __init__(
+            self,
+            num_queries,
+            embed_dim,
+            num_heads,
+            kv_dim=None,
+            norm_layer=partial(nn.LayerNorm, eps=1e-6),
+            adaptive=False,
+            max_size=(70, 70),
+    ):
+        super().__init__()
+        self.num_queries = num_queries
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.adaptive = adaptive
+        self.max_size = max_size
+
+        self.query = nn.Parameter(torch.zeros(self.num_queries, embed_dim))
+
+        if kv_dim is not None and kv_dim != embed_dim:
+            self.kv_proj = nn.Linear(kv_dim, embed_dim, bias=False)
+        else:
+            self.kv_proj = nn.Identity()
+
+        # Change to nn.MultiheadAttention instead of MultiheadAttention in this file.
+        self.attn = nn.MultiheadAttention(embed_dim, num_heads)
+        self.ln_q = norm_layer(embed_dim)
+        self.ln_kv = norm_layer(embed_dim)
+
+        self.ln_post = norm_layer(embed_dim)
+        self.proj = nn.Parameter((embed_dim ** -0.5) * torch.randn(embed_dim, embed_dim))
+
+        self._set_2d_pos_cache(self.max_size)
+
+    def _set_2d_pos_cache(self, max_size, device='cpu'):
+        if is_deepspeed_zero3_enabled():
+            device='cuda'
+        pos_embed = torch.from_numpy(get_2d_sincos_pos_embed(self.embed_dim, max_size)).float().to(device)
+        self.register_buffer("pos_embed", pos_embed, persistent=False)
+
+    def _adjust_pos_cache(self, tgt_sizes, device):
+        max_h = torch.max(tgt_sizes[:, 0])
+        max_w = torch.max(tgt_sizes[:, 1])
+        if max_h > self.max_size[0] or max_w > self.max_size[1]:
+            self.max_size = [max(max_h, self.max_size[0]), max(max_w, self.max_size[1])]
+            self._set_2d_pos_cache(self.max_size, device)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x, tgt_sizes=None):
+        assert x.shape[0] == tgt_sizes.shape[0]
+        bs = x.shape[0]
+
+        device = x.device
+        dtype = x.dtype
+
+        patch_len = tgt_sizes[:, 0] * tgt_sizes[:, 1]
+
+        self._adjust_pos_cache(tgt_sizes, device=device)
+
+        max_patch_len = torch.max(patch_len)
+        key_padding_mask = torch.zeros((bs, max_patch_len), dtype=torch.bool, device=device)
+
+        pos_embed = []
+        for i in range(bs):
+            tgt_h, tgt_w = tgt_sizes[i]
+            pos_embed.append(self.pos_embed[:tgt_h, :tgt_w, :].reshape((tgt_h * tgt_w, -1)).to(dtype))  # patches * D
+            key_padding_mask[i, patch_len[i]:] = True
+
+        pos_embed = torch.nn.utils.rnn.pad_sequence(
+            pos_embed, batch_first=True, padding_value=0.0).permute(1, 0, 2)  # BLD => L * B * D
+
+        x = self.kv_proj(x)  # B * L * D
+        x = self.ln_kv(x).permute(1, 0, 2)  # L * B * D
+
+        q = self.ln_q(self.query)  # Q * D
+
+        out = self.attn(
+            self._repeat(q, bs),  # Q * B * D
+            x + pos_embed,  # L * B * D +  L * B * D
+            x,
+            key_padding_mask=key_padding_mask)[0]
+        #  out: Q * B * D
+        x = out.permute(1, 0, 2)  # B * Q * D
+
+        x = self.ln_post(x)
+        x = x @ self.proj
+        return x
+
+    def _repeat(self, query, N: int):
+        return query.unsqueeze(1).repeat(1, N, 1)
+
+
+class MultiheadAttention(nn.MultiheadAttention):
+    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, 
+                 add_zero_attn=False, kdim=None, vdim=None, batch_first=False, device=None, dtype=None):
+        super().__init__(embed_dim, num_heads, dropout, bias, add_bias_kv, add_zero_attn, kdim, vdim, batch_first, device, dtype)
+
+        # rewrite out_proj layer，with nn.Linear
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias, device=device, dtype=dtype)
+
+    def forward(
+                self,
+                query: Tensor,
+                key: Tensor,
+                value: Tensor,
+                key_padding_mask: Optional[Tensor] = None,
+                need_weights: bool = True,
+                attn_mask: Optional[Tensor] = None,
+                average_attn_weights: bool = True,
+                is_causal : bool = False) -> Tuple[Tensor, Optional[Tensor]]:
+        why_not_fast_path = ''
+        if ((attn_mask is not None and torch.is_floating_point(attn_mask))
+           or (key_padding_mask is not None) and torch.is_floating_point(key_padding_mask)):
+            why_not_fast_path = "floating-point masks are not supported for fast path."
+
+        is_batched = query.dim() == 3
+
+        key_padding_mask = _canonical_mask(
+            mask=key_padding_mask,
+            mask_name="key_padding_mask",
+            other_type=F._none_or_dtype(attn_mask),
+            other_name="attn_mask",
+            target_type=query.dtype
+        )
+
+        attn_mask = _canonical_mask(
+            mask=attn_mask,
+            mask_name="attn_mask",
+            other_type=None,
+            other_name="",
+            target_type=query.dtype,
+            check_other=False,
+        )
+
+
+        if not is_batched:
+            why_not_fast_path = f"input not batched; expected query.dim() of 3 but got {query.dim()}"
+        elif query is not key or key is not value:
+            # When lifting this restriction, don't forget to either
+            # enforce that the dtypes all match or test cases where
+            # they don't!
+            why_not_fast_path = "non-self attention was used (query, key, and value are not the same Tensor)"
+        elif self.in_proj_bias is not None and query.dtype != self.in_proj_bias.dtype:
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_bias ({self.in_proj_bias.dtype}) don't match"
+        elif self.in_proj_weight is None:
+            why_not_fast_path = "in_proj_weight was None"
+        elif query.dtype != self.in_proj_weight.dtype:
+            # this case will fail anyway, but at least they'll get a useful error message.
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_weight ({self.in_proj_weight.dtype}) don't match"
+        elif self.training:
+            why_not_fast_path = "training is enabled"
+        elif (self.num_heads % 2) != 0:
+            why_not_fast_path = "self.num_heads is not even"
+        elif not self.batch_first:
+            why_not_fast_path = "batch_first was not True"
+        elif self.bias_k is not None:
+            why_not_fast_path = "self.bias_k was not None"
+        elif self.bias_v is not None:
+            why_not_fast_path = "self.bias_v was not None"
+        elif self.add_zero_attn:
+            why_not_fast_path = "add_zero_attn was enabled"
+        elif not self._qkv_same_embed_dim:
+            why_not_fast_path = "_qkv_same_embed_dim was not True"
+        elif query.is_nested and (key_padding_mask is not None or attn_mask is not None):
+            why_not_fast_path = "supplying both src_key_padding_mask and src_mask at the same time \
+                                 is not supported with NestedTensor input"
+        elif torch.is_autocast_enabled():
+            why_not_fast_path = "autocast is enabled"
+
+        if not why_not_fast_path:
+            tensor_args = (
+                query,
+                key,
+                value,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.out_proj.weight,
+                self.out_proj.bias,
+            )
+            # We have to use list comprehensions below because TorchScript does not support
+            # generator expressions.
+            if torch.overrides.has_torch_function(tensor_args):
+                why_not_fast_path = "some Tensor argument has_torch_function"
+            elif _is_make_fx_tracing():
+                why_not_fast_path = "we are running make_fx tracing"
+            elif not all(_check_arg_device(x) for x in tensor_args):
+                why_not_fast_path = ("some Tensor argument's device is neither one of "
+                                     f"cpu, cuda or {torch.utils.backend_registration._privateuse1_backend_name}")
+            elif torch.is_grad_enabled() and any(_arg_requires_grad(x) for x in tensor_args):
+                why_not_fast_path = ("grad is enabled and at least one of query or the "
+                                     "input/output projection weights or biases requires_grad")
+            if not why_not_fast_path:
+                merged_mask, mask_type = self.merge_masks(attn_mask, key_padding_mask, query)
+
+                if self.in_proj_bias is not None and self.in_proj_weight is not None:
+                    return torch._native_multi_head_attention(
+                        query,
+                        key,
+                        value,
+                        self.embed_dim,
+                        self.num_heads,
+                        self.in_proj_weight,
+                        self.in_proj_bias,
+                        self.out_proj.weight,
+                        self.out_proj.bias,
+                        merged_mask,
+                        need_weights,
+                        average_attn_weights,
+                        mask_type)
+
+        any_nested = query.is_nested or key.is_nested or value.is_nested
+        assert not any_nested, ("MultiheadAttention does not support NestedTensor outside of its fast path. " +
+                                f"The fast path was not hit because {why_not_fast_path}")
+
+        if self.batch_first and is_batched:
+            # make sure that the transpose op does not affect the "is" property
+            if key is value:
+                if query is key:
+                    query = key = value = query.transpose(1, 0)
+                else:
+                    query, key = (x.transpose(1, 0) for x in (query, key))
+                    value = key
+            else:
+                query, key, value = (x.transpose(1, 0) for x in (query, key, value))
+        
+        if not self._qkv_same_embed_dim:
+            attn_output, attn_output_weights = self.multi_head_attention_forward(
+                query, key, value, self.embed_dim, self.num_heads,
+                self.in_proj_weight, self.in_proj_bias,
+                self.bias_k, self.bias_v, self.add_zero_attn,
+                self.dropout, self.out_proj.weight, self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask, need_weights=need_weights,
+                attn_mask=attn_mask,
+                use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
+                v_proj_weight=self.v_proj_weight,
+                average_attn_weights=average_attn_weights,
+                is_causal=is_causal)
+        else:
+            attn_output, attn_output_weights = self.multi_head_attention_forward(
+                query, key, value, self.embed_dim, self.num_heads,
+                self.in_proj_weight, self.in_proj_bias,
+                self.bias_k, self.bias_v, self.add_zero_attn,
+                self.dropout, self.out_proj.weight, self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask,
+                need_weights=need_weights,
+                attn_mask=attn_mask,
+                average_attn_weights=average_attn_weights,
+                is_causal=is_causal)
+        if self.batch_first and is_batched:
+            return attn_output.transpose(1, 0), attn_output_weights
+        else:
+            return attn_output, attn_output_weights
+            
+    def multi_head_attention_forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        embed_dim_to_check: int,
+        num_heads: int,
+        in_proj_weight: Optional[Tensor],
+        in_proj_bias: Optional[Tensor],
+        bias_k: Optional[Tensor],
+        bias_v: Optional[Tensor],
+        add_zero_attn: bool,
+        dropout_p: float,
+        out_proj_weight: Tensor,
+        out_proj_bias: Optional[Tensor],
+        training: bool = True,
+        key_padding_mask: Optional[Tensor] = None,
+        need_weights: bool = True,
+        attn_mask: Optional[Tensor] = None,
+        use_separate_proj_weight: bool = False,
+        q_proj_weight: Optional[Tensor] = None,
+        k_proj_weight: Optional[Tensor] = None,
+        v_proj_weight: Optional[Tensor] = None,
+        static_k: Optional[Tensor] = None,
+        static_v: Optional[Tensor] = None,
+        average_attn_weights: bool = True,
+        is_causal: bool = False,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v, out_proj_weight, out_proj_bias)
+    
+        is_batched = _mha_shape_check(query, key, value, key_padding_mask, attn_mask, num_heads)
+    
+        # For unbatched input, we unsqueeze at the expected batch-dim to pretend that the input
+        # is batched, run the computation and before returning squeeze the
+        # batch dimension so that the output doesn't carry this temporary batch dimension.
+        if not is_batched:
+            # unsqueeze if the input is unbatched
+            query = query.unsqueeze(1)
+            key = key.unsqueeze(1)
+            value = value.unsqueeze(1)
+            if key_padding_mask is not None:
+                key_padding_mask = key_padding_mask.unsqueeze(0)
+    
+        # set up shape vars
+        tgt_len, bsz, embed_dim = query.shape
+        src_len, _, _ = key.shape
+    
+        key_padding_mask = _canonical_mask(
+            mask=key_padding_mask,
+            mask_name="key_padding_mask",
+            other_type=_none_or_dtype(attn_mask),
+            other_name="attn_mask",
+            target_type=query.dtype
+        )
+    
+        if is_causal and attn_mask is None:
+            raise RuntimeError(
+                "Need attn_mask if specifying the is_causal hint. "
+                "You may use the Transformer module method "
+                "`generate_square_subsequent_mask` to create this mask."
+            )
+    
+        if is_causal and key_padding_mask is None and not need_weights:
+            # when we have a kpm or need weights, we need attn_mask
+            # Otherwise, we use the is_causal hint go as is_causal
+            # indicator to SDPA.
+            attn_mask = None
+        else:
+            attn_mask = _canonical_mask(
+                mask=attn_mask,
+                mask_name="attn_mask",
+                other_type=None,
+                other_name="",
+                target_type=query.dtype,
+                check_other=False,
+            )
+    
+            if key_padding_mask is not None:
+                # We have the attn_mask, and use that to merge kpm into it.
+                # Turn off use of is_causal hint, as the merged mask is no
+                # longer causal.
+                is_causal = False
+    
+        assert embed_dim == embed_dim_to_check, \
+            f"was expecting embedding dimension of {embed_dim_to_check}, but got {embed_dim}"
+        if isinstance(embed_dim, torch.Tensor):
+            # embed_dim can be a tensor when JIT tracing
+            head_dim = embed_dim.div(num_heads, rounding_mode='trunc')
+        else:
+            head_dim = embed_dim // num_heads
+        assert head_dim * num_heads == embed_dim, f"embed_dim {embed_dim} not divisible by num_heads {num_heads}"
+        if use_separate_proj_weight:
+            # allow MHA to have different embedding dimensions when separate projection weights are used
+            assert key.shape[:2] == value.shape[:2], \
+                f"key's sequence and batch dims {key.shape[:2]} do not match value's {value.shape[:2]}"
+        else:
+            assert key.shape == value.shape, f"key shape {key.shape} does not match value shape {value.shape}"
+    
+        #
+        # compute in-projection
+        #
+        if not use_separate_proj_weight:
+            assert in_proj_weight is not None, "use_separate_proj_weight is False but in_proj_weight is None"
+            q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)
+        else:
+            assert q_proj_weight is not None, "use_separate_proj_weight is True but q_proj_weight is None"
+            assert k_proj_weight is not None, "use_separate_proj_weight is True but k_proj_weight is None"
+            assert v_proj_weight is not None, "use_separate_proj_weight is True but v_proj_weight is None"
+            if in_proj_bias is None:
+                b_q = b_k = b_v = None
+            else:
+                b_q, b_k, b_v = in_proj_bias.chunk(3)
+            q, k, v = _in_projection(query, key, value, q_proj_weight, k_proj_weight, v_proj_weight, b_q, b_k, b_v)
+    
+        # prep attention mask
+    
+        if attn_mask is not None:
+            # ensure attn_mask's dim is 3
+            if attn_mask.dim() == 2:
+                correct_2d_size = (tgt_len, src_len)
+                if attn_mask.shape != correct_2d_size:
+                    raise RuntimeError(f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
+                attn_mask = attn_mask.unsqueeze(0)
+            elif attn_mask.dim() == 3:
+                correct_3d_size = (bsz * num_heads, tgt_len, src_len)
+                if attn_mask.shape != correct_3d_size:
+                    raise RuntimeError(f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.")
+            else:
+                raise RuntimeError(f"attn_mask's dimension {attn_mask.dim()} is not supported")
+    
+        # add bias along batch dimension (currently second)
+        if bias_k is not None and bias_v is not None:
+            assert static_k is None, "bias cannot be added to static key."
+            assert static_v is None, "bias cannot be added to static value."
+            k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
+            v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
+            if attn_mask is not None:
+                attn_mask = pad(attn_mask, (0, 1))
+            if key_padding_mask is not None:
+                key_padding_mask = pad(key_padding_mask, (0, 1))
+        else:
+            assert bias_k is None
+            assert bias_v is None
+    
+        #
+        # reshape q, k, v for multihead attention and make em batch first
+        #
+        q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+        if static_k is None:
+            k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
+        else:
+            # TODO finish disentangling control flow so we don't do in-projections when statics are passed
+            assert static_k.size(0) == bsz * num_heads, \
+                f"expecting static_k.size(0) of {bsz * num_heads}, but got {static_k.size(0)}"
+            assert static_k.size(2) == head_dim, \
+                f"expecting static_k.size(2) of {head_dim}, but got {static_k.size(2)}"
+            k = static_k
+        if static_v is None:
+            v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
+        else:
+            # TODO finish disentangling control flow so we don't do in-projections when statics are passed
+            assert static_v.size(0) == bsz * num_heads, \
+                f"expecting static_v.size(0) of {bsz * num_heads}, but got {static_v.size(0)}"
+            assert static_v.size(2) == head_dim, \
+                f"expecting static_v.size(2) of {head_dim}, but got {static_v.size(2)}"
+            v = static_v
+    
+        # add zero attention along batch dimension (now first)
+        if add_zero_attn:
+            zero_attn_shape = (bsz * num_heads, 1, head_dim)
+            k = torch.cat([k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=1)
+            v = torch.cat([v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=1)
+            if attn_mask is not None:
+                attn_mask = pad(attn_mask, (0, 1))
+            if key_padding_mask is not None:
+                key_padding_mask = pad(key_padding_mask, (0, 1))
+    
+        # update source sequence length after adjustments
+        src_len = k.size(1)
+    
+        # merge key padding and attention masks
+        if key_padding_mask is not None:
+            assert key_padding_mask.shape == (bsz, src_len), \
+                f"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}"
+            key_padding_mask = key_padding_mask.view(bsz, 1, 1, src_len).   \
+                expand(-1, num_heads, -1, -1).reshape(bsz * num_heads, 1, src_len)
+            if attn_mask is None:
+                attn_mask = key_padding_mask
+            else:
+                attn_mask = attn_mask + key_padding_mask
+    
+        # adjust dropout probability
+        if not training:
+            dropout_p = 0.0
+    
+        #
+        # (deep breath) calculate attention and out projection
+        #
+    
+        if need_weights:
+            B, Nt, E = q.shape
+            q_scaled = q / math.sqrt(E)
+    
+            assert not (is_causal and attn_mask is None), "FIXME: is_causal not implemented for need_weights"
+    
+            if attn_mask is not None:
+                attn_output_weights = torch.baddbmm(attn_mask, q_scaled, k.transpose(-2, -1))
+            else:
+                attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1))
+            attn_output_weights = softmax(attn_output_weights, dim=-1)
+            if dropout_p > 0.0:
+                attn_output_weights = dropout(attn_output_weights, p=dropout_p)
+    
+            attn_output = torch.bmm(attn_output_weights, v)
+    
+            attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)
+            attn_output = self.out_proj(attn_output)
+            attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+    
+            # optionally average attention weights over heads
+            attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+            if average_attn_weights:
+                attn_output_weights = attn_output_weights.mean(dim=1)
+    
+            if not is_batched:
+                # squeeze the output if input was unbatched
+                attn_output = attn_output.squeeze(1)
+                attn_output_weights = attn_output_weights.squeeze(0)
+            return attn_output, attn_output_weights
+        else:
+            # attn_mask can be either (L,S) or (N*num_heads, L, S)
+            # if attn_mask's shape is (1, L, S) we need to unsqueeze to (1, 1, L, S)
+            # in order to match the input for SDPA of (N, num_heads, L, S)
+            if attn_mask is not None:
+                if attn_mask.size(0) == 1 and attn_mask.dim() == 3:
+                    attn_mask = attn_mask.unsqueeze(0)
+                else:
+                    attn_mask = attn_mask.view(bsz, num_heads, -1, src_len)
+    
+            q = q.view(bsz, num_heads, tgt_len, head_dim)
+            k = k.view(bsz, num_heads, src_len, head_dim)
+            v = v.view(bsz, num_heads, src_len, head_dim)
+    
+            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)
+            attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)
+    
+            attn_output = self.out_proj(attn_output)
+            attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+            if not is_batched:
+                # squeeze the output if input was unbatched
+                attn_output = attn_output.squeeze(1)
+            return attn_output, None
+
+
+def _mha_shape_check(query: Tensor, key: Tensor, value: Tensor,
+                     key_padding_mask: Optional[Tensor], attn_mask: Optional[Tensor], num_heads: int):
+    # Verifies the expected shape for `query, `key`, `value`, `key_padding_mask` and `attn_mask`
+    # and returns if the input is batched or not.
+    # Raises an error if `query` is not 2-D (unbatched) or 3-D (batched) tensor.
+
+    # Shape check.
+    if query.dim() == 3:
+        # Batched Inputs
+        is_batched = True
+        assert key.dim() == 3 and value.dim() == 3, \
+            ("For batched (3-D) `query`, expected `key` and `value` to be 3-D"
+             f" but found {key.dim()}-D and {value.dim()}-D tensors respectively")
+        if key_padding_mask is not None:
+            assert key_padding_mask.dim() == 2, \
+                ("For batched (3-D) `query`, expected `key_padding_mask` to be `None` or 2-D"
+                 f" but found {key_padding_mask.dim()}-D tensor instead")
+        if attn_mask is not None:
+            assert attn_mask.dim() in (2, 3), \
+                ("For batched (3-D) `query`, expected `attn_mask` to be `None`, 2-D or 3-D"
+                 f" but found {attn_mask.dim()}-D tensor instead")
+    elif query.dim() == 2:
+        # Unbatched Inputs
+        is_batched = False
+        assert key.dim() == 2 and value.dim() == 2, \
+            ("For unbatched (2-D) `query`, expected `key` and `value` to be 2-D"
+             f" but found {key.dim()}-D and {value.dim()}-D tensors respectively")
+
+        if key_padding_mask is not None:
+            assert key_padding_mask.dim() == 1, \
+                ("For unbatched (2-D) `query`, expected `key_padding_mask` to be `None` or 1-D"
+                 f" but found {key_padding_mask.dim()}-D tensor instead")
+
+        if attn_mask is not None:
+            assert attn_mask.dim() in (2, 3), \
+                ("For unbatched (2-D) `query`, expected `attn_mask` to be `None`, 2-D or 3-D"
+                 f" but found {attn_mask.dim()}-D tensor instead")
+            if attn_mask.dim() == 3:
+                expected_shape = (num_heads, query.shape[0], key.shape[0])
+                assert attn_mask.shape == expected_shape, \
+                    (f"Expected `attn_mask` shape to be {expected_shape} but got {attn_mask.shape}")
+    else:
+        raise AssertionError(
+            f"query should be unbatched 2D or batched 3D tensor but received {query.dim()}-D query tensor")
+
+    return is_batched
+
+
+def _canonical_mask(
+        mask: Optional[Tensor],
+        mask_name: str,
+        other_type: Optional[DType],
+        other_name: str,
+        target_type: DType,
+        check_other: bool = True,
+) -> Optional[Tensor]:
+
+    if mask is not None:
+        _mask_dtype = mask.dtype
+        _mask_is_float = torch.is_floating_point(mask)
+        if _mask_dtype != torch.bool and not _mask_is_float:
+            raise AssertionError(
+                f"only bool and floating types of {mask_name} are supported")
+        if check_other and other_type is not None:
+            if _mask_dtype != other_type:
+                warnings.warn(
+                    f"Support for mismatched {mask_name} and {other_name} "
+                    "is deprecated. Use same type for both instead."
+                )
+        if not _mask_is_float:
+            mask = (
+                torch.zeros_like(mask, dtype=target_type)
+                .masked_fill_(mask, float("-inf"))
+            )
+    return mask
+
+
+def _none_or_dtype(input: Optional[Tensor]) -> Optional[DType]:
+    if input is None:
+        return None
+    elif isinstance(input, torch.Tensor):
+        return input.dtype
+    raise RuntimeError("input to _none_or_dtype() must be None or torch.Tensor")
+
+def _in_projection_packed(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    w: Tensor,
+    b: Optional[Tensor] = None,
+) -> List[Tensor]:
+    r"""
+    Performs the in-projection step of the attention operation, using packed weights.
+    Output is a triple containing projection tensors for query, key and value.
+    Args:
+        q, k, v: query, key and value tensors to be projected. For self-attention,
+            these are typically the same tensor; for encoder-decoder attention,
+            k and v are typically the same tensor. (We take advantage of these
+            identities for performance if they are present.) Regardless, q, k and v
+            must share a common embedding dimension; otherwise their shapes may vary.
+        w: projection weights for q, k and v, packed into a single tensor. Weights
+            are packed along dimension 0, in q, k, v order.
+        b: optional projection biases for q, k and v, packed into a single tensor
+            in q, k, v order.
+    Shape:
+        Inputs:
+        - q: :math:`(..., E)` where E is the embedding dimension
+        - k: :math:`(..., E)` where E is the embedding dimension
+        - v: :math:`(..., E)` where E is the embedding dimension
+        - w: :math:`(E * 3, E)` where E is the embedding dimension
+        - b: :math:`E * 3` where E is the embedding dimension
+        Output:
+        - in output list :math:`[q', k', v']`, each output tensor will have the
+            same shape as the corresponding input tensor.
+    """
+    E = q.size(-1)
+    if k is v:
+        if q is k:
+            # self-attention
+            proj = linear(q, w, b)
+            # reshape to 3, E and not E, 3 is deliberate for better memory coalescing and keeping same order as chunk()
+            proj = proj.unflatten(-1, (3, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
+            return proj[0], proj[1], proj[2]
+        else:
+            # encoder-decoder attention
+            w_q, w_kv = w.split([E, E * 2])
+            if b is None:
+                b_q = b_kv = None
+            else:
+                b_q, b_kv = b.split([E, E * 2])
+            q_proj = linear(q, w_q, b_q)
+            kv_proj = linear(k, w_kv, b_kv)
+            # reshape to 2, E and not E, 2 is deliberate for better memory coalescing and keeping same order as chunk()
+            kv_proj = kv_proj.unflatten(-1, (2, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
+            return (q_proj, kv_proj[0], kv_proj[1])
+    else:
+        w_q, w_k, w_v = w.chunk(3)
+        if b is None:
+            b_q = b_k = b_v = None
+        else:
+            b_q, b_k, b_v = b.chunk(3)
+        return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
+
+
+def _in_projection(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    w_q: Tensor,
+    w_k: Tensor,
+    w_v: Tensor,
+    b_q: Optional[Tensor] = None,
+    b_k: Optional[Tensor] = None,
+    b_v: Optional[Tensor] = None,
+) -> Tuple[Tensor, Tensor, Tensor]:
+    r"""
+    Performs the in-projection step of the attention operation. This is simply
+    a triple of linear projections, with shape constraints on the weights which
+    ensure embedding dimension uniformity in the projected outputs.
+    Output is a triple containing projection tensors for query, key and value.
+    Args:
+        q, k, v: query, key and value tensors to be projected.
+        w_q, w_k, w_v: weights for q, k and v, respectively.
+        b_q, b_k, b_v: optional biases for q, k and v, respectively.
+    Shape:
+        Inputs:
+        - q: :math:`(Qdims..., Eq)` where Eq is the query embedding dimension and Qdims are any
+            number of leading dimensions.
+        - k: :math:`(Kdims..., Ek)` where Ek is the key embedding dimension and Kdims are any
+            number of leading dimensions.
+        - v: :math:`(Vdims..., Ev)` where Ev is the value embedding dimension and Vdims are any
+            number of leading dimensions.
+        - w_q: :math:`(Eq, Eq)`
+        - w_k: :math:`(Eq, Ek)`
+        - w_v: :math:`(Eq, Ev)`
+        - b_q: :math:`(Eq)`
+        - b_k: :math:`(Eq)`
+        - b_v: :math:`(Eq)`
+        Output: in output triple :math:`(q', k', v')`,
+         - q': :math:`[Qdims..., Eq]`
+         - k': :math:`[Kdims..., Eq]`
+         - v': :math:`[Vdims..., Eq]`
+    """
+    Eq, Ek, Ev = q.size(-1), k.size(-1), v.size(-1)
+    assert w_q.shape == (Eq, Eq), f"expecting query weights shape of {(Eq, Eq)}, but got {w_q.shape}"
+    assert w_k.shape == (Eq, Ek), f"expecting key weights shape of {(Eq, Ek)}, but got {w_k.shape}"
+    assert w_v.shape == (Eq, Ev), f"expecting value weights shape of {(Eq, Ev)}, but got {w_v.shape}"
+    assert b_q is None or b_q.shape == (Eq,), f"expecting query bias shape of {(Eq,)}, but got {b_q.shape}"
+    assert b_k is None or b_k.shape == (Eq,), f"expecting key bias shape of {(Eq,)}, but got {b_k.shape}"
+    assert b_v is None or b_v.shape == (Eq,), f"expecting value bias shape of {(Eq,)}, but got {b_v.shape}"
+    return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v) 

special_tokens_map.json

@@ -0,0 +1,10 @@
+{
+  "eos_token": "<|turn_end|>",
+  "unk_token": {
+    "content": "<|unk|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,257 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "120000": {
+      "content": "<|eos|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120001": {
+      "content": "<|unk|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120002": {
+      "content": "<|pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120003": {
+      "content": "<|role_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120004": {
+      "content": "<|role_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120005": {
+      "content": "<|turn_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120006": {
+      "content": "<|code_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120007": {
+      "content": "<|code_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120008": {
+      "content": "<|commit_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120009": {
+      "content": "<|commit_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120010": {
+      "content": "<|diff_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120011": {
+      "content": "<|diff_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120012": {
+      "content": "<|code_execution_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120013": {
+      "content": "<|code_execution_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120014": {
+      "content": "<|image_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120015": {
+      "content": "<|image_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120016": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120017": {
+      "content": "<|video_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120018": {
+      "content": "<|video_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120019": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120020": {
+      "content": "<|audio_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120021": {
+      "content": "<|audio_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120022": {
+      "content": "<|audio_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120023": {
+      "content": "<|function_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120024": {
+      "content": "<|function_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120025": {
+      "content": "<|slice_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120026": {
+      "content": "<|slice_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120027": {
+      "content": "<|image_id_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "120028": {
+      "content": "<|image_id_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoProcessor": "processing_megrezo.MegrezOProcessor",
+    "AutoTokenizer": [
+      "tokenizer_wrapper.LlamaTokenizerWrapper",
+      null
+    ]
+  },
+  "bos_token": null,
+  "chat_template": "{% set audio_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|role_start|>system<|role_end|>你是Megrez-3B-Instruct，将针对用户的问题给出详细的、积极的回答。<|turn_end|>{% endif %}<|role_start|>{{ message['role'] }}<|role_end|>{% if message['content'] is string %}{{ message['content'] }}{% else %}{% if 'image' in message['content'] %}{% if message['content']['image'] is sequence and message['content']['image'] is not string %}{% for image in message['content']['image'] %}(<image>./</image>)\n{% endfor %}{% else %}(<image>./</image>)\n{% endif %}{% endif %}{% if 'audio' in message['content'] %}{% if message['content']['audio'] is sequence and message['content']['audio'] is not string %}{% for audio in message['content']['audio'] %}{% set audio_count.value = audio_count.value + 1 %}Audio {{ audio_count.value }}: (<audio>./</audio>)\n{% endfor %}{% else %}{% set audio_count.value = audio_count.value + 1 %}Audio {{ audio_count.value }}: (<audio>./</audio>)\n{% endif %}{% endif %}{% if 'text' in message['content'] and message['content']['text'] %}{{ message['content']['text'] }}{% endif %}{% endif %}<|turn_end|>{% endfor %}{% if add_generation_prompt %}<|role_start|>assistant<|role_end|>{% endif %}",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|turn_end|>",
+  "legacy": true,
+  "model_max_length": 4096,
+  "processor_class": "MegrezOProcessor",
+  "tokenizer_class": "LlamaTokenizerWrapper",
+  "unk_token": "<|unk|>",
+  "pad_token": "<|pad|>",
+  "use_default_system_prompt": false
+}

tokenizer_wrapper.py

@@ -0,0 +1,63 @@
+from transformers import LlamaTokenizerFast
+
+
+class LlamaTokenizerWrapper(LlamaTokenizerFast):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        self.im_start = "<|image_start|>"
+        self.im_end = "<|image_end|>"
+        self.ref_start = "<|ref_start>|"
+        self.ref_end = "<|ref_end|>"
+        self.box_start = "<|box_start|>"
+        self.box_end = "<|box_end|>"
+        self.quad_start = "<|quad_start>"
+        self.quad_end = "<|quad_end|>"
+        self.point_start = "<|point_start|>"
+        self.point_end = "<|point_end|>"
+        self.slice_start = "<|slice_start|>"
+        self.slice_end = "<|slice_end|>"
+        self.audio_start = "<|audio_start|>"
+        self.audio_end = "<|audio_end|>"
+        self.eos_token = "<|turn_end|>"
+        self.pad_token = "<|pad|>"
+
+    @property
+    def eos_id(self):
+        return self.eos_token_id
+
+    @property
+    def unk_id(self):
+        return self.unk_token_id
+
+    @property
+    def im_start_id(self):
+        return self.encode(self.im_start, add_special_tokens=False)[0]
+
+    @property
+    def im_end_id(self):
+        return self.encode(self.im_end, add_special_tokens=False)[0]
+
+    @property
+    def slice_start_id(self):
+        return self.encode(self.slice_start, add_special_tokens=False)[0]
+
+    @property
+    def slice_end_id(self):
+        return self.encode(self.slice_end, add_special_tokens=False)[0]
+
+    @property
+    def audio_start_id(self):
+        return self.encode(self.audio_start, add_special_tokens=False)[0]
+
+    @property
+    def audio_end_id(self):
+        return self.encode(self.audio_end, add_special_tokens=False)[0]
+
+    @property
+    def eos_token_id(self):
+        return self.encode(self.eos_token, add_special_tokens=False)[0]
+
+    @property
+    def pad_token_id(self):
+        return self.encode(self.eos_token, add_special_tokens=False)[0]