smthemex
@@ -0,0 +1,78 @@
|
||||
# VLM
|
||||
We follow [InternVL2](https://internvl.readthedocs.io/en/latest/internvl2.0/evaluation.html) to evaluate the performance on MME, MMBench, MMMU, MMVet, MathVista and MMVP.
|
||||
|
||||
## Data prepration
|
||||
Please follow the [InternVL2](https://internvl.readthedocs.io/en/latest/get_started/eval_data_preparation.html) to prepare the corresponding data. And the link the data under `vlm`.
|
||||
|
||||
The final directory structure is:
|
||||
```shell
|
||||
data
|
||||
├── MathVista
|
||||
├── mmbench
|
||||
├── mme
|
||||
├── MMMU
|
||||
├── mm-vet
|
||||
└── MMVP
|
||||
```
|
||||
|
||||
## Evaluation
|
||||
|
||||
Directly run `scripts/eval/run_eval_vlm.sh` to evaluate different benchmarks. The output will be saved in `$output_path`.
|
||||
- Set `$model_path` and `$output_path` for the path for checkpoint and log.
|
||||
- Increase `GPUS` if you want to run faster.
|
||||
- For MMBench, please use the official [evaluation server](https://mmbench.opencompass.org.cn/mmbench-submission).
|
||||
- For MMVet, please use the official [evaluation server](https://huggingface.co/spaces/whyu/MM-Vet_Evaluator).
|
||||
- For MathVista, please set `$openai_api_key` in `scripts/eval/run_eval_vlm.sh` and `your_api_url` in `eval/vlm/eval/mathvista/utilities.py`. The default GPT version is `gpt-4o-2024-11-20`.
|
||||
- For MMMU, we use CoT in the report, which improve the accuracy by about 2%. For evaluation of the oprn-ended answer, we use GPT-4o for judgement.
|
||||
|
||||
|
||||
# GenEval
|
||||
We modify the code in [GenEval](https://github.com/djghosh13/geneval/tree/main) for faster evaluation.
|
||||
|
||||
## Setup
|
||||
Install the following dependencies:
|
||||
```shell
|
||||
pip install open-clip-torch
|
||||
pip install clip-benchmark
|
||||
pip install --upgrade setuptools
|
||||
|
||||
sudo pip install -U openmim
|
||||
sudo mim install mmengine mmcv-full==1.7.2
|
||||
|
||||
git clone https://github.com/open-mmlab/mmdetection.git
|
||||
cd mmdetection; git checkout 2.x
|
||||
pip install -v -e .
|
||||
```
|
||||
|
||||
Download Detector:
|
||||
```shell
|
||||
cd ./eval/gen/geneval
|
||||
mkdir model
|
||||
|
||||
bash ./evaluation/download_models.sh ./model
|
||||
```
|
||||
|
||||
## Evaluation
|
||||
Directly run `scripts/eval/run_geneval.sh` to evaluate GenEVAL. The output will be saved in `$output_path`.
|
||||
- Set `$model_path` and `$output_path` for the path for checkpoint and log.
|
||||
- Set `metadata_file` to `./eval/gen/geneval/prompts/evaluation_metadata.jsonl` for original GenEval prompts.
|
||||
|
||||
|
||||
# WISE
|
||||
We modify the code in [WISE](https://github.com/PKU-YuanGroup/WISE/tree/main) for faster evaluation.
|
||||
|
||||
|
||||
## Evaluation
|
||||
Directly run `scripts/eval/run_wise.sh` to evaluate WISE. The output will be saved in `$output_path`.
|
||||
- Set `$model_path` and `$output_path` for the path for checkpoint and log.
|
||||
- Set `$openai_api_key` in `scripts/eval/run_wise.sh` and `your_api_url` in `eval/gen/wise/gpt_eval_mp.py`. The default GPT version is `gpt-4o-2024-11-20`.
|
||||
- Use `think` for thinking mode.
|
||||
|
||||
|
||||
|
||||
# GEdit-Bench
|
||||
Please follow [GEdit-Bench](https://github.com/stepfun-ai/Step1X-Edit/blob/main/GEdit-Bench/EVAL.md) for evaluation.
|
||||
|
||||
|
||||
# IntelligentBench
|
||||
TBD
|
||||
+201
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
+199
@@ -0,0 +1,199 @@
|
||||
<p align="center">
|
||||
<img src="https://lf3-static.bytednsdoc.com/obj/eden-cn/nuhojubrps/banner.png" alt="BAGEL" width="480"/>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="https://bagel-ai.org/">
|
||||
<img
|
||||
src="https://img.shields.io/badge/BAGEL-Website-0A66C2?logo=safari&logoColor=white"
|
||||
alt="BAGEL Website"
|
||||
/>
|
||||
</a>
|
||||
<a href="https://arxiv.org/abs/2505.14683">
|
||||
<img
|
||||
src="https://img.shields.io/badge/BAGEL-Paper-red?logo=arxiv&logoColor=red"
|
||||
alt="BAGEL Paper on arXiv"
|
||||
/>
|
||||
</a>
|
||||
<a href="https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT">
|
||||
<img
|
||||
src="https://img.shields.io/badge/BAGEL-Model-yellow?logo=huggingface&logoColor=yellow"
|
||||
alt="BAGEL Model"
|
||||
/>
|
||||
</a>
|
||||
<a href="https://demo.bagel-ai.org/">
|
||||
<img
|
||||
src="https://img.shields.io/badge/BAGEL-Demo-blue?logo=googleplay&logoColor=blue"
|
||||
alt="BAGEL Demo"
|
||||
/>
|
||||
</a>
|
||||
<a href="https://huggingface.co/spaces/ByteDance-Seed/BAGEL">
|
||||
<img
|
||||
src="https://img.shields.io/badge/BAGEL-Space-orange?logo=huggingface&logoColor=yellow"
|
||||
alt="BAGEL Model"
|
||||
/>
|
||||
</a>
|
||||
<a href="https://discord.gg/Z836xxzy">
|
||||
<img
|
||||
src="https://img.shields.io/badge/BAGEL-Discord-5865F2?logo=discord&logoColor=purple"
|
||||
alt="BAGEL Discord"
|
||||
/>
|
||||
</a>
|
||||
<a href="mailto:bagel@bytedance.com">
|
||||
<img
|
||||
src="https://img.shields.io/badge/BAGEL-Email-D14836?logo=gmail&logoColor=red"
|
||||
alt="BAGEL Email"
|
||||
/>
|
||||
</a>
|
||||
</p>
|
||||
|
||||
# Unified Model for Multimodal Understanding and Generation
|
||||
> [Chaorui Deng*](https://scholar.google.com/citations?hl=en&user=k0TWfBoAAAAJ), [Deyao Zhu*](https://tsutikgiau.github.io/), [Kunchang Li*](https://andy1621.github.io/), [Chenhui Gou*](https://www.linkedin.com/in/chenhui-gou-9201081a1/?originalSubdomain=au), [Feng Li*](https://fengli-ust.github.io/), [Zeyu Wang](https://zw615.github.io/), Shu Zhong, [Weihao Yu](https://whyu.me/), [Xiaonan Nie](https://codecaution.github.io/), [Ziang Song](https://www.linkedin.com/in/ziang-song-43b0ab8a/), Guang Shi :email: , [Haoqi Fan* :tophat: ](https://haoqifan.github.io/)
|
||||
>
|
||||
> contact: shiguang.sg@bytedance.com
|
||||
>
|
||||
> We present **BAGEL**, an open‑source multimodal foundation model with 7B active parameters (14B total) trained on large‑scale interleaved multimodal data. BAGEL outperforms the current top‑tier open‑source VLMs like Qwen2.5-VL and InternVL-2.5 on standard multimodal understanding leaderboards, and delivers text‑to‑image quality that is competitive with strong specialist generators such as SD3.
|
||||
Moreover, BAGEL demonstrates superior qualitative results in classical image‑editing scenarios than the leading open-source models. More importantly, it extends to free-form visual manipulation, multiview synthesis, and world navigation, capabilities that constitute "world-modeling" tasks beyond the scope of previous image-editing models.
|
||||
The figure below showcases BAGEL's qualitative performance.
|
||||
|
||||
<p align="center"><img src="assets/teaser.webp" width="95%"></p>
|
||||
|
||||
|
||||
<!-- ## 🧠 Method
|
||||
BAGEL adopts a Mixture-of-Transformer-Experts (MoT) architecture to maximize the model’s capacity to learn from richly diverse multimodal information. Following the same principle of capacity maximization, it utilizes two separate encoders to capture pixel-level and semantic-level features of an image. The overall framework follows a Next Group of Token Prediction paradigm, where the model is trained to predict the next group of language or visual tokens as a compression target.
|
||||
|
||||
BAGEL scales MoT’s capacity through Pre-training, Continued Training, and Supervised Finetuning on trillions of interleaved multimodal tokens spanning language, image, video, and web data. It surpasses open models on standard understanding and generation benchmarks and demonstrates advanced in-context multimodal abilities like free-form image editing, future frame prediction, 3D manipulation, world navigation, and sequential reasoning.
|
||||
|
||||
<p align="center"><img src="assets/arch.png" width="95%"></p>
|
||||
|
||||
|
||||
## 🌱 Emerging Properties
|
||||
<p align="center"><img src="assets/emerging_curves.png" width="95%"></p>
|
||||
|
||||
As we scale up BAGEL’s pretraining with more multimodal tokens, we observe consistent performance gains across understanding, generation, and editing tasks. Different capabilities emerge at distinct training stages—multimodal understanding and generation appear early, followed by basic editing, while complex, intelligent editing emerges later. This staged progression suggests an emergent pattern, where advanced multimodal reasoning builds on well-formed foundational skills. Ablation studies further show that combining VAE and ViT features significantly improves intelligent editing, underscoring the importance of visual-semantic context in enabling complex multimodal reasoning and further supporting its role in the emergence of advanced capabilities. -->
|
||||
|
||||
## 📢 News
|
||||
|
||||
We sincerely thank all contributors from the open community for their valuable support.
|
||||
|
||||
- **May 26, 2025:** Thanks to [@neverbiasu](https://github.com/neverbiasu) for contributing [ComfyUI](https://github.com/neverbiasu/ComfyUI-BAGEL).
|
||||
- **May 25, 2025:** Special thanks to [@LeanModels](https://github.com/LeanModels) for providing the [DF11-compressed version](https://huggingface.co/DFloat11/BAGEL-7B-MoT-DF11), and to [@Gapeleon](https://huggingface.co/Gapeleon) for the [INT8-compressed version](https://huggingface.co/Gapeleon/bytedance_BAGEL-7B-MoT-INT8). We also appreciate [@gluttony-10](https://github.com/gluttony-10) for contributions to the [Windows package](https://github.com/ByteDance-Seed/Bagel/issues/51).
|
||||
- **May 24, 2025:** Together with [@wangwei1237](https://github.com/wangwei1237), [@gluttony-10](https://github.com/gluttony-10), and [@KingNish24](https://github.com/KingNish24), we built a Gradio [app](app.py) and launched a [Hugging Face Space](https://huggingface.co/spaces/ByteDance-Seed/BAGEL).
|
||||
- **May 23, 2025:** We have provided a training guideline in [TRAIN](./TRAIN.md).
|
||||
- **May 20, 2025:** We released the official [website](https://bagel-ai.org/), [demo](https://demo.bagel-ai.org/), [model](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT), and [report](https://arxiv.org/abs/2505.14683) for BAGEL.
|
||||
|
||||
|
||||
## 📮 Notice
|
||||
**Call for Bad Cases:** If you have encountered any cases where the model performs poorly, we would greatly appreciate it if you could share them in the [issue#11](https://github.com/ByteDance-Seed/Bagel/issues/11) or [Discord](https://discord.gg/Z836xxzy).
|
||||
|
||||
**About Inference Hyperparameters:**
|
||||
- **`cfg_text_scale`:** Controls how strongly the model follows the text prompt. `1.0` disables text guidance. Typical range: `4.0–8.0`.
|
||||
- **`cfg_image_scale`:** Controls how much the model preserves input image details. `1.0` disables image guidance. Typical range: `1.0–2.0`.
|
||||
- **`cfg_interval`:** Fraction of denoising steps where CFG is applied. Later steps can skip CFG to reduce computation. Typical: `[0.4, 1.0]`.
|
||||
- **`timestep_shift`:** Shifts the distribution of denoising steps. Higher values allocate more steps at the start (affects layout); lower values allocate more at the end (improves details).
|
||||
- **`num_timesteps`:** Total denoising steps. Typical: `50`.
|
||||
- **`cfg_renorm_min`:** Minimum value for CFG-Renorm. `1.0` disables renorm. Typical: `0`.
|
||||
- **`cfg_renorm_type`:** CFG-Renorm method:
|
||||
- `global`: Normalize over all tokens and channels (default for T2I).
|
||||
- `channel`: Normalize across channels for each token.
|
||||
- `text_channel`: Like `channel`, but only applies to text condition (good for editing, may cause blur).
|
||||
- **If edited images appear blurry, try `global` CFG-Renorm, decrease `cfg_renorm_min` or decrease `cfg_scale`.**
|
||||
|
||||
|
||||
## 🔥 Quick Start
|
||||
|
||||
1️⃣ Set up environment
|
||||
```bash
|
||||
git clone https://github.com/bytedance-seed/BAGEL.git
|
||||
cd BAGEL
|
||||
conda create -n bagel python=3.10 -y
|
||||
conda activate bagel
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
2️⃣ Download pretrained checkpoint
|
||||
```python
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
save_dir = "/path/to/save/BAGEL-7B-MoT"
|
||||
repo_id = "ByteDance-Seed/BAGEL-7B-MoT"
|
||||
cache_dir = save_dir + "/cache"
|
||||
|
||||
snapshot_download(cache_dir=cache_dir,
|
||||
local_dir=save_dir,
|
||||
repo_id=repo_id,
|
||||
local_dir_use_symlinks=False,
|
||||
resume_download=True,
|
||||
allow_patterns=["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
|
||||
)
|
||||
|
||||
```
|
||||
|
||||
3️⃣ Go to [`inference.ipynb`](inference.ipynb) to start playing with BAGEL!
|
||||
|
||||
4️⃣ Use Gradio WebUI to start playing with BAGEL!
|
||||
```bash
|
||||
pip install gradio
|
||||
python app.py
|
||||
```
|
||||
|
||||
## 🔥 Train & Eval
|
||||
|
||||
### Train
|
||||
|
||||
```bash
|
||||
bash scripts/train.sh
|
||||
```
|
||||
|
||||
You can replace the variables in the script with your own before running.
|
||||
See [TRAIN](TRAIN.md) for more details.
|
||||
|
||||
### Eval
|
||||
We provide the scripts for evaluating VLM, T2I and Editing benchmarks.
|
||||
Please See [EVAL](EVAL.md) for more details.
|
||||
|
||||
|
||||
## 📊 Benchmarks
|
||||
|
||||
### 1. Visual Understanding
|
||||
|
||||
| Model | MME ↑ | MMBench ↑ | MMMU ↑ | MM-Vet ↑ | MathVista ↑ |
|
||||
| ------------------- | ----------: | ----------: | -------: | -------: | ----------: |
|
||||
| Janus-Pro-7B | - | 79.2 | 41.0 | 50.0 | – |
|
||||
| Qwen2.5-VL-7B | 2347 | 83.5 | **58.6** | 67.1 | 68.2 |
|
||||
| **BAGEL** | **2388** | **85.0** | 55.3 | **67.2** | **73.1** |
|
||||
|
||||
### 2. Text-to-Image Generation
|
||||
|
||||
| Model | GenEval ↑ | WISE ↑|
|
||||
| ------------ | --------- | --------- |
|
||||
| Janus-Pro-7B | 0.80 | 0.35 |
|
||||
| SD3-Medium | 0.74 | - |
|
||||
| FLUX-1-dev | 0.82 | 0.50 |
|
||||
| **BAGEL** | - | **0.52** |
|
||||
| **BAGEL + CoT** | **0.88** | **0.70** |
|
||||
|
||||
### 3. Image Editing
|
||||
|
||||
| Model | GEdit-Bench-EN (SC) ↑ | GEdit-Bench-EN (PQ) ↑ | GEdit-Bench-EN (O) ↑ | IntelligentBench ↑ |
|
||||
| ------------- | --------------------- | --------------------- | ------------------- | ------------------ |
|
||||
| Step1X-Edit | 7.09 | 6.76 | **6.70** | 14.9 |
|
||||
| Gemini-2-exp. | 6.73 | 6.61 | 6.32 | **57.6** |
|
||||
| **BAGEL** | **7.36** | **6.83** | 6.52 | 44.0 |
|
||||
| **BAGEL+CoT** | – | – | – | 55.3 |
|
||||
|
||||
|
||||
## ✍️ Citation
|
||||
|
||||
```bibtex
|
||||
@article{deng2025bagel,
|
||||
title = {Emerging Properties in Unified Multimodal Pretraining},
|
||||
author = {Deng, Chaorui and Zhu, Deyao and Li, Kunchang and Gou, Chenhui and Li, Feng and Wang, Zeyu and Zhong, Shu and Yu, Weihao and Nie, Xiaonan and Song, Ziang and Shi, Guang and Fan, Haoqi},
|
||||
journal = {arXiv preprint arXiv:2505.14683},
|
||||
year = {2025}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## 📜 License
|
||||
BAGEL is licensed under the Apache 2.0.
|
||||
+168
@@ -0,0 +1,168 @@
|
||||
# Data prepration
|
||||
|
||||
We provide data examples for **T2I**, **Editing**, and **VLM** tasks. The T2I dataset is generated using [FLUX.1‑dev](https://huggingface.co/black-forest-labs/FLUX.1-dev); the editing examples are randomly sampled from [SEED‑Data‑Edit‑Part3](https://huggingface.co/datasets/AILab-CVC/SEED-Data-Edit-Part2-3); and the VLM set is sourced from [LLaVA‑OneVision‑Data](https://huggingface.co/datasets/lmms-lab/LLaVA-OneVision-Data).
|
||||
|
||||
We offer examples in both raw-image folder and parquet shard formats. For other data formats, you can use our dataset code as a template and extend it as needed.
|
||||
|
||||
|
||||
1. **Download the sample dataset**
|
||||
|
||||
```bash
|
||||
wget -O bagel_example.zip \
|
||||
https://lf3-static.bytednsdoc.com/obj/eden-cn/nuhojubrps/bagel_example.zip
|
||||
unzip bagel_example.zip -d /data
|
||||
```
|
||||
2. **Expected hierarchy**
|
||||
|
||||
```text
|
||||
bagel_example
|
||||
├── t2i/ # text-to-image (parquet)
|
||||
├── editing/ # image editing (parquet)
|
||||
│ ├── seedxedit_multi/
|
||||
│ └── parquet_info/
|
||||
└── vlm/
|
||||
├── images/ # JPEG / PNG frames
|
||||
└── llava_ov_si.jsonl # vision‑language SFT conversations
|
||||
```
|
||||
3. Edit every `your_data_path` placeholder in **`data/dataset_info.py`**.
|
||||
4. *(Optional)* Extend `DATASET_INFO` with your own parquet shards or JSONL files to mix extra data.
|
||||
|
||||
---
|
||||
|
||||
# Training
|
||||
|
||||
The baseline training recipe looks like this (replace environment variables with real paths or values):
|
||||
|
||||
```shell
|
||||
# Pre-training
|
||||
torchrun \
|
||||
--nnodes=$num_nodes \
|
||||
--node_rank=$node_rank \
|
||||
--nproc_per_node=8 \
|
||||
--master_addr=$master_addr \
|
||||
--master_port=$master_port \
|
||||
train/pretrain_unified_navit.py \
|
||||
--dataset_config_file ./data/configs/example.yaml \
|
||||
--llm_path $llm_path \
|
||||
--vae_path $vae_path \
|
||||
--vit_path $vit_path \
|
||||
--layer_module Qwen2MoTDecoderLayer \
|
||||
--use_flex True \
|
||||
--resume_from $resume_from \
|
||||
--results_dir $output_path \
|
||||
--checkpoint_dir $ckpt_path \
|
||||
--max_latent_size 64 # 32 for low-resolution pre-training
|
||||
|
||||
# Fine-tuning
|
||||
torchrun \
|
||||
--nnodes=$num_nodes \
|
||||
--node_rank=$node_rank \
|
||||
--nproc_per_node=8 \
|
||||
--master_addr=$master_addr \
|
||||
--master_port=$master_port \
|
||||
train/pretrain_unified_navit.py \
|
||||
--dataset_config_file ./data/configs/example.yaml \
|
||||
--model_path $model_path \
|
||||
--layer_module Qwen2MoTDecoderLayer \
|
||||
--max_latent_size 64 \
|
||||
--resume-from $model_path \
|
||||
--finetune_from_hf True \
|
||||
--auto_resume True \
|
||||
--resume-model-only True \
|
||||
--finetune-from-ema True \
|
||||
--log_every 1 \
|
||||
--lr 2e-5 \
|
||||
--num_worker 1 \
|
||||
--expected_num_tokens 10240 \
|
||||
--max_num_tokens 11520 \
|
||||
--max_num_tokens_per_sample 10240
|
||||
```
|
||||
|
||||
- **When fine-tuning BAGEL, set `max_latent_size=64` to ensure the correct pretrained weights are loaded.** If this is not set, an out-of-bounds error may occur.
|
||||
- The total value of `num_used_data` should be greater than `NUM_GPUS × NUM_WORKERS`. (For toy data, use `num_worker=1`.)
|
||||
- For T2I-only fine-tuning, set `visual_und=False`. For VLM-only fine-tuning, set `visual_gen=False`.
|
||||
- For debugging purposes, use smaller values for `expected_num_tokens`, `max_num_tokens`, and `max_num_tokens_per_sample`.
|
||||
- When fine-tuning on toy data, the loss behaves as follows:
|
||||
```shell
|
||||
[2025-05-25 17:01:37] (step=0000000) Train Loss mse: 0.4063, Train Loss ce: 0.5504, Train Steps/Sec: 0.01,
|
||||
[2025-05-25 17:01:40] (step=0000001) Train Loss mse: 0.4121, Train Loss ce: 0.8152, Train Steps/Sec: 0.44,
|
||||
[2025-05-25 17:01:42] (step=0000002) Train Loss mse: 0.3876, Train Loss ce: 1.3411, Train Steps/Sec: 0.40,
|
||||
[2025-05-25 17:01:45] (step=0000003) Train Loss mse: 0.3825, Train Loss ce: 0.7360, Train Steps/Sec: 0.44,
|
||||
```
|
||||
|
||||
|
||||
You are encouraged to adjust any of these hyperparameters to fit your GPU budget and the scale of your dataset. If you encounter any issues, please open an issue for assistance. 🎉
|
||||
|
||||
|
||||
## Model config
|
||||
|
||||
|
||||
| Argument | Default | Description |
|
||||
| ---------------------------- | ------------------------------------------- | --------------------------------------------------------------- |
|
||||
| `llm_path` | `hf/Qwen2.5-0.5B-Instruct` | Language‑model backbone (HuggingFace repo or local folder). |
|
||||
| `vae_path` | `flux/vae/ae.safetensors` | Pre‑trained VAE checkpoint for latent diffusion. |
|
||||
| `vit_path` | `hf/siglip-so400m-14-980-flash-attn2-navit` | SigLIP ViT used for image understanding. |
|
||||
| `max_latent_size` | `32` | Maximum latent grid side; defines highest generable resolution. |
|
||||
| `latent_patch_size` | `2` | VAE pixels represented by one latent patch. |
|
||||
| `vit_max_num_patch_per_side` | `70` | Max ViT patches per image side after resizing. |
|
||||
| `text_cond_dropout_prob` | `0.1` | Probability to drop text conditioning while training. |
|
||||
| `vae_cond_dropout_prob` | `0.3` | Dropout on VAE latent inputs. |
|
||||
| `vit_cond_dropout_prob` | `0.3` | Dropout on visual features. |
|
||||
|
||||
*(See `ModelArguments` for many more options.)*
|
||||
|
||||
|
||||
## Data config
|
||||
|
||||
|
||||
| Argument | Default | Description |
|
||||
| --------------------------- | --------------------------- | --------------------------------------------------------- |
|
||||
| `dataset_config_file` | `data/configs/example.yaml` | YAML that groups datasets and assigns sampling weights. |
|
||||
| `num_workers` | `4` | Background workers per rank for the PyTorch `DataLoader`. |
|
||||
| `prefetch_factor` | `2` | Batches pre‑fetched by each worker. |
|
||||
| `max_num_tokens_per_sample` | `16384` | Skip raw samples longer than this. |
|
||||
| `max_num_tokens` | `36864` | Hard cap for a packed batch (prevents OOM). |
|
||||
| `max_buffer_size` | `50` | Overflow buffer length for oversized samples. |
|
||||
| `data_seed` | `42` | Seed for reproducible shuffling and sampling. |
|
||||
|
||||
|
||||
## Training config
|
||||
|
||||
| Argument | Default | Description |
|
||||
| -------------------------------------- | ---------------------- | ------------------------------------------------------ |
|
||||
| `total_steps` | `500_000` | Optimiser steps to run. |
|
||||
| `lr` | `1e-4` | Peak learning rate after warm‑up. |
|
||||
| `lr_scheduler` | `constant` | Learning‑rate schedule (`constant` or `cosine`). |
|
||||
| `warmup_steps` | `2000` | Linear warm‑up duration. |
|
||||
| `ema` | `0.9999` | Exponential moving‑average decay for model weights. |
|
||||
| `max_grad_norm` | `1.0` | Gradient‑clipping threshold. |
|
||||
| `save_every` | `2000` | Checkpoint frequency (steps). |
|
||||
| `visual_gen / visual_und` | `True` | Enable image generation / understanding branches. |
|
||||
| `freeze_llm / freeze_vit / freeze_vae` | `False / False / True` | Freeze selected modules to save VRAM or for ablations. |
|
||||
| `use_flex` | `True` (in example) | Enable FLEX packing for higher GPU utilisation. |
|
||||
| `sharding_strategy` | `HYBRID_SHARD` | FSDP sharding mode. |
|
||||
| `num_shard` | `8` | Parameter shards per rank in HYBRID mode. |
|
||||
|
||||
**Distributed‑launch environment variables**
|
||||
|
||||
| Var | Meaning |
|
||||
| ----------------------------- | --------------------------------- |
|
||||
| `num_nodes` / `node_rank` | Multi‑node orchestration indices. |
|
||||
| `nproc_per_node` | Number of GPUs per node. |
|
||||
| `master_addr` / `master_port` | NCCL rendezvous endpoint. |
|
||||
|
||||
|
||||
## Logging config
|
||||
|
||||
|
||||
| Argument | Default | Description |
|
||||
| ---------------- | --------------------- | ---------------------------------------------------- |
|
||||
| `results_dir` | `results` | Root directory for logs and metrics. |
|
||||
| `checkpoint_dir` | `results/checkpoints` | Checkpoints are saved here. |
|
||||
| `log_every` | `10` | Steps between console / W\&B logs. |
|
||||
| `wandb_project` | `bagel` | Weights & Biases project name. |
|
||||
| `wandb_name` | `run` | Run name inside the project. |
|
||||
| `wandb_offline` | `False` | Switch to offline mode (logs locally, sync later). |
|
||||
| `wandb_resume` | `allow` | Resumption policy if an existing run ID is detected. |
|
||||
|
||||
> **Tip** Export `WANDB_API_KEY` before launching if you want online dashboards.
|
||||
+256
@@ -0,0 +1,256 @@
|
||||
|
||||
import numpy as np
|
||||
import os
|
||||
import torch
|
||||
import random
|
||||
|
||||
from accelerate import infer_auto_device_map, load_checkpoint_and_dispatch, init_empty_weights
|
||||
from accelerate.utils import BnbQuantizationConfig, load_and_quantize_model
|
||||
from PIL import Image
|
||||
|
||||
from .data.data_utils import add_special_tokens, pil_img2rgb
|
||||
from .data.transforms import ImageTransform
|
||||
from .inferencer import InterleaveInferencer
|
||||
from .modeling.autoencoder import load_ae
|
||||
from .modeling.bagel.qwen2_navit import NaiveCache
|
||||
from .modeling.bagel import (
|
||||
BagelConfig, Bagel, Qwen2Config, Qwen2ForCausalLM,
|
||||
SiglipVisionConfig, SiglipVisionModel
|
||||
)
|
||||
from .modeling.qwen2 import Qwen2Tokenizer
|
||||
|
||||
|
||||
|
||||
def load_bagel_model(model_path,quantize_mode,max_mem_per_gpu = "24GiB"):
|
||||
print ("Loading model .... ")
|
||||
llm_config = Qwen2Config.from_json_file(os.path.join(model_path, "llm_config.json"))
|
||||
llm_config.qk_norm = True
|
||||
llm_config.tie_word_embeddings = False
|
||||
llm_config.layer_module = "Qwen2MoTDecoderLayer"
|
||||
vit_config = SiglipVisionConfig.from_json_file(os.path.join(model_path, "vit_config.json"))
|
||||
vit_config.rope = False
|
||||
vit_config.num_hidden_layers -= 1
|
||||
vae_model, vae_config = load_ae(local_path=os.path.join(model_path, "ae.safetensors"))
|
||||
config = BagelConfig(
|
||||
visual_gen=True,
|
||||
visual_und=True,
|
||||
llm_config=llm_config,
|
||||
vit_config=vit_config,
|
||||
vae_config=vae_config,
|
||||
vit_max_num_patch_per_side=70,
|
||||
connector_act='gelu_pytorch_tanh',
|
||||
latent_patch_size=2,
|
||||
max_latent_size=64,
|
||||
)
|
||||
with init_empty_weights():
|
||||
language_model = Qwen2ForCausalLM(llm_config)
|
||||
vit_model = SiglipVisionModel(vit_config)
|
||||
model = Bagel(language_model, vit_model, config)
|
||||
model.vit_model.vision_model.embeddings.convert_conv2d_to_linear(vit_config, meta=True)
|
||||
|
||||
tokenizer = Qwen2Tokenizer.from_pretrained(model_path)
|
||||
tokenizer, new_token_ids, _ = add_special_tokens(tokenizer)
|
||||
|
||||
vae_transform = ImageTransform(1024, 512, 16)
|
||||
vit_transform = ImageTransform(980, 224, 14)
|
||||
# Model Loading and Multi GPU Infernece Preparing
|
||||
|
||||
|
||||
#max_mem_per_gpu = "24GiB" # Modify it according to your GPU setting. On an A100, 80 GiB is sufficient to load on a single GPU.
|
||||
|
||||
device_map = infer_auto_device_map(
|
||||
model,
|
||||
max_memory={0: max_mem_per_gpu},
|
||||
no_split_module_classes=["Bagel", "Qwen2MoTDecoderLayer", "SiglipVisionModel"],
|
||||
)
|
||||
#print(f"device_map is {device_map}")
|
||||
|
||||
same_device_modules = [
|
||||
'language_model.model.embed_tokens',
|
||||
'time_embedder',
|
||||
'latent_pos_embed',
|
||||
'vae2llm',
|
||||
'llm2vae',
|
||||
'connector',
|
||||
'vit_pos_embed'
|
||||
]
|
||||
|
||||
if torch.cuda.device_count() == 1:
|
||||
first_device = device_map.get(same_device_modules[0], "cuda:0")
|
||||
for k in same_device_modules:
|
||||
if k in device_map:
|
||||
device_map[k] = first_device
|
||||
else:
|
||||
device_map[k] = "cuda:0"
|
||||
else:
|
||||
first_device = device_map.get(same_device_modules[0])
|
||||
for k in same_device_modules:
|
||||
if k in device_map:
|
||||
device_map[k] = first_device
|
||||
|
||||
if quantize_mode == "fp16":
|
||||
model = load_checkpoint_and_dispatch(
|
||||
model,
|
||||
checkpoint=os.path.join(model_path, "ema.safetensors"),
|
||||
device_map=device_map,
|
||||
offload_buffers=True,
|
||||
offload_folder="offload",
|
||||
dtype=torch.bfloat16,
|
||||
force_hooks=True,
|
||||
).eval()
|
||||
elif quantize_mode == "nf4":
|
||||
print("Quantizing to NF4")
|
||||
bnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=False, bnb_4bit_quant_type="nf4")
|
||||
model = load_and_quantize_model(
|
||||
model,
|
||||
weights_location=os.path.join(model_path, "ema.safetensors"),
|
||||
bnb_quantization_config=bnb_quantization_config,
|
||||
device_map=device_map,
|
||||
offload_folder="offload",
|
||||
).eval()
|
||||
else:#args.mode == 3: # INT8
|
||||
print("Quantizing to INT8/fp8")
|
||||
bnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, torch_dtype=torch.float32)
|
||||
model = load_and_quantize_model(
|
||||
model,
|
||||
weights_location=os.path.join(model_path, "ema.safetensors"),
|
||||
bnb_quantization_config=bnb_quantization_config,
|
||||
device_map=device_map,
|
||||
offload_folder="offload",
|
||||
).eval()
|
||||
|
||||
model = model.eval()
|
||||
print('Model loaded is done!')
|
||||
inferencer = InterleaveInferencer(
|
||||
model=model,
|
||||
vae_model=vae_model,
|
||||
tokenizer=tokenizer,
|
||||
vae_transform=vae_transform,
|
||||
vit_transform=vit_transform,
|
||||
new_token_ids=new_token_ids
|
||||
)
|
||||
return inferencer
|
||||
|
||||
|
||||
|
||||
def set_seed(seed):
|
||||
"""Set random seeds for reproducibility"""
|
||||
if seed > 0:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
return seed
|
||||
|
||||
# Text to Image function with thinking option and hyperparameters
|
||||
def text_to_image(inferencer,prompt, show_thinking=False, cfg_text_scale=4.0, cfg_interval=0.4,
|
||||
timestep_shift=3.0, num_timesteps=50,
|
||||
cfg_renorm_min=0.0, cfg_renorm_type="global",
|
||||
max_think_token_n=1024, do_sample=False, text_temperature=0.3,
|
||||
seed=0, image_ratio="1:1"):
|
||||
# Set seed for reproducibility
|
||||
set_seed(seed)
|
||||
|
||||
if image_ratio == "1:1":
|
||||
image_shapes = (1024, 1024)
|
||||
elif image_ratio == "4:3":
|
||||
image_shapes = (768, 1024)
|
||||
elif image_ratio == "3:4":
|
||||
image_shapes = (1024, 768)
|
||||
elif image_ratio == "16:9":
|
||||
image_shapes = (576, 1024)
|
||||
elif image_ratio == "9:16":
|
||||
image_shapes = (1024, 576)
|
||||
|
||||
# Set hyperparameters
|
||||
inference_hyper = dict(
|
||||
max_think_token_n=max_think_token_n if show_thinking else 1024,
|
||||
do_sample=do_sample if show_thinking else False,
|
||||
text_temperature=text_temperature if show_thinking else 0.3,
|
||||
cfg_text_scale=cfg_text_scale,
|
||||
cfg_interval=[cfg_interval, 1.0], # End fixed at 1.0
|
||||
timestep_shift=timestep_shift,
|
||||
num_timesteps=num_timesteps,
|
||||
cfg_renorm_min=cfg_renorm_min,
|
||||
cfg_renorm_type=cfg_renorm_type,
|
||||
image_shapes=image_shapes,
|
||||
)
|
||||
|
||||
# Call inferencer with or without think parameter based on user choice
|
||||
result = inferencer(text=prompt, think=show_thinking, **inference_hyper)
|
||||
return result["image"], result.get("text", "")
|
||||
|
||||
|
||||
# Image Understanding function with thinking option and hyperparameters
|
||||
def image_understanding(inferencer,image: Image.Image, prompt: str, show_thinking=False,
|
||||
do_sample=False, text_temperature=0.3, max_new_tokens=512):
|
||||
if image is None:
|
||||
return "Please upload an image."
|
||||
|
||||
if isinstance(image, np.ndarray):
|
||||
image = Image.fromarray(image)
|
||||
|
||||
image = pil_img2rgb(image)
|
||||
|
||||
# Set hyperparameters
|
||||
inference_hyper = dict(
|
||||
do_sample=do_sample,
|
||||
text_temperature=text_temperature,
|
||||
max_think_token_n=max_new_tokens, # Set max_length
|
||||
)
|
||||
|
||||
# Use show_thinking parameter to control thinking process
|
||||
result = inferencer(image=image, text=prompt, think=show_thinking,
|
||||
understanding_output=True, **inference_hyper)
|
||||
return result["text"]
|
||||
|
||||
|
||||
# Image Editing function with thinking option and hyperparameters
|
||||
def edit_image(inferencer,image: Image.Image, prompt: str, show_thinking=False, cfg_text_scale=4.0,
|
||||
cfg_img_scale=2.0, cfg_interval=0.0,
|
||||
timestep_shift=3.0, num_timesteps=50, cfg_renorm_min=0.0,
|
||||
cfg_renorm_type="text_channel", max_think_token_n=1024,
|
||||
do_sample=False, text_temperature=0.3, seed=0):
|
||||
# Set seed for reproducibility
|
||||
set_seed(seed)
|
||||
|
||||
if image is None:
|
||||
return "Please upload an image.", ""
|
||||
|
||||
if isinstance(image, np.ndarray):
|
||||
image = Image.fromarray(image)
|
||||
|
||||
image = pil_img2rgb(image)
|
||||
|
||||
# Set hyperparameters
|
||||
inference_hyper = dict(
|
||||
max_think_token_n=max_think_token_n if show_thinking else 1024,
|
||||
do_sample=do_sample if show_thinking else False,
|
||||
text_temperature=text_temperature if show_thinking else 0.3,
|
||||
cfg_text_scale=cfg_text_scale,
|
||||
cfg_img_scale=cfg_img_scale,
|
||||
cfg_interval=[cfg_interval, 1.0], # End fixed at 1.0
|
||||
timestep_shift=timestep_shift,
|
||||
num_timesteps=num_timesteps,
|
||||
cfg_renorm_min=cfg_renorm_min,
|
||||
cfg_renorm_type=cfg_renorm_type,
|
||||
)
|
||||
|
||||
# Include thinking parameter based on user choice
|
||||
result = inferencer(image=image, text=prompt, think=show_thinking, **inference_hyper)
|
||||
return result["image"], result.get("text", "")
|
||||
|
||||
|
||||
# Helper function to load example images
|
||||
def load_example_image(image_path):
|
||||
try:
|
||||
return Image.open(image_path)
|
||||
except Exception as e:
|
||||
print(f"Error loading example image: {e}")
|
||||
return None
|
||||
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 164 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 364 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
@@ -0,0 +1,2 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
@@ -0,0 +1,45 @@
|
||||
t2i_pretrain:
|
||||
dataset_names:
|
||||
- t2i
|
||||
image_transform_args:
|
||||
image_stride: 16
|
||||
max_image_size: 1024
|
||||
min_image_size: 512
|
||||
is_mandatory: true
|
||||
num_used_data: # The sum should be larger that NUM_GPUS x NUM_WORKERS
|
||||
- 10
|
||||
weight: 1
|
||||
|
||||
unified_edit:
|
||||
dataset_names:
|
||||
- seedxedit_multi
|
||||
image_transform_args:
|
||||
image_stride: 16
|
||||
max_image_size: 1024
|
||||
min_image_size: 512
|
||||
vit_image_transform_args:
|
||||
image_stride: 14
|
||||
max_image_size: 518
|
||||
min_image_size: 224
|
||||
is_mandatory: false
|
||||
num_used_data:
|
||||
- 10
|
||||
weight: 1
|
||||
|
||||
vlm_sft:
|
||||
dataset_names:
|
||||
- llava_ov
|
||||
image_transform_args:
|
||||
image_stride: 14
|
||||
max_image_size: 980
|
||||
min_image_size: 378
|
||||
max_pixels: 2_007_040
|
||||
frame_sampler_args:
|
||||
max_num_frames: 12
|
||||
min_num_frames: 8
|
||||
is_mandatory: true
|
||||
shuffle_lines: True
|
||||
shuffle_seed: 0
|
||||
num_used_data:
|
||||
- 1000
|
||||
weight: 1
|
||||
@@ -0,0 +1,177 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
|
||||
import math
|
||||
import random
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
from torch.nn.attention.flex_attention import or_masks, and_masks
|
||||
|
||||
|
||||
def create_sparse_mask(document_lens, split_lens, attn_modes, device):
|
||||
def causal_mask(b, h, q_idx, kv_idx):
|
||||
return q_idx >= kv_idx
|
||||
|
||||
def full_and_noise_mask(b, h, q_idx, kv_idx):
|
||||
return (full_and_noise_seq_id[q_idx] == full_and_noise_seq_id[kv_idx]) & (full_and_noise_seq_id[q_idx] >= 0)
|
||||
|
||||
def remove_noise_mask(b, h, q_idx, kv_idx):
|
||||
return (~((noise_seq_id[kv_idx] >= 0) & (noise_seq_id[q_idx] != noise_seq_id[kv_idx])))
|
||||
|
||||
def sample_mask(b, h, q_idx, kv_idx):
|
||||
return document_id[q_idx] == document_id[kv_idx]
|
||||
|
||||
full_and_noise_tmp = []
|
||||
noise_tmp = []
|
||||
|
||||
for i, (length, model) in enumerate(zip(split_lens, attn_modes)):
|
||||
value = i if model in ['full', 'noise'] else -1
|
||||
full_and_noise_tmp.extend([value] * length)
|
||||
value_noise = i if model == 'noise' else -1
|
||||
noise_tmp.extend([value_noise] * length)
|
||||
|
||||
full_and_noise_seq_id = torch.Tensor(full_and_noise_tmp).to(device)
|
||||
noise_seq_id = torch.Tensor(noise_tmp).to(device)
|
||||
|
||||
document_id = torch.cat([torch.full((l,), i) for i, l in enumerate(document_lens, start=1)]).to(device)
|
||||
|
||||
return and_masks(or_masks(causal_mask, full_and_noise_mask), remove_noise_mask, sample_mask)
|
||||
|
||||
|
||||
def patchify(image, patch_size):
|
||||
p = patch_size
|
||||
c, h, w = image.shape
|
||||
assert h % p == 0 and w % p == 0
|
||||
image = image.reshape(c, h // p, p, w // p, p)
|
||||
image = torch.einsum("chpwq->hwpqc", image)
|
||||
image = image.reshape(-1, p**2 * c)
|
||||
return image
|
||||
|
||||
|
||||
def get_flattened_position_ids_extrapolate(img_h, img_w, patch_size, max_num_patches_per_side):
|
||||
num_patches_h, num_patches_w = img_h // patch_size, img_w // patch_size
|
||||
coords_h = torch.arange(0, num_patches_h)
|
||||
coords_w = torch.arange(0, num_patches_w)
|
||||
pos_ids = (coords_h[:, None] * max_num_patches_per_side + coords_w).flatten()
|
||||
return pos_ids
|
||||
|
||||
|
||||
def get_flattened_position_ids_interpolate(img_h, img_w, patch_size, max_num_patches_per_side):
|
||||
num_patches_h, num_patches_w = img_h // patch_size, img_w // patch_size
|
||||
boundaries = torch.arange(1 / max_num_patches_per_side, 1.0, 1 / max_num_patches_per_side)
|
||||
fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / num_patches_h)
|
||||
fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / num_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] * max_num_patches_per_side + bucket_coords_w).flatten()
|
||||
return pos_ids
|
||||
|
||||
|
||||
def prepare_attention_mask_per_sample(split_lens, attn_modes, device="cpu"):
|
||||
"""
|
||||
nested_split_lens: A list of N lists of ints. Each int indicates the length of a split within
|
||||
a sample, where each sample contains multiple splits with different attn modes.
|
||||
nested_attn_modes: whether to use full attn in each split.
|
||||
"""
|
||||
sample_len = sum(split_lens)
|
||||
attention_mask = torch.zeros((sample_len, sample_len), dtype=torch.bool, device=device)
|
||||
|
||||
csum = 0
|
||||
for s, attn_mode in zip(split_lens, attn_modes):
|
||||
assert attn_mode in ['causal', 'full', 'noise']
|
||||
if attn_mode == "causal":
|
||||
attention_mask[csum:csum + s, csum:csum + s] = torch.ones((s, s), device=device).tril()
|
||||
attention_mask[csum:csum + s, :csum] = 1
|
||||
else:
|
||||
attention_mask[csum:csum + s, csum:csum + s] = torch.ones((s, s))
|
||||
attention_mask[csum:csum + s, :csum] = 1
|
||||
csum += s
|
||||
|
||||
csum = 0
|
||||
for s, attn_mode in zip(split_lens, attn_modes):
|
||||
if attn_mode == "noise":
|
||||
attention_mask[:, csum : csum + s] = torch.zeros((sample_len, s))
|
||||
attention_mask[csum : csum + s, csum : csum + s] = torch.ones((s, s))
|
||||
csum += s
|
||||
|
||||
attention_mask = torch.zeros_like(attention_mask, dtype=torch.float).masked_fill_(
|
||||
~attention_mask, float("-inf")
|
||||
)
|
||||
|
||||
return attention_mask
|
||||
|
||||
|
||||
def split_integer_exp_decay(S, ng_sample_decay=1.0):
|
||||
if ng_sample_decay == 1.0:
|
||||
N = random.randint(1, S)
|
||||
else:
|
||||
base = (1 - ng_sample_decay) / (1 - math.pow(ng_sample_decay, S))
|
||||
p = [base * math.pow(ng_sample_decay, i) for i in range(S)]
|
||||
N = random.choices(list(range(1, S + 1)), p, k=1)[0]
|
||||
cumsum = [0] + sorted(random.sample(range(1, S), N - 1)) + [S]
|
||||
result = [cumsum[i+1] - cumsum[i] for i in range(len(cumsum) - 1)]
|
||||
return result, cumsum
|
||||
|
||||
|
||||
def pil_img2rgb(image):
|
||||
if image.mode == "RGBA" or image.info.get("transparency", None) is not None:
|
||||
image = image.convert("RGBA")
|
||||
white = Image.new(mode="RGB", size=image.size, color=(255, 255, 255))
|
||||
white.paste(image, mask=image.split()[3])
|
||||
image = white
|
||||
else:
|
||||
image = image.convert("RGB")
|
||||
|
||||
return image
|
||||
|
||||
|
||||
def add_special_tokens(tokenizer):
|
||||
all_special_tokens = []
|
||||
for k, v in tokenizer.special_tokens_map.items():
|
||||
if isinstance(v, str):
|
||||
all_special_tokens.append(v)
|
||||
elif isinstance(v, list):
|
||||
all_special_tokens += v
|
||||
|
||||
new_tokens = []
|
||||
|
||||
if '<|im_start|>' not in all_special_tokens:
|
||||
new_tokens.append('<|im_start|>')
|
||||
|
||||
if '<|im_end|>' not in all_special_tokens:
|
||||
new_tokens.append('<|im_end|>')
|
||||
|
||||
if '<|vision_start|>' not in all_special_tokens:
|
||||
new_tokens.append('<|vision_start|>')
|
||||
|
||||
if '<|vision_end|>' not in all_special_tokens:
|
||||
new_tokens.append('<|vision_end|>')
|
||||
|
||||
num_new_tokens = tokenizer.add_tokens(new_tokens)
|
||||
bos_token_id = tokenizer.convert_tokens_to_ids('<|im_start|>')
|
||||
eos_token_id = tokenizer.convert_tokens_to_ids('<|im_end|>')
|
||||
start_of_image = tokenizer.convert_tokens_to_ids('<|vision_start|>')
|
||||
end_of_image = tokenizer.convert_tokens_to_ids('<|vision_end|>')
|
||||
|
||||
new_token_ids = dict(
|
||||
bos_token_id=bos_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
start_of_image=start_of_image,
|
||||
end_of_image=end_of_image,
|
||||
)
|
||||
|
||||
return tokenizer, new_token_ids, num_new_tokens
|
||||
|
||||
|
||||
def len2weight(x, loss_reduction='square'):
|
||||
if x == 0:
|
||||
return x
|
||||
if loss_reduction == 'token':
|
||||
return 1
|
||||
if loss_reduction == 'sample':
|
||||
return 1 / x
|
||||
if loss_reduction == 'square':
|
||||
return 1 / (x ** 0.5)
|
||||
raise NotImplementedError(loss_reduction)
|
||||
@@ -0,0 +1,620 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
|
||||
import random
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .data_utils import (
|
||||
get_flattened_position_ids_interpolate,
|
||||
get_flattened_position_ids_extrapolate,
|
||||
len2weight,
|
||||
patchify,
|
||||
prepare_attention_mask_per_sample,
|
||||
)
|
||||
from .dataset_info import DATASET_INFO, DATASET_REGISTRY
|
||||
from .transforms import ImageTransform
|
||||
from .video_utils import FrameSampler
|
||||
|
||||
|
||||
class DataConfig:
|
||||
def __init__(
|
||||
self,
|
||||
grouped_datasets,
|
||||
text_cond_dropout_prob=0.1,
|
||||
vit_cond_dropout_prob=0.4,
|
||||
vae_cond_dropout_prob=0.1,
|
||||
vae_image_downsample=16,
|
||||
max_latent_size=32,
|
||||
vit_patch_size=14,
|
||||
max_num_patch_per_side=70,
|
||||
):
|
||||
self.grouped_datasets = grouped_datasets
|
||||
self.text_cond_dropout_prob = text_cond_dropout_prob
|
||||
self.vit_cond_dropout_prob = vit_cond_dropout_prob
|
||||
self.vit_patch_size = vit_patch_size
|
||||
self.max_num_patch_per_side = max_num_patch_per_side
|
||||
self.vae_cond_dropout_prob = vae_cond_dropout_prob
|
||||
self.vae_image_downsample = vae_image_downsample
|
||||
self.max_latent_size = max_latent_size
|
||||
|
||||
|
||||
class PackedDataset(torch.utils.data.IterableDataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_config,
|
||||
tokenizer,
|
||||
special_tokens,
|
||||
local_rank,
|
||||
world_size,
|
||||
num_workers,
|
||||
expected_num_tokens=32768,
|
||||
max_num_tokens_per_sample=16384,
|
||||
max_num_tokens=36864,
|
||||
prefer_buffer_before=16384,
|
||||
max_buffer_size=50,
|
||||
interpolate_pos=False,
|
||||
use_flex=False,
|
||||
data_status=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.expected_num_tokens = expected_num_tokens
|
||||
self.max_num_tokens_per_sample = max_num_tokens_per_sample
|
||||
self.prefer_buffer_before = prefer_buffer_before
|
||||
self.max_num_tokens = max_num_tokens
|
||||
self.max_buffer_size = max_buffer_size
|
||||
self.tokenizer = tokenizer
|
||||
self.local_rank = local_rank
|
||||
self.world_size = world_size
|
||||
self.num_workers = num_workers
|
||||
self.use_flex = use_flex
|
||||
for k, v in special_tokens.items():
|
||||
setattr(self, k, v)
|
||||
|
||||
grouped_datasets, is_mandatory, grouped_weights = self.build_datasets(
|
||||
data_config.grouped_datasets, data_status
|
||||
)
|
||||
self.grouped_datasets = grouped_datasets
|
||||
self.dataset_iters = [iter(dataset) for dataset in grouped_datasets]
|
||||
self.is_mandatory = is_mandatory
|
||||
self.grouped_weights = grouped_weights
|
||||
self.data_config = data_config
|
||||
self.interpolate_pos = interpolate_pos
|
||||
if self.interpolate_pos:
|
||||
self.get_flattened_position_ids = get_flattened_position_ids_interpolate
|
||||
else:
|
||||
self.get_flattened_position_ids = get_flattened_position_ids_extrapolate
|
||||
|
||||
def build_datasets(self, datasets_metainfo, data_status):
|
||||
datasets = []
|
||||
is_mandatory = []
|
||||
grouped_weights = []
|
||||
for grouped_dataset_name, dataset_args in datasets_metainfo.items():
|
||||
is_mandatory.append(dataset_args.pop('is_mandatory', False))
|
||||
grouped_weights.append(dataset_args.pop('weight', 0.0))
|
||||
|
||||
if 'frame_sampler_args' in dataset_args.keys():
|
||||
frame_sampler = FrameSampler(**dataset_args.pop('frame_sampler_args'))
|
||||
dataset_args['frame_sampler'] = frame_sampler
|
||||
if 'image_transform_args' in dataset_args.keys():
|
||||
transform = ImageTransform(**dataset_args.pop('image_transform_args'))
|
||||
dataset_args['transform'] = transform
|
||||
if 'vit_image_transform_args' in dataset_args.keys():
|
||||
vit_transform = ImageTransform(**dataset_args.pop('vit_image_transform_args'))
|
||||
dataset_args['vit_transform'] = vit_transform
|
||||
|
||||
assert 'dataset_names' in dataset_args.keys()
|
||||
dataset_names = dataset_args.pop('dataset_names')
|
||||
dataset_args['data_dir_list'] = []
|
||||
for item in dataset_names:
|
||||
if self.local_rank == 0:
|
||||
print(f'Preparing Dataset {grouped_dataset_name}/{item}')
|
||||
meta_info = DATASET_INFO[grouped_dataset_name][item]
|
||||
dataset_args['data_dir_list'].append(meta_info['data_dir'])
|
||||
|
||||
if "parquet_info_path" in meta_info.keys():
|
||||
if 'parquet_info' not in dataset_args.keys():
|
||||
dataset_args['parquet_info'] = {}
|
||||
with open(meta_info['parquet_info_path'], 'r') as f:
|
||||
parquet_info = json.load(f)
|
||||
dataset_args['parquet_info'].update(parquet_info)
|
||||
|
||||
if 'json_dir' in meta_info.keys():
|
||||
# parquet/tar with json
|
||||
if 'json_dir_list' not in dataset_args.keys():
|
||||
dataset_args['json_dir_list'] = [meta_info['json_dir']]
|
||||
else:
|
||||
dataset_args['json_dir_list'].append(meta_info['json_dir'])
|
||||
|
||||
if 'jsonl_path' in meta_info.keys():
|
||||
# jsonl with jpeg
|
||||
if 'jsonl_path_list' not in dataset_args.keys():
|
||||
dataset_args['jsonl_path_list'] = [meta_info['jsonl_path']]
|
||||
else:
|
||||
dataset_args['jsonl_path_list'].append(meta_info['jsonl_path'])
|
||||
|
||||
resume_data_status = dataset_args.pop('resume_data_status', True)
|
||||
if data_status is not None and grouped_dataset_name in data_status.keys() and resume_data_status:
|
||||
data_status_per_group = data_status[grouped_dataset_name]
|
||||
else:
|
||||
data_status_per_group = None
|
||||
dataset = DATASET_REGISTRY[grouped_dataset_name](
|
||||
dataset_name=grouped_dataset_name,
|
||||
tokenizer=self.tokenizer,
|
||||
local_rank=self.local_rank,
|
||||
world_size=self.world_size,
|
||||
num_workers=self.num_workers,
|
||||
data_status=data_status_per_group,
|
||||
**dataset_args
|
||||
)
|
||||
datasets.append(dataset)
|
||||
|
||||
return datasets, is_mandatory, grouped_weights
|
||||
|
||||
def set_epoch(self, seed):
|
||||
for dataset in self.grouped_datasets:
|
||||
dataset.set_epoch(seed)
|
||||
|
||||
def set_sequence_status(self):
|
||||
sequence_status = dict(
|
||||
curr = 0,
|
||||
sample_lens = list(),
|
||||
packed_position_ids = list(),
|
||||
nested_attention_masks = list(),
|
||||
split_lens = list(),
|
||||
attn_modes = list(),
|
||||
packed_text_ids = list(),
|
||||
packed_text_indexes = list(),
|
||||
packed_label_ids = list(),
|
||||
ce_loss_indexes = list(),
|
||||
ce_loss_weights = list(),
|
||||
vae_image_tensors = list(),
|
||||
packed_latent_position_ids = list(),
|
||||
vae_latent_shapes = list(),
|
||||
packed_vae_token_indexes = list(),
|
||||
packed_timesteps = list(),
|
||||
mse_loss_indexes = list(),
|
||||
packed_vit_tokens = list(),
|
||||
vit_token_seqlens = list(),
|
||||
packed_vit_position_ids = list(),
|
||||
packed_vit_token_indexes = list(),
|
||||
)
|
||||
return sequence_status
|
||||
|
||||
def to_tensor(self, sequence_status):
|
||||
data = dict(
|
||||
sequence_length=sum(sequence_status['sample_lens']),
|
||||
sample_lens=sequence_status['sample_lens'],
|
||||
packed_text_ids=torch.tensor(sequence_status['packed_text_ids']),
|
||||
packed_text_indexes=torch.tensor(sequence_status['packed_text_indexes']),
|
||||
packed_position_ids=torch.tensor(sequence_status['packed_position_ids']),
|
||||
)
|
||||
if not self.use_flex:
|
||||
data['nested_attention_masks'] = sequence_status['nested_attention_masks']
|
||||
else:
|
||||
sequence_len = data['sequence_length']
|
||||
pad_len = self.max_num_tokens - sequence_len
|
||||
data['split_lens'] = sequence_status['split_lens'] + [pad_len]
|
||||
data['attn_modes'] = sequence_status['attn_modes'] + ['causal']
|
||||
data['sample_lens'] += [pad_len]
|
||||
|
||||
# if the model has a convnet vae (e.g., as visual tokenizer)
|
||||
if len(sequence_status['vae_image_tensors']) > 0:
|
||||
image_tensors = sequence_status.pop('vae_image_tensors')
|
||||
image_sizes = [item.shape for item in image_tensors]
|
||||
max_image_size = [max(item) for item in list(zip(*image_sizes))]
|
||||
padded_images = torch.zeros(size=(len(image_tensors), *max_image_size))
|
||||
for i, image_tensor in enumerate(image_tensors):
|
||||
padded_images[i, :, :image_tensor.shape[1], :image_tensor.shape[2]] = image_tensor
|
||||
|
||||
data['padded_images'] = padded_images
|
||||
data['patchified_vae_latent_shapes'] = sequence_status['vae_latent_shapes']
|
||||
data['packed_latent_position_ids'] = torch.cat(sequence_status['packed_latent_position_ids'], dim=0)
|
||||
data['packed_vae_token_indexes'] = torch.tensor(sequence_status['packed_vae_token_indexes'])
|
||||
|
||||
# if the model has a vit (e.g., as visual tokenizer)
|
||||
if len(sequence_status['packed_vit_tokens']) > 0:
|
||||
data['packed_vit_tokens'] = torch.cat(sequence_status['packed_vit_tokens'], dim=0)
|
||||
data['packed_vit_position_ids'] = torch.cat(sequence_status['packed_vit_position_ids'], dim=0)
|
||||
data['packed_vit_token_indexes'] = torch.tensor(sequence_status['packed_vit_token_indexes'])
|
||||
data['vit_token_seqlens'] = torch.tensor(sequence_status['vit_token_seqlens'])
|
||||
|
||||
# if the model is required to perform visual generation
|
||||
if len(sequence_status['packed_timesteps']) > 0:
|
||||
data['packed_timesteps'] = torch.tensor(sequence_status['packed_timesteps'])
|
||||
data['mse_loss_indexes'] = torch.tensor(sequence_status['mse_loss_indexes'])
|
||||
|
||||
# if the model is required to perform text generation
|
||||
if len(sequence_status['packed_label_ids']) > 0:
|
||||
data['packed_label_ids'] = torch.tensor(sequence_status['packed_label_ids'])
|
||||
data['ce_loss_indexes'] = torch.tensor(sequence_status['ce_loss_indexes'])
|
||||
data['ce_loss_weights'] = torch.tensor(sequence_status['ce_loss_weights'])
|
||||
|
||||
return data
|
||||
|
||||
def __iter__(self):
|
||||
total_weights = sum(self.grouped_weights)
|
||||
assert total_weights > 0.0
|
||||
group_cumprobs = [sum(self.grouped_weights[:i + 1]) / total_weights
|
||||
for i in range(len(self.grouped_weights))]
|
||||
sequence_status = self.set_sequence_status()
|
||||
batch_data_indexes = []
|
||||
|
||||
buffer = []
|
||||
while True:
|
||||
# Ensure at least one sample from each group
|
||||
if sequence_status['curr'] == 0:
|
||||
for group_index, group_iter in enumerate(self.dataset_iters):
|
||||
if self.is_mandatory[group_index]:
|
||||
while True:
|
||||
sample = next(group_iter)
|
||||
# if a sample is too long, skip it
|
||||
num_tokens = sample['num_tokens'] + 2 * len(sample['sequence_plan'])
|
||||
if num_tokens < self.max_num_tokens_per_sample:
|
||||
sequence_status = self.pack_sequence(sample, sequence_status)
|
||||
batch_data_indexes.append(sample['data_indexes'])
|
||||
break
|
||||
else:
|
||||
print(f"skip a sample with length {num_tokens}")
|
||||
continue
|
||||
|
||||
if sequence_status['curr'] < self.prefer_buffer_before and len(buffer) > 0:
|
||||
sample = buffer.pop(0)
|
||||
sample_from_buffer = True
|
||||
else:
|
||||
# sample normally across all groups
|
||||
n = random.random()
|
||||
group_index = 0
|
||||
for i, cumprob in enumerate(group_cumprobs):
|
||||
if n < cumprob:
|
||||
group_index = i
|
||||
break
|
||||
sample = next(self.dataset_iters[group_index])
|
||||
sample_from_buffer = False
|
||||
|
||||
# if a sample is too long, skip it
|
||||
num_tokens = sample['num_tokens'] + 2 * len(sample['sequence_plan'])
|
||||
if num_tokens > self.max_num_tokens_per_sample:
|
||||
print(f"skip a sample with length {num_tokens}")
|
||||
continue
|
||||
|
||||
if sequence_status['curr'] + num_tokens > self.max_num_tokens:
|
||||
if len(buffer) < self.max_buffer_size and not sample_from_buffer:
|
||||
buffer.append(sample)
|
||||
else:
|
||||
print(f"Yielding data with length {sum(sequence_status['sample_lens'])}")
|
||||
data = self.to_tensor(sequence_status)
|
||||
data['batch_data_indexes'] = batch_data_indexes
|
||||
yield data
|
||||
sequence_status = self.set_sequence_status()
|
||||
batch_data_indexes = []
|
||||
continue
|
||||
|
||||
sequence_status = self.pack_sequence(sample, sequence_status)
|
||||
batch_data_indexes.append(sample['data_indexes'])
|
||||
|
||||
if sequence_status['curr'] >= self.expected_num_tokens:
|
||||
data = self.to_tensor(sequence_status)
|
||||
data['batch_data_indexes'] = batch_data_indexes
|
||||
yield data
|
||||
sequence_status = self.set_sequence_status()
|
||||
batch_data_indexes = []
|
||||
|
||||
def pack_sequence(self, sample, sequence_status):
|
||||
image_tensor_list = sample['image_tensor_list']
|
||||
text_ids_list = sample['text_ids_list']
|
||||
sequence_plan = sample['sequence_plan']
|
||||
|
||||
split_lens, attn_modes = list(), list()
|
||||
curr = sequence_status['curr']
|
||||
curr_rope_id = 0
|
||||
sample_lens = 0
|
||||
|
||||
for item in sequence_plan:
|
||||
split_start = item.get('split_start', True)
|
||||
if split_start:
|
||||
curr_split_len = 0
|
||||
|
||||
if item['type'] == 'text':
|
||||
text_ids = text_ids_list.pop(0)
|
||||
if item['enable_cfg'] == 1 and random.random() < self.data_config.text_cond_dropout_prob:
|
||||
continue
|
||||
|
||||
shifted_text_ids = [self.bos_token_id] + text_ids
|
||||
sequence_status['packed_text_ids'].extend(shifted_text_ids)
|
||||
sequence_status['packed_text_indexes'].extend(range(curr, curr + len(shifted_text_ids)))
|
||||
if item['loss'] == 1:
|
||||
sequence_status['ce_loss_indexes'].extend(range(curr, curr + len(shifted_text_ids)))
|
||||
sequence_status['ce_loss_weights'].extend(
|
||||
[len2weight(len(shifted_text_ids))] * len(shifted_text_ids)
|
||||
)
|
||||
sequence_status['packed_label_ids'].extend(text_ids + [self.eos_token_id])
|
||||
curr += len(shifted_text_ids)
|
||||
curr_split_len += len(shifted_text_ids)
|
||||
|
||||
# add a <|im_end|> token
|
||||
sequence_status['packed_text_ids'].append(self.eos_token_id)
|
||||
sequence_status['packed_text_indexes'].append(curr)
|
||||
if item['special_token_loss'] == 1: # <|im_end|> may have loss
|
||||
sequence_status['ce_loss_indexes'].append(curr)
|
||||
sequence_status['ce_loss_weights'].append(1.0)
|
||||
sequence_status['packed_label_ids'].append(item['special_token_label'])
|
||||
curr += 1
|
||||
curr_split_len += 1
|
||||
|
||||
# update sequence status
|
||||
attn_modes.append("causal")
|
||||
sequence_status['packed_position_ids'].extend(range(curr_rope_id, curr_rope_id + curr_split_len))
|
||||
curr_rope_id += curr_split_len
|
||||
|
||||
elif item['type'] == 'vit_image':
|
||||
image_tensor = image_tensor_list.pop(0)
|
||||
if item['enable_cfg'] == 1 and random.random() < self.data_config.vit_cond_dropout_prob:
|
||||
curr_rope_id += 1
|
||||
continue
|
||||
|
||||
# add a <|startofimage|> token
|
||||
sequence_status['packed_text_ids'].append(self.start_of_image)
|
||||
sequence_status['packed_text_indexes'].append(curr)
|
||||
curr += 1
|
||||
curr_split_len += 1
|
||||
|
||||
# preprocess image
|
||||
vit_tokens = patchify(image_tensor, self.data_config.vit_patch_size)
|
||||
num_img_tokens = vit_tokens.shape[0]
|
||||
sequence_status['packed_vit_token_indexes'].extend(range(curr, curr + num_img_tokens))
|
||||
curr += num_img_tokens
|
||||
curr_split_len += num_img_tokens
|
||||
|
||||
sequence_status['packed_vit_tokens'].append(vit_tokens)
|
||||
sequence_status['vit_token_seqlens'].append(num_img_tokens)
|
||||
sequence_status['packed_vit_position_ids'].append(
|
||||
self.get_flattened_position_ids(
|
||||
image_tensor.size(1), image_tensor.size(2),
|
||||
self.data_config.vit_patch_size,
|
||||
max_num_patches_per_side=self.data_config.max_num_patch_per_side
|
||||
)
|
||||
)
|
||||
|
||||
# add a <|endofimage|> token
|
||||
sequence_status['packed_text_ids'].append(self.end_of_image)
|
||||
sequence_status['packed_text_indexes'].append(curr)
|
||||
if item['special_token_loss'] == 1: # <|endofimage|> may have loss
|
||||
sequence_status['ce_loss_indexes'].append(curr)
|
||||
sequence_status['ce_loss_weights'].append(1.0)
|
||||
sequence_status['packed_label_ids'].append(item['special_token_label'])
|
||||
curr += 1
|
||||
curr_split_len += 1
|
||||
|
||||
# update sequence status
|
||||
attn_modes.append("full")
|
||||
sequence_status['packed_position_ids'].extend([curr_rope_id] * curr_split_len)
|
||||
curr_rope_id += 1
|
||||
|
||||
elif item['type'] == 'vae_image':
|
||||
image_tensor = image_tensor_list.pop(0)
|
||||
if item['enable_cfg'] == 1 and random.random() < self.data_config.vae_cond_dropout_prob:
|
||||
# FIXME fix vae dropout in video2video setting.
|
||||
curr_rope_id += 1
|
||||
continue
|
||||
|
||||
# add a <|startofimage|> token
|
||||
sequence_status['packed_text_ids'].append(self.start_of_image)
|
||||
sequence_status['packed_text_indexes'].append(curr)
|
||||
curr += 1
|
||||
curr_split_len += 1
|
||||
|
||||
# preprocess image
|
||||
sequence_status['vae_image_tensors'].append(image_tensor)
|
||||
sequence_status['packed_latent_position_ids'].append(
|
||||
self.get_flattened_position_ids(
|
||||
image_tensor.size(1), image_tensor.size(2),
|
||||
self.data_config.vae_image_downsample,
|
||||
max_num_patches_per_side=self.data_config.max_latent_size
|
||||
)
|
||||
)
|
||||
H, W = image_tensor.shape[1:]
|
||||
h = H // self.data_config.vae_image_downsample
|
||||
w = W // self.data_config.vae_image_downsample
|
||||
sequence_status['vae_latent_shapes'].append((h, w))
|
||||
|
||||
num_img_tokens = w * h
|
||||
sequence_status['packed_vae_token_indexes'].extend(range(curr, curr + num_img_tokens))
|
||||
if item['loss'] == 1:
|
||||
sequence_status['mse_loss_indexes'].extend(range(curr, curr + num_img_tokens))
|
||||
if split_start:
|
||||
timestep = np.random.randn()
|
||||
else:
|
||||
timestep = float('-inf')
|
||||
|
||||
sequence_status['packed_timesteps'].extend([timestep] * num_img_tokens)
|
||||
curr += num_img_tokens
|
||||
curr_split_len += num_img_tokens
|
||||
|
||||
# add a <|endofimage|> token
|
||||
sequence_status['packed_text_ids'].append(self.end_of_image)
|
||||
sequence_status['packed_text_indexes'].append(curr)
|
||||
# <|endofimage|> may have loss
|
||||
if item['special_token_loss'] == 1:
|
||||
sequence_status['ce_loss_indexes'].append(curr)
|
||||
sequence_status['ce_loss_weights'].append(1.0)
|
||||
sequence_status['packed_label_ids'].append(item['special_token_label'])
|
||||
curr += 1
|
||||
curr_split_len += 1
|
||||
|
||||
# update sequence status
|
||||
if split_start:
|
||||
if item['loss'] == 1 and 'frame_delta' not in item.keys():
|
||||
attn_modes.append("noise")
|
||||
else:
|
||||
attn_modes.append("full")
|
||||
sequence_status['packed_position_ids'].extend([curr_rope_id] * (num_img_tokens + 2))
|
||||
if 'frame_delta' in item.keys():
|
||||
curr_rope_id += item['frame_delta']
|
||||
elif item['loss'] == 0:
|
||||
curr_rope_id += 1
|
||||
|
||||
if item.get('split_end', True):
|
||||
split_lens.append(curr_split_len)
|
||||
sample_lens += curr_split_len
|
||||
|
||||
sequence_status['curr'] = curr
|
||||
sequence_status['sample_lens'].append(sample_lens)
|
||||
# prepare attention mask
|
||||
if not self.use_flex:
|
||||
sequence_status['nested_attention_masks'].append(
|
||||
prepare_attention_mask_per_sample(split_lens, attn_modes)
|
||||
)
|
||||
else:
|
||||
sequence_status['split_lens'].extend(split_lens)
|
||||
sequence_status['attn_modes'].extend(attn_modes)
|
||||
|
||||
return sequence_status
|
||||
|
||||
|
||||
class SimpleCustomBatch:
|
||||
def __init__(self, batch):
|
||||
data = batch[0]
|
||||
self.batch_data_indexes = data['batch_data_indexes']
|
||||
self.sequence_length = data["sequence_length"]
|
||||
self.sample_lens = data["sample_lens"]
|
||||
self.packed_text_ids = data["packed_text_ids"]
|
||||
self.packed_text_indexes = data["packed_text_indexes"]
|
||||
self.packed_position_ids = data["packed_position_ids"]
|
||||
|
||||
self.use_flex = "nested_attention_masks" not in data.keys()
|
||||
|
||||
if self.use_flex:
|
||||
self.split_lens = data["split_lens"]
|
||||
self.attn_modes = data["attn_modes"]
|
||||
else:
|
||||
self.nested_attention_masks = data["nested_attention_masks"]
|
||||
|
||||
if "padded_images" in data.keys():
|
||||
self.padded_images = data["padded_images"]
|
||||
self.patchified_vae_latent_shapes = data["patchified_vae_latent_shapes"]
|
||||
self.packed_latent_position_ids = data["packed_latent_position_ids"]
|
||||
self.packed_vae_token_indexes = data["packed_vae_token_indexes"]
|
||||
|
||||
if "packed_vit_tokens" in data.keys():
|
||||
self.packed_vit_tokens = data["packed_vit_tokens"]
|
||||
self.packed_vit_position_ids = data["packed_vit_position_ids"]
|
||||
self.packed_vit_token_indexes = data["packed_vit_token_indexes"]
|
||||
self.vit_token_seqlens = data["vit_token_seqlens"]
|
||||
|
||||
if "packed_timesteps" in data.keys():
|
||||
self.packed_timesteps = data["packed_timesteps"]
|
||||
self.mse_loss_indexes = data["mse_loss_indexes"]
|
||||
|
||||
if "packed_label_ids" in data.keys():
|
||||
self.packed_label_ids = data["packed_label_ids"]
|
||||
self.ce_loss_indexes = data["ce_loss_indexes"]
|
||||
self.ce_loss_weights = data["ce_loss_weights"]
|
||||
|
||||
def pin_memory(self):
|
||||
self.packed_text_ids = self.packed_text_ids.pin_memory()
|
||||
self.packed_text_indexes = self.packed_text_indexes.pin_memory()
|
||||
self.packed_position_ids = self.packed_position_ids.pin_memory()
|
||||
|
||||
if not self.use_flex:
|
||||
self.nested_attention_masks = [item.pin_memory() for item in self.nested_attention_masks]
|
||||
|
||||
if hasattr(self, 'padded_images'):
|
||||
self.padded_images = self.padded_images.pin_memory()
|
||||
self.packed_vae_token_indexes = self.packed_vae_token_indexes.pin_memory()
|
||||
self.packed_latent_position_ids = self.packed_latent_position_ids.pin_memory()
|
||||
|
||||
if hasattr(self, 'packed_timesteps'):
|
||||
self.packed_timesteps = self.packed_timesteps.pin_memory()
|
||||
self.mse_loss_indexes = self.mse_loss_indexes.pin_memory()
|
||||
|
||||
if hasattr(self, 'packed_vit_tokens'):
|
||||
self.packed_vit_tokens = self.packed_vit_tokens.pin_memory()
|
||||
self.packed_vit_position_ids = self.packed_vit_position_ids.pin_memory()
|
||||
self.packed_vit_token_indexes = self.packed_vit_token_indexes.pin_memory()
|
||||
self.vit_token_seqlens = self.vit_token_seqlens.pin_memory()
|
||||
|
||||
if hasattr(self, 'packed_label_ids'):
|
||||
self.packed_label_ids = self.packed_label_ids.pin_memory()
|
||||
self.ce_loss_indexes = self.ce_loss_indexes.pin_memory()
|
||||
self.ce_loss_weights = self.ce_loss_weights.pin_memory()
|
||||
|
||||
return self
|
||||
|
||||
def cuda(self, device):
|
||||
self.packed_text_ids = self.packed_text_ids.to(device)
|
||||
self.packed_text_indexes = self.packed_text_indexes.to(device)
|
||||
self.packed_position_ids = self.packed_position_ids.to(device)
|
||||
|
||||
if not self.use_flex:
|
||||
self.nested_attention_masks = [item.to(device) for item in self.nested_attention_masks]
|
||||
|
||||
if hasattr(self, 'padded_images'):
|
||||
self.padded_images = self.padded_images.to(device)
|
||||
self.packed_vae_token_indexes = self.packed_vae_token_indexes.to(device)
|
||||
self.packed_latent_position_ids = self.packed_latent_position_ids.to(device)
|
||||
|
||||
if hasattr(self, 'packed_timesteps'):
|
||||
self.packed_timesteps = self.packed_timesteps.to(device)
|
||||
self.mse_loss_indexes = self.mse_loss_indexes.to(device)
|
||||
|
||||
if hasattr(self, 'packed_vit_tokens'):
|
||||
self.packed_vit_tokens = self.packed_vit_tokens.to(device)
|
||||
self.packed_vit_position_ids = self.packed_vit_position_ids.to(device)
|
||||
self.packed_vit_token_indexes = self.packed_vit_token_indexes.to(device)
|
||||
self.vit_token_seqlens = self.vit_token_seqlens.to(device)
|
||||
|
||||
if hasattr(self, 'packed_label_ids'):
|
||||
self.packed_label_ids = self.packed_label_ids.to(device)
|
||||
self.ce_loss_indexes = self.ce_loss_indexes.to(device)
|
||||
self.ce_loss_weights = self.ce_loss_weights.to(device)
|
||||
|
||||
return self
|
||||
|
||||
def to_dict(self):
|
||||
data = dict(
|
||||
sequence_length = self.sequence_length,
|
||||
sample_lens = self.sample_lens,
|
||||
packed_text_ids = self.packed_text_ids,
|
||||
packed_text_indexes = self.packed_text_indexes,
|
||||
packed_position_ids = self.packed_position_ids,
|
||||
batch_data_indexes = self.batch_data_indexes,
|
||||
)
|
||||
|
||||
if not self.use_flex:
|
||||
data['nested_attention_masks'] = self.nested_attention_masks
|
||||
else:
|
||||
data['split_lens'] = self.split_lens
|
||||
data['attn_modes'] = self.attn_modes
|
||||
|
||||
if hasattr(self, 'padded_images'):
|
||||
data['padded_images'] = self.padded_images
|
||||
data['patchified_vae_latent_shapes'] = self.patchified_vae_latent_shapes
|
||||
data['packed_latent_position_ids'] = self.packed_latent_position_ids
|
||||
data['packed_vae_token_indexes'] = self.packed_vae_token_indexes
|
||||
|
||||
if hasattr(self, 'packed_vit_tokens'):
|
||||
data['packed_vit_tokens'] = self.packed_vit_tokens
|
||||
data['packed_vit_position_ids'] = self.packed_vit_position_ids
|
||||
data['packed_vit_token_indexes'] = self.packed_vit_token_indexes
|
||||
data['vit_token_seqlens'] = self.vit_token_seqlens
|
||||
|
||||
if hasattr(self, 'packed_timesteps'):
|
||||
data['packed_timesteps'] = self.packed_timesteps
|
||||
data['mse_loss_indexes'] = self.mse_loss_indexes
|
||||
|
||||
if hasattr(self, 'packed_label_ids'):
|
||||
data['packed_label_ids'] = self.packed_label_ids
|
||||
data['ce_loss_indexes'] = self.ce_loss_indexes
|
||||
data['ce_loss_weights'] = self.ce_loss_weights
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def collate_wrapper():
|
||||
def collate_fn(batch):
|
||||
return SimpleCustomBatch(batch)
|
||||
return collate_fn
|
||||
@@ -0,0 +1,39 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from .interleave_datasets import UnifiedEditIterableDataset
|
||||
from .t2i_dataset import T2IIterableDataset
|
||||
from .vlm_dataset import SftJSONLIterableDataset
|
||||
|
||||
|
||||
DATASET_REGISTRY = {
|
||||
't2i_pretrain': T2IIterableDataset,
|
||||
'vlm_sft': SftJSONLIterableDataset,
|
||||
'unified_edit': UnifiedEditIterableDataset,
|
||||
}
|
||||
|
||||
|
||||
DATASET_INFO = {
|
||||
't2i_pretrain': {
|
||||
't2i': {
|
||||
'data_dir': 'your_data_path/bagel_example/t2i', # path of the parquet files
|
||||
'num_files': 10, # number of data units to be sharded across all ranks and workers
|
||||
'num_total_samples': 1000, # number of total samples in the dataset
|
||||
},
|
||||
},
|
||||
'unified_edit':{
|
||||
'seedxedit_multi': {
|
||||
'data_dir': 'your_data_path/bagel_example/editing/seedxedit_multi',
|
||||
'num_files': 10,
|
||||
'num_total_samples': 1000,
|
||||
"parquet_info_path": 'your_data_path/bagel_example/editing/parquet_info/seedxedit_multi_nas.json', # information of the parquet files
|
||||
},
|
||||
},
|
||||
'vlm_sft': {
|
||||
'llava_ov': {
|
||||
'data_dir': 'your_data_path/bagel_example/vlm/images',
|
||||
'jsonl_path': 'your_data_path/bagel_example/vlm/llava_ov_si.jsonl',
|
||||
'num_total_samples': 1000
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import random
|
||||
import torch
|
||||
|
||||
|
||||
class DistributedIterableDataset(torch.utils.data.IterableDataset):
|
||||
def __init__(self, dataset_name, local_rank=0, world_size=1, num_workers=8):
|
||||
self.dataset_name = dataset_name
|
||||
self.local_rank = local_rank
|
||||
self.world_size = world_size
|
||||
self.num_workers = num_workers
|
||||
self.rng = random.Random()
|
||||
self.data_paths = None
|
||||
|
||||
def get_data_paths(self, *args, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
def set_epoch(self, seed=42):
|
||||
if self.data_paths is None:
|
||||
return
|
||||
|
||||
if isinstance(self.data_paths[0], tuple):
|
||||
data_paths = sorted(self.data_paths, key=lambda x: (x[0], x[1]))
|
||||
elif isinstance(self.data_paths[0], str):
|
||||
data_paths = sorted(self.data_paths)
|
||||
else:
|
||||
raise ValueError(f"Unknown data_paths type: {type(self.data_paths[0])}")
|
||||
|
||||
self.rng.seed(seed)
|
||||
self.rng.shuffle(data_paths)
|
||||
|
||||
num_files_per_rank = len(data_paths) // self.world_size
|
||||
local_start = self.local_rank * num_files_per_rank
|
||||
local_end = (self.local_rank + 1) * num_files_per_rank
|
||||
self.num_files_per_rank = num_files_per_rank
|
||||
self.data_paths_per_rank = data_paths[local_start:local_end]
|
||||
|
||||
def get_data_paths_per_worker(self):
|
||||
if self.data_paths is None:
|
||||
return None
|
||||
|
||||
info = torch.utils.data.get_worker_info()
|
||||
if info is None:
|
||||
# Single worker: Use all files assigned to the rank
|
||||
return self.data_paths_per_rank, 0
|
||||
|
||||
worker_id = info.id
|
||||
num_files_per_worker = self.num_files_per_rank // info.num_workers
|
||||
start = num_files_per_worker * worker_id
|
||||
end = num_files_per_worker * (worker_id + 1)
|
||||
data_paths_per_worker = self.data_paths_per_rank[start:end]
|
||||
|
||||
return data_paths_per_worker[::-1], worker_id
|
||||
|
||||
def __iter__(self):
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from .edit_dataset import UnifiedEditIterableDataset
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import io
|
||||
import random
|
||||
from PIL import Image, ImageFile, PngImagePlugin
|
||||
|
||||
from .interleave_t2i_dataset import InterleavedBaseIterableDataset, ParquetStandardIterableDataset
|
||||
from ..data_utils import pil_img2rgb
|
||||
|
||||
|
||||
Image.MAX_IMAGE_PIXELS = 200000000
|
||||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||||
MaximumDecompressedSize = 1024
|
||||
MegaByte = 2 ** 20
|
||||
PngImagePlugin.MAX_TEXT_CHUNK = MaximumDecompressedSize * MegaByte
|
||||
|
||||
|
||||
class UnifiedEditIterableDataset(InterleavedBaseIterableDataset, ParquetStandardIterableDataset):
|
||||
|
||||
def parse_row(self, row):
|
||||
image_num = len(row["image_list"])
|
||||
# randomly choose start and end, return [0, 1] when only two images
|
||||
start_idx = random.choice(range(image_num - 1))
|
||||
max_end = min(start_idx + 3, image_num)
|
||||
end_idx = random.choice(range(start_idx + 1, max_end))
|
||||
|
||||
data = self._init_data()
|
||||
data = self._add_image(
|
||||
data,
|
||||
pil_img2rgb(Image.open(io.BytesIO(row["image_list"][start_idx]))),
|
||||
need_loss=False,
|
||||
need_vae=True,
|
||||
need_vit=True,
|
||||
)
|
||||
|
||||
if end_idx - start_idx > 1 and random.random() < 0.5: # concat multiple insturction
|
||||
if end_idx == image_num - 1:
|
||||
end_idx -= 1
|
||||
|
||||
instruction = ""
|
||||
for idx in range(start_idx + 1, end_idx + 1):
|
||||
instruction += random.choice(row["instruction_list"][idx-1]) + ". "
|
||||
data = self._add_text(data, instruction.rstrip(), need_loss=False)
|
||||
data = self._add_image(
|
||||
data,
|
||||
pil_img2rgb(Image.open(io.BytesIO(row["image_list"][end_idx]))),
|
||||
need_loss=True,
|
||||
need_vae=False,
|
||||
need_vit=False,
|
||||
)
|
||||
else:
|
||||
for idx in range(start_idx + 1, end_idx + 1):
|
||||
instruction = random.choice(row["instruction_list"][idx-1])
|
||||
data = self._add_text(data, instruction, need_loss=False)
|
||||
if idx != end_idx:
|
||||
data = self._add_image(
|
||||
data,
|
||||
pil_img2rgb(Image.open(io.BytesIO(row["image_list"][idx]))),
|
||||
need_loss=True,
|
||||
need_vae=True,
|
||||
need_vit=True,
|
||||
)
|
||||
else:
|
||||
data = self._add_image(
|
||||
data,
|
||||
pil_img2rgb(Image.open(io.BytesIO(row["image_list"][idx]))),
|
||||
need_loss=True,
|
||||
need_vae=False,
|
||||
need_vit=False,
|
||||
)
|
||||
return data
|
||||
@@ -0,0 +1,212 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pyarrow.parquet as pq
|
||||
|
||||
from ..distributed_iterable_dataset import DistributedIterableDataset
|
||||
from ..parquet_utils import get_parquet_data_paths, init_arrow_pf_fs
|
||||
|
||||
|
||||
class InterleavedBaseIterableDataset(DistributedIterableDataset):
|
||||
|
||||
def _init_data(self):
|
||||
data = {
|
||||
'sequence_plan': [],
|
||||
'text_ids_list': [],
|
||||
'image_tensor_list': [],
|
||||
'num_tokens': 0,
|
||||
}
|
||||
return data
|
||||
|
||||
def _add_text(self, data, text, need_loss, enable_cfg=True):
|
||||
text_ids = self.tokenizer.encode(text)
|
||||
data['num_tokens'] += len(text_ids)
|
||||
data['text_ids_list'].append(text_ids)
|
||||
data['sequence_plan'].append(
|
||||
{
|
||||
'type': 'text',
|
||||
'enable_cfg': int(enable_cfg),
|
||||
'loss': int(need_loss),
|
||||
'special_token_loss': 0,
|
||||
'special_token_label': None,
|
||||
}
|
||||
)
|
||||
return data
|
||||
|
||||
def _add_image(self, data, image, need_loss, need_vae, need_vit, enable_cfg=True):
|
||||
assert need_loss or need_vae or need_vit
|
||||
|
||||
if need_loss:
|
||||
data['sequence_plan'].append(
|
||||
{
|
||||
'type': 'vae_image',
|
||||
'enable_cfg': 0,
|
||||
'loss': 1,
|
||||
'special_token_loss': 0,
|
||||
'special_token_label': None,
|
||||
}
|
||||
)
|
||||
|
||||
image_tensor = self.transform(image)
|
||||
height, width = image_tensor.shape[1:]
|
||||
data['num_tokens'] += width * height // self.transform.stride ** 2
|
||||
data['image_tensor_list'].append(image_tensor)
|
||||
|
||||
if need_vae:
|
||||
data['sequence_plan'].append(
|
||||
{
|
||||
'type': 'vae_image',
|
||||
'enable_cfg': int(enable_cfg),
|
||||
'loss': 0,
|
||||
'special_token_loss': 0,
|
||||
'special_token_label': None,
|
||||
}
|
||||
)
|
||||
|
||||
image_tensor = self.transform(image)
|
||||
height, width = image_tensor.shape[1:]
|
||||
data['num_tokens'] += width * height // self.transform.stride ** 2
|
||||
data['image_tensor_list'].append(image_tensor.clone())
|
||||
|
||||
if need_vit:
|
||||
data['sequence_plan'].append(
|
||||
{
|
||||
'type': 'vit_image',
|
||||
'enable_cfg': int(enable_cfg),
|
||||
'loss': 0,
|
||||
'special_token_loss': 0,
|
||||
'special_token_label': None,
|
||||
},
|
||||
)
|
||||
vit_image_tensor = self.vit_transform(image)
|
||||
height, width = vit_image_tensor.shape[1:]
|
||||
data['num_tokens'] += width * height // self.vit_transform.stride ** 2
|
||||
data['image_tensor_list'].append(vit_image_tensor)
|
||||
|
||||
return data
|
||||
|
||||
def _add_video(self, data, frames, frame_indexes, need_loss, need_vae, enable_cfg=True):
|
||||
assert int(need_loss) + int(need_vae) == 1
|
||||
|
||||
if need_loss:
|
||||
for idx, (image, frame_idx) in enumerate(zip(frames, frame_indexes)):
|
||||
current_sequence_plan = {
|
||||
'type': 'vae_image',
|
||||
'enable_cfg': 0,
|
||||
'loss': 1,
|
||||
'special_token_loss': 0,
|
||||
'special_token_label': None,
|
||||
'split_start': idx == 0,
|
||||
'split_end': idx == len(frames) - 1,
|
||||
}
|
||||
if idx < len(frame_indexes) - 1:
|
||||
current_sequence_plan['frame_delta'] = frame_indexes[idx + 1] - frame_idx
|
||||
data['sequence_plan'].append(current_sequence_plan)
|
||||
image_tensor = self.transform(image)
|
||||
height, width = image_tensor.shape[1:]
|
||||
data['image_tensor_list'].append(image_tensor)
|
||||
data['num_tokens'] += width * height // self.transform.stride ** 2
|
||||
|
||||
elif need_vae:
|
||||
for idx, (image, frame_idx) in enumerate(zip(frames, frame_indexes)):
|
||||
current_sequence_plan = {
|
||||
'type': 'vae_image',
|
||||
'enable_cfg': int(enable_cfg),
|
||||
'loss': 0,
|
||||
'special_token_loss': 0,
|
||||
'special_token_label': None,
|
||||
'split_start': idx == 0,
|
||||
'split_end': idx == len(frames) - 1,
|
||||
}
|
||||
if idx < len(frame_indexes) - 1:
|
||||
current_sequence_plan['frame_delta'] = frame_indexes[idx + 1] - frame_idx
|
||||
data['sequence_plan'].append(current_sequence_plan)
|
||||
image_tensor = self.transform(image)
|
||||
height, width = image_tensor.shape[1:]
|
||||
data['image_tensor_list'].append(image_tensor)
|
||||
data['num_tokens'] += width * height // self.transform.stride ** 2
|
||||
|
||||
return data
|
||||
|
||||
|
||||
class ParquetStandardIterableDataset(DistributedIterableDataset):
|
||||
|
||||
def __init__(
|
||||
self, dataset_name, transform, tokenizer, vit_transform,
|
||||
data_dir_list, num_used_data, parquet_info,
|
||||
local_rank=0, world_size=1, num_workers=8, data_status=None,
|
||||
):
|
||||
"""
|
||||
data_dir_list: list of data directories contains parquet files
|
||||
num_used_data: list of number of sampled data paths for each data directory
|
||||
vit_transform: input transform for vit model.
|
||||
"""
|
||||
super().__init__(dataset_name, local_rank, world_size, num_workers)
|
||||
self.transform = transform
|
||||
self.vit_transform = vit_transform
|
||||
self.tokenizer = tokenizer
|
||||
self.data_status = data_status
|
||||
self.data_paths = self.get_data_paths(data_dir_list, num_used_data, parquet_info)
|
||||
self.set_epoch()
|
||||
|
||||
def get_data_paths(self, data_dir_list, num_used_data, parquet_info):
|
||||
row_groups = []
|
||||
for data_dir, num_data_path in zip(data_dir_list, num_used_data):
|
||||
data_paths = get_parquet_data_paths([data_dir], [num_data_path])
|
||||
for data_path in data_paths:
|
||||
if data_path in parquet_info.keys():
|
||||
num_row_groups = parquet_info[data_path]['num_row_groups']
|
||||
for rg_idx in range(num_row_groups):
|
||||
row_groups.append((data_path, rg_idx))
|
||||
return row_groups
|
||||
|
||||
def parse_row(self, row):
|
||||
raise NotImplementedError
|
||||
|
||||
def __iter__(self):
|
||||
file_paths_per_worker, worker_id = self.get_data_paths_per_worker()
|
||||
if self.data_status is not None:
|
||||
global_row_group_start_id = self.data_status[worker_id][0]
|
||||
row_start_id = self.data_status[worker_id][1] + 1
|
||||
else:
|
||||
global_row_group_start_id = 0
|
||||
row_start_id = 0
|
||||
|
||||
print(
|
||||
f"rank-{self.local_rank} worker-{worker_id} dataset-{self.dataset_name}: "
|
||||
f"resuming data at global_rg#{global_row_group_start_id}, row#{row_start_id}"
|
||||
)
|
||||
|
||||
while True:
|
||||
file_paths_per_worker_ = file_paths_per_worker[global_row_group_start_id:]
|
||||
for global_row_group_idx, (parquet_file_path, row_group_id) in enumerate(
|
||||
file_paths_per_worker_, start=global_row_group_start_id
|
||||
):
|
||||
fs = init_arrow_pf_fs(parquet_file_path)
|
||||
with fs.open_input_file(parquet_file_path) as f:
|
||||
try:
|
||||
fr = pq.ParquetFile(f)
|
||||
df = fr.read_row_group(row_group_id).to_pandas()
|
||||
df = df.iloc[row_start_id:]
|
||||
except Exception as e:
|
||||
print(f'Error {e} in rg#{row_group_id}, {parquet_file_path}')
|
||||
continue
|
||||
|
||||
for row_idx, row in df.iterrows():
|
||||
try:
|
||||
data = self.parse_row(row)
|
||||
if len(data) == 0:
|
||||
continue
|
||||
data['data_indexes'] = {
|
||||
"data_indexes": [global_row_group_idx, row_idx],
|
||||
"worker_id": worker_id,
|
||||
"dataset_name": self.dataset_name,
|
||||
}
|
||||
except Exception as e:
|
||||
print(f'Error {e} in rg#{row_group_id}, {parquet_file_path}')
|
||||
continue
|
||||
yield data
|
||||
|
||||
row_start_id = 0
|
||||
global_row_group_start_id = 0
|
||||
print(f"{self.dataset_name} repeat in rank-{self.local_rank} worker-{worker_id}")
|
||||
@@ -0,0 +1,89 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import logging
|
||||
|
||||
import pyarrow.fs as pf
|
||||
import torch.distributed as dist
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_parquet_data_paths(data_dir_list, num_sampled_data_paths, rank=0, world_size=1):
|
||||
num_data_dirs = len(data_dir_list)
|
||||
if world_size > 1:
|
||||
chunk_size = (num_data_dirs + world_size - 1) // world_size
|
||||
start_idx = rank * chunk_size
|
||||
end_idx = min(start_idx + chunk_size, num_data_dirs)
|
||||
local_data_dir_list = data_dir_list[start_idx:end_idx]
|
||||
local_num_sampled_data_paths = num_sampled_data_paths[start_idx:end_idx]
|
||||
else:
|
||||
local_data_dir_list = data_dir_list
|
||||
local_num_sampled_data_paths = num_sampled_data_paths
|
||||
|
||||
local_data_paths = []
|
||||
for data_dir, num_data_path in zip(local_data_dir_list, local_num_sampled_data_paths):
|
||||
if data_dir.startswith("hdfs://"):
|
||||
files = hdfs_ls_cmd(data_dir)
|
||||
data_paths_per_dir = [
|
||||
file for file in files if file.endswith(".parquet")
|
||||
]
|
||||
else:
|
||||
files = os.listdir(data_dir)
|
||||
data_paths_per_dir = [
|
||||
os.path.join(data_dir, name)
|
||||
for name in files
|
||||
if name.endswith(".parquet")
|
||||
]
|
||||
repeat = num_data_path // len(data_paths_per_dir)
|
||||
data_paths_per_dir = data_paths_per_dir * (repeat + 1)
|
||||
local_data_paths.extend(data_paths_per_dir[:num_data_path])
|
||||
|
||||
if world_size > 1:
|
||||
gather_list = [None] * world_size
|
||||
dist.all_gather_object(gather_list, local_data_paths)
|
||||
|
||||
combined_chunks = []
|
||||
for chunk_list in gather_list:
|
||||
if chunk_list is not None:
|
||||
combined_chunks.extend(chunk_list)
|
||||
else:
|
||||
combined_chunks = local_data_paths
|
||||
|
||||
return combined_chunks
|
||||
|
||||
|
||||
# NOTE: cumtomize this function for your cluster
|
||||
def get_hdfs_host():
|
||||
return "hdfs://xxx"
|
||||
|
||||
|
||||
# NOTE: cumtomize this function for your cluster
|
||||
def get_hdfs_block_size():
|
||||
return 134217728
|
||||
|
||||
|
||||
# NOTE: cumtomize this function for your cluster
|
||||
def get_hdfs_extra_conf():
|
||||
return None
|
||||
|
||||
|
||||
def init_arrow_pf_fs(parquet_file_path):
|
||||
if parquet_file_path.startswith("hdfs://"):
|
||||
fs = pf.HadoopFileSystem(
|
||||
host=get_hdfs_host(),
|
||||
port=0,
|
||||
buffer_size=get_hdfs_block_size(),
|
||||
extra_conf=get_hdfs_extra_conf(),
|
||||
)
|
||||
else:
|
||||
fs = pf.LocalFileSystem()
|
||||
return fs
|
||||
|
||||
|
||||
def hdfs_ls_cmd(dir):
|
||||
result = subprocess.run(["hdfs", "dfs", "ls", dir], capture_output=True, text=True).stdout
|
||||
return ['hdfs://' + i.split('hdfs://')[-1].strip() for i in result.split('\n') if 'hdfs://' in i]
|
||||
@@ -0,0 +1,128 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import io
|
||||
import json
|
||||
import pyarrow.parquet as pq
|
||||
import random
|
||||
from PIL import Image
|
||||
|
||||
from .data_utils import pil_img2rgb
|
||||
from .distributed_iterable_dataset import DistributedIterableDataset
|
||||
from .parquet_utils import get_parquet_data_paths, init_arrow_pf_fs
|
||||
|
||||
Image.MAX_IMAGE_PIXELS = 20_000_000
|
||||
|
||||
|
||||
class T2IIterableDataset(DistributedIterableDataset):
|
||||
def __init__(
|
||||
self, dataset_name, transform, tokenizer, data_dir_list, num_used_data,
|
||||
local_rank=0, world_size=1, num_workers=8, data_status=None,
|
||||
):
|
||||
"""
|
||||
data_dir_list: list of data directories contains parquet files
|
||||
num_used_data: list of number of sampled data paths for each data directory
|
||||
"""
|
||||
super().__init__(dataset_name, local_rank, world_size, num_workers)
|
||||
self.transform = transform
|
||||
self.tokenizer = tokenizer
|
||||
self.data_status = data_status
|
||||
self.data_paths = self.get_data_paths(data_dir_list, num_used_data)
|
||||
self.set_epoch()
|
||||
|
||||
def get_data_paths(self, data_dir_list, num_used_data):
|
||||
return get_parquet_data_paths(data_dir_list, num_used_data)
|
||||
|
||||
def __iter__(self):
|
||||
data_paths_per_worker, worker_id = self.get_data_paths_per_worker()
|
||||
if self.data_status is not None:
|
||||
parquet_start_id = self.data_status[worker_id][0]
|
||||
row_group_start_id = self.data_status[worker_id][1]
|
||||
row_start_id = self.data_status[worker_id][2] + 1
|
||||
else:
|
||||
parquet_start_id = 0
|
||||
row_group_start_id = 0
|
||||
row_start_id = 0
|
||||
transform_stride = self.transform.stride
|
||||
|
||||
print(
|
||||
f"rank-{self.local_rank} worker-{worker_id} dataset-{self.dataset_name}: "
|
||||
f"resuming data at parquet#{parquet_start_id}, rg#{row_group_start_id}, row#{row_start_id}"
|
||||
)
|
||||
|
||||
while True:
|
||||
data_paths_per_worker_ = data_paths_per_worker[parquet_start_id:]
|
||||
for parquet_idx, parquet_file_path in enumerate(data_paths_per_worker_, start=parquet_start_id):
|
||||
fs = init_arrow_pf_fs(parquet_file_path)
|
||||
with fs.open_input_file(parquet_file_path) as f:
|
||||
fr = pq.ParquetFile(f)
|
||||
row_group_ids = list(range(fr.num_row_groups))
|
||||
row_group_ids_ = row_group_ids[row_group_start_id:]
|
||||
|
||||
for row_group_id in row_group_ids_:
|
||||
df = fr.read_row_group(row_group_id).to_pandas()
|
||||
df = df.iloc[row_start_id:]
|
||||
|
||||
for row_idx, row in df.iterrows():
|
||||
num_tokens = 0
|
||||
try:
|
||||
image_byte = row['image']
|
||||
image = pil_img2rgb(Image.open(io.BytesIO(image_byte)))
|
||||
except Exception as e:
|
||||
print(f'Error: {e} in rg#{row_group_id}, {parquet_file_path}')
|
||||
continue
|
||||
image_tensor = self.transform(image)
|
||||
height, width = image_tensor.shape[1:]
|
||||
num_tokens += width * height // transform_stride ** 2
|
||||
|
||||
try:
|
||||
caption_dict = row['captions']
|
||||
caption_dict = json.loads(caption_dict)
|
||||
except Exception as e:
|
||||
print(f'Error: {e} in rg#{row_group_id}, {parquet_file_path}')
|
||||
continue
|
||||
|
||||
caps_token = [self.tokenizer.encode(v) for _, v in caption_dict.items()]
|
||||
if len(caps_token) == 0:
|
||||
print(f'no caption in rg#{row_group_id}, {parquet_file_path}')
|
||||
caption_token = self.tokenizer.encode(' ')
|
||||
else:
|
||||
caption_token = random.choice(caps_token)
|
||||
|
||||
sequence_plan, text_ids_list = [], []
|
||||
text_ids = caption_token
|
||||
num_tokens += len(caption_token)
|
||||
text_ids_list.append(text_ids)
|
||||
sequence_plan.append({
|
||||
'type': 'text',
|
||||
'enable_cfg': 1,
|
||||
'loss': 0,
|
||||
'special_token_loss': 0,
|
||||
'special_token_label': None,
|
||||
})
|
||||
|
||||
sequence_plan.append({
|
||||
'type': 'vae_image',
|
||||
'enable_cfg': 0,
|
||||
'loss': 1,
|
||||
'special_token_loss': 0,
|
||||
'special_token_label': None,
|
||||
})
|
||||
|
||||
sample = dict(
|
||||
image_tensor_list=[image_tensor],
|
||||
text_ids_list=text_ids_list,
|
||||
num_tokens=num_tokens,
|
||||
sequence_plan=sequence_plan,
|
||||
data_indexes={
|
||||
"data_indexes": [parquet_idx, row_group_id, row_idx],
|
||||
"worker_id": worker_id,
|
||||
"dataset_name": self.dataset_name,
|
||||
}
|
||||
)
|
||||
yield sample
|
||||
|
||||
row_start_id = 0
|
||||
row_group_start_id = 0
|
||||
parquet_start_id = 0
|
||||
print(f"{self.dataset_name} repeat in rank-{self.local_rank} worker-{worker_id}")
|
||||
@@ -0,0 +1,287 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import random
|
||||
from PIL import Image
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
from torchvision.transforms import functional as F
|
||||
from torchvision.transforms import InterpolationMode
|
||||
|
||||
|
||||
class MaxLongEdgeMinShortEdgeResize(torch.nn.Module):
|
||||
"""Resize the input image so that its longest side and shortest side are within a specified range,
|
||||
ensuring that both sides are divisible by a specified stride.
|
||||
|
||||
Args:
|
||||
max_size (int): Maximum size for the longest edge of the image.
|
||||
min_size (int): Minimum size for the shortest edge of the image.
|
||||
stride (int): Value by which the height and width of the image must be divisible.
|
||||
max_pixels (int): Maximum pixels for the full image.
|
||||
interpolation (InterpolationMode): Desired interpolation enum defined by
|
||||
:class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``.
|
||||
If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.NEAREST_EXACT``,
|
||||
``InterpolationMode.BILINEAR``, and ``InterpolationMode.BICUBIC`` are supported.
|
||||
The corresponding Pillow integer constants, e.g., ``PIL.Image.BILINEAR`` are also accepted.
|
||||
antialias (bool, optional): Whether to apply antialiasing (default is True).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_size: int,
|
||||
min_size: int,
|
||||
stride: int,
|
||||
max_pixels: int,
|
||||
interpolation=InterpolationMode.BICUBIC,
|
||||
antialias=True
|
||||
):
|
||||
super().__init__()
|
||||
self.max_size = max_size
|
||||
self.min_size = min_size
|
||||
self.stride = stride
|
||||
self.max_pixels = max_pixels
|
||||
self.interpolation = interpolation
|
||||
self.antialias = antialias
|
||||
|
||||
def _make_divisible(self, value, stride):
|
||||
"""Ensure the value is divisible by the stride."""
|
||||
return max(stride, int(round(value / stride) * stride))
|
||||
|
||||
def _apply_scale(self, width, height, scale):
|
||||
new_width = round(width * scale)
|
||||
new_height = round(height * scale)
|
||||
new_width = self._make_divisible(new_width, self.stride)
|
||||
new_height = self._make_divisible(new_height, self.stride)
|
||||
return new_width, new_height
|
||||
|
||||
def forward(self, img, img_num=1):
|
||||
"""
|
||||
Args:
|
||||
img (PIL Image): Image to be resized.
|
||||
img_num (int): Number of images, used to change max_tokens.
|
||||
Returns:
|
||||
PIL Image or Tensor: Rescaled image with divisible dimensions.
|
||||
"""
|
||||
if isinstance(img, torch.Tensor):
|
||||
height, width = img.shape[-2:]
|
||||
else:
|
||||
width, height = img.size
|
||||
|
||||
scale = min(self.max_size / max(width, height), 1.0)
|
||||
scale = max(scale, self.min_size / min(width, height))
|
||||
new_width, new_height = self._apply_scale(width, height, scale)
|
||||
|
||||
# Ensure the number of pixels does not exceed max_pixels
|
||||
if new_width * new_height > self.max_pixels / img_num:
|
||||
scale = self.max_pixels / img_num / (new_width * new_height)
|
||||
new_width, new_height = self._apply_scale(new_width, new_height, scale)
|
||||
|
||||
# Ensure longest edge does not exceed max_size
|
||||
if max(new_width, new_height) > self.max_size:
|
||||
scale = self.max_size / max(new_width, new_height)
|
||||
new_width, new_height = self._apply_scale(new_width, new_height, scale)
|
||||
|
||||
return F.resize(img, (new_height, new_width), self.interpolation, antialias=self.antialias)
|
||||
|
||||
|
||||
class ImageTransform:
|
||||
def __init__(
|
||||
self,
|
||||
max_image_size,
|
||||
min_image_size,
|
||||
image_stride,
|
||||
max_pixels=14*14*9*1024,
|
||||
image_mean=[0.5, 0.5, 0.5],
|
||||
image_std=[0.5, 0.5, 0.5]
|
||||
):
|
||||
self.stride = image_stride
|
||||
|
||||
self.resize_transform = MaxLongEdgeMinShortEdgeResize(
|
||||
max_size=max_image_size,
|
||||
min_size=min_image_size,
|
||||
stride=image_stride,
|
||||
max_pixels=max_pixels,
|
||||
)
|
||||
self.to_tensor_transform = transforms.ToTensor()
|
||||
self.normalize_transform = transforms.Normalize(mean=image_mean, std=image_std, inplace=True)
|
||||
|
||||
def __call__(self, img, img_num=1):
|
||||
img = self.resize_transform(img, img_num=img_num)
|
||||
img = self.to_tensor_transform(img)
|
||||
img = self.normalize_transform(img)
|
||||
return img
|
||||
|
||||
|
||||
def decolorization(image):
|
||||
gray_image = image.convert('L')
|
||||
return Image.merge(image.mode, [gray_image] * 3) if image.mode in ('RGB', 'L') else gray_image
|
||||
|
||||
|
||||
def downscale(image, scale_factor):
|
||||
new_width = int(round(image.width * scale_factor))
|
||||
new_height = int(round(image.height * scale_factor))
|
||||
new_width = max(1, new_width)
|
||||
new_height = max(1, new_height)
|
||||
return image.resize((new_width, new_height), resample=Image.BICUBIC)
|
||||
|
||||
|
||||
def crop(image, crop_factors):
|
||||
target_h, target_w = crop_factors
|
||||
img_w, img_h = image.size
|
||||
|
||||
if target_h > img_h or target_w > img_w:
|
||||
raise ValueError("Crop size exceeds image dimensions")
|
||||
|
||||
x = random.randint(0, img_w - target_w)
|
||||
y = random.randint(0, img_h - target_h)
|
||||
|
||||
return image.crop((x, y, x + target_w, y + target_h)), [[x, y], [x + target_w, y + target_h]]
|
||||
|
||||
|
||||
def motion_blur_opencv(image, kernel_size=15, angle=0):
|
||||
# 线性核
|
||||
kernel = np.zeros((kernel_size, kernel_size), dtype=np.float32)
|
||||
kernel[kernel_size // 2, :] = np.ones(kernel_size, dtype=np.float32)
|
||||
|
||||
# 旋转核
|
||||
center = (kernel_size / 2 - 0.5, kernel_size / 2 - 0.5)
|
||||
M = cv2.getRotationMatrix2D(center, angle, 1)
|
||||
rotated_kernel = cv2.warpAffine(kernel, M, (kernel_size, kernel_size))
|
||||
|
||||
# 归一化核
|
||||
rotated_kernel /= rotated_kernel.sum() if rotated_kernel.sum() != 0 else 1
|
||||
|
||||
img = np.array(image)
|
||||
if img.ndim == 2:
|
||||
blurred = cv2.filter2D(img, -1, rotated_kernel, borderType=cv2.BORDER_REFLECT)
|
||||
else:
|
||||
# 对于彩色图像，各通道独立卷积
|
||||
blurred = np.zeros_like(img)
|
||||
for c in range(img.shape[2]):
|
||||
blurred[..., c] = cv2.filter2D(img[..., c], -1, rotated_kernel, borderType=cv2.BORDER_REFLECT)
|
||||
|
||||
return Image.fromarray(blurred.astype(np.uint8))
|
||||
|
||||
|
||||
def shuffle_patch(image, num_splits, gap_size=2):
|
||||
"""将图像分割为块（允许尺寸不整除），随机打乱后拼接，块间保留间隙"""
|
||||
h_splits, w_splits = num_splits
|
||||
img_w, img_h = image.size
|
||||
|
||||
base_patch_h = img_h // h_splits
|
||||
patch_heights = [base_patch_h] * (h_splits - 1)
|
||||
patch_heights.append(img_h - sum(patch_heights))
|
||||
|
||||
base_patch_w = img_w // w_splits
|
||||
patch_widths = [base_patch_w] * (w_splits - 1)
|
||||
patch_widths.append(img_w - sum(patch_widths))
|
||||
|
||||
patches = []
|
||||
current_y = 0
|
||||
for i in range(h_splits):
|
||||
current_x = 0
|
||||
patch_h = patch_heights[i]
|
||||
for j in range(w_splits):
|
||||
patch_w = patch_widths[j]
|
||||
patch = image.crop((current_x, current_y, current_x + patch_w, current_y + patch_h))
|
||||
patches.append(patch)
|
||||
current_x += patch_w
|
||||
current_y += patch_h
|
||||
|
||||
random.shuffle(patches)
|
||||
|
||||
total_width = sum(patch_widths) + (w_splits - 1) * gap_size
|
||||
total_height = sum(patch_heights) + (h_splits - 1) * gap_size
|
||||
new_image = Image.new(image.mode, (total_width, total_height), color=(255, 255, 255))
|
||||
|
||||
current_y = 0 # 当前行的起始 Y 坐标
|
||||
patch_idx = 0 # 当前处理的块索引
|
||||
for i in range(h_splits):
|
||||
current_x = 0 # 当前列的起始 X 坐标
|
||||
patch_h = patch_heights[i] # 当前行块的高度
|
||||
for j in range(w_splits):
|
||||
# 取出打乱后的块
|
||||
patch = patches[patch_idx]
|
||||
patch_w = patch_widths[j] # 当前列块的宽度
|
||||
# 粘贴块（左上角坐标为 (current_x, current_y)）
|
||||
new_image.paste(patch, (current_x, current_y))
|
||||
# 更新 X 坐标（下一个块的起始位置 = 当前块宽度 + 间隙）
|
||||
current_x += patch_w + gap_size
|
||||
patch_idx += 1
|
||||
# 更新 Y 坐标（下一行的起始位置 = 当前行高度 + 间隙）
|
||||
current_y += patch_h + gap_size
|
||||
|
||||
return new_image
|
||||
|
||||
|
||||
def inpainting(image, num_splits, blank_ratio=0.3, blank_color=(255, 255, 255)):
|
||||
"""
|
||||
图像分割后随机空白部分patch，用于inpainting任务
|
||||
|
||||
参数：
|
||||
image: PIL.Image 输入图像（RGB模式）
|
||||
h_splits: int 行分割数（垂直方向分割块数）
|
||||
w_splits: int 列分割数（水平方向分割块数）
|
||||
blank_ratio: float 空白patch的比例（0~1）
|
||||
blank_color: tuple 空白区域的颜色（RGB，如白色(255,255,255)）
|
||||
|
||||
返回：
|
||||
PIL.Image 处理后拼接的图像
|
||||
"""
|
||||
h_splits, w_splits = num_splits
|
||||
img_w, img_h = image.size
|
||||
|
||||
base_patch_h = img_h // h_splits
|
||||
patch_heights = [base_patch_h] * (h_splits - 1)
|
||||
patch_heights.append(img_h - sum(patch_heights))
|
||||
|
||||
base_patch_w = img_w // w_splits
|
||||
patch_widths = [base_patch_w] * (w_splits - 1)
|
||||
patch_widths.append(img_w - sum(patch_widths))
|
||||
|
||||
patches = []
|
||||
current_y = 0
|
||||
for i in range(h_splits):
|
||||
current_x = 0
|
||||
patch_h = patch_heights[i]
|
||||
for j in range(w_splits):
|
||||
patch_w = patch_widths[j]
|
||||
patch = image.crop((current_x, current_y, current_x + patch_w, current_y + patch_h))
|
||||
patches.append(patch)
|
||||
current_x += patch_w
|
||||
current_y += patch_h
|
||||
|
||||
total_patches = h_splits * w_splits
|
||||
num_blank = int(total_patches * blank_ratio)
|
||||
num_blank = max(0, min(num_blank, total_patches))
|
||||
blank_indices = random.sample(range(total_patches), num_blank)
|
||||
|
||||
processed_patches = []
|
||||
for idx, patch in enumerate(patches):
|
||||
if idx in blank_indices:
|
||||
blank_patch = Image.new("RGB", patch.size, color=blank_color)
|
||||
processed_patches.append(blank_patch)
|
||||
else:
|
||||
processed_patches.append(patch)
|
||||
|
||||
# 创建结果图像（尺寸与原图一致）
|
||||
result_image = Image.new("RGB", (img_w, img_h))
|
||||
current_y = 0
|
||||
patch_idx = 0
|
||||
for i in range(h_splits):
|
||||
current_x = 0
|
||||
patch_h = patch_heights[i]
|
||||
for j in range(w_splits):
|
||||
# 取出处理后的patch
|
||||
patch = processed_patches[patch_idx]
|
||||
patch_w = patch_widths[j]
|
||||
# 粘贴到原位置
|
||||
result_image.paste(patch, (current_x, current_y))
|
||||
current_x += patch_w
|
||||
patch_idx += 1
|
||||
current_y += patch_h
|
||||
|
||||
return result_image
|
||||
@@ -0,0 +1,165 @@
|
||||
# Copyright (c) 2023 OpenGVLab
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: MIT
|
||||
#
|
||||
# This file has been modified by ByteDance Ltd. and/or its affiliates. on 2025-05-20.
|
||||
#
|
||||
# Original file was released under MIT, with the full license text
|
||||
# available at https://github.com/OpenGVLab/InternVL/blob/main/LICENSE.
|
||||
#
|
||||
# This modified file is released under the same license.
|
||||
|
||||
|
||||
import io
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
|
||||
import numpy as np
|
||||
import decord
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def get_frame_indices(num_frames, vlen, sample='rand', fix_start=None, input_fps=1, max_num_frames=-1):
|
||||
if sample in ['rand', 'middle']: # uniform sampling
|
||||
acc_samples = min(num_frames, vlen)
|
||||
# split the video into `acc_samples` intervals, and sample from each interval.
|
||||
intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
|
||||
ranges = []
|
||||
for idx, interv in enumerate(intervals[:-1]):
|
||||
ranges.append((interv, intervals[idx + 1] - 1))
|
||||
if sample == 'rand':
|
||||
try:
|
||||
frame_indices = [random.choice(range(x[0], x[1])) for x in ranges]
|
||||
except:
|
||||
frame_indices = np.random.permutation(vlen)[:acc_samples]
|
||||
frame_indices.sort()
|
||||
frame_indices = list(frame_indices)
|
||||
elif fix_start is not None:
|
||||
frame_indices = [x[0] + fix_start for x in ranges]
|
||||
elif sample == 'middle':
|
||||
frame_indices = [(x[0] + x[1]) // 2 for x in ranges]
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
if len(frame_indices) < num_frames: # padded with last frame
|
||||
padded_frame_indices = [frame_indices[-1]] * num_frames
|
||||
padded_frame_indices[:len(frame_indices)] = frame_indices
|
||||
frame_indices = padded_frame_indices
|
||||
elif 'fps' in sample: # fps0.5, sequentially sample frames at 0.5 fps
|
||||
output_fps = float(sample[3:])
|
||||
duration = float(vlen) / input_fps
|
||||
delta = 1 / output_fps # gap between frames, this is also the clip length each frame represents
|
||||
frame_seconds = np.arange(0 + delta / 2, duration + delta / 2, delta)
|
||||
frame_indices = np.around(frame_seconds * input_fps).astype(int)
|
||||
frame_indices = [e for e in frame_indices if e < vlen]
|
||||
if max_num_frames > 0 and len(frame_indices) > max_num_frames:
|
||||
frame_indices = frame_indices[:max_num_frames]
|
||||
else:
|
||||
raise ValueError
|
||||
return frame_indices
|
||||
|
||||
|
||||
def read_frames_decord(video_path, num_frames, sample='rand', fix_start=None, clip=None, min_num_frames=4):
|
||||
video_reader = decord.VideoReader(video_path, num_threads=1)
|
||||
vlen = len(video_reader)
|
||||
fps = video_reader.get_avg_fps()
|
||||
duration = vlen / float(fps)
|
||||
if clip:
|
||||
start, end = clip
|
||||
duration = end - start
|
||||
vlen = int(duration * fps)
|
||||
start_index = int(start * fps)
|
||||
|
||||
t_num_frames = np.random.randint(min_num_frames, num_frames + 1)
|
||||
|
||||
frame_indices = get_frame_indices(
|
||||
t_num_frames, vlen, sample=sample, fix_start=fix_start,
|
||||
input_fps=fps
|
||||
)
|
||||
if clip:
|
||||
frame_indices = [f + start_index for f in frame_indices]
|
||||
frames = video_reader.get_batch(frame_indices).asnumpy() # (T, H, W, C), np.uint8
|
||||
frames = [Image.fromarray(frames[i]) for i in range(frames.shape[0])]
|
||||
return frames
|
||||
|
||||
|
||||
def extract_frame_number(filename):
|
||||
# Extract the numeric part from the filename using regular expressions
|
||||
match = re.search(r'_(\d+).jpg$', filename)
|
||||
return int(match.group(1)) if match else -1
|
||||
|
||||
|
||||
def sort_frames(frame_paths):
|
||||
# Extract filenames from each path and sort by their numeric part
|
||||
return sorted(frame_paths, key=lambda x: extract_frame_number(os.path.basename(x)))
|
||||
|
||||
|
||||
def read_frames_folder(video_path, num_frames, sample='rand', fix_start=None, min_num_frames=4):
|
||||
image_list = sort_frames(list(os.listdir(video_path)))
|
||||
frames = []
|
||||
for image in image_list:
|
||||
fp = os.path.join(video_path, image)
|
||||
frame = Image.open(fp).convert('RGB')
|
||||
frames.append(frame)
|
||||
vlen = len(frames)
|
||||
|
||||
t_num_frames = np.random.randint(min_num_frames, num_frames + 1)
|
||||
|
||||
if vlen > t_num_frames:
|
||||
frame_indices = get_frame_indices(
|
||||
t_num_frames, vlen, sample=sample, fix_start=fix_start
|
||||
)
|
||||
frames = [frames[i] for i in frame_indices]
|
||||
return frames
|
||||
|
||||
|
||||
class FrameSampler:
|
||||
def __init__(self, max_num_frames=-1, min_num_frames=8, sample='rand'):
|
||||
self.max_num_frames = max_num_frames
|
||||
self.min_num_frames = min_num_frames
|
||||
self.sample = sample
|
||||
|
||||
def __call__(self, file_name):
|
||||
fn = read_frames_folder if file_name.endswith('/') else read_frames_decord
|
||||
frames = fn(file_name, num_frames=self.max_num_frames, min_num_frames=self.min_num_frames, sample=self.sample)
|
||||
return frames
|
||||
|
||||
|
||||
def decode_video_byte(video_bytes):
|
||||
video_stream = io.BytesIO(video_bytes)
|
||||
vr = decord.VideoReader(video_stream)
|
||||
return vr
|
||||
|
||||
|
||||
def sample_mp4_frames(mp4_p, n_frames=None, fps=None, return_frame_indices=False, random_sample=False):
|
||||
if isinstance(mp4_p, str):
|
||||
vr = decord.VideoReader(mp4_p, num_threads=1)
|
||||
elif isinstance(mp4_p, decord.video_reader.VideoReader):
|
||||
vr = mp4_p
|
||||
video_fps = vr.get_avg_fps() # 获取视频的帧率
|
||||
video_duration = len(vr) / video_fps
|
||||
if n_frames is not None:
|
||||
if random_sample:
|
||||
frame_indices = sorted(random.sample(range(len(vr)), n_frames))
|
||||
else:
|
||||
frame_indices = np.linspace(0, len(vr)-1, n_frames, dtype=int).tolist()
|
||||
else:
|
||||
frame_indices = [int(i) for i in np.arange(0, len(vr)-1, video_fps/fps)]
|
||||
frames = vr.get_batch(frame_indices).asnumpy() # 转换为 numpy 数组
|
||||
frames = [Image.fromarray(frame).convert("RGB") for frame in frames]
|
||||
if not return_frame_indices:
|
||||
return frames, video_duration
|
||||
else:
|
||||
return frames, video_duration, frame_indices
|
||||
|
||||
|
||||
def sample_mp4_frames_by_indices(mp4_p, frame_indices: list):
|
||||
if isinstance(mp4_p, str):
|
||||
vr = decord.VideoReader(mp4_p, num_threads=1)
|
||||
elif isinstance(mp4_p, decord.video_reader.VideoReader):
|
||||
vr = mp4_p
|
||||
# sample the frames in frame_indices
|
||||
frames = vr.get_batch(frame_indices).asnumpy() # 转换为 numpy 数组
|
||||
frames = [Image.fromarray(frame).convert("RGB") for frame in frames]
|
||||
return frames
|
||||
@@ -0,0 +1,195 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
import os
|
||||
import traceback
|
||||
from PIL import Image, ImageFile, PngImagePlugin
|
||||
|
||||
from .data_utils import pil_img2rgb
|
||||
from .distributed_iterable_dataset import DistributedIterableDataset
|
||||
|
||||
|
||||
Image.MAX_IMAGE_PIXELS = 200000000
|
||||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||||
MaximumDecompressedSize = 1024
|
||||
MegaByte = 2 ** 20
|
||||
PngImagePlugin.MAX_TEXT_CHUNK = MaximumDecompressedSize * MegaByte
|
||||
|
||||
|
||||
class SftJSONLIterableDataset(DistributedIterableDataset):
|
||||
def __init__(
|
||||
self, dataset_name, transform, tokenizer, frame_sampler,
|
||||
jsonl_path_list, data_dir_list, num_used_data,
|
||||
local_rank=0, world_size=1, num_workers=8, data_status=None,
|
||||
shuffle_lines=False, shuffle_seed=0,
|
||||
):
|
||||
"""
|
||||
jsonl_path_list: list of jsonl file paths
|
||||
data_dir_list: list of image directories containing the images of each jsonl file
|
||||
num_used_data: list of number of sampled data points for each jsonl
|
||||
"""
|
||||
super().__init__(dataset_name, local_rank, world_size, num_workers)
|
||||
self.transform = transform
|
||||
self.tokenizer = tokenizer
|
||||
self.frame_sampler = frame_sampler
|
||||
self.data_status = data_status
|
||||
self.data_paths = self.get_data_paths(
|
||||
jsonl_path_list,
|
||||
data_dir_list,
|
||||
num_used_data,
|
||||
shuffle_lines,
|
||||
shuffle_seed,
|
||||
)
|
||||
self.set_epoch()
|
||||
|
||||
def get_data_paths(
|
||||
self,
|
||||
jsonl_path_list,
|
||||
data_dir_list,
|
||||
num_used_data,
|
||||
shuffle_lines,
|
||||
shuffle_seed,
|
||||
):
|
||||
data_paths = []
|
||||
for jsonl_path, image_dir, num_data_point in zip(
|
||||
jsonl_path_list, data_dir_list, num_used_data
|
||||
):
|
||||
with open(jsonl_path, 'r') as f:
|
||||
raw_data = f.readlines()
|
||||
if shuffle_lines:
|
||||
self.rng.seed(shuffle_seed)
|
||||
self.rng.shuffle(raw_data)
|
||||
raw_data = raw_data[:num_data_point]
|
||||
data_paths.extend([(json_data, image_dir) for json_data in raw_data])
|
||||
return data_paths
|
||||
|
||||
def change_format(self, data, num_images):
|
||||
elements = []
|
||||
for conversation in data['conversations']:
|
||||
if conversation['from'] == 'human':
|
||||
if '<image>' not in conversation['value']:
|
||||
elements.append({
|
||||
'type': 'text',
|
||||
'has_loss': 0,
|
||||
'text': conversation['value'],
|
||||
})
|
||||
else:
|
||||
text_list = conversation['value'].split('<image>')
|
||||
for idx, text in enumerate(text_list):
|
||||
if text.strip() != '':
|
||||
elements.append({
|
||||
'type': 'text',
|
||||
'has_loss': 0,
|
||||
'text': text.strip(),
|
||||
})
|
||||
if (idx != len(text_list) - 1) and (idx < num_images):
|
||||
elements.append({'type': 'image',})
|
||||
elif conversation['from'] == 'gpt':
|
||||
elements.append({
|
||||
'type': 'text',
|
||||
'has_loss': 1,
|
||||
'text': conversation['value'],
|
||||
})
|
||||
return elements
|
||||
|
||||
def __iter__(self):
|
||||
data_paths_per_worker, worker_id = self.get_data_paths_per_worker()
|
||||
if self.data_status is not None:
|
||||
row_start_id = self.data_status[worker_id] + 1
|
||||
else:
|
||||
row_start_id = 0
|
||||
transform_stride = self.transform.stride
|
||||
|
||||
print(
|
||||
f"rank-{self.local_rank} worker-{worker_id} dataset-{self.dataset_name}: "
|
||||
f"resuming data at row#{row_start_id}"
|
||||
)
|
||||
|
||||
while True:
|
||||
data_paths_per_worker_ = data_paths_per_worker[row_start_id:]
|
||||
for row_idx, (data, image_dir) in enumerate(data_paths_per_worker_, start=row_start_id):
|
||||
num_tokens = 0
|
||||
image_tensor_list = []
|
||||
text_ids_list = []
|
||||
sequence_plan = []
|
||||
|
||||
try:
|
||||
data_item = json.loads(data)
|
||||
raw_images = None
|
||||
if 'image' in data_item:
|
||||
if type(data_item['image']) == list:
|
||||
raw_images = [
|
||||
pil_img2rgb(Image.open(os.path.join(image_dir, image)))
|
||||
for image in data_item['image']
|
||||
]
|
||||
else:
|
||||
raw_images = [
|
||||
pil_img2rgb(Image.open(os.path.join(image_dir, data_item['image'])))
|
||||
]
|
||||
elif 'video' in data_item:
|
||||
raw_images = self.frame_sampler(os.path.join(image_dir, data_item['video']))
|
||||
special_tokens = '<image>' * len(raw_images)
|
||||
for item in data_item['conversations']:
|
||||
if '<video>' in item['value']:
|
||||
item['value'] = item['value'].replace('<video>', special_tokens)
|
||||
break
|
||||
else:
|
||||
raise ValueError("Cannot find <video> in the conversation!")
|
||||
except:
|
||||
traceback.print_exc()
|
||||
continue
|
||||
|
||||
if raw_images:
|
||||
for raw_image in raw_images:
|
||||
image_tensor = self.transform(raw_image, img_num=len(raw_images))
|
||||
image_tensor_list.append(image_tensor)
|
||||
height, width = image_tensor.shape[1:]
|
||||
num_tokens += width * height // transform_stride ** 2
|
||||
|
||||
elements = self.change_format(data_item, len(image_tensor_list))
|
||||
|
||||
for item in elements:
|
||||
if item['type'] == 'text':
|
||||
text_data = item['text']
|
||||
text_ids = self.tokenizer.encode(text_data)
|
||||
if len(text_ids) > 0:
|
||||
text_ids_list.append(text_ids)
|
||||
num_tokens += len(text_ids)
|
||||
current_plan = {
|
||||
'type': 'text',
|
||||
'enable_cfg': 0,
|
||||
'loss': item['has_loss'],
|
||||
'special_token_loss': 0,
|
||||
'special_token_label': None,
|
||||
}
|
||||
sequence_plan.append(current_plan)
|
||||
elif item['type'] == 'image':
|
||||
current_plan = {
|
||||
'type': 'vit_image',
|
||||
'enable_cfg': 0,
|
||||
'loss': 0,
|
||||
'special_token_loss': 0,
|
||||
'special_token_label': None,
|
||||
}
|
||||
sequence_plan.append(current_plan)
|
||||
|
||||
has_loss = [item['loss'] for item in sequence_plan]
|
||||
if sum(has_loss) == 0:
|
||||
print(f'No loss defined, skipped.')
|
||||
continue
|
||||
|
||||
yield dict(
|
||||
image_tensor_list=image_tensor_list,
|
||||
text_ids_list=text_ids_list,
|
||||
sequence_plan=sequence_plan,
|
||||
num_tokens=num_tokens,
|
||||
data_indexes={
|
||||
"data_indexes": row_idx,
|
||||
"worker_id": worker_id,
|
||||
"dataset_name": self.dataset_name,
|
||||
}
|
||||
)
|
||||
|
||||
row_start_id = 0
|
||||
print(f"{self.dataset_name} repeat in rank-{self.local_rank} worker-{worker_id}")
|
||||
@@ -0,0 +1,535 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Copyright 2025 Bytedance Ltd. and/or its affiliates.\n",
|
||||
"# SPDX-License-Identifier: Apache-2.0"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%load_ext autoreload\n",
|
||||
"%autoreload 2"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"from copy import deepcopy\n",
|
||||
"from typing import (\n",
|
||||
" Any,\n",
|
||||
" AsyncIterable,\n",
|
||||
" Callable,\n",
|
||||
" Dict,\n",
|
||||
" Generator,\n",
|
||||
" List,\n",
|
||||
" NamedTuple,\n",
|
||||
" Optional,\n",
|
||||
" Tuple,\n",
|
||||
" Union,\n",
|
||||
")\n",
|
||||
"import requests\n",
|
||||
"from io import BytesIO\n",
|
||||
"\n",
|
||||
"from PIL import Image\n",
|
||||
"import torch\n",
|
||||
"from accelerate import infer_auto_device_map, load_checkpoint_and_dispatch, init_empty_weights\n",
|
||||
"\n",
|
||||
"from data.transforms import ImageTransform\n",
|
||||
"from data.data_utils import pil_img2rgb, add_special_tokens\n",
|
||||
"from modeling.bagel import (\n",
|
||||
" BagelConfig, Bagel, Qwen2Config, Qwen2ForCausalLM, SiglipVisionConfig, SiglipVisionModel\n",
|
||||
")\n",
|
||||
"from modeling.qwen2 import Qwen2Tokenizer\n",
|
||||
"from modeling.bagel.qwen2_navit import NaiveCache\n",
|
||||
"from modeling.autoencoder import load_ae\n",
|
||||
"from safetensors.torch import load_file"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Model Initialization"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"model_path = \"/path/to/BAGEL-7B-MoT/weights\" # Download from https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT\n",
|
||||
"\n",
|
||||
"# LLM config preparing\n",
|
||||
"llm_config = Qwen2Config.from_json_file(os.path.join(model_path, \"llm_config.json\"))\n",
|
||||
"llm_config.qk_norm = True\n",
|
||||
"llm_config.tie_word_embeddings = False\n",
|
||||
"llm_config.layer_module = \"Qwen2MoTDecoderLayer\"\n",
|
||||
"\n",
|
||||
"# ViT config preparing\n",
|
||||
"vit_config = SiglipVisionConfig.from_json_file(os.path.join(model_path, \"vit_config.json\"))\n",
|
||||
"vit_config.rope = False\n",
|
||||
"vit_config.num_hidden_layers = vit_config.num_hidden_layers - 1\n",
|
||||
"\n",
|
||||
"# VAE loading\n",
|
||||
"vae_model, vae_config = load_ae(local_path=os.path.join(model_path, \"ae.safetensors\"))\n",
|
||||
"\n",
|
||||
"# Bagel config preparing\n",
|
||||
"config = BagelConfig(\n",
|
||||
" visual_gen=True,\n",
|
||||
" visual_und=True,\n",
|
||||
" llm_config=llm_config, \n",
|
||||
" vit_config=vit_config,\n",
|
||||
" vae_config=vae_config,\n",
|
||||
" vit_max_num_patch_per_side=70,\n",
|
||||
" connector_act='gelu_pytorch_tanh',\n",
|
||||
" latent_patch_size=2,\n",
|
||||
" max_latent_size=64,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"with init_empty_weights():\n",
|
||||
" language_model = Qwen2ForCausalLM(llm_config)\n",
|
||||
" vit_model = SiglipVisionModel(vit_config)\n",
|
||||
" model = Bagel(language_model, vit_model, config)\n",
|
||||
" model.vit_model.vision_model.embeddings.convert_conv2d_to_linear(vit_config, meta=True)\n",
|
||||
"\n",
|
||||
"# Tokenizer Preparing\n",
|
||||
"tokenizer = Qwen2Tokenizer.from_pretrained(model_path)\n",
|
||||
"tokenizer, new_token_ids, _ = add_special_tokens(tokenizer)\n",
|
||||
"\n",
|
||||
"# Image Transform Preparing\n",
|
||||
"vae_transform = ImageTransform(1024, 512, 16)\n",
|
||||
"vit_transform = ImageTransform(980, 224, 14)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Model Loading and Multi GPU Infernece Preparing"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"max_mem_per_gpu = \"40GiB\" # Modify it according to your GPU setting. On an A100, 80 GiB is sufficient to load on a single GPU.\n",
|
||||
"\n",
|
||||
"device_map = infer_auto_device_map(\n",
|
||||
" model,\n",
|
||||
" max_memory={i: max_mem_per_gpu for i in range(torch.cuda.device_count())},\n",
|
||||
" no_split_module_classes=[\"Bagel\", \"Qwen2MoTDecoderLayer\"],\n",
|
||||
")\n",
|
||||
"print(device_map)\n",
|
||||
"\n",
|
||||
"same_device_modules = [\n",
|
||||
" 'language_model.model.embed_tokens',\n",
|
||||
" 'time_embedder',\n",
|
||||
" 'latent_pos_embed',\n",
|
||||
" 'vae2llm',\n",
|
||||
" 'llm2vae',\n",
|
||||
" 'connector',\n",
|
||||
" 'vit_pos_embed'\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"if torch.cuda.device_count() == 1:\n",
|
||||
" first_device = device_map.get(same_device_modules[0], \"cuda:0\")\n",
|
||||
" for k in same_device_modules:\n",
|
||||
" if k in device_map:\n",
|
||||
" device_map[k] = first_device\n",
|
||||
" else:\n",
|
||||
" device_map[k] = \"cuda:0\"\n",
|
||||
"else:\n",
|
||||
" first_device = device_map.get(same_device_modules[0])\n",
|
||||
" for k in same_device_modules:\n",
|
||||
" if k in device_map:\n",
|
||||
" device_map[k] = first_device\n",
|
||||
"\n",
|
||||
"# Thanks @onion-liu: https://github.com/ByteDance-Seed/Bagel/pull/8\n",
|
||||
"model = load_checkpoint_and_dispatch(\n",
|
||||
" model,\n",
|
||||
" checkpoint=os.path.join(model_path, \"ema.safetensors\"),\n",
|
||||
" device_map=device_map,\n",
|
||||
" offload_buffers=True,\n",
|
||||
" dtype=torch.bfloat16,\n",
|
||||
" force_hooks=True,\n",
|
||||
" offload_folder=\"/tmp/offload\"\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"model = model.eval()\n",
|
||||
"print('Model loaded')"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Inferencer Preparing "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from inferencer import InterleaveInferencer\n",
|
||||
"\n",
|
||||
"inferencer = InterleaveInferencer(\n",
|
||||
" model=model, \n",
|
||||
" vae_model=vae_model, \n",
|
||||
" tokenizer=tokenizer, \n",
|
||||
" vae_transform=vae_transform, \n",
|
||||
" vit_transform=vit_transform, \n",
|
||||
" new_token_ids=new_token_ids\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import random\n",
|
||||
"import numpy as np\n",
|
||||
"\n",
|
||||
"seed = 42\n",
|
||||
"random.seed(seed)\n",
|
||||
"np.random.seed(seed)\n",
|
||||
"torch.manual_seed(seed)\n",
|
||||
"if torch.cuda.is_available():\n",
|
||||
" torch.cuda.manual_seed(seed)\n",
|
||||
" torch.cuda.manual_seed_all(seed)\n",
|
||||
"torch.backends.cudnn.deterministic = True\n",
|
||||
"torch.backends.cudnn.benchmark = False"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"**About Inference Hyperparameters:**\n",
|
||||
"- **`cfg_text_scale`:** Controls how strongly the model follows the text prompt. `1.0` disables text guidance. Typical range: `4.0–8.0`.\n",
|
||||
"- **`cfg_image_scale`:** Controls how much the model preserves input image details. `1.0` disables image guidance. Typical range: `1.0–2.0`.\n",
|
||||
"- **`cfg_interval`:** Fraction of denoising steps where CFG is applied. Later steps can skip CFG to reduce computation. Typical: `[0.4, 1.0]`.\n",
|
||||
"- **`timestep_shift`:** Shifts the distribution of denoising steps. Higher values allocate more steps at the start (affects layout); lower values allocate more at the end (improves details).\n",
|
||||
"- **`num_timesteps`:** Total denoising steps. Typical: `50`.\n",
|
||||
"- **`cfg_renorm_min`:** Minimum value for CFG-Renorm. `1.0` disables renorm. Typical: `0`.\n",
|
||||
"- **`cfg_renorm_type`:** CFG-Renorm method: \n",
|
||||
" - `global`: Normalize over all tokens and channels (default for T2I).\n",
|
||||
" - `channel`: Normalize across channels for each token.\n",
|
||||
" - `text_channel`: Like `channel`, but only applies to text condition (good for editing, may cause blur).\n",
|
||||
"- **If edited images appear blurry, try `global` CFG-Renorm, decrease `cfg_renorm_min` or decrease `cfg_scale`.**\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Image Generation"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"inference_hyper=dict(\n",
|
||||
" cfg_text_scale=4.0,\n",
|
||||
" cfg_img_scale=1.0,\n",
|
||||
" cfg_interval=[0.4, 1.0],\n",
|
||||
" timestep_shift=3.0,\n",
|
||||
" num_timesteps=50,\n",
|
||||
" cfg_renorm_min=1.0,\n",
|
||||
" cfg_renorm_type=\"global\",\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"prompt = \"A female cosplayer portraying an ethereal fairy or elf, wearing a flowing dress made of delicate fabrics in soft, mystical colors like emerald green and silver. She has pointed ears, a gentle, enchanting expression, and her outfit is adorned with sparkling jewels and intricate patterns. The background is a magical forest with glowing plants, mystical creatures, and a serene atmosphere.\"\n",
|
||||
"\n",
|
||||
"print(prompt)\n",
|
||||
"print('-' * 10)\n",
|
||||
"output_dict = inferencer(text=prompt, **inference_hyper)\n",
|
||||
"display(output_dict['image'])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"source": [
|
||||
"## Image Generation with Think"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"inference_hyper=dict(\n",
|
||||
" max_think_token_n=1000,\n",
|
||||
" do_sample=False,\n",
|
||||
" # text_temperature=0.3,\n",
|
||||
" cfg_text_scale=4.0,\n",
|
||||
" cfg_img_scale=1.0,\n",
|
||||
" cfg_interval=[0.4, 1.0],\n",
|
||||
" timestep_shift=3.0,\n",
|
||||
" num_timesteps=50,\n",
|
||||
" cfg_renorm_min=1.0,\n",
|
||||
" cfg_renorm_type=\"global\",\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"prompt = 'a car made of small cars'\n",
|
||||
"\n",
|
||||
"print(prompt)\n",
|
||||
"print('-' * 10)\n",
|
||||
"output_dict = inferencer(text=prompt, think=True, **inference_hyper)\n",
|
||||
"print(output_dict['text'])\n",
|
||||
"display(output_dict['image'])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Editing"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"inference_hyper=dict(\n",
|
||||
" cfg_text_scale=4.0,\n",
|
||||
" cfg_img_scale=2.0,\n",
|
||||
" cfg_interval=[0.0, 1.0],\n",
|
||||
" timestep_shift=3.0,\n",
|
||||
" num_timesteps=50,\n",
|
||||
" cfg_renorm_min=1.0,\n",
|
||||
" cfg_renorm_type=\"text_channel\",\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"image = Image.open('test_images/women.jpg')\n",
|
||||
"prompt = 'She boards a modern subway, quietly reading a folded newspaper, wearing the same clothes.'\n",
|
||||
"\n",
|
||||
"display(image)\n",
|
||||
"print(prompt)\n",
|
||||
"print('-'*10)\n",
|
||||
"output_dict = inferencer(image=image, text=prompt, **inference_hyper)\n",
|
||||
"display(output_dict['image'])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Edit with Think"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"inference_hyper=dict(\n",
|
||||
" max_think_token_n=1000,\n",
|
||||
" do_sample=False,\n",
|
||||
" # text_temperature=0.3,\n",
|
||||
" cfg_text_scale=4.0,\n",
|
||||
" cfg_img_scale=2.0,\n",
|
||||
" cfg_interval=[0.0, 1.0],\n",
|
||||
" timestep_shift=3.0,\n",
|
||||
" num_timesteps=50,\n",
|
||||
" cfg_renorm_min=0.0,\n",
|
||||
" cfg_renorm_type=\"text_channel\",\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"image = Image.open('test_images/octupusy.jpg')\n",
|
||||
"prompt = 'Could you display the sculpture that takes after this design?'\n",
|
||||
"\n",
|
||||
"display(image)\n",
|
||||
"print('-'*10)\n",
|
||||
"output_dict = inferencer(image=image, text=prompt, think=True, **inference_hyper)\n",
|
||||
"print(output_dict['text'])\n",
|
||||
"display(output_dict['image'])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Understanding"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"inference_hyper=dict(\n",
|
||||
" max_think_token_n=1000,\n",
|
||||
" do_sample=False,\n",
|
||||
" # text_temperature=0.3,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"image = Image.open('test_images/meme.jpg')\n",
|
||||
"prompt = \"Can someone explain what’s funny about this meme??\"\n",
|
||||
"\n",
|
||||
"display(image)\n",
|
||||
"print(prompt)\n",
|
||||
"print('-'*10)\n",
|
||||
"output_dict = inferencer(image=image, text=prompt, understanding_output=True, **inference_hyper)\n",
|
||||
"print(output_dict['text'])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"fileId": "1bfaa82d-51b0-4c13-9e4c-295ba28bcd8a",
|
||||
"filePath": "/mnt/bn/seed-aws-va/chaorui/code/cdt-hf/notebooks/chat.ipynb",
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -0,0 +1,309 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from copy import deepcopy
|
||||
from typing import List, Dict, Optional, Union, Any
|
||||
|
||||
from PIL import Image
|
||||
import torch
|
||||
|
||||
from .data.data_utils import pil_img2rgb
|
||||
from .modeling.bagel.qwen2_navit import NaiveCache
|
||||
|
||||
|
||||
|
||||
VLM_THINK_SYSTEM_PROMPT = '''You should first think about the reasoning process in the mind and then provide the user with the answer.
|
||||
The reasoning process is enclosed within <think> </think> tags, i.e. <think> reasoning process here </think> answer here'''
|
||||
|
||||
GEN_THINK_SYSTEM_PROMPT = '''You should first think about the planning process in the mind and then generate the image.
|
||||
The planning process is enclosed within <think> </think> tags, i.e. <think> planning process here </think> image here'''
|
||||
|
||||
|
||||
class InterleaveInferencer:
|
||||
def __init__(self, model, vae_model, tokenizer, vae_transform, vit_transform, new_token_ids):
|
||||
self.model = model
|
||||
self.vae_model = vae_model
|
||||
self.tokenizer = tokenizer
|
||||
self.vae_transform = vae_transform
|
||||
self.vit_transform = vit_transform
|
||||
self.new_token_ids = new_token_ids
|
||||
|
||||
def init_gen_context(self):
|
||||
gen_context = {
|
||||
'kv_lens': [0],
|
||||
'ropes': [0],
|
||||
'past_key_values': NaiveCache(self.model.config.llm_config.num_hidden_layers),
|
||||
}
|
||||
return gen_context
|
||||
|
||||
@torch.no_grad()
|
||||
def update_context_text(self, text, gen_context):
|
||||
# used for interleave data, currently only support 1 data inference,
|
||||
|
||||
past_key_values = gen_context['past_key_values']
|
||||
kv_lens = gen_context['kv_lens']
|
||||
ropes = gen_context['ropes']
|
||||
generation_input, kv_lens, ropes = self.model.prepare_prompts(
|
||||
curr_kvlens=kv_lens,
|
||||
curr_rope=ropes,
|
||||
prompts=[text],
|
||||
tokenizer=self.tokenizer,
|
||||
new_token_ids=self.new_token_ids,
|
||||
)
|
||||
|
||||
past_key_values = self.model.forward_cache_update_text(past_key_values, **generation_input)
|
||||
gen_context['kv_lens'] = kv_lens
|
||||
gen_context['ropes'] = ropes
|
||||
gen_context['past_key_values'] = past_key_values
|
||||
|
||||
return gen_context
|
||||
|
||||
@torch.no_grad()
|
||||
def update_context_image(self, image, gen_context, vae=True, vit=True):
|
||||
# used for interleave data, currently only support 1 data inference,
|
||||
|
||||
assert vae or vit
|
||||
past_key_values = gen_context['past_key_values']
|
||||
kv_lens = gen_context['kv_lens']
|
||||
ropes = gen_context['ropes']
|
||||
|
||||
if vae:
|
||||
## update vae
|
||||
generation_input, kv_lens, ropes = self.model.prepare_vae_images(
|
||||
curr_kvlens=kv_lens,
|
||||
curr_rope=ropes,
|
||||
images=[image],
|
||||
transforms=self.vae_transform,
|
||||
new_token_ids=self.new_token_ids,
|
||||
)
|
||||
past_key_values = self.model.forward_cache_update_vae(self.vae_model, past_key_values, **generation_input)
|
||||
|
||||
if vit:
|
||||
## update vit
|
||||
generation_input, kv_lens, ropes = self.model.prepare_vit_images(
|
||||
curr_kvlens=kv_lens,
|
||||
curr_rope=ropes,
|
||||
images=[image],
|
||||
transforms=self.vit_transform,
|
||||
new_token_ids=self.new_token_ids,
|
||||
)
|
||||
past_key_values = self.model.forward_cache_update_vit(past_key_values, **generation_input)
|
||||
|
||||
gen_context['kv_lens'] = kv_lens
|
||||
gen_context['ropes'] = ropes
|
||||
gen_context['past_key_values'] = past_key_values
|
||||
|
||||
return gen_context
|
||||
|
||||
@torch.no_grad()
|
||||
def gen_image(
|
||||
self,
|
||||
image_shape,
|
||||
gen_context,
|
||||
cfg_text_scale=4.0,
|
||||
cfg_img_scale=1.5,
|
||||
|
||||
cfg_text_precontext=None,
|
||||
cfg_img_precontext=None,
|
||||
cfg_interval=(0.4, 1.0),
|
||||
cfg_renorm_min=0.0,
|
||||
cfg_renorm_type="global",
|
||||
|
||||
num_timesteps=50,
|
||||
timestep_shift=3.0
|
||||
):
|
||||
# print(cfg_renorm_type)
|
||||
past_key_values = gen_context['past_key_values']
|
||||
kv_lens = gen_context['kv_lens']
|
||||
ropes = gen_context['ropes']
|
||||
generation_input = self.model.prepare_vae_latent(
|
||||
curr_kvlens=kv_lens,
|
||||
curr_rope=ropes,
|
||||
image_sizes=[image_shape],
|
||||
new_token_ids=self.new_token_ids,
|
||||
)
|
||||
|
||||
# text cfg
|
||||
cfg_text_past_key_values = cfg_text_precontext['past_key_values']
|
||||
kv_lens_cfg = cfg_text_precontext['kv_lens']
|
||||
ropes_cfg = cfg_text_precontext['ropes']
|
||||
generation_input_cfg_text = self.model.prepare_vae_latent_cfg(
|
||||
curr_kvlens=kv_lens_cfg,
|
||||
curr_rope=ropes_cfg,
|
||||
image_sizes=[image_shape],
|
||||
)
|
||||
|
||||
# img cfg
|
||||
cfg_img_past_key_values = cfg_img_precontext['past_key_values']
|
||||
kv_lens_cfg = cfg_img_precontext['kv_lens']
|
||||
ropes_cfg = cfg_img_precontext['ropes']
|
||||
generation_input_cfg_img = self.model.prepare_vae_latent_cfg(
|
||||
curr_kvlens=kv_lens_cfg,
|
||||
curr_rope=ropes_cfg,
|
||||
image_sizes=[image_shape],
|
||||
)
|
||||
|
||||
unpacked_latent = self.model.generate_image(
|
||||
past_key_values=past_key_values,
|
||||
cfg_text_past_key_values=cfg_text_past_key_values,
|
||||
cfg_img_past_key_values=cfg_img_past_key_values,
|
||||
num_timesteps=num_timesteps,
|
||||
cfg_text_scale=cfg_text_scale,
|
||||
cfg_img_scale=cfg_img_scale,
|
||||
cfg_interval=cfg_interval,
|
||||
cfg_renorm_min=cfg_renorm_min,
|
||||
cfg_renorm_type=cfg_renorm_type,
|
||||
timestep_shift=timestep_shift,
|
||||
**generation_input,
|
||||
cfg_text_packed_position_ids=generation_input_cfg_text['cfg_packed_position_ids'],
|
||||
cfg_text_packed_query_indexes=generation_input_cfg_text['cfg_packed_query_indexes'],
|
||||
cfg_text_key_values_lens=generation_input_cfg_text['cfg_key_values_lens'],
|
||||
cfg_text_packed_key_value_indexes=generation_input_cfg_text['cfg_packed_key_value_indexes'],
|
||||
cfg_img_packed_position_ids=generation_input_cfg_img['cfg_packed_position_ids'],
|
||||
cfg_img_packed_query_indexes=generation_input_cfg_img['cfg_packed_query_indexes'],
|
||||
cfg_img_key_values_lens=generation_input_cfg_img['cfg_key_values_lens'],
|
||||
cfg_img_packed_key_value_indexes=generation_input_cfg_img['cfg_packed_key_value_indexes'],
|
||||
)
|
||||
|
||||
image = self.decode_image(unpacked_latent[0], image_shape)
|
||||
return image
|
||||
|
||||
|
||||
def decode_image(self, latent, image_shape):
|
||||
H, W = image_shape
|
||||
h, w = H // self.model.latent_downsample, W // self.model.latent_downsample
|
||||
|
||||
latent = latent.reshape(1, h, w, self.model.latent_patch_size, self.model.latent_patch_size, self.model.latent_channel)
|
||||
latent = torch.einsum("nhwpqc->nchpwq", latent)
|
||||
latent = latent.reshape(1, self.model.latent_channel, h * self.model.latent_patch_size, w * self.model.latent_patch_size)
|
||||
image = self.vae_model.decode(latent)
|
||||
image = (image * 0.5 + 0.5).clamp(0, 1)[0].permute(1, 2, 0) * 255
|
||||
image = Image.fromarray((image).to(torch.uint8).cpu().numpy())
|
||||
|
||||
return image
|
||||
|
||||
@torch.no_grad()
|
||||
def gen_text(self, gen_context, max_length: int = 500, do_sample: bool = True, temperature: float = 1.0):
|
||||
gen_context = deepcopy(gen_context)
|
||||
past_key_values = gen_context['past_key_values']
|
||||
kv_lens = gen_context['kv_lens']
|
||||
ropes = gen_context['ropes']
|
||||
|
||||
generation_input = self.model.prepare_start_tokens(kv_lens, ropes, self.new_token_ids)
|
||||
unpacked_latent = self.model.generate_text(
|
||||
past_key_values=past_key_values,
|
||||
max_length=max_length,
|
||||
do_sample=do_sample,
|
||||
temperature=temperature,
|
||||
end_token_id=self.new_token_ids['eos_token_id'],
|
||||
**generation_input,
|
||||
)
|
||||
output = self.tokenizer.decode(unpacked_latent[:,0])
|
||||
output = output.split('<|im_end|>')[0].split('<|im_start|>')[1]
|
||||
return output
|
||||
|
||||
@torch.no_grad()
|
||||
def interleave_inference(
|
||||
self,
|
||||
input_lists: List[Union[str, Image.Image]],
|
||||
think=False,
|
||||
understanding_output=False,
|
||||
|
||||
max_think_token_n=1000,
|
||||
do_sample=False,
|
||||
text_temperature=0.3,
|
||||
cfg_text_scale=3.0,
|
||||
cfg_img_scale=1.5,
|
||||
cfg_interval=[0.4, 1.0],
|
||||
timestep_shift=3.0,
|
||||
num_timesteps=50,
|
||||
cfg_renorm_min=0.0,
|
||||
cfg_renorm_type="global",
|
||||
image_shapes=(1024, 1024),
|
||||
) -> List[Union[str, Image.Image]]:
|
||||
|
||||
output_list = []
|
||||
gen_context = self.init_gen_context()
|
||||
cfg_text_context = deepcopy(gen_context)
|
||||
cfg_img_context = deepcopy(gen_context)
|
||||
|
||||
with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):
|
||||
if think:
|
||||
if understanding_output:
|
||||
system_prompt = VLM_THINK_SYSTEM_PROMPT
|
||||
else:
|
||||
system_prompt = GEN_THINK_SYSTEM_PROMPT
|
||||
gen_context = self.update_context_text(system_prompt, gen_context)
|
||||
cfg_img_context = self.update_context_text(system_prompt, cfg_img_context)
|
||||
|
||||
for input_term in input_lists:
|
||||
if isinstance(input_term, str):
|
||||
cfg_text_context = deepcopy(gen_context)
|
||||
gen_context = self.update_context_text(input_term, gen_context)
|
||||
cfg_img_context = self.update_context_text(input_term, cfg_img_context)
|
||||
|
||||
elif isinstance(input_term, Image.Image):
|
||||
input_term = self.vae_transform.resize_transform(pil_img2rgb(input_term))
|
||||
gen_context = self.update_context_image(input_term, gen_context, vae=not understanding_output)
|
||||
|
||||
image_shapes = input_term.size[::-1]
|
||||
cfg_text_context = deepcopy(gen_context)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported input type: {type(input_term)}")
|
||||
|
||||
if understanding_output:
|
||||
gen_text = self.gen_text(gen_context, do_sample=do_sample, temperature=text_temperature, max_length=max_think_token_n)
|
||||
output_list.append(gen_text)
|
||||
|
||||
else:
|
||||
if think:
|
||||
gen_text = self.gen_text(gen_context, do_sample=do_sample, temperature=text_temperature, max_length=max_think_token_n)
|
||||
gen_context = self.update_context_text(gen_text, gen_context)
|
||||
output_list.append(gen_text)
|
||||
|
||||
img = self.gen_image(
|
||||
image_shapes,
|
||||
gen_context,
|
||||
cfg_text_precontext=cfg_text_context,
|
||||
cfg_img_precontext=cfg_img_context,
|
||||
|
||||
cfg_text_scale=cfg_text_scale,
|
||||
cfg_img_scale=cfg_img_scale,
|
||||
cfg_interval=cfg_interval,
|
||||
timestep_shift=timestep_shift,
|
||||
num_timesteps=num_timesteps,
|
||||
cfg_renorm_min=cfg_renorm_min,
|
||||
cfg_renorm_type=cfg_renorm_type,
|
||||
)
|
||||
|
||||
output_list.append(img)
|
||||
|
||||
return output_list
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
image: Optional[Image.Image] = None,
|
||||
text: Optional[str] = None,
|
||||
**kargs
|
||||
) -> Dict[str, Any]:
|
||||
output_dict = {'image': None, 'text': None}
|
||||
|
||||
if image is None and text is None:
|
||||
print('Please provide at least one input: either an image or text.')
|
||||
return output_dict
|
||||
|
||||
input_list = []
|
||||
if image is not None:
|
||||
input_list.append(image)
|
||||
if text is not None:
|
||||
input_list.append(text)
|
||||
|
||||
output_list = self.interleave_inference(input_list, **kargs)
|
||||
|
||||
for i in output_list:
|
||||
if isinstance(i, Image.Image):
|
||||
output_dict['image'] = i
|
||||
elif isinstance(i, str):
|
||||
output_dict['text'] = i
|
||||
return output_dict
|
||||
@@ -0,0 +1,4 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from . import bagel, qwen2, siglip, autoencoder
|
||||
@@ -0,0 +1,360 @@
|
||||
# Copyright (c) 2024 Black Forest Labs.
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# This file has been modified by ByteDance Ltd. and/or its affiliates. on 2025-05-20.
|
||||
#
|
||||
# Original file was released under Apache-2.0, with the full license text
|
||||
# available at https://github.com/black-forest-labs/flux/blob/main/LICENSE.
|
||||
#
|
||||
# This modified file is released under the same license.
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor, nn
|
||||
from safetensors.torch import load_file as load_sft
|
||||
|
||||
|
||||
@dataclass
|
||||
class AutoEncoderParams:
|
||||
resolution: int
|
||||
in_channels: int
|
||||
downsample: int
|
||||
ch: int
|
||||
out_ch: int
|
||||
ch_mult: list[int]
|
||||
num_res_blocks: int
|
||||
z_channels: int
|
||||
scale_factor: float
|
||||
shift_factor: float
|
||||
|
||||
|
||||
def swish(x: Tensor) -> Tensor:
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
class AttnBlock(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
|
||||
self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
|
||||
def attention(self, h_: Tensor) -> Tensor:
|
||||
h_ = self.norm(h_)
|
||||
q = self.q(h_)
|
||||
k = self.k(h_)
|
||||
v = self.v(h_)
|
||||
|
||||
b, c, h, w = q.shape
|
||||
q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
|
||||
k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
|
||||
v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
|
||||
h_ = nn.functional.scaled_dot_product_attention(q, k, v)
|
||||
|
||||
return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return x + self.proj_out(self.attention(x))
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
|
||||
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
if self.in_channels != self.out_channels:
|
||||
self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
h = x
|
||||
h = self.norm1(h)
|
||||
h = swish(h)
|
||||
h = self.conv1(h)
|
||||
|
||||
h = self.norm2(h)
|
||||
h = swish(h)
|
||||
h = self.conv2(h)
|
||||
|
||||
if self.in_channels != self.out_channels:
|
||||
x = self.nin_shortcut(x)
|
||||
|
||||
return x + h
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
# no asymmetric padding in torch conv, must do it ourselves
|
||||
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
pad = (0, 1, 0, 1)
|
||||
x = nn.functional.pad(x, pad, mode="constant", value=0)
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
resolution: int,
|
||||
in_channels: int,
|
||||
ch: int,
|
||||
ch_mult: list[int],
|
||||
num_res_blocks: int,
|
||||
z_channels: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
# downsampling
|
||||
self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
curr_res = resolution
|
||||
in_ch_mult = (1,) + tuple(ch_mult)
|
||||
self.in_ch_mult = in_ch_mult
|
||||
self.down = nn.ModuleList()
|
||||
block_in = self.ch
|
||||
for i_level in range(self.num_resolutions):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_in = ch * in_ch_mult[i_level]
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for _ in range(self.num_res_blocks):
|
||||
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
|
||||
block_in = block_out
|
||||
down = nn.Module()
|
||||
down.block = block
|
||||
down.attn = attn
|
||||
if i_level != self.num_resolutions - 1:
|
||||
down.downsample = Downsample(block_in)
|
||||
curr_res = curr_res // 2
|
||||
self.down.append(down)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
|
||||
# end
|
||||
self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
|
||||
self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
# downsampling
|
||||
hs = [self.conv_in(x)]
|
||||
for i_level in range(self.num_resolutions):
|
||||
for i_block in range(self.num_res_blocks):
|
||||
h = self.down[i_level].block[i_block](hs[-1])
|
||||
if len(self.down[i_level].attn) > 0:
|
||||
h = self.down[i_level].attn[i_block](h)
|
||||
hs.append(h)
|
||||
if i_level != self.num_resolutions - 1:
|
||||
hs.append(self.down[i_level].downsample(hs[-1]))
|
||||
|
||||
# middle
|
||||
h = hs[-1]
|
||||
h = self.mid.block_1(h)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h)
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = swish(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
ch: int,
|
||||
out_ch: int,
|
||||
ch_mult: list[int],
|
||||
num_res_blocks: int,
|
||||
in_channels: int,
|
||||
resolution: int,
|
||||
z_channels: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
self.ffactor = 2 ** (self.num_resolutions - 1)
|
||||
|
||||
# compute in_ch_mult, block_in and curr_res at lowest res
|
||||
block_in = ch * ch_mult[self.num_resolutions - 1]
|
||||
curr_res = resolution // 2 ** (self.num_resolutions - 1)
|
||||
self.z_shape = (1, z_channels, curr_res, curr_res)
|
||||
|
||||
# z to block_in
|
||||
self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
|
||||
# upsampling
|
||||
self.up = nn.ModuleList()
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for _ in range(self.num_res_blocks + 1):
|
||||
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
|
||||
block_in = block_out
|
||||
up = nn.Module()
|
||||
up.block = block
|
||||
up.attn = attn
|
||||
if i_level != 0:
|
||||
up.upsample = Upsample(block_in)
|
||||
curr_res = curr_res * 2
|
||||
self.up.insert(0, up) # prepend to get consistent order
|
||||
|
||||
# end
|
||||
self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
|
||||
self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, z: Tensor) -> Tensor:
|
||||
# z to block_in
|
||||
h = self.conv_in(z)
|
||||
|
||||
# middle
|
||||
h = self.mid.block_1(h)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h)
|
||||
|
||||
# upsampling
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
for i_block in range(self.num_res_blocks + 1):
|
||||
h = self.up[i_level].block[i_block](h)
|
||||
if len(self.up[i_level].attn) > 0:
|
||||
h = self.up[i_level].attn[i_block](h)
|
||||
if i_level != 0:
|
||||
h = self.up[i_level].upsample(h)
|
||||
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = swish(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class DiagonalGaussian(nn.Module):
|
||||
def __init__(self, sample: bool = True, chunk_dim: int = 1):
|
||||
super().__init__()
|
||||
self.sample = sample
|
||||
self.chunk_dim = chunk_dim
|
||||
|
||||
def forward(self, z: Tensor) -> Tensor:
|
||||
mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
|
||||
if self.sample:
|
||||
std = torch.exp(0.5 * logvar)
|
||||
return mean + std * torch.randn_like(mean)
|
||||
else:
|
||||
return mean
|
||||
|
||||
|
||||
class AutoEncoder(nn.Module):
|
||||
def __init__(self, params: AutoEncoderParams):
|
||||
super().__init__()
|
||||
self.encoder = Encoder(
|
||||
resolution=params.resolution,
|
||||
in_channels=params.in_channels,
|
||||
ch=params.ch,
|
||||
ch_mult=params.ch_mult,
|
||||
num_res_blocks=params.num_res_blocks,
|
||||
z_channels=params.z_channels,
|
||||
)
|
||||
self.decoder = Decoder(
|
||||
resolution=params.resolution,
|
||||
in_channels=params.in_channels,
|
||||
ch=params.ch,
|
||||
out_ch=params.out_ch,
|
||||
ch_mult=params.ch_mult,
|
||||
num_res_blocks=params.num_res_blocks,
|
||||
z_channels=params.z_channels,
|
||||
)
|
||||
self.reg = DiagonalGaussian()
|
||||
|
||||
self.scale_factor = params.scale_factor
|
||||
self.shift_factor = params.shift_factor
|
||||
|
||||
def encode(self, x: Tensor) -> Tensor:
|
||||
z = self.reg(self.encoder(x))
|
||||
z = self.scale_factor * (z - self.shift_factor)
|
||||
return z
|
||||
|
||||
def decode(self, z: Tensor) -> Tensor:
|
||||
z = z / self.scale_factor + self.shift_factor
|
||||
return self.decoder(z)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.decode(self.encode(x))
|
||||
|
||||
|
||||
def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
|
||||
if len(missing) > 0 and len(unexpected) > 0:
|
||||
print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
|
||||
print("\n" + "-" * 79 + "\n")
|
||||
print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
|
||||
elif len(missing) > 0:
|
||||
print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
|
||||
elif len(unexpected) > 0:
|
||||
print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
|
||||
|
||||
|
||||
def load_ae(local_path: str) -> AutoEncoder:
|
||||
ae_params = AutoEncoderParams(
|
||||
resolution=256,
|
||||
in_channels=3,
|
||||
downsample=8,
|
||||
ch=128,
|
||||
out_ch=3,
|
||||
ch_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
z_channels=16,
|
||||
scale_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
)
|
||||
|
||||
# Loading the autoencoder
|
||||
ae = AutoEncoder(ae_params)
|
||||
|
||||
if local_path is not None:
|
||||
sd = load_sft(local_path)
|
||||
missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
|
||||
print_load_warning(missing, unexpected)
|
||||
return ae, ae_params
|
||||
@@ -0,0 +1,18 @@
|
||||
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
|
||||
from .bagel import BagelConfig, Bagel
|
||||
from .qwen2_navit import Qwen2Config, Qwen2Model, Qwen2ForCausalLM
|
||||
from .siglip_navit import SiglipVisionConfig, SiglipVisionModel
|
||||
|
||||
|
||||
__all__ = [
|
||||
'BagelConfig',
|
||||
'Bagel',
|
||||
'Qwen2Config',
|
||||
'Qwen2Model',
|
||||
'Qwen2ForCausalLM',
|
||||
'SiglipVisionConfig',
|
||||
'SiglipVisionModel',
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,144 @@
|
||||
# Copyright (c) 2022 Facebook, Inc. and its affiliates.
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: CC BY-NC 4.0
|
||||
#
|
||||
# This file has been modified by ByteDance Ltd. and/or its affiliates. on 2025-05-20.
|
||||
#
|
||||
# Original file was released under CC BY-NC 4.0, with the full license text
|
||||
# available at https://github.com/facebookresearch/DiT/blob/main/LICENSE.txt.
|
||||
#
|
||||
# This modified file is released under the same license.
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers.activations import ACT2FN
|
||||
|
||||
# --------------------------------------------------------
|
||||
# 2D sine-cosine position embedding
|
||||
# References:
|
||||
# DiT: https://github.com/facebookresearch/DiT/blob/main/models.py
|
||||
# --------------------------------------------------------
|
||||
def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
|
||||
grid_h = np.arange(grid_size, dtype=np.float32)
|
||||
grid_w = np.arange(grid_size, dtype=np.float32)
|
||||
grid = np.meshgrid(grid_w, grid_h) # here w goes first
|
||||
grid = np.stack(grid, axis=0)
|
||||
|
||||
grid = grid.reshape([2, 1, grid_size, grid_size])
|
||||
pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
|
||||
if cls_token and extra_tokens > 0:
|
||||
pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
|
||||
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(embed_dim // 2, grid[0]) # (H*W, D/2)
|
||||
emb_w = get_1d_sincos_pos_embed_from_grid(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(embed_dim, pos):
|
||||
"""
|
||||
embed_dim: output dimension for each position
|
||||
pos: a list of positions to be encoded: size (M,)
|
||||
out: (M, D)
|
||||
"""
|
||||
assert embed_dim % 2 == 0
|
||||
omega = np.arange(embed_dim // 2, dtype=np.float64)
|
||||
omega /= embed_dim / 2.
|
||||
omega = 1. / 10000**omega # (D/2,)
|
||||
|
||||
pos = pos.reshape(-1) # (M,)
|
||||
out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product
|
||||
|
||||
emb_sin = np.sin(out) # (M, D/2)
|
||||
emb_cos = np.cos(out) # (M, D/2)
|
||||
|
||||
emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
|
||||
return emb
|
||||
|
||||
|
||||
# --------------------------------------------------------
|
||||
# TimestepEmbedder
|
||||
# Reference:
|
||||
# DiT: https://github.com/facebookresearch/DiT/blob/main/models.py
|
||||
# --------------------------------------------------------
|
||||
class TimestepEmbedder(nn.Module):
|
||||
"""
|
||||
Embeds scalar timesteps into vector representations.
|
||||
"""
|
||||
def __init__(self, hidden_size, frequency_embedding_size=256):
|
||||
super().__init__()
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||||
)
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
|
||||
@staticmethod
|
||||
def timestep_embedding(t, dim, max_period=10000):
|
||||
"""
|
||||
Create sinusoidal timestep embeddings.
|
||||
:param t: a 1-D Tensor of N indices, one per batch element.
|
||||
These may be fractional.
|
||||
:param dim: the dimension of the output.
|
||||
:param max_period: controls the minimum frequency of the embeddings.
|
||||
:return: an (N, D) Tensor of positional embeddings.
|
||||
"""
|
||||
half = dim // 2
|
||||
freqs = torch.exp(
|
||||
-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
|
||||
).to(device=t.device)
|
||||
args = t[:, None].float() * freqs[None]
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
||||
return embedding
|
||||
|
||||
def forward(self, t):
|
||||
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
|
||||
t_emb = self.mlp(t_freq)
|
||||
return t_emb
|
||||
|
||||
|
||||
class MLPconnector(nn.Module):
|
||||
def __init__(self, in_dim: int, out_dim: int, hidden_act: str):
|
||||
super().__init__()
|
||||
self.activation_fn = ACT2FN[hidden_act]
|
||||
self.fc1 = nn.Linear(in_dim, out_dim)
|
||||
self.fc2 = nn.Linear(out_dim, out_dim)
|
||||
|
||||
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
|
||||
|
||||
|
||||
class PositionEmbedding(nn.Module):
|
||||
def __init__(self, max_num_patch_per_side, hidden_size):
|
||||
super().__init__()
|
||||
self.max_num_patch_per_side = max_num_patch_per_side
|
||||
self.hidden_size = hidden_size
|
||||
self.pos_embed = nn.Parameter(
|
||||
torch.zeros(max_num_patch_per_side ** 2, hidden_size),
|
||||
requires_grad=False
|
||||
)
|
||||
self._init_weights()
|
||||
|
||||
def _init_weights(self):
|
||||
# Initialize (and freeze) pos_embed by sin-cos embedding:
|
||||
pos_embed = get_2d_sincos_pos_embed(self.hidden_size, self.max_num_patch_per_side)
|
||||
self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float())
|
||||
|
||||
def forward(self, position_ids):
|
||||
return self.pos_embed[position_ids]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,402 @@
|
||||
# Copyright (c) 2024 The HuggingFace Inc. team.
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# This file has been modified by ByteDance Ltd. and/or its affiliates. on 2025-05-20.
|
||||
#
|
||||
# Original file was released under Apache-2.0, with the full license text
|
||||
# available at https://github.com/huggingface/transformers/blob/main/LICENSE.
|
||||
#
|
||||
# This modified file is released under the same license.
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from transformers.activations import ACT2FN
|
||||
from ...modeling.siglip.configuration_siglip import SiglipVisionConfig as _SiglipVisionConfig
|
||||
from ...modeling.siglip.modeling_siglip import SiglipAttention, SiglipPreTrainedModel
|
||||
from flash_attn import flash_attn_varlen_func
|
||||
|
||||
|
||||
class SiglipVisionConfig(_SiglipVisionConfig):
|
||||
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,
|
||||
rope=True,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
num_hidden_layers=num_hidden_layers,
|
||||
num_attention_heads=num_attention_heads,
|
||||
num_channels=num_channels,
|
||||
image_size=image_size,
|
||||
patch_size=patch_size,
|
||||
hidden_act=hidden_act,
|
||||
layer_norm_eps=layer_norm_eps,
|
||||
attention_dropout=attention_dropout,
|
||||
**kwargs)
|
||||
|
||||
self.rope = rope
|
||||
|
||||
|
||||
class RotaryEmbedding2D(torch.nn.Module):
|
||||
def __init__(self, dim, max_h, max_w, base=10000):
|
||||
super().__init__()
|
||||
freq = torch.arange(0, dim, 2, dtype=torch.int64).float() / dim
|
||||
inv_freq = 1.0 / (base ** freq)
|
||||
|
||||
grid_h = torch.arange(0, max_h)
|
||||
grid_h = grid_h.to(inv_freq.dtype)
|
||||
grid_h = grid_h[:, None].repeat(1, max_w)
|
||||
|
||||
grid_w = torch.arange(0, max_w)
|
||||
grid_w = grid_w.to(inv_freq.dtype)
|
||||
grid_w = grid_w[None, :].repeat(max_h, 1)
|
||||
|
||||
cos_h, sin_h = self._forward_one_side(grid_h, inv_freq)
|
||||
cos_w, sin_w = self._forward_one_side(grid_w, inv_freq)
|
||||
|
||||
self.register_buffer("cos_h", cos_h)
|
||||
self.register_buffer("sin_h", sin_h)
|
||||
self.register_buffer("cos_w", cos_w)
|
||||
self.register_buffer("sin_w", sin_w)
|
||||
|
||||
def _forward_one_side(self, grid, inv_freq):
|
||||
freqs = grid[..., None] * inv_freq[None, None, :]
|
||||
emb = torch.cat((freqs, freqs), dim=-1).flatten(0, 1)
|
||||
return emb.cos(), emb.sin()
|
||||
|
||||
|
||||
def rotate_half(x):
|
||||
x1 = x[..., : x.shape[-1] // 2]
|
||||
x2 = x[..., x.shape[-1] // 2 :]
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def apply_rotary_pos_emb(q, k, cos, sin):
|
||||
# unsqueeze due to the head dimension
|
||||
cos = cos.unsqueeze(1)
|
||||
sin = sin.unsqueeze(1)
|
||||
q_embed = (q * cos) + (rotate_half(q) * sin)
|
||||
k_embed = (k * cos) + (rotate_half(k) * sin)
|
||||
return q_embed, k_embed
|
||||
|
||||
|
||||
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
|
||||
if not config.rope:
|
||||
self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
|
||||
|
||||
def convert_conv2d_to_linear(self, config, meta=False):
|
||||
if meta:
|
||||
linear_patch_embedding = nn.Linear(
|
||||
config.num_channels * self.patch_size ** 2, self.embed_dim, bias=True, device='meta'
|
||||
)
|
||||
else:
|
||||
linear_patch_embedding = nn.Linear(
|
||||
config.num_channels * self.patch_size ** 2, self.embed_dim, bias=True
|
||||
)
|
||||
W = self.patch_embedding.weight.permute(0, 2, 3, 1).reshape(
|
||||
self.embed_dim, config.num_channels * self.patch_size ** 2
|
||||
)
|
||||
linear_patch_embedding.weight.data = W
|
||||
linear_patch_embedding.bias.data = self.patch_embedding.bias.data
|
||||
del self.patch_embedding
|
||||
self.patch_embedding = linear_patch_embedding
|
||||
|
||||
def forward(
|
||||
self,
|
||||
packed_pixel_values: torch.FloatTensor,
|
||||
packed_flattened_position_ids: torch.LongTensor
|
||||
) -> torch.Tensor:
|
||||
|
||||
patch_embeds = self.patch_embedding(packed_pixel_values)
|
||||
if not self.config.rope:
|
||||
embeddings = patch_embeds + self.position_embedding(packed_flattened_position_ids)
|
||||
else:
|
||||
embeddings = patch_embeds
|
||||
return embeddings
|
||||
|
||||
|
||||
class SiglipFlashAttention2(SiglipAttention):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
cu_seqlens: torch.IntTensor,
|
||||
max_seqlen: int,
|
||||
cos_h: torch.Tensor = None,
|
||||
sin_h: torch.Tensor = None,
|
||||
cos_w: torch.Tensor = None,
|
||||
sin_w: torch.Tensor = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
|
||||
total_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(total_q_len, self.num_heads, self.head_dim)
|
||||
key_states = key_states.view(total_q_len, self.num_heads, self.head_dim)
|
||||
value_states = value_states.view(total_q_len, self.num_heads, self.head_dim)
|
||||
|
||||
if self.config.rope:
|
||||
qh, qw = query_states[:, :, :self.head_dim // 2], query_states[:, :, self.head_dim // 2:]
|
||||
kh, kw = key_states[:, :, :self.head_dim // 2], key_states[:, :, self.head_dim // 2:]
|
||||
qh, kh = apply_rotary_pos_emb(qh, kh, cos_h, sin_h)
|
||||
qw, kw = apply_rotary_pos_emb(qw, kw, cos_w, sin_w)
|
||||
query_states = torch.cat([qh, qw], dim=-1)
|
||||
key_states = torch.cat([kh, kw], dim=-1)
|
||||
|
||||
attn_output = flash_attn_varlen_func(
|
||||
query_states.to(torch.bfloat16),
|
||||
key_states.to(torch.bfloat16),
|
||||
value_states.to(torch.bfloat16),
|
||||
cu_seqlens_q=cu_seqlens,
|
||||
cu_seqlens_k=cu_seqlens,
|
||||
max_seqlen_q=max_seqlen,
|
||||
max_seqlen_k=max_seqlen,
|
||||
causal=False,
|
||||
)
|
||||
|
||||
attn_output = self.out_proj(attn_output.reshape(total_q_len, -1))
|
||||
return attn_output
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
class SiglipEncoderLayer(nn.Module):
|
||||
def __init__(self, config: SiglipVisionConfig):
|
||||
super().__init__()
|
||||
self.embed_dim = config.hidden_size
|
||||
self.self_attn = 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,
|
||||
cu_seqlens: torch.IntTensor,
|
||||
max_seqlen: int,
|
||||
cos_h: torch.Tensor = None,
|
||||
sin_h: torch.Tensor = None,
|
||||
cos_w: torch.Tensor = None,
|
||||
sin_w: torch.Tensor = None
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
|
||||
hidden_states = self.layer_norm1(hidden_states)
|
||||
hidden_states = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
cos_h=cos_h,
|
||||
sin_h=sin_h,
|
||||
cos_w=cos_w,
|
||||
sin_w=sin_w
|
||||
)
|
||||
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
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class SiglipEncoder(nn.Module):
|
||||
def __init__(self, config: SiglipVisionConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.layers = nn.ModuleList(
|
||||
[SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)]
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
inputs_embeds: torch.Tensor,
|
||||
cu_seqlens: torch.IntTensor,
|
||||
max_seqlen: int,
|
||||
cos_h: torch.Tensor = None,
|
||||
sin_h: torch.Tensor = None,
|
||||
cos_w: torch.Tensor = None,
|
||||
sin_w: torch.Tensor = None,
|
||||
) -> torch.Tensor:
|
||||
|
||||
hidden_states = inputs_embeds
|
||||
for encoder_layer in self.layers:
|
||||
hidden_states = encoder_layer(hidden_states, cu_seqlens, max_seqlen,
|
||||
cos_h=cos_h, sin_h=sin_h, cos_w=cos_w, sin_w=sin_w)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class SiglipVisionTransformer(nn.Module):
|
||||
def __init__(self, config: SiglipVisionConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
embed_dim = config.hidden_size
|
||||
|
||||
self.embeddings = SiglipVisionEmbeddings(config)
|
||||
if config.rope:
|
||||
max_size = config.image_size // config.patch_size
|
||||
dim_head = config.hidden_size // config.num_attention_heads
|
||||
self.rope = RotaryEmbedding2D(dim_head // 2, max_size, max_size)
|
||||
|
||||
self.encoder = SiglipEncoder(config)
|
||||
self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
packed_pixel_values: torch.Tensor,
|
||||
packed_flattened_position_ids: torch.LongTensor,
|
||||
cu_seqlens: torch.IntTensor,
|
||||
max_seqlen: int,
|
||||
) -> torch.Tensor:
|
||||
hidden_states = self.embeddings(
|
||||
packed_pixel_values=packed_pixel_values,
|
||||
packed_flattened_position_ids=packed_flattened_position_ids
|
||||
)
|
||||
|
||||
extra_inputs = {}
|
||||
if self.config.rope:
|
||||
extra_inputs.update(
|
||||
cos_h = self.rope.cos_h[packed_flattened_position_ids],
|
||||
sin_h = self.rope.sin_h[packed_flattened_position_ids],
|
||||
cos_w = self.rope.cos_w[packed_flattened_position_ids],
|
||||
sin_w = self.rope.sin_w[packed_flattened_position_ids]
|
||||
)
|
||||
|
||||
last_hidden_state = self.encoder(
|
||||
inputs_embeds=hidden_states, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
|
||||
**extra_inputs
|
||||
)
|
||||
last_hidden_state = self.post_layernorm(last_hidden_state)
|
||||
return last_hidden_state
|
||||
|
||||
|
||||
class SiglipVisionModel(SiglipPreTrainedModel):
|
||||
config_class = SiglipVisionConfig
|
||||
main_input_name = "packed_pixel_values"
|
||||
|
||||
def __init__(self, config: SiglipVisionConfig):
|
||||
super().__init__(config)
|
||||
|
||||
self.vision_model = SiglipVisionTransformer(config)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def get_input_embeddings(self) -> nn.Module:
|
||||
return self.vision_model.embeddings.patch_embedding
|
||||
|
||||
def forward(
|
||||
self,
|
||||
packed_pixel_values: torch.Tensor,
|
||||
packed_flattened_position_ids: torch.LongTensor,
|
||||
cu_seqlens: torch.IntTensor,
|
||||
max_seqlen: int,
|
||||
) -> torch.Tensor:
|
||||
|
||||
return self.vision_model(
|
||||
packed_pixel_values=packed_pixel_values,
|
||||
packed_flattened_position_ids=packed_flattened_position_ids,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
)
|
||||
@@ -0,0 +1,68 @@
|
||||
# Copyright 2024 The Qwen Team and The HuggingFace Inc. team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from transformers.utils import (
|
||||
OptionalDependencyNotAvailable,
|
||||
_LazyModule,
|
||||
is_tokenizers_available,
|
||||
is_torch_available,
|
||||
)
|
||||
|
||||
|
||||
_import_structure = {
|
||||
"configuration_qwen2": ["Qwen2Config"],
|
||||
"tokenization_qwen2": ["Qwen2Tokenizer"],
|
||||
}
|
||||
|
||||
try:
|
||||
if not is_tokenizers_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
_import_structure["tokenization_qwen2_fast"] = ["Qwen2TokenizerFast"]
|
||||
|
||||
try:
|
||||
if not is_torch_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
_import_structure["modeling_qwen2"] = [
|
||||
"Qwen2ForCausalLM",
|
||||
"Qwen2Model",
|
||||
"Qwen2PreTrainedModel",
|
||||
]
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .configuration_qwen2 import Qwen2Config
|
||||
from .tokenization_qwen2 import Qwen2Tokenizer
|
||||
|
||||
try:
|
||||
if not is_tokenizers_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
from .tokenization_qwen2_fast import Qwen2TokenizerFast
|
||||
|
||||
try:
|
||||
if not is_torch_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
from .modeling_qwen2 import (
|
||||
Qwen2ForCausalLM,
|
||||
Qwen2Model,
|
||||
Qwen2PreTrainedModel,
|
||||
)
|
||||
|
||||
|
||||
else:
|
||||
import sys
|
||||
|
||||
sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
|
||||
@@ -0,0 +1,179 @@
|
||||
# Copyright 2024 The Qwen Team and The HuggingFace Inc. team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Qwen2 model configuration"""
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.modeling_rope_utils import rope_config_validation
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
class Qwen2Config(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`Qwen2Model`]. It is used to instantiate a
|
||||
Qwen2 model according to the specified arguments, defining the model architecture. Instantiating a configuration
|
||||
with the defaults will yield a similar configuration to that of
|
||||
Qwen2-7B-beta [Qwen/Qwen2-7B-beta](https://huggingface.co/Qwen/Qwen2-7B-beta).
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
|
||||
Args:
|
||||
vocab_size (`int`, *optional*, defaults to 151936):
|
||||
Vocabulary size of the Qwen2 model. Defines the number of different tokens that can be represented by the
|
||||
`inputs_ids` passed when calling [`Qwen2Model`]
|
||||
hidden_size (`int`, *optional*, defaults to 4096):
|
||||
Dimension of the hidden representations.
|
||||
intermediate_size (`int`, *optional*, defaults to 22016):
|
||||
Dimension of the MLP representations.
|
||||
num_hidden_layers (`int`, *optional*, defaults to 32):
|
||||
Number of hidden layers in the Transformer encoder.
|
||||
num_attention_heads (`int`, *optional*, defaults to 32):
|
||||
Number of attention heads for each attention layer in the Transformer encoder.
|
||||
num_key_value_heads (`int`, *optional*, defaults to 32):
|
||||
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
|
||||
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
|
||||
`num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
|
||||
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
|
||||
by meanpooling all the original heads within that group. For more details checkout [this
|
||||
paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
|
||||
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
|
||||
The non-linear activation function (function or string) in the decoder.
|
||||
max_position_embeddings (`int`, *optional*, defaults to 32768):
|
||||
The maximum sequence length that this model might ever be used with.
|
||||
initializer_range (`float`, *optional*, defaults to 0.02):
|
||||
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||
rms_norm_eps (`float`, *optional*, defaults to 1e-06):
|
||||
The epsilon used by the rms normalization layers.
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models). Only
|
||||
relevant if `config.is_decoder=True`.
|
||||
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
|
||||
Whether the model's input and output word embeddings should be tied.
|
||||
rope_theta (`float`, *optional*, defaults to 10000.0):
|
||||
The base period of the RoPE embeddings.
|
||||
rope_scaling (`Dict`, *optional*):
|
||||
Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
|
||||
and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
|
||||
accordingly.
|
||||
Expected contents:
|
||||
`rope_type` (`str`):
|
||||
The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
|
||||
'llama3'], with 'default' being the original RoPE implementation.
|
||||
`factor` (`float`, *optional*):
|
||||
Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
|
||||
most scaling types, a `factor` of x will enable the model to handle sequences of length x *
|
||||
original maximum pre-trained length.
|
||||
`original_max_position_embeddings` (`int`, *optional*):
|
||||
Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
|
||||
pretraining.
|
||||
`attention_factor` (`float`, *optional*):
|
||||
Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
|
||||
computation. If unspecified, it defaults to value recommended by the implementation, using the
|
||||
`factor` field to infer the suggested value.
|
||||
`beta_fast` (`float`, *optional*):
|
||||
Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
|
||||
ramp function. If unspecified, it defaults to 32.
|
||||
`beta_slow` (`float`, *optional*):
|
||||
Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
|
||||
ramp function. If unspecified, it defaults to 1.
|
||||
`short_factor` (`List[float]`, *optional*):
|
||||
Only used with 'longrope'. The scaling factor to be applied to short contexts (<
|
||||
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||
size divided by the number of attention heads divided by 2
|
||||
`long_factor` (`List[float]`, *optional*):
|
||||
Only used with 'longrope'. The scaling factor to be applied to long contexts (<
|
||||
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||
size divided by the number of attention heads divided by 2
|
||||
`low_freq_factor` (`float`, *optional*):
|
||||
Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
|
||||
`high_freq_factor` (`float`, *optional*):
|
||||
Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
|
||||
use_sliding_window (`bool`, *optional*, defaults to `False`):
|
||||
Whether to use sliding window attention.
|
||||
sliding_window (`int`, *optional*, defaults to 4096):
|
||||
Sliding window attention (SWA) window size. If not specified, will default to `4096`.
|
||||
max_window_layers (`int`, *optional*, defaults to 28):
|
||||
The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
|
||||
attention_dropout (`float`, *optional*, defaults to 0.0):
|
||||
The dropout ratio for the attention probabilities.
|
||||
|
||||
```python
|
||||
>>> from transformers import Qwen2Model, Qwen2Config
|
||||
|
||||
>>> # Initializing a Qwen2 style configuration
|
||||
>>> configuration = Qwen2Config()
|
||||
|
||||
>>> # Initializing a model from the Qwen2-7B style configuration
|
||||
>>> model = Qwen2Model(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "qwen2"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=151936,
|
||||
hidden_size=4096,
|
||||
intermediate_size=22016,
|
||||
num_hidden_layers=32,
|
||||
num_attention_heads=32,
|
||||
num_key_value_heads=32,
|
||||
hidden_act="silu",
|
||||
max_position_embeddings=32768,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-6,
|
||||
use_cache=True,
|
||||
tie_word_embeddings=False,
|
||||
rope_theta=10000.0,
|
||||
rope_scaling=None,
|
||||
use_sliding_window=False,
|
||||
sliding_window=4096,
|
||||
max_window_layers=28,
|
||||
attention_dropout=0.0,
|
||||
is_causal=True,
|
||||
_attn_implementation="flash_attention_2",
|
||||
**kwargs,
|
||||
):
|
||||
self.vocab_size = vocab_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.use_sliding_window = use_sliding_window
|
||||
self.sliding_window = sliding_window if use_sliding_window else None
|
||||
self.max_window_layers = max_window_layers
|
||||
|
||||
# for backward compatibility
|
||||
if num_key_value_heads is None:
|
||||
num_key_value_heads = num_attention_heads
|
||||
|
||||
self.num_key_value_heads = num_key_value_heads
|
||||
self.hidden_act = hidden_act
|
||||
self.initializer_range = initializer_range
|
||||
self.rms_norm_eps = rms_norm_eps
|
||||
self.use_cache = use_cache
|
||||
self.rope_theta = rope_theta
|
||||
self.rope_scaling = rope_scaling
|
||||
self.attention_dropout = attention_dropout
|
||||
self.is_causal = is_causal
|
||||
self._attn_implementation = _attn_implementation
|
||||
|
||||
# Validate the correctness of rotary position embeddings parameters
|
||||
# BC: if there is a 'type' field, move it to 'rope_type'.
|
||||
if self.rope_scaling is not None and "type" in self.rope_scaling:
|
||||
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
|
||||
rope_config_validation(self)
|
||||
|
||||
super().__init__(
|
||||
tie_word_embeddings=tie_word_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,929 @@
|
||||
# Copyright 2024 The Qwen Team and The HuggingFace Inc. team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""PyTorch Qwen2 model."""
|
||||
|
||||
import math
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from torch import nn
|
||||
|
||||
from transformers.activations import ACT2FN
|
||||
from transformers.cache_utils import Cache, DynamicCache
|
||||
from transformers.generation import GenerationMixin
|
||||
from transformers.modeling_outputs import (
|
||||
BaseModelOutputWithPast,
|
||||
CausalLMOutputWithPast,
|
||||
)
|
||||
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
|
||||
from transformers.modeling_utils import PreTrainedModel
|
||||
from transformers.utils import (
|
||||
add_start_docstrings,
|
||||
add_start_docstrings_to_model_forward,
|
||||
is_flash_attn_2_available,
|
||||
is_flash_attn_greater_or_equal_2_10,
|
||||
logging,
|
||||
replace_return_docstrings,
|
||||
)
|
||||
from .configuration_qwen2 import Qwen2Config
|
||||
|
||||
|
||||
if is_flash_attn_2_available():
|
||||
from transformers.modeling_flash_attention_utils import _flash_attention_forward
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
_CHECKPOINT_FOR_DOC = "Qwen/Qwen2-7B"
|
||||
_CONFIG_FOR_DOC = "Qwen2Config"
|
||||
|
||||
|
||||
# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Qwen2
|
||||
class Qwen2RMSNorm(nn.Module):
|
||||
def __init__(self, hidden_size, eps=1e-6):
|
||||
"""
|
||||
Qwen2RMSNorm is equivalent to T5LayerNorm
|
||||
"""
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
self.variance_epsilon = eps
|
||||
|
||||
def forward(self, hidden_states):
|
||||
input_dtype = hidden_states.dtype
|
||||
hidden_states = hidden_states.to(torch.float32)
|
||||
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
||||
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
||||
return self.weight * hidden_states.to(input_dtype)
|
||||
|
||||
def extra_repr(self):
|
||||
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
|
||||
|
||||
|
||||
# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Qwen2
|
||||
class Qwen2RotaryEmbedding(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim=None,
|
||||
max_position_embeddings=2048,
|
||||
base=10000,
|
||||
device=None,
|
||||
scaling_factor=1.0,
|
||||
rope_type="default",
|
||||
config: Optional[Qwen2Config] = None,
|
||||
):
|
||||
super().__init__()
|
||||
# TODO (joao): remove the `if` below, only used for BC
|
||||
self.rope_kwargs = {}
|
||||
if config is None:
|
||||
logger.warning_once(
|
||||
"`Qwen2RotaryEmbedding` can now be fully parameterized by passing the model config through the "
|
||||
"`config` argument. All other arguments will be removed in v4.46"
|
||||
)
|
||||
self.rope_kwargs = {
|
||||
"rope_type": rope_type,
|
||||
"factor": scaling_factor,
|
||||
"dim": dim,
|
||||
"base": base,
|
||||
"max_position_embeddings": max_position_embeddings,
|
||||
}
|
||||
self.rope_type = rope_type
|
||||
self.max_seq_len_cached = max_position_embeddings
|
||||
self.original_max_seq_len = max_position_embeddings
|
||||
else:
|
||||
# BC: "rope_type" was originally "type"
|
||||
if config.rope_scaling is not None:
|
||||
self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
|
||||
else:
|
||||
self.rope_type = "default"
|
||||
self.max_seq_len_cached = config.max_position_embeddings
|
||||
self.original_max_seq_len = config.max_position_embeddings
|
||||
|
||||
self.config = config
|
||||
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
|
||||
|
||||
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, **self.rope_kwargs)
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
self.original_inv_freq = self.inv_freq
|
||||
|
||||
def _dynamic_frequency_update(self, position_ids, device):
|
||||
"""
|
||||
dynamic RoPE layers should recompute `inv_freq` in the following situations:
|
||||
1 - growing beyond the cached sequence length (allow scaling)
|
||||
2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
|
||||
"""
|
||||
seq_len = torch.max(position_ids) + 1
|
||||
if seq_len > self.max_seq_len_cached: # growth
|
||||
inv_freq, self.attention_scaling = self.rope_init_fn(
|
||||
self.config, device, seq_len=seq_len, **self.rope_kwargs
|
||||
)
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation
|
||||
self.max_seq_len_cached = seq_len
|
||||
|
||||
if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset
|
||||
self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
|
||||
self.max_seq_len_cached = self.original_max_seq_len
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, x, position_ids):
|
||||
if "dynamic" in self.rope_type:
|
||||
self._dynamic_frequency_update(position_ids, device=x.device)
|
||||
|
||||
# Core RoPE block
|
||||
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
|
||||
position_ids_expanded = position_ids[:, None, :].float()
|
||||
# Force float32 (see https://github.com/huggingface/transformers/pull/29285)
|
||||
device_type = x.device.type
|
||||
device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
|
||||
with torch.autocast(device_type=device_type, enabled=False):
|
||||
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
cos = emb.cos()
|
||||
sin = emb.sin()
|
||||
|
||||
# Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
|
||||
cos = cos * self.attention_scaling
|
||||
sin = sin * self.attention_scaling
|
||||
|
||||
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
||||
|
||||
|
||||
# Copied from transformers.models.llama.modeling_llama.rotate_half
|
||||
def rotate_half(x):
|
||||
"""Rotates half the hidden dims of the input."""
|
||||
x1 = x[..., : x.shape[-1] // 2]
|
||||
x2 = x[..., x.shape[-1] // 2 :]
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
|
||||
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
|
||||
"""Applies Rotary Position Embedding to the query and key tensors.
|
||||
|
||||
Args:
|
||||
q (`torch.Tensor`): The query tensor.
|
||||
k (`torch.Tensor`): The key tensor.
|
||||
cos (`torch.Tensor`): The cosine part of the rotary embedding.
|
||||
sin (`torch.Tensor`): The sine part of the rotary embedding.
|
||||
position_ids (`torch.Tensor`, *optional*):
|
||||
Deprecated and unused.
|
||||
unsqueeze_dim (`int`, *optional*, defaults to 1):
|
||||
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
|
||||
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
|
||||
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
|
||||
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
|
||||
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
|
||||
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
|
||||
Returns:
|
||||
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
|
||||
"""
|
||||
cos = cos.unsqueeze(unsqueeze_dim)
|
||||
sin = sin.unsqueeze(unsqueeze_dim)
|
||||
q_embed = (q * cos) + (rotate_half(q) * sin)
|
||||
k_embed = (k * cos) + (rotate_half(k) * sin)
|
||||
return q_embed, k_embed
|
||||
|
||||
|
||||
# Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Qwen2
|
||||
class Qwen2MLP(nn.Module):
|
||||
def __init__(self, config):
|
||||
super().__init__()
|
||||
self.hidden_size = config.hidden_size
|
||||
self.intermediate_size = config.intermediate_size
|
||||
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
||||
self.act_fn = ACT2FN[config.hidden_act]
|
||||
|
||||
def forward(self, hidden_state):
|
||||
return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
|
||||
|
||||
|
||||
# Copied from transformers.models.llama.modeling_llama.repeat_kv
|
||||
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||
"""
|
||||
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
||||
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
||||
"""
|
||||
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
||||
if n_rep == 1:
|
||||
return hidden_states
|
||||
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
||||
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
||||
|
||||
|
||||
class Qwen2Attention(nn.Module):
|
||||
"""
|
||||
Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
|
||||
and "Generating Long Sequences with Sparse Transformers".
|
||||
"""
|
||||
|
||||
def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.layer_idx = layer_idx
|
||||
if layer_idx is None:
|
||||
logger.warning_once(
|
||||
f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
|
||||
"to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
|
||||
"when creating this class."
|
||||
)
|
||||
|
||||
self.hidden_size = config.hidden_size
|
||||
self.num_heads = config.num_attention_heads
|
||||
self.head_dim = self.hidden_size // self.num_heads
|
||||
self.num_key_value_heads = config.num_key_value_heads
|
||||
self.num_key_value_groups = self.num_heads // self.num_key_value_heads
|
||||
self.max_position_embeddings = config.max_position_embeddings
|
||||
self.rope_theta = config.rope_theta
|
||||
self.is_causal = config.is_causal
|
||||
self.attention_dropout = config.attention_dropout
|
||||
|
||||
if (self.head_dim * self.num_heads) != self.hidden_size:
|
||||
raise ValueError(
|
||||
f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
|
||||
f" and `num_heads`: {self.num_heads})."
|
||||
)
|
||||
self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
|
||||
self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
|
||||
self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
|
||||
self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_value: Optional[Cache] = None,
|
||||
output_attentions: bool = False,
|
||||
use_cache: bool = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46
|
||||
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
|
||||
bsz, q_len, _ = hidden_states.size()
|
||||
|
||||
query_states = self.q_proj(hidden_states)
|
||||
key_states = self.k_proj(hidden_states)
|
||||
value_states = self.v_proj(hidden_states)
|
||||
|
||||
query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
|
||||
value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
|
||||
|
||||
if position_embeddings is None:
|
||||
logger.warning_once(
|
||||
"The attention layers in this model are transitioning from computing the RoPE embeddings internally "
|
||||
"through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
|
||||
"`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
|
||||
"removed and `position_embeddings` will be mandatory."
|
||||
)
|
||||
cos, sin = self.rotary_emb(value_states, position_ids)
|
||||
else:
|
||||
cos, sin = position_embeddings
|
||||
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
||||
|
||||
if past_key_value is not None:
|
||||
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models
|
||||
key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
|
||||
|
||||
# repeat k/v heads if n_kv_heads < n_heads
|
||||
key_states = repeat_kv(key_states, self.num_key_value_groups)
|
||||
value_states = repeat_kv(value_states, self.num_key_value_groups)
|
||||
|
||||
attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
|
||||
if attention_mask is not None: # no matter the length, we just slice it
|
||||
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
|
||||
attn_weights = attn_weights + causal_mask
|
||||
|
||||
# upcast attention to fp32
|
||||
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
|
||||
attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
|
||||
attn_output = torch.matmul(attn_weights, value_states)
|
||||
|
||||
if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
|
||||
raise ValueError(
|
||||
f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
|
||||
f" {attn_output.size()}"
|
||||
)
|
||||
|
||||
attn_output = attn_output.transpose(1, 2).contiguous()
|
||||
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
|
||||
|
||||
attn_output = self.o_proj(attn_output)
|
||||
|
||||
if not output_attentions:
|
||||
attn_weights = None
|
||||
|
||||
return attn_output, attn_weights, past_key_value
|
||||
|
||||
|
||||
class Qwen2FlashAttention2(Qwen2Attention):
|
||||
"""
|
||||
Qwen2 flash attention module, following Qwen2 attention module. This module inherits from `Qwen2Attention`
|
||||
as the weights of the module stays untouched. The only required change would be on the forward pass
|
||||
where it needs to correctly call the public API of flash attention and deal with padding tokens
|
||||
in case the input contains any of them. Additionally, for sliding window attention, we apply SWA only to the bottom
|
||||
config.max_window_layers layers.
|
||||
"""
|
||||
|
||||
# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
# TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
|
||||
# flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
|
||||
# Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
|
||||
self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_value: Optional[Cache] = None,
|
||||
output_attentions: bool = False,
|
||||
use_cache: bool = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46
|
||||
):
|
||||
bsz, q_len, _ = hidden_states.size()
|
||||
|
||||
query_states = self.q_proj(hidden_states)
|
||||
key_states = self.k_proj(hidden_states)
|
||||
value_states = self.v_proj(hidden_states)
|
||||
|
||||
query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
|
||||
value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
|
||||
|
||||
if position_embeddings is None:
|
||||
logger.warning_once(
|
||||
"The attention layers in this model are transitioning from computing the RoPE embeddings internally "
|
||||
"through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
|
||||
"`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
|
||||
"removed and `position_embeddings` will be mandatory."
|
||||
)
|
||||
cos, sin = self.rotary_emb(value_states, position_ids)
|
||||
else:
|
||||
cos, sin = position_embeddings
|
||||
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
||||
|
||||
if past_key_value is not None:
|
||||
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models
|
||||
key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
|
||||
|
||||
# repeat k/v heads if n_kv_heads < n_heads
|
||||
key_states = repeat_kv(key_states, self.num_key_value_groups)
|
||||
value_states = repeat_kv(value_states, self.num_key_value_groups)
|
||||
dropout_rate = 0.0 if not self.training else self.attention_dropout
|
||||
|
||||
# In PEFT, usually we cast the layer norms in float32 for training stability reasons
|
||||
# therefore the input hidden states gets silently casted in float32. Hence, we need
|
||||
# cast them back in float16 just to be sure everything works as expected.
|
||||
input_dtype = query_states.dtype
|
||||
if input_dtype == torch.float32:
|
||||
if torch.is_autocast_enabled():
|
||||
target_dtype = torch.get_autocast_gpu_dtype()
|
||||
# Handle the case where the model is quantized
|
||||
elif hasattr(self.config, "_pre_quantization_dtype"):
|
||||
target_dtype = self.config._pre_quantization_dtype
|
||||
else:
|
||||
target_dtype = self.q_proj.weight.dtype
|
||||
|
||||
logger.warning_once(
|
||||
f"The input hidden states seems to be silently casted in float32, this might be related to"
|
||||
f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
|
||||
f" {target_dtype}."
|
||||
)
|
||||
|
||||
query_states = query_states.to(target_dtype)
|
||||
key_states = key_states.to(target_dtype)
|
||||
value_states = value_states.to(target_dtype)
|
||||
|
||||
# Reashape to the expected shape for Flash Attention
|
||||
query_states = query_states.transpose(1, 2)
|
||||
key_states = key_states.transpose(1, 2)
|
||||
value_states = value_states.transpose(1, 2)
|
||||
|
||||
if (
|
||||
self.config.use_sliding_window
|
||||
and getattr(self.config, "sliding_window", None) is not None
|
||||
and self.layer_idx >= self.config.max_window_layers
|
||||
):
|
||||
sliding_window = self.config.sliding_window
|
||||
else:
|
||||
sliding_window = None
|
||||
|
||||
attn_output = _flash_attention_forward(
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
attention_mask,
|
||||
q_len,
|
||||
position_ids=position_ids,
|
||||
dropout=dropout_rate,
|
||||
sliding_window=sliding_window,
|
||||
is_causal=self.is_causal,
|
||||
use_top_left_mask=self._flash_attn_uses_top_left_mask,
|
||||
)
|
||||
|
||||
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
|
||||
attn_output = self.o_proj(attn_output)
|
||||
|
||||
if not output_attentions:
|
||||
attn_weights = None
|
||||
|
||||
return attn_output, attn_weights, past_key_value
|
||||
|
||||
|
||||
QWEN2_ATTENTION_CLASSES = {
|
||||
"eager": Qwen2Attention,
|
||||
"flash_attention_2": Qwen2FlashAttention2,
|
||||
}
|
||||
|
||||
|
||||
class Qwen2DecoderLayer(nn.Module):
|
||||
def __init__(self, config: Qwen2Config, layer_idx: int):
|
||||
super().__init__()
|
||||
self.hidden_size = config.hidden_size
|
||||
|
||||
if config.sliding_window and config._attn_implementation != "flash_attention_2":
|
||||
logger.warning_once(
|
||||
f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "
|
||||
"unexpected results may be encountered."
|
||||
)
|
||||
self.self_attn = QWEN2_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
|
||||
|
||||
self.mlp = Qwen2MLP(config)
|
||||
self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||||
output_attentions: Optional[bool] = False,
|
||||
use_cache: Optional[bool] = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46
|
||||
**kwargs,
|
||||
) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
|
||||
"""
|
||||
Args:
|
||||
hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
|
||||
attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
|
||||
`(batch, sequence_length)` where padding elements are indicated by 0.
|
||||
output_attentions (`bool`, *optional*):
|
||||
Whether or not to return the attentions tensors of all attention layers. See `attentions` under
|
||||
returned tensors for more detail.
|
||||
use_cache (`bool`, *optional*):
|
||||
If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
|
||||
(see `past_key_values`).
|
||||
past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
|
||||
cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
|
||||
Indices depicting the position of the input sequence tokens in the sequence.
|
||||
position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
|
||||
Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
|
||||
with `head_dim` being the embedding dimension of each attention head.
|
||||
kwargs (`dict`, *optional*):
|
||||
Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
|
||||
into the model
|
||||
"""
|
||||
|
||||
residual = hidden_states
|
||||
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
|
||||
# Self Attention
|
||||
hidden_states, self_attn_weights, present_key_value = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_value=past_key_value,
|
||||
output_attentions=output_attentions,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
position_embeddings=position_embeddings,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
# Fully Connected
|
||||
residual = hidden_states
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
outputs = (hidden_states,)
|
||||
|
||||
if output_attentions:
|
||||
outputs += (self_attn_weights,)
|
||||
|
||||
if use_cache:
|
||||
outputs += (present_key_value,)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
QWEN2_START_DOCSTRING = r"""
|
||||
This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
|
||||
library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
|
||||
etc.)
|
||||
|
||||
This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
|
||||
Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
|
||||
and behavior.
|
||||
|
||||
Parameters:
|
||||
config ([`Qwen2Config`]):
|
||||
Model configuration class with all the parameters of the model. Initializing with a config file does not
|
||||
load the weights associated with the model, only the configuration. Check out the
|
||||
[`~PreTrainedModel.from_pretrained`] method to load the model weights.
|
||||
"""
|
||||
|
||||
|
||||
@add_start_docstrings(
|
||||
"The bare Qwen2 Model outputting raw hidden-states without any specific head on top.",
|
||||
QWEN2_START_DOCSTRING,
|
||||
)
|
||||
class Qwen2PreTrainedModel(PreTrainedModel):
|
||||
config_class = Qwen2Config
|
||||
base_model_prefix = "model"
|
||||
supports_gradient_checkpointing = True
|
||||
_no_split_modules = ["Qwen2DecoderLayer"]
|
||||
_skip_keys_device_placement = "past_key_values"
|
||||
_supports_flash_attn_2 = True
|
||||
_supports_cache_class = True
|
||||
_supports_quantized_cache = True
|
||||
_supports_static_cache = True
|
||||
|
||||
def _init_weights(self, module):
|
||||
std = self.config.initializer_range
|
||||
if isinstance(module, nn.Linear):
|
||||
module.weight.data.normal_(mean=0.0, std=std)
|
||||
if module.bias is not None:
|
||||
module.bias.data.zero_()
|
||||
elif isinstance(module, nn.Embedding):
|
||||
module.weight.data.normal_(mean=0.0, std=std)
|
||||
if module.padding_idx is not None:
|
||||
module.weight.data[module.padding_idx].zero_()
|
||||
|
||||
|
||||
QWEN2_INPUTS_DOCSTRING = r"""
|
||||
Args:
|
||||
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
|
||||
Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
|
||||
it.
|
||||
|
||||
Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
|
||||
[`PreTrainedTokenizer.__call__`] for details.
|
||||
|
||||
[What are input IDs?](../glossary#input-ids)
|
||||
attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
|
||||
Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
|
||||
|
||||
- 1 for tokens that are **not masked**,
|
||||
- 0 for tokens that are **masked**.
|
||||
|
||||
[What are attention masks?](../glossary#attention-mask)
|
||||
|
||||
Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
|
||||
[`PreTrainedTokenizer.__call__`] for details.
|
||||
|
||||
If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
|
||||
`past_key_values`).
|
||||
|
||||
If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
|
||||
and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
|
||||
information on the default strategy.
|
||||
|
||||
- 1 indicates the head is **not masked**,
|
||||
- 0 indicates the head is **masked**.
|
||||
position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
||||
Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
|
||||
config.n_positions - 1]`.
|
||||
|
||||
[What are position IDs?](../glossary#position-ids)
|
||||
past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
|
||||
Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
|
||||
blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
|
||||
returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
|
||||
|
||||
Two formats are allowed:
|
||||
- a [`~cache_utils.Cache`] instance, see our
|
||||
[kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache);
|
||||
- Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
|
||||
shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
|
||||
cache format.
|
||||
|
||||
The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
|
||||
legacy cache format will be returned.
|
||||
|
||||
If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
|
||||
have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
|
||||
of shape `(batch_size, sequence_length)`.
|
||||
inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
|
||||
Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
|
||||
is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
|
||||
model's internal embedding lookup matrix.
|
||||
use_cache (`bool`, *optional*):
|
||||
If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
|
||||
`past_key_values`).
|
||||
output_attentions (`bool`, *optional*):
|
||||
Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
|
||||
tensors for more detail.
|
||||
output_hidden_states (`bool`, *optional*):
|
||||
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
|
||||
more detail.
|
||||
return_dict (`bool`, *optional*):
|
||||
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
||||
cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
|
||||
Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
|
||||
this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
|
||||
the complete sequence length.
|
||||
"""
|
||||
|
||||
|
||||
@add_start_docstrings(
|
||||
"The bare Qwen2 Model outputting raw hidden-states without any specific head on top.",
|
||||
QWEN2_START_DOCSTRING,
|
||||
)
|
||||
class Qwen2Model(Qwen2PreTrainedModel):
|
||||
"""
|
||||
Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2DecoderLayer`]
|
||||
|
||||
Args:
|
||||
config: Qwen2Config
|
||||
"""
|
||||
|
||||
def __init__(self, config: Qwen2Config):
|
||||
super().__init__(config)
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
||||
self.layers = nn.ModuleList(
|
||||
[Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
||||
)
|
||||
self._attn_implementation = config._attn_implementation
|
||||
self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.rotary_emb = Qwen2RotaryEmbedding(config=config)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def get_input_embeddings(self):
|
||||
return self.embed_tokens
|
||||
|
||||
def set_input_embeddings(self, value):
|
||||
self.embed_tokens = value
|
||||
|
||||
@add_start_docstrings_to_model_forward(QWEN2_INPUTS_DOCSTRING)
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
) -> Union[Tuple, BaseModelOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
||||
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
if (input_ids is None) ^ (inputs_embeds is not None):
|
||||
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
||||
|
||||
if self.gradient_checkpointing and self.training:
|
||||
if use_cache:
|
||||
logger.warning_once(
|
||||
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
|
||||
)
|
||||
use_cache = False
|
||||
|
||||
# kept for BC (non `Cache` `past_key_values` inputs)
|
||||
return_legacy_cache = False
|
||||
if use_cache and not isinstance(past_key_values, Cache):
|
||||
return_legacy_cache = True
|
||||
if past_key_values is None:
|
||||
past_key_values = DynamicCache()
|
||||
else:
|
||||
past_key_values = DynamicCache.from_legacy_cache(past_key_values)
|
||||
logger.warning_once(
|
||||
"We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
|
||||
"will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
|
||||
"(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"
|
||||
)
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embed_tokens(input_ids)
|
||||
|
||||
if cache_position is None:
|
||||
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
||||
cache_position = torch.arange(
|
||||
past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
|
||||
)
|
||||
if position_ids is None:
|
||||
position_ids = cache_position.unsqueeze(0)
|
||||
|
||||
if attention_mask is not None and 0.0 in attention_mask:
|
||||
causal_mask = attention_mask
|
||||
else:
|
||||
causal_mask = None
|
||||
|
||||
hidden_states = inputs_embeds
|
||||
# create position embeddings to be shared across the decoder layers
|
||||
position_embeddings = self.rotary_emb(hidden_states, position_ids)
|
||||
|
||||
# decoder layers
|
||||
all_hidden_states = () if output_hidden_states else None
|
||||
all_self_attns = () if output_attentions else None
|
||||
next_decoder_cache = None
|
||||
|
||||
for decoder_layer in self.layers:
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
if self.gradient_checkpointing and self.training:
|
||||
layer_outputs = self._gradient_checkpointing_func(
|
||||
decoder_layer.__call__,
|
||||
hidden_states,
|
||||
causal_mask,
|
||||
position_ids,
|
||||
past_key_values,
|
||||
output_attentions,
|
||||
use_cache,
|
||||
cache_position,
|
||||
position_embeddings,
|
||||
)
|
||||
else:
|
||||
layer_outputs = decoder_layer(
|
||||
hidden_states,
|
||||
attention_mask=causal_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_value=past_key_values,
|
||||
output_attentions=output_attentions,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
position_embeddings=position_embeddings,
|
||||
)
|
||||
|
||||
hidden_states = layer_outputs[0]
|
||||
|
||||
if use_cache:
|
||||
next_decoder_cache = layer_outputs[2 if output_attentions else 1]
|
||||
|
||||
if output_attentions:
|
||||
all_self_attns += (layer_outputs[1],)
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
|
||||
# add hidden states from the last decoder layer
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
next_cache = next_decoder_cache if use_cache else None
|
||||
if return_legacy_cache:
|
||||
next_cache = next_cache.to_legacy_cache()
|
||||
|
||||
if not return_dict:
|
||||
return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=hidden_states,
|
||||
past_key_values=next_cache,
|
||||
hidden_states=all_hidden_states,
|
||||
attentions=all_self_attns,
|
||||
)
|
||||
|
||||
|
||||
class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin):
|
||||
_tied_weights_keys = ["lm_head.weight"]
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self.model = Qwen2Model(config)
|
||||
self.vocab_size = config.vocab_size
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def get_input_embeddings(self):
|
||||
return self.model.embed_tokens
|
||||
|
||||
def set_input_embeddings(self, value):
|
||||
self.model.embed_tokens = value
|
||||
|
||||
def get_output_embeddings(self):
|
||||
return self.lm_head
|
||||
|
||||
def set_output_embeddings(self, new_embeddings):
|
||||
self.lm_head = new_embeddings
|
||||
|
||||
def set_decoder(self, decoder):
|
||||
self.model = decoder
|
||||
|
||||
def get_decoder(self):
|
||||
return self.model
|
||||
|
||||
@add_start_docstrings_to_model_forward(QWEN2_INPUTS_DOCSTRING)
|
||||
@replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
num_logits_to_keep: int = 0,
|
||||
**loss_kwargs,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
r"""
|
||||
Args:
|
||||
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
||||
Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
|
||||
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
|
||||
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
|
||||
|
||||
num_logits_to_keep (`int`, *optional*):
|
||||
Calculate logits for the last `num_logits_to_keep` tokens. If `0`, calculate logits for all
|
||||
`input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
|
||||
token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
|
||||
|
||||
Returns:
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import AutoTokenizer, Qwen2ForCausalLM
|
||||
|
||||
>>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
|
||||
>>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
|
||||
|
||||
>>> prompt = "Hey, are you conscious? Can you talk to me?"
|
||||
>>> inputs = tokenizer(prompt, return_tensors="pt")
|
||||
|
||||
>>> # Generate
|
||||
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
|
||||
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
|
||||
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
|
||||
```"""
|
||||
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
|
||||
logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
@@ -0,0 +1,328 @@
|
||||
# Copyright 2024 The Qwen Team and The HuggingFace Inc. team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Tokenization classes for Qwen2."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import unicodedata
|
||||
from functools import lru_cache
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import regex as re
|
||||
|
||||
from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
VOCAB_FILES_NAMES = {
|
||||
"vocab_file": "vocab.json",
|
||||
"merges_file": "merges.txt",
|
||||
}
|
||||
|
||||
|
||||
MAX_MODEL_INPUT_SIZES = {"qwen/qwen-tokenizer": 32768}
|
||||
|
||||
PRETOKENIZE_REGEX = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
|
||||
|
||||
|
||||
@lru_cache()
|
||||
# Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode
|
||||
def bytes_to_unicode():
|
||||
"""
|
||||
Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
|
||||
characters the bpe code barfs on.
|
||||
|
||||
The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
|
||||
if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
|
||||
decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
|
||||
tables between utf-8 bytes and unicode strings.
|
||||
"""
|
||||
bs = (
|
||||
list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
|
||||
)
|
||||
cs = bs[:]
|
||||
n = 0
|
||||
for b in range(2**8):
|
||||
if b not in bs:
|
||||
bs.append(b)
|
||||
cs.append(2**8 + n)
|
||||
n += 1
|
||||
cs = [chr(n) for n in cs]
|
||||
return dict(zip(bs, cs))
|
||||
|
||||
|
||||
# Copied from transformers.models.gpt2.tokenization_gpt2.get_pairs
|
||||
def get_pairs(word):
|
||||
"""
|
||||
Return set of symbol pairs in a word.
|
||||
|
||||
Word is represented as tuple of symbols (symbols being variable-length strings).
|
||||
"""
|
||||
pairs = set()
|
||||
prev_char = word[0]
|
||||
for char in word[1:]:
|
||||
pairs.add((prev_char, char))
|
||||
prev_char = char
|
||||
return pairs
|
||||
|
||||
|
||||
class Qwen2Tokenizer(PreTrainedTokenizer):
|
||||
"""
|
||||
Construct a Qwen2 tokenizer. Based on byte-level Byte-Pair-Encoding.
|
||||
|
||||
Same with GPT2Tokenizer, this tokenizer has been trained to treat spaces like parts of the tokens so a word will
|
||||
be encoded differently whether it is at the beginning of the sentence (without space) or not:
|
||||
|
||||
```python
|
||||
>>> from transformers import Qwen2Tokenizer
|
||||
|
||||
>>> tokenizer = Qwen2Tokenizer.from_pretrained("Qwen/Qwen-tokenizer")
|
||||
>>> tokenizer("Hello world")["input_ids"]
|
||||
[9707, 1879]
|
||||
|
||||
>>> tokenizer(" Hello world")["input_ids"]
|
||||
[21927, 1879]
|
||||
```
|
||||
This is expected.
|
||||
|
||||
You should not use GPT2Tokenizer instead, because of the different pretokenization rules.
|
||||
|
||||
This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
|
||||
this superclass for more information regarding those methods.
|
||||
|
||||
Args:
|
||||
vocab_file (`str`):
|
||||
Path to the vocabulary file.
|
||||
merges_file (`str`):
|
||||
Path to the merges file.
|
||||
errors (`str`, *optional*, defaults to `"replace"`):
|
||||
Paradigm to follow when decoding bytes to UTF-8. See
|
||||
[bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
|
||||
unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
|
||||
The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
|
||||
token instead.
|
||||
bos_token (`str`, *optional*):
|
||||
The beginning of sequence token. Not applicable for this tokenizer.
|
||||
eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
|
||||
The end of sequence token.
|
||||
pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
|
||||
The token used for padding, for example when batching sequences of different lengths.
|
||||
clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not the model should cleanup the spaces that were added when splitting the input text during the
|
||||
tokenization process. Not applicable to this tokenizer, since tokenization does not add spaces.
|
||||
split_special_tokens (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not the special tokens should be split during the tokenization process. The default behavior is
|
||||
to not split special tokens. This means that if `<|endoftext|>` is the `eos_token`, then `tokenizer.tokenize("<|endoftext|>") =
|
||||
['<|endoftext|>`]. Otherwise, if `split_special_tokens=True`, then `tokenizer.tokenize("<|endoftext|>")` will be give `['<',
|
||||
'|', 'endo', 'ft', 'ext', '|', '>']`. This argument is only supported for `slow` tokenizers for the moment.
|
||||
"""
|
||||
|
||||
vocab_files_names = VOCAB_FILES_NAMES
|
||||
model_input_names = ["input_ids", "attention_mask"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_file,
|
||||
merges_file,
|
||||
errors="replace",
|
||||
unk_token="<|endoftext|>",
|
||||
bos_token=None,
|
||||
eos_token="<|endoftext|>",
|
||||
pad_token="<|endoftext|>",
|
||||
clean_up_tokenization_spaces=False,
|
||||
split_special_tokens=False,
|
||||
**kwargs,
|
||||
):
|
||||
# Qwen vocab does not contain control tokens; added tokens need to be special
|
||||
bos_token = (
|
||||
AddedToken(bos_token, lstrip=False, rstrip=False, special=True, normalized=False)
|
||||
if isinstance(bos_token, str)
|
||||
else bos_token
|
||||
)
|
||||
eos_token = (
|
||||
AddedToken(eos_token, lstrip=False, rstrip=False, special=True, normalized=False)
|
||||
if isinstance(eos_token, str)
|
||||
else eos_token
|
||||
)
|
||||
unk_token = (
|
||||
AddedToken(unk_token, lstrip=False, rstrip=False, special=True, normalized=False)
|
||||
if isinstance(unk_token, str)
|
||||
else unk_token
|
||||
)
|
||||
pad_token = (
|
||||
AddedToken(pad_token, lstrip=False, rstrip=False, special=True, normalized=False)
|
||||
if isinstance(pad_token, str)
|
||||
else pad_token
|
||||
)
|
||||
|
||||
with open(vocab_file, encoding="utf-8") as vocab_handle:
|
||||
self.encoder = json.load(vocab_handle)
|
||||
self.decoder = {v: k for k, v in self.encoder.items()}
|
||||
self.errors = errors # how to handle errors in decoding
|
||||
self.byte_encoder = bytes_to_unicode()
|
||||
self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
|
||||
bpe_merges = []
|
||||
with open(merges_file, encoding="utf-8") as merges_handle:
|
||||
for i, line in enumerate(merges_handle):
|
||||
line = line.strip()
|
||||
if (i == 0 and line.startswith("#version:")) or not line:
|
||||
continue
|
||||
bpe_merges.append(tuple(line.split()))
|
||||
self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
|
||||
# NOTE: the cache can grow without bound and will get really large for long running processes
|
||||
# (esp. for texts of language that do not use space between word, e.g. Chinese); technically
|
||||
# not a memory leak but appears as one.
|
||||
# GPT2Tokenizer has the same problem, so let's be consistent.
|
||||
self.cache = {}
|
||||
|
||||
self.pat = re.compile(PRETOKENIZE_REGEX)
|
||||
|
||||
if kwargs.get("add_prefix_space", False):
|
||||
logger.warning_once(
|
||||
f"{self.__class__.__name} does not support `add_prefix_space`, setting it to True has no effect."
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
errors=errors,
|
||||
bos_token=bos_token,
|
||||
eos_token=eos_token,
|
||||
pad_token=pad_token,
|
||||
unk_token=unk_token,
|
||||
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
|
||||
split_special_tokens=split_special_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return len(self.encoder)
|
||||
|
||||
# Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.get_vocab
|
||||
def get_vocab(self):
|
||||
return dict(self.encoder, **self.added_tokens_encoder)
|
||||
|
||||
# Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.bpe
|
||||
def bpe(self, token):
|
||||
if token in self.cache:
|
||||
return self.cache[token]
|
||||
word = tuple(token)
|
||||
pairs = get_pairs(word)
|
||||
|
||||
if not pairs:
|
||||
return token
|
||||
|
||||
while True:
|
||||
bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
|
||||
if bigram not in self.bpe_ranks:
|
||||
break
|
||||
first, second = bigram
|
||||
new_word = []
|
||||
i = 0
|
||||
while i < len(word):
|
||||
try:
|
||||
j = word.index(first, i)
|
||||
except ValueError:
|
||||
new_word.extend(word[i:])
|
||||
break
|
||||
else:
|
||||
new_word.extend(word[i:j])
|
||||
i = j
|
||||
|
||||
if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
|
||||
new_word.append(first + second)
|
||||
i += 2
|
||||
else:
|
||||
new_word.append(word[i])
|
||||
i += 1
|
||||
new_word = tuple(new_word)
|
||||
word = new_word
|
||||
if len(word) == 1:
|
||||
break
|
||||
else:
|
||||
pairs = get_pairs(word)
|
||||
word = " ".join(word)
|
||||
self.cache[token] = word
|
||||
return word
|
||||
|
||||
# Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._tokenize
|
||||
def _tokenize(self, text):
|
||||
"""Tokenize a string."""
|
||||
bpe_tokens = []
|
||||
for token in re.findall(self.pat, text):
|
||||
token = "".join(
|
||||
self.byte_encoder[b] for b in token.encode("utf-8")
|
||||
) # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
|
||||
bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
|
||||
return bpe_tokens
|
||||
|
||||
# Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_token_to_id
|
||||
def _convert_token_to_id(self, token):
|
||||
"""Converts a token (str) in an id using the vocab."""
|
||||
return self.encoder.get(token, self.encoder.get(self.unk_token))
|
||||
|
||||
# Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_id_to_token
|
||||
def _convert_id_to_token(self, index):
|
||||
"""Converts an index (integer) in a token (str) using the vocab."""
|
||||
return self.decoder.get(index)
|
||||
|
||||
# Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.convert_tokens_to_string
|
||||
def convert_tokens_to_string(self, tokens):
|
||||
"""Converts a sequence of tokens (string) in a single string."""
|
||||
text = "".join(tokens)
|
||||
text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
|
||||
return text
|
||||
|
||||
def decode(
|
||||
self,
|
||||
token_ids,
|
||||
skip_special_tokens: bool = False,
|
||||
clean_up_tokenization_spaces: Optional[bool] = False,
|
||||
spaces_between_special_tokens: bool = False,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
# `spaces_between_special_tokens` defaults to True for _decode in slow tokenizers
|
||||
# and cannot be configured elsewhere, but it should default to False for Qwen2Tokenizer
|
||||
return super().decode(
|
||||
token_ids,
|
||||
skip_special_tokens=skip_special_tokens,
|
||||
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
|
||||
spaces_between_special_tokens=spaces_between_special_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.save_vocabulary
|
||||
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
|
||||
if not os.path.isdir(save_directory):
|
||||
logger.error(f"Vocabulary path ({save_directory}) should be a directory")
|
||||
return
|
||||
vocab_file = os.path.join(
|
||||
save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
|
||||
)
|
||||
merge_file = os.path.join(
|
||||
save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
|
||||
)
|
||||
|
||||
with open(vocab_file, "w", encoding="utf-8") as f:
|
||||
f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
|
||||
|
||||
index = 0
|
||||
with open(merge_file, "w", encoding="utf-8") as writer:
|
||||
writer.write("#version: 0.2\n")
|
||||
for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
|
||||
if index != token_index:
|
||||
logger.warning(
|
||||
f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
|
||||
" Please check that the tokenizer is not corrupted!"
|
||||
)
|
||||
index = token_index
|
||||
writer.write(" ".join(bpe_tokens) + "\n")
|
||||
index += 1
|
||||
|
||||
return vocab_file, merge_file
|
||||
|
||||
def prepare_for_tokenization(self, text, **kwargs):
|
||||
text = unicodedata.normalize("NFC", text)
|
||||
return (text, kwargs)
|
||||
@@ -0,0 +1,123 @@
|
||||
# Copyright 2024 The Qwen Team and The HuggingFace Inc. team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Tokenization classes for Qwen2."""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from transformers.tokenization_utils import AddedToken
|
||||
from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
|
||||
from transformers.utils import logging
|
||||
from .tokenization_qwen2 import Qwen2Tokenizer
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
VOCAB_FILES_NAMES = {
|
||||
"vocab_file": "vocab.json",
|
||||
"merges_file": "merges.txt",
|
||||
"tokenizer_file": "tokenizer.json",
|
||||
}
|
||||
|
||||
|
||||
MAX_MODEL_INPUT_SIZES = {"qwen/qwen-tokenizer": 32768}
|
||||
|
||||
|
||||
class Qwen2TokenizerFast(PreTrainedTokenizerFast):
|
||||
"""
|
||||
Construct a "fast" Qwen2 tokenizer (backed by HuggingFace's *tokenizers* library). Based on byte-level
|
||||
Byte-Pair-Encoding.
|
||||
|
||||
Same with GPT2Tokenizer, this tokenizer has been trained to treat spaces like parts of the tokens so a word will
|
||||
be encoded differently whether it is at the beginning of the sentence (without space) or not:
|
||||
|
||||
```python
|
||||
>>> from transformers import Qwen2TokenizerFast
|
||||
|
||||
>>> tokenizer = Qwen2TokenizerFast.from_pretrained("Qwen/Qwen-tokenizer")
|
||||
>>> tokenizer("Hello world")["input_ids"]
|
||||
[9707, 1879]
|
||||
|
||||
>>> tokenizer(" Hello world")["input_ids"]
|
||||
[21927, 1879]
|
||||
```
|
||||
This is expected.
|
||||
|
||||
This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
|
||||
refer to this superclass for more information regarding those methods.
|
||||
|
||||
Args:
|
||||
vocab_file (`str`, *optional*):
|
||||
Path to the vocabulary file.
|
||||
merges_file (`str`, *optional*):
|
||||
Path to the merges file.
|
||||
tokenizer_file (`str`, *optional*):
|
||||
Path to [tokenizers](https://github.com/huggingface/tokenizers) file (generally has a .json extension) that
|
||||
contains everything needed to load the tokenizer.
|
||||
unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
|
||||
The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
|
||||
token instead. Not applicable to this tokenizer.
|
||||
bos_token (`str`, *optional*):
|
||||
The beginning of sequence token. Not applicable for this tokenizer.
|
||||
eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
|
||||
The end of sequence token.
|
||||
pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
|
||||
The token used for padding, for example when batching sequences of different lengths.
|
||||
"""
|
||||
|
||||
vocab_files_names = VOCAB_FILES_NAMES
|
||||
model_input_names = ["input_ids", "attention_mask"]
|
||||
slow_tokenizer_class = Qwen2Tokenizer
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_file=None,
|
||||
merges_file=None,
|
||||
tokenizer_file=None,
|
||||
unk_token="<|endoftext|>",
|
||||
bos_token=None,
|
||||
eos_token="<|endoftext|>",
|
||||
pad_token="<|endoftext|>",
|
||||
**kwargs,
|
||||
):
|
||||
# We need to at least pass vocab_file and merges_file to base class
|
||||
# in case a slow tokenizer needs to be initialized; other can be
|
||||
# configured through files.
|
||||
# following GPT2TokenizerFast, also adding unk_token, bos_token, and eos_token
|
||||
|
||||
bos_token = (
|
||||
AddedToken(bos_token, lstrip=False, rstrip=False, special=True, normalized=False)
|
||||
if isinstance(bos_token, str)
|
||||
else bos_token
|
||||
)
|
||||
eos_token = (
|
||||
AddedToken(eos_token, lstrip=False, rstrip=False, special=True, normalized=False)
|
||||
if isinstance(eos_token, str)
|
||||
else eos_token
|
||||
)
|
||||
unk_token = (
|
||||
AddedToken(unk_token, lstrip=False, rstrip=False, special=True, normalized=False)
|
||||
if isinstance(unk_token, str)
|
||||
else unk_token
|
||||
)
|
||||
pad_token = (
|
||||
AddedToken(pad_token, lstrip=False, rstrip=False, special=True, normalized=False)
|
||||
if isinstance(pad_token, str)
|
||||
else pad_token
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
vocab_file=vocab_file,
|
||||
merges_file=merges_file,
|
||||
tokenizer_file=tokenizer_file,
|
||||
unk_token=unk_token,
|
||||
bos_token=bos_token,
|
||||
eos_token=eos_token,
|
||||
pad_token=pad_token,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Copied from transformers.models.gpt2.tokenization_gpt2_fast.GPT2TokenizerFast.save_vocabulary
|
||||
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
|
||||
files = self._tokenizer.model.save(save_directory, name=filename_prefix)
|
||||
return tuple(files)
|
||||
@@ -0,0 +1,98 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from transformers.utils import (
|
||||
OptionalDependencyNotAvailable,
|
||||
_LazyModule,
|
||||
is_sentencepiece_available,
|
||||
is_torch_available,
|
||||
is_vision_available,
|
||||
)
|
||||
|
||||
|
||||
_import_structure = {
|
||||
"configuration_siglip": [
|
||||
"SiglipConfig",
|
||||
"SiglipTextConfig",
|
||||
"SiglipVisionConfig",
|
||||
],
|
||||
"processing_siglip": ["SiglipProcessor"],
|
||||
}
|
||||
|
||||
try:
|
||||
if not is_sentencepiece_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
_import_structure["tokenization_siglip"] = ["SiglipTokenizer"]
|
||||
|
||||
|
||||
try:
|
||||
if not is_vision_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
_import_structure["image_processing_siglip"] = ["SiglipImageProcessor"]
|
||||
|
||||
try:
|
||||
if not is_torch_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
_import_structure["modeling_siglip"] = [
|
||||
"SiglipModel",
|
||||
"SiglipPreTrainedModel",
|
||||
"SiglipTextModel",
|
||||
"SiglipVisionModel",
|
||||
"SiglipForImageClassification",
|
||||
]
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .configuration_siglip import (
|
||||
SiglipConfig,
|
||||
SiglipTextConfig,
|
||||
SiglipVisionConfig,
|
||||
)
|
||||
from .processing_siglip import SiglipProcessor
|
||||
|
||||
try:
|
||||
if not is_sentencepiece_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
from .tokenization_siglip import SiglipTokenizer
|
||||
|
||||
try:
|
||||
if not is_vision_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
from .image_processing_siglip import SiglipImageProcessor
|
||||
|
||||
try:
|
||||
if not is_torch_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
from .modeling_siglip import (
|
||||
SiglipForImageClassification,
|
||||
SiglipModel,
|
||||
SiglipPreTrainedModel,
|
||||
SiglipTextModel,
|
||||
SiglipVisionModel,
|
||||
)
|
||||
|
||||
|
||||
else:
|
||||
import sys
|
||||
|
||||
sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
|
||||
@@ -0,0 +1,287 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Siglip model configuration"""
|
||||
|
||||
import os
|
||||
from typing import Union
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
class SiglipTextConfig(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`SiglipTextModel`]. It is used to instantiate a
|
||||
Siglip text 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 text 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:
|
||||
vocab_size (`int`, *optional*, defaults to 32000):
|
||||
Vocabulary size of the Siglip text model. Defines the number of different tokens that can be represented by
|
||||
the `inputs_ids` passed when calling [`SiglipModel`].
|
||||
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.
|
||||
max_position_embeddings (`int`, *optional*, defaults to 64):
|
||||
The maximum sequence length that this model might ever be used with. Typically set this to something large
|
||||
just in case (e.g., 512 or 1024 or 2048).
|
||||
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.
|
||||
pad_token_id (`int`, *optional*, defaults to 1):
|
||||
The id of the padding token in the vocabulary.
|
||||
bos_token_id (`int`, *optional*, defaults to 49406):
|
||||
The id of the beginning-of-sequence token in the vocabulary.
|
||||
eos_token_id (`int`, *optional*, defaults to 49407):
|
||||
The id of the end-of-sequence token in the vocabulary.
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import SiglipTextConfig, SiglipTextModel
|
||||
|
||||
>>> # Initializing a SiglipTextConfig with google/siglip-base-patch16-224 style configuration
|
||||
>>> configuration = SiglipTextConfig()
|
||||
|
||||
>>> # Initializing a SiglipTextModel (with random weights) from the google/siglip-base-patch16-224 style configuration
|
||||
>>> model = SiglipTextModel(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "siglip_text_model"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=32000,
|
||||
hidden_size=768,
|
||||
intermediate_size=3072,
|
||||
num_hidden_layers=12,
|
||||
num_attention_heads=12,
|
||||
max_position_embeddings=64,
|
||||
hidden_act="gelu_pytorch_tanh",
|
||||
layer_norm_eps=1e-6,
|
||||
attention_dropout=0.0,
|
||||
# This differs from `CLIPTokenizer`'s default and from openai/siglip
|
||||
# See https://github.com/huggingface/transformers/pull/24773#issuecomment-1632287538
|
||||
pad_token_id=1,
|
||||
bos_token_id=49406,
|
||||
eos_token_id=49407,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
|
||||
|
||||
self.vocab_size = vocab_size
|
||||
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.max_position_embeddings = max_position_embeddings
|
||||
self.layer_norm_eps = layer_norm_eps
|
||||
self.hidden_act = hidden_act
|
||||
self.attention_dropout = attention_dropout
|
||||
|
||||
@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 text config dict if we are loading from SiglipConfig
|
||||
if config_dict.get("model_type") == "siglip":
|
||||
config_dict = config_dict["text_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)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
class SiglipConfig(PretrainedConfig):
|
||||
r"""
|
||||
[`SiglipConfig`] is the configuration class to store the configuration of a [`SiglipModel`]. It is used to
|
||||
instantiate a Siglip model according to the specified arguments, defining the text model and vision model configs.
|
||||
Instantiating a configuration with the defaults will yield a similar configuration to that 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:
|
||||
text_config (`dict`, *optional*):
|
||||
Dictionary of configuration options used to initialize [`SiglipTextConfig`].
|
||||
vision_config (`dict`, *optional*):
|
||||
Dictionary of configuration options used to initialize [`SiglipVisionConfig`].
|
||||
kwargs (*optional*):
|
||||
Dictionary of keyword arguments.
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import SiglipConfig, SiglipModel
|
||||
|
||||
>>> # Initializing a SiglipConfig with google/siglip-base-patch16-224 style configuration
|
||||
>>> configuration = SiglipConfig()
|
||||
|
||||
>>> # Initializing a SiglipModel (with random weights) from the google/siglip-base-patch16-224 style configuration
|
||||
>>> model = SiglipModel(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
|
||||
>>> # We can also initialize a SiglipConfig from a SiglipTextConfig and a SiglipVisionConfig
|
||||
>>> from transformers import SiglipTextConfig, SiglipVisionConfig
|
||||
|
||||
>>> # Initializing a SiglipText and SiglipVision configuration
|
||||
>>> config_text = SiglipTextConfig()
|
||||
>>> config_vision = SiglipVisionConfig()
|
||||
|
||||
>>> config = SiglipConfig.from_text_vision_configs(config_text, config_vision)
|
||||
```"""
|
||||
|
||||
model_type = "siglip"
|
||||
|
||||
def __init__(self, text_config=None, vision_config=None, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
if text_config is None:
|
||||
text_config = {}
|
||||
logger.info("`text_config` is `None`. Initializing the `SiglipTextConfig` with default values.")
|
||||
|
||||
if vision_config is None:
|
||||
vision_config = {}
|
||||
logger.info("`vision_config` is `None`. initializing the `SiglipVisionConfig` with default values.")
|
||||
|
||||
self.text_config = SiglipTextConfig(**text_config)
|
||||
self.vision_config = SiglipVisionConfig(**vision_config)
|
||||
|
||||
self.initializer_factor = 1.0
|
||||
|
||||
@classmethod
|
||||
def from_text_vision_configs(cls, text_config: SiglipTextConfig, vision_config: SiglipVisionConfig, **kwargs):
|
||||
r"""
|
||||
Instantiate a [`SiglipConfig`] (or a derived class) from siglip text model configuration and siglip vision
|
||||
model configuration.
|
||||
|
||||
Returns:
|
||||
[`SiglipConfig`]: An instance of a configuration object
|
||||
"""
|
||||
|
||||
return cls(text_config=text_config.to_dict(), vision_config=vision_config.to_dict(), **kwargs)
|
||||
@@ -0,0 +1,401 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Convert SigLIP checkpoints from the original repository.
|
||||
|
||||
URL: https://github.com/google-research/big_vision/tree/main
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import collections
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
from numpy import load
|
||||
from PIL import Image
|
||||
|
||||
from transformers import SiglipConfig, SiglipImageProcessor, SiglipModel, SiglipProcessor, SiglipTokenizer
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logging.set_verbosity_info()
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
model_name_to_checkpoint = {
|
||||
# base checkpoints
|
||||
"siglip-base-patch16-224": "/Users/nielsrogge/Documents/SigLIP/webli_en_b16_224_63724782.npz",
|
||||
"siglip-base-patch16-256": "/Users/nielsrogge/Documents/SigLIP/webli_en_b16_256_60500360.npz",
|
||||
"siglip-base-patch16-384": "/Users/nielsrogge/Documents/SigLIP/webli_en_b16_384_68578854.npz",
|
||||
"siglip-base-patch16-512": "/Users/nielsrogge/Documents/SigLIP/webli_en_b16_512_68580893.npz",
|
||||
# large checkpoints
|
||||
"siglip-large-patch16-256": "/Users/nielsrogge/Documents/SigLIP/webli_en_l16_256_60552751.npz",
|
||||
"siglip-large-patch16-384": "/Users/nielsrogge/Documents/SigLIP/webli_en_l16_384_63634585.npz",
|
||||
# multilingual checkpoint
|
||||
"siglip-base-patch16-256-i18n": "/Users/nielsrogge/Documents/SigLIP/webli_i18n_b16_256_66117334.npz",
|
||||
# so400m checkpoints
|
||||
"siglip-so400m-patch14-384": "/Users/nielsrogge/Documents/SigLIP/webli_en_so400m_384_58765454.npz",
|
||||
}
|
||||
|
||||
model_name_to_image_size = {
|
||||
"siglip-base-patch16-224": 224,
|
||||
"siglip-base-patch16-256": 256,
|
||||
"siglip-base-patch16-384": 384,
|
||||
"siglip-base-patch16-512": 512,
|
||||
"siglip-large-patch16-256": 256,
|
||||
"siglip-large-patch16-384": 384,
|
||||
"siglip-base-patch16-256-i18n": 256,
|
||||
"siglip-so400m-patch14-384": 384,
|
||||
}
|
||||
|
||||
|
||||
def get_siglip_config(model_name):
|
||||
config = SiglipConfig()
|
||||
|
||||
vocab_size = 250000 if "i18n" in model_name else 32000
|
||||
image_size = model_name_to_image_size[model_name]
|
||||
patch_size = 16 if "patch16" in model_name else 14
|
||||
|
||||
# size of the architecture
|
||||
config.vision_config.image_size = image_size
|
||||
config.vision_config.patch_size = patch_size
|
||||
config.text_config.vocab_size = vocab_size
|
||||
|
||||
if "base" in model_name:
|
||||
pass
|
||||
elif "large" in model_name:
|
||||
config.text_config.hidden_size = 1024
|
||||
config.text_config.intermediate_size = 4096
|
||||
config.text_config.num_hidden_layers = 24
|
||||
config.text_config.num_attention_heads = 16
|
||||
config.vision_config.hidden_size = 1024
|
||||
config.vision_config.intermediate_size = 4096
|
||||
config.vision_config.num_hidden_layers = 24
|
||||
config.vision_config.num_attention_heads = 16
|
||||
elif "so400m" in model_name:
|
||||
config.text_config.hidden_size = 1152
|
||||
config.text_config.intermediate_size = 4304
|
||||
config.text_config.num_hidden_layers = 27
|
||||
config.text_config.num_attention_heads = 16
|
||||
config.vision_config.hidden_size = 1152
|
||||
config.vision_config.intermediate_size = 4304
|
||||
config.vision_config.num_hidden_layers = 27
|
||||
config.vision_config.num_attention_heads = 16
|
||||
else:
|
||||
raise ValueError("Model not supported")
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def create_rename_keys(config):
|
||||
rename_keys = []
|
||||
# fmt: off
|
||||
|
||||
# vision encoder
|
||||
|
||||
rename_keys.append(("params/img/embedding/kernel", "vision_model.embeddings.patch_embedding.weight"))
|
||||
rename_keys.append(("params/img/embedding/bias", "vision_model.embeddings.patch_embedding.bias"))
|
||||
rename_keys.append(("params/img/pos_embedding", "vision_model.embeddings.position_embedding.weight"))
|
||||
|
||||
for i in range(config.vision_config.num_hidden_layers):
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/LayerNorm_0/scale", f"vision_model.encoder.layers.{i}.layer_norm1.weight"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/LayerNorm_0/bias", f"vision_model.encoder.layers.{i}.layer_norm1.bias"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/LayerNorm_1/scale", f"vision_model.encoder.layers.{i}.layer_norm2.weight"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/LayerNorm_1/bias", f"vision_model.encoder.layers.{i}.layer_norm2.bias"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MlpBlock_0/Dense_0/kernel", f"vision_model.encoder.layers.{i}.mlp.fc1.weight"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MlpBlock_0/Dense_0/bias", f"vision_model.encoder.layers.{i}.mlp.fc1.bias"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MlpBlock_0/Dense_1/kernel", f"vision_model.encoder.layers.{i}.mlp.fc2.weight"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MlpBlock_0/Dense_1/bias", f"vision_model.encoder.layers.{i}.mlp.fc2.bias"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/key/kernel", f"vision_model.encoder.layers.{i}.self_attn.k_proj.weight"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/key/bias", f"vision_model.encoder.layers.{i}.self_attn.k_proj.bias"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/value/kernel", f"vision_model.encoder.layers.{i}.self_attn.v_proj.weight"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/value/bias", f"vision_model.encoder.layers.{i}.self_attn.v_proj.bias"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/query/kernel", f"vision_model.encoder.layers.{i}.self_attn.q_proj.weight"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/query/bias", f"vision_model.encoder.layers.{i}.self_attn.q_proj.bias"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/out/kernel", f"vision_model.encoder.layers.{i}.self_attn.out_proj.weight"))
|
||||
rename_keys.append((f"params/img/Transformer/encoderblock_{i}/MultiHeadDotProductAttention_0/out/bias", f"vision_model.encoder.layers.{i}.self_attn.out_proj.bias"))
|
||||
|
||||
rename_keys.append(("params/img/Transformer/encoder_norm/scale", "vision_model.post_layernorm.weight"))
|
||||
rename_keys.append(("params/img/Transformer/encoder_norm/bias", "vision_model.post_layernorm.bias"))
|
||||
|
||||
rename_keys.append(("params/img/MAPHead_0/probe", "vision_model.head.probe"))
|
||||
rename_keys.append(("params/img/MAPHead_0/LayerNorm_0/scale", "vision_model.head.layernorm.weight"))
|
||||
rename_keys.append(("params/img/MAPHead_0/LayerNorm_0/bias", "vision_model.head.layernorm.bias"))
|
||||
rename_keys.append(("params/img/MAPHead_0/MlpBlock_0/Dense_0/kernel", "vision_model.head.mlp.fc1.weight"))
|
||||
rename_keys.append(("params/img/MAPHead_0/MlpBlock_0/Dense_0/bias", "vision_model.head.mlp.fc1.bias"))
|
||||
rename_keys.append(("params/img/MAPHead_0/MlpBlock_0/Dense_1/kernel", "vision_model.head.mlp.fc2.weight"))
|
||||
rename_keys.append(("params/img/MAPHead_0/MlpBlock_0/Dense_1/bias", "vision_model.head.mlp.fc2.bias"))
|
||||
rename_keys.append(("params/img/MAPHead_0/MultiHeadDotProductAttention_0/out/kernel", "vision_model.head.attention.out_proj.weight"))
|
||||
rename_keys.append(("params/img/MAPHead_0/MultiHeadDotProductAttention_0/out/bias", "vision_model.head.attention.out_proj.bias"))
|
||||
|
||||
# text encoder
|
||||
|
||||
rename_keys.append(("params/txt/Embed_0/embedding", "text_model.embeddings.token_embedding.weight"))
|
||||
rename_keys.append(("params/txt/pos_embedding", "text_model.embeddings.position_embedding.weight"))
|
||||
|
||||
for i in range(config.text_config.num_hidden_layers):
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/LayerNorm_0/scale", f"text_model.encoder.layers.{i}.layer_norm1.weight"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/LayerNorm_0/bias", f"text_model.encoder.layers.{i}.layer_norm1.bias"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/LayerNorm_1/scale", f"text_model.encoder.layers.{i}.layer_norm2.weight"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/LayerNorm_1/bias", f"text_model.encoder.layers.{i}.layer_norm2.bias"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MlpBlock_0/Dense_0/kernel", f"text_model.encoder.layers.{i}.mlp.fc1.weight"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MlpBlock_0/Dense_0/bias", f"text_model.encoder.layers.{i}.mlp.fc1.bias"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MlpBlock_0/Dense_1/kernel", f"text_model.encoder.layers.{i}.mlp.fc2.weight"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MlpBlock_0/Dense_1/bias", f"text_model.encoder.layers.{i}.mlp.fc2.bias"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/key/kernel", f"text_model.encoder.layers.{i}.self_attn.k_proj.weight"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/key/bias", f"text_model.encoder.layers.{i}.self_attn.k_proj.bias"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/value/kernel", f"text_model.encoder.layers.{i}.self_attn.v_proj.weight"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/value/bias", f"text_model.encoder.layers.{i}.self_attn.v_proj.bias"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/query/kernel", f"text_model.encoder.layers.{i}.self_attn.q_proj.weight"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/query/bias", f"text_model.encoder.layers.{i}.self_attn.q_proj.bias"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/out/kernel", f"text_model.encoder.layers.{i}.self_attn.out_proj.weight"))
|
||||
rename_keys.append((f"params/txt/Encoder_0/encoderblock_{i}/MultiHeadDotProductAttention_0/out/bias", f"text_model.encoder.layers.{i}.self_attn.out_proj.bias"))
|
||||
|
||||
rename_keys.append(("params/txt/Encoder_0/encoder_norm/scale", "text_model.final_layer_norm.weight"))
|
||||
rename_keys.append(("params/txt/Encoder_0/encoder_norm/bias", "text_model.final_layer_norm.bias"))
|
||||
rename_keys.append(("params/txt/head/kernel", "text_model.head.weight"))
|
||||
rename_keys.append(("params/txt/head/bias", "text_model.head.bias"))
|
||||
|
||||
# learned temperature and bias
|
||||
rename_keys.append(("params/t", "logit_scale"))
|
||||
rename_keys.append(("params/b", "logit_bias"))
|
||||
|
||||
# fmt: on
|
||||
return rename_keys
|
||||
|
||||
|
||||
def rename_key(dct, old, new, config):
|
||||
val = dct.pop(old)
|
||||
|
||||
if ("out_proj" in new or "v_proj" in new or "k_proj" in new or "q_proj" in new) and "vision" in new:
|
||||
val = val.reshape(-1, config.vision_config.hidden_size)
|
||||
if ("out_proj" in new or "v_proj" in new or "k_proj" in new or "q_proj" in new) and "text" in new:
|
||||
val = val.reshape(-1, config.text_config.hidden_size)
|
||||
|
||||
if "patch_embedding.weight" in new:
|
||||
val = val.transpose(3, 2, 0, 1)
|
||||
elif new.endswith("weight") and "position_embedding" not in new and "token_embedding" not in new:
|
||||
val = val.T
|
||||
|
||||
if "position_embedding" in new and "vision" in new:
|
||||
val = val.reshape(-1, config.vision_config.hidden_size)
|
||||
if "position_embedding" in new and "text" in new:
|
||||
val = val.reshape(-1, config.text_config.hidden_size)
|
||||
|
||||
if new.endswith("bias"):
|
||||
val = val.reshape(-1)
|
||||
|
||||
dct[new] = torch.from_numpy(val)
|
||||
|
||||
|
||||
def read_in_q_k_v_head(state_dict, config):
|
||||
# read in individual input projection layers
|
||||
key_proj_weight = (
|
||||
state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/key/kernel")
|
||||
.reshape(-1, config.vision_config.hidden_size)
|
||||
.T
|
||||
)
|
||||
key_proj_bias = state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/key/bias").reshape(-1)
|
||||
value_proj_weight = (
|
||||
state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/value/kernel")
|
||||
.reshape(-1, config.vision_config.hidden_size)
|
||||
.T
|
||||
)
|
||||
value_proj_bias = state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/value/bias").reshape(-1)
|
||||
query_proj_weight = (
|
||||
state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/query/kernel")
|
||||
.reshape(-1, config.vision_config.hidden_size)
|
||||
.T
|
||||
)
|
||||
query_proj_bias = state_dict.pop("params/img/MAPHead_0/MultiHeadDotProductAttention_0/query/bias").reshape(-1)
|
||||
|
||||
# next, add them to the state dict as a single matrix + vector
|
||||
state_dict["vision_model.head.attention.in_proj_weight"] = torch.from_numpy(
|
||||
np.concatenate([query_proj_weight, key_proj_weight, value_proj_weight], axis=0)
|
||||
)
|
||||
state_dict["vision_model.head.attention.in_proj_bias"] = torch.from_numpy(
|
||||
np.concatenate([query_proj_bias, key_proj_bias, value_proj_bias], axis=0)
|
||||
)
|
||||
|
||||
|
||||
# We will verify our results on an image of cute cats
|
||||
def prepare_img():
|
||||
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
|
||||
image = Image.open(requests.get(url, stream=True).raw)
|
||||
return image
|
||||
|
||||
|
||||
def flatten_nested_dict(params, parent_key="", sep="/"):
|
||||
items = []
|
||||
|
||||
for k, v in params.items():
|
||||
new_key = parent_key + sep + k if parent_key else k
|
||||
|
||||
if isinstance(v, collections.abc.MutableMapping):
|
||||
items.extend(flatten_nested_dict(v, new_key, sep=sep).items())
|
||||
else:
|
||||
items.append((new_key, v))
|
||||
return dict(items)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def convert_siglip_checkpoint(model_name, pytorch_dump_folder_path, verify_logits=True, push_to_hub=False):
|
||||
"""
|
||||
Copy/paste/tweak model's weights to our SigLIP structure.
|
||||
"""
|
||||
|
||||
# define default SigLIP configuration
|
||||
config = get_siglip_config(model_name)
|
||||
|
||||
# get checkpoint
|
||||
checkpoint = model_name_to_checkpoint[model_name]
|
||||
|
||||
# get vocab file
|
||||
if "i18n" in model_name:
|
||||
vocab_file = "/Users/nielsrogge/Documents/SigLIP/multilingual_vocab/sentencepiece.model"
|
||||
else:
|
||||
vocab_file = "/Users/nielsrogge/Documents/SigLIP/english_vocab/sentencepiece.model"
|
||||
|
||||
# load original state dict
|
||||
data = load(checkpoint)
|
||||
state_dict = flatten_nested_dict(data)
|
||||
|
||||
# remove and rename some keys
|
||||
rename_keys = create_rename_keys(config)
|
||||
for src, dest in rename_keys:
|
||||
rename_key(state_dict, src, dest, config)
|
||||
|
||||
# qkv matrices of attention pooling head need special treatment
|
||||
read_in_q_k_v_head(state_dict, config)
|
||||
|
||||
# load HuggingFace model
|
||||
model = SiglipModel(config).eval()
|
||||
model.load_state_dict(state_dict)
|
||||
|
||||
# create processor
|
||||
# important: make tokenizer not return attention_mask since original one doesn't require it
|
||||
image_size = config.vision_config.image_size
|
||||
size = {"height": image_size, "width": image_size}
|
||||
image_processor = SiglipImageProcessor(size=size)
|
||||
tokenizer = SiglipTokenizer(vocab_file=vocab_file, model_input_names=["input_ids"])
|
||||
processor = SiglipProcessor(image_processor=image_processor, tokenizer=tokenizer)
|
||||
|
||||
# verify on dummy images and texts
|
||||
url_1 = "https://cdn.openai.com/multimodal-neurons/assets/apple/apple-ipod.jpg"
|
||||
image_1 = Image.open(requests.get(url_1, stream=True).raw).convert("RGB")
|
||||
url_2 = "https://cdn.openai.com/multimodal-neurons/assets/apple/apple-blank.jpg"
|
||||
image_2 = Image.open(requests.get(url_2, stream=True).raw).convert("RGB")
|
||||
texts = ["an apple", "a picture of an apple"]
|
||||
|
||||
inputs = processor(images=[image_1, image_2], text=texts, return_tensors="pt", padding="max_length")
|
||||
|
||||
# verify input_ids against original ones
|
||||
if image_size == 224:
|
||||
filename = "siglip_pixel_values.pt"
|
||||
elif image_size == 256:
|
||||
filename = "siglip_pixel_values_256.pt"
|
||||
elif image_size == 384:
|
||||
filename = "siglip_pixel_values_384.pt"
|
||||
elif image_size == 512:
|
||||
filename = "siglip_pixel_values_512.pt"
|
||||
else:
|
||||
raise ValueError("Image size not supported")
|
||||
|
||||
filepath = hf_hub_download(repo_id="nielsr/test-image", filename=filename, repo_type="dataset")
|
||||
original_pixel_values = torch.load(filepath)
|
||||
filepath = hf_hub_download(repo_id="nielsr/test-image", filename="siglip_input_ids.pt", repo_type="dataset")
|
||||
original_input_ids = torch.load(filepath)
|
||||
|
||||
if "i18n" not in model_name:
|
||||
assert inputs.input_ids.tolist() == original_input_ids.tolist()
|
||||
|
||||
print("Mean of original pixel values:", original_pixel_values.mean())
|
||||
print("Mean of new pixel values:", inputs.pixel_values.mean())
|
||||
|
||||
# note: we're testing with original pixel values here since we don't have exact pixel values
|
||||
with torch.no_grad():
|
||||
outputs = model(input_ids=inputs.input_ids, pixel_values=original_pixel_values)
|
||||
|
||||
# with torch.no_grad():
|
||||
# outputs = model(input_ids=inputs.input_ids, pixel_values=inputs.pixel_values)
|
||||
|
||||
print(outputs.logits_per_image[:3, :3])
|
||||
|
||||
probs = torch.sigmoid(outputs.logits_per_image) # these are the probabilities
|
||||
print(f"{probs[0][0]:.1%} that image 0 is '{texts[0]}'")
|
||||
print(f"{probs[0][1]:.1%} that image 0 is '{texts[1]}'")
|
||||
|
||||
if verify_logits:
|
||||
if model_name == "siglip-base-patch16-224":
|
||||
expected_slice = torch.tensor(
|
||||
[[-2.9621, -2.1672], [-0.2713, 0.2910]],
|
||||
)
|
||||
elif model_name == "siglip-base-patch16-256":
|
||||
expected_slice = torch.tensor(
|
||||
[[-3.1146, -1.9894], [-0.7312, 0.6387]],
|
||||
)
|
||||
elif model_name == "siglip-base-patch16-384":
|
||||
expected_slice = torch.tensor(
|
||||
[[-2.8098, -2.1891], [-0.4242, 0.4102]],
|
||||
)
|
||||
elif model_name == "siglip-base-patch16-512":
|
||||
expected_slice = torch.tensor(
|
||||
[[-2.7899, -2.2668], [-0.4295, -0.0735]],
|
||||
)
|
||||
elif model_name == "siglip-large-patch16-256":
|
||||
expected_slice = torch.tensor(
|
||||
[[-1.5827, -0.5801], [-0.9153, 0.1363]],
|
||||
)
|
||||
elif model_name == "siglip-large-patch16-384":
|
||||
expected_slice = torch.tensor(
|
||||
[[-2.1523, -0.2899], [-0.2959, 0.7884]],
|
||||
)
|
||||
elif model_name == "siglip-so400m-patch14-384":
|
||||
expected_slice = torch.tensor([[-1.2441, -0.6649], [-0.7060, 0.7374]])
|
||||
elif model_name == "siglip-base-patch16-256-i18n":
|
||||
expected_slice = torch.tensor(
|
||||
[[-0.9064, 0.1073], [-0.0299, 0.5304]],
|
||||
)
|
||||
|
||||
assert torch.allclose(outputs.logits_per_image[:3, :3], expected_slice, atol=1e-4)
|
||||
print("Looks ok!")
|
||||
|
||||
if pytorch_dump_folder_path is not None:
|
||||
Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
|
||||
print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
|
||||
model.save_pretrained(pytorch_dump_folder_path)
|
||||
print(f"Saving processor to {pytorch_dump_folder_path}")
|
||||
processor.save_pretrained(pytorch_dump_folder_path)
|
||||
|
||||
if push_to_hub:
|
||||
model.push_to_hub(f"nielsr/{model_name}")
|
||||
processor.push_to_hub(f"nielsr/{model_name}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
# Required parameters
|
||||
parser.add_argument(
|
||||
"--model_name",
|
||||
default="siglip-base-patch16-224",
|
||||
type=str,
|
||||
choices=model_name_to_checkpoint.keys(),
|
||||
help="Name of the model you'd like to convert.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verify_logits",
|
||||
action="store_false",
|
||||
help="Whether to verify logits against the original implementation.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
convert_siglip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.verify_logits, args.push_to_hub)
|
||||
@@ -0,0 +1,230 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Image processor class for SigLIP."""
|
||||
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
|
||||
from transformers.image_transforms import (
|
||||
convert_to_rgb,
|
||||
resize,
|
||||
to_channel_dimension_format,
|
||||
)
|
||||
from transformers.image_utils import (
|
||||
IMAGENET_STANDARD_MEAN,
|
||||
IMAGENET_STANDARD_STD,
|
||||
ChannelDimension,
|
||||
ImageInput,
|
||||
PILImageResampling,
|
||||
infer_channel_dimension_format,
|
||||
is_scaled_image,
|
||||
make_list_of_images,
|
||||
to_numpy_array,
|
||||
valid_images,
|
||||
validate_preprocess_arguments,
|
||||
)
|
||||
from transformers.utils import TensorType, filter_out_non_signature_kwargs, is_vision_available, logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
if is_vision_available():
|
||||
import PIL
|
||||
|
||||
|
||||
class SiglipImageProcessor(BaseImageProcessor):
|
||||
r"""
|
||||
Constructs a SigLIP image processor.
|
||||
|
||||
Args:
|
||||
do_resize (`bool`, *optional*, defaults to `True`):
|
||||
Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by
|
||||
`do_resize` in the `preprocess` method.
|
||||
size (`Dict[str, int]` *optional*, defaults to `{"height": 224, "width": 224}`):
|
||||
Size of the image after resizing. Can be overridden by `size` in the `preprocess` method.
|
||||
resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):
|
||||
Resampling filter to use if resizing the image. Can be overridden by `resample` in the `preprocess` method.
|
||||
do_rescale (`bool`, *optional*, defaults to `True`):
|
||||
Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by `do_rescale` in
|
||||
the `preprocess` method.
|
||||
rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
|
||||
Scale factor to use if rescaling the image. Can be overridden by `rescale_factor` in the `preprocess`
|
||||
method.
|
||||
do_normalize (`bool`, *optional*, defaults to `True`):
|
||||
Whether to normalize the image by the specified mean and standard deviation. Can be overridden by
|
||||
`do_normalize` in the `preprocess` method.
|
||||
image_mean (`float` or `List[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
|
||||
Mean to use if normalizing the image. This is a float or list of floats the length of the number of
|
||||
channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
|
||||
image_std (`float` or `List[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
|
||||
Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
|
||||
number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
|
||||
Can be overridden by the `image_std` parameter in the `preprocess` method.
|
||||
do_convert_rgb (`bool`, *optional*, defaults to `True`):
|
||||
Whether to convert the image to RGB.
|
||||
"""
|
||||
|
||||
model_input_names = ["pixel_values"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
do_resize: bool = True,
|
||||
size: Dict[str, int] = None,
|
||||
resample: PILImageResampling = PILImageResampling.BICUBIC,
|
||||
do_rescale: bool = True,
|
||||
rescale_factor: Union[int, float] = 1 / 255,
|
||||
do_normalize: bool = True,
|
||||
image_mean: Optional[Union[float, List[float]]] = None,
|
||||
image_std: Optional[Union[float, List[float]]] = None,
|
||||
do_convert_rgb: bool = None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__(**kwargs)
|
||||
size = size if size is not None else {"height": 224, "width": 224}
|
||||
image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
|
||||
image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD
|
||||
|
||||
self.do_resize = do_resize
|
||||
self.size = size
|
||||
self.resample = resample
|
||||
self.do_rescale = do_rescale
|
||||
self.rescale_factor = rescale_factor
|
||||
self.do_normalize = do_normalize
|
||||
self.image_mean = image_mean
|
||||
self.image_std = image_std
|
||||
self.do_convert_rgb = do_convert_rgb
|
||||
|
||||
@filter_out_non_signature_kwargs()
|
||||
def preprocess(
|
||||
self,
|
||||
images: ImageInput,
|
||||
do_resize: bool = None,
|
||||
size: Dict[str, int] = None,
|
||||
resample: PILImageResampling = None,
|
||||
do_rescale: bool = None,
|
||||
rescale_factor: float = None,
|
||||
do_normalize: bool = None,
|
||||
image_mean: Optional[Union[float, List[float]]] = None,
|
||||
image_std: Optional[Union[float, List[float]]] = None,
|
||||
return_tensors: Optional[Union[str, TensorType]] = None,
|
||||
data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
|
||||
input_data_format: Optional[Union[str, ChannelDimension]] = None,
|
||||
do_convert_rgb: bool = None,
|
||||
) -> PIL.Image.Image:
|
||||
"""
|
||||
Preprocess an image or batch of images.
|
||||
|
||||
Args:
|
||||
images (`ImageInput`):
|
||||
Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
|
||||
passing in images with pixel values between 0 and 1, set `do_rescale=False`.
|
||||
do_resize (`bool`, *optional*, defaults to `self.do_resize`):
|
||||
Whether to resize the image.
|
||||
size (`Dict[str, int]`, *optional*, defaults to `self.size`):
|
||||
Size of the image after resizing.
|
||||
resample (`int`, *optional*, defaults to `self.resample`):
|
||||
Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only
|
||||
has an effect if `do_resize` is set to `True`.
|
||||
do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
|
||||
Whether to rescale the image.
|
||||
rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
|
||||
Rescale factor to rescale the image by if `do_rescale` is set to `True`.
|
||||
do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
|
||||
Whether to normalize the image.
|
||||
image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
|
||||
Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
|
||||
image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
|
||||
Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
|
||||
`True`.
|
||||
return_tensors (`str` or `TensorType`, *optional*):
|
||||
The type of tensors to return. Can be one of:
|
||||
- Unset: Return a list of `np.ndarray`.
|
||||
- `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
|
||||
- `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
|
||||
- `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
|
||||
- `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
|
||||
data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
|
||||
The channel dimension format for the output image. Can be one of:
|
||||
- `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
|
||||
- `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
|
||||
- Unset: Use the channel dimension format of the input image.
|
||||
input_data_format (`ChannelDimension` or `str`, *optional*):
|
||||
The channel dimension format for the input image. If unset, the channel dimension format is inferred
|
||||
from the input image. Can be one of:
|
||||
- `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
|
||||
- `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
|
||||
- `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
|
||||
do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
|
||||
Whether to convert the image to RGB.
|
||||
"""
|
||||
do_resize = do_resize if do_resize is not None else self.do_resize
|
||||
size = size if size is not None else self.size
|
||||
size = get_size_dict(size, param_name="size", default_to_square=False)
|
||||
resample = resample if resample is not None else self.resample
|
||||
do_rescale = do_rescale if do_rescale is not None else self.do_rescale
|
||||
rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
|
||||
do_normalize = do_normalize if do_normalize is not None else self.do_normalize
|
||||
image_mean = image_mean if image_mean is not None else self.image_mean
|
||||
image_std = image_std if image_std is not None else self.image_std
|
||||
do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
|
||||
|
||||
images = make_list_of_images(images)
|
||||
|
||||
if not valid_images(images):
|
||||
raise ValueError(
|
||||
"Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
|
||||
"torch.Tensor, tf.Tensor or jax.ndarray."
|
||||
)
|
||||
validate_preprocess_arguments(
|
||||
do_rescale=do_rescale,
|
||||
rescale_factor=rescale_factor,
|
||||
do_normalize=do_normalize,
|
||||
image_mean=image_mean,
|
||||
image_std=image_std,
|
||||
do_resize=do_resize,
|
||||
size=size,
|
||||
resample=resample,
|
||||
)
|
||||
# All transformations expect numpy arrays.
|
||||
images = [to_numpy_array(image) for image in images]
|
||||
|
||||
if do_convert_rgb:
|
||||
images = [convert_to_rgb(image) for image in images]
|
||||
|
||||
if is_scaled_image(images[0]) and do_rescale:
|
||||
logger.warning_once(
|
||||
"It looks like you are trying to rescale already rescaled images. If the input"
|
||||
" images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
|
||||
)
|
||||
|
||||
if input_data_format is None:
|
||||
# We assume that all images have the same channel dimension format.
|
||||
input_data_format = infer_channel_dimension_format(images[0])
|
||||
|
||||
if do_resize:
|
||||
height, width = size["height"], size["width"]
|
||||
images = [
|
||||
resize(image=image, size=(height, width), resample=resample, input_data_format=input_data_format)
|
||||
for image in images
|
||||
]
|
||||
|
||||
if do_rescale:
|
||||
images = [
|
||||
self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
|
||||
for image in images
|
||||
]
|
||||
|
||||
if do_normalize:
|
||||
images = [
|
||||
self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
|
||||
for image in images
|
||||
]
|
||||
|
||||
images = [
|
||||
to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
|
||||
]
|
||||
|
||||
data = {"pixel_values": images}
|
||||
return BatchFeature(data=data, tensor_type=return_tensors)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,131 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""
|
||||
Image/Text processor class for SigLIP.
|
||||
"""
|
||||
|
||||
from typing import List, Optional, Union
|
||||
|
||||
from transformers.feature_extraction_utils import BatchFeature
|
||||
from transformers.image_utils import ImageInput
|
||||
from transformers.processing_utils import ProcessorMixin
|
||||
from transformers.tokenization_utils_base import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
|
||||
from transformers.utils import TensorType
|
||||
|
||||
|
||||
class SiglipProcessor(ProcessorMixin):
|
||||
r"""
|
||||
Constructs a Siglip processor which wraps a Siglip image processor and a Siglip tokenizer into a single processor.
|
||||
|
||||
[`SiglipProcessor`] offers all the functionalities of [`SiglipImageProcessor`] and [`SiglipTokenizer`]. See the
|
||||
[`~SiglipProcessor.__call__`] and [`~SiglipProcessor.decode`] for more information.
|
||||
|
||||
Args:
|
||||
image_processor ([`SiglipImageProcessor`]):
|
||||
The image processor is a required input.
|
||||
tokenizer ([`SiglipTokenizer`]):
|
||||
The tokenizer is a required input.
|
||||
"""
|
||||
|
||||
attributes = ["image_processor", "tokenizer"]
|
||||
image_processor_class = "SiglipImageProcessor"
|
||||
tokenizer_class = "SiglipTokenizer"
|
||||
|
||||
def __init__(self, image_processor, tokenizer):
|
||||
super().__init__(image_processor, tokenizer)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
|
||||
images: ImageInput = None,
|
||||
padding: Union[bool, str, PaddingStrategy] = False,
|
||||
truncation: Union[bool, str, TruncationStrategy] = None,
|
||||
max_length: int = None,
|
||||
return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH,
|
||||
) -> BatchFeature:
|
||||
"""
|
||||
Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
|
||||
and `kwargs` arguments to SiglipTokenizer's [`~SiglipTokenizer.__call__`] if `text` is not `None` to encode
|
||||
the text. To prepare the image(s), this method forwards the `images` argument to
|
||||
SiglipImageProcessor's [`~SiglipImageProcessor.__call__`] if `images` is not `None`. Please refer to the doctsring
|
||||
of the above two methods for more information.
|
||||
|
||||
Args:
|
||||
text (`str`, `List[str]`, `List[List[str]]`):
|
||||
The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
|
||||
(pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
|
||||
`is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
|
||||
images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
|
||||
The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
|
||||
tensor. Both channels-first and channels-last formats are supported.
|
||||
padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):
|
||||
Select a strategy to pad the returned sequences (according to the model's padding side and padding
|
||||
index) among:
|
||||
- `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
|
||||
sequence if provided).
|
||||
- `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
|
||||
acceptable input length for the model if that argument is not provided.
|
||||
- `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
|
||||
lengths).
|
||||
max_length (`int`, *optional*):
|
||||
Maximum length of the returned list and optionally padding length (see above).
|
||||
truncation (`bool`, *optional*):
|
||||
Activates truncation to cut input sequences longer than `max_length` to `max_length`.
|
||||
return_tensors (`str` or [`~utils.TensorType`], *optional*):
|
||||
If set, will return tensors of a particular framework. Acceptable values are:
|
||||
|
||||
- `'tf'`: Return TensorFlow `tf.constant` objects.
|
||||
- `'pt'`: Return PyTorch `torch.Tensor` objects.
|
||||
- `'np'`: Return NumPy `np.ndarray` objects.
|
||||
- `'jax'`: Return JAX `jnp.ndarray` objects.
|
||||
|
||||
Returns:
|
||||
[`BatchFeature`]: A [`BatchFeature`] with the following fields:
|
||||
|
||||
- **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
|
||||
- **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
|
||||
`return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
|
||||
`None`).
|
||||
- **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
|
||||
"""
|
||||
|
||||
if text is None and images is None:
|
||||
raise ValueError("You have to specify either text or images. Both cannot be none.")
|
||||
|
||||
if text is not None:
|
||||
encoding = self.tokenizer(
|
||||
text, return_tensors=return_tensors, padding=padding, truncation=truncation, max_length=max_length
|
||||
)
|
||||
|
||||
if images is not None:
|
||||
image_features = self.image_processor(images, return_tensors=return_tensors)
|
||||
|
||||
if text is not None and images is not None:
|
||||
encoding["pixel_values"] = image_features.pixel_values
|
||||
return encoding
|
||||
elif text is not None:
|
||||
return encoding
|
||||
else:
|
||||
return BatchFeature(data=dict(**image_features), tensor_type=return_tensors)
|
||||
|
||||
def decode(self, *args, **kwargs):
|
||||
"""
|
||||
This method forwards all its arguments to SiglipTokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to
|
||||
the docstring of this method for more information.
|
||||
"""
|
||||
return self.tokenizer.decode(*args, **kwargs)
|
||||
|
||||
def batch_decode(self, *args, **kwargs):
|
||||
"""
|
||||
This method forwards all its arguments to SiglipTokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please
|
||||
refer to the docstring of this method for more information.
|
||||
"""
|
||||
return self.tokenizer.batch_decode(*args, **kwargs)
|
||||
|
||||
@property
|
||||
# Copied from transformers.models.clip.processing_clip.CLIPProcessor.model_input_names with CLIP->Siglip, T5->Siglip
|
||||
def model_input_names(self):
|
||||
tokenizer_input_names = self.tokenizer.model_input_names
|
||||
image_processor_input_names = self.image_processor.model_input_names
|
||||
return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
|
||||
@@ -0,0 +1,364 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Tokenization class for SigLIP model."""
|
||||
|
||||
import os
|
||||
import re
|
||||
import string
|
||||
import warnings
|
||||
from shutil import copyfile
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
||||
|
||||
import sentencepiece as spm
|
||||
|
||||
from transformers.convert_slow_tokenizer import import_protobuf
|
||||
from transformers.tokenization_utils import PreTrainedTokenizer
|
||||
from transformers.tokenization_utils_base import AddedToken
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from transformers.tokenization_utils_base import TextInput
|
||||
from transformers.utils import logging, requires_backends
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
VOCAB_FILES_NAMES = {"vocab_file": "spiece.model"}
|
||||
|
||||
|
||||
SPIECE_UNDERLINE = "▁"
|
||||
|
||||
|
||||
class SiglipTokenizer(PreTrainedTokenizer):
|
||||
"""
|
||||
Construct a Siglip tokenizer. Based on [SentencePiece](https://github.com/google/sentencepiece).
|
||||
|
||||
This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
|
||||
this superclass for more information regarding those methods.
|
||||
|
||||
Args:
|
||||
vocab_file (`str`):
|
||||
[SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that
|
||||
contains the vocabulary necessary to instantiate a tokenizer.
|
||||
eos_token (`str`, *optional*, defaults to `"</s>"`):
|
||||
The end of sequence token.
|
||||
unk_token (`str`, *optional*, defaults to `"<unk>"`):
|
||||
The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
|
||||
token instead.
|
||||
pad_token (`str`, *optional*, defaults to `"</s>"`):
|
||||
The token used for padding, for example when batching sequences of different lengths.
|
||||
additional_special_tokens (`List[str]`, *optional*):
|
||||
Additional special tokens used by the tokenizer.
|
||||
sp_model_kwargs (`dict`, *optional*):
|
||||
Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
|
||||
SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
|
||||
to set:
|
||||
|
||||
- `enable_sampling`: Enable subword regularization.
|
||||
- `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
|
||||
|
||||
- `nbest_size = {0,1}`: No sampling is performed.
|
||||
- `nbest_size > 1`: samples from the nbest_size results.
|
||||
- `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
|
||||
using forward-filtering-and-backward-sampling algorithm.
|
||||
|
||||
- `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
|
||||
BPE-dropout.
|
||||
model_max_length (`int`, *optional*, defaults to 64):
|
||||
The maximum length (in number of tokens) for model inputs.
|
||||
do_lower_case (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to lowercase the input when tokenizing.
|
||||
"""
|
||||
|
||||
vocab_files_names = VOCAB_FILES_NAMES
|
||||
model_input_names = ["input_ids", "attention_mask"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_file,
|
||||
eos_token="</s>",
|
||||
unk_token="<unk>",
|
||||
pad_token="</s>",
|
||||
additional_special_tokens=None,
|
||||
sp_model_kwargs: Optional[Dict[str, Any]] = None,
|
||||
model_max_length=64,
|
||||
do_lower_case=True,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
requires_backends(self, "protobuf")
|
||||
|
||||
pad_token = (
|
||||
AddedToken(pad_token, rstrip=True, lstrip=True, normalized=False, special=True)
|
||||
if isinstance(pad_token, str)
|
||||
else pad_token
|
||||
)
|
||||
unk_token = (
|
||||
AddedToken(unk_token, rstrip=True, lstrip=True, normalized=False, special=True)
|
||||
if isinstance(unk_token, str)
|
||||
else unk_token
|
||||
)
|
||||
eos_token = (
|
||||
AddedToken(eos_token, rstrip=True, lstrip=True, normalized=False, special=True)
|
||||
if isinstance(eos_token, str)
|
||||
else eos_token
|
||||
)
|
||||
|
||||
self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
|
||||
|
||||
self.do_lower_case = do_lower_case
|
||||
self.vocab_file = vocab_file
|
||||
|
||||
self.sp_model = self.get_spm_processor()
|
||||
self.vocab_file = vocab_file
|
||||
|
||||
super().__init__(
|
||||
eos_token=eos_token,
|
||||
unk_token=unk_token,
|
||||
pad_token=pad_token,
|
||||
additional_special_tokens=additional_special_tokens,
|
||||
sp_model_kwargs=self.sp_model_kwargs,
|
||||
model_max_length=model_max_length,
|
||||
do_lower_case=do_lower_case,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def get_spm_processor(self):
|
||||
tokenizer = spm.SentencePieceProcessor(**self.sp_model_kwargs)
|
||||
with open(self.vocab_file, "rb") as f:
|
||||
sp_model = f.read()
|
||||
model_pb2 = import_protobuf()
|
||||
model = model_pb2.ModelProto.FromString(sp_model)
|
||||
normalizer_spec = model_pb2.NormalizerSpec()
|
||||
normalizer_spec.add_dummy_prefix = False
|
||||
model.normalizer_spec.MergeFrom(normalizer_spec)
|
||||
sp_model = model.SerializeToString()
|
||||
tokenizer.LoadFromSerializedProto(sp_model)
|
||||
return tokenizer
|
||||
|
||||
@property
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.vocab_size
|
||||
def vocab_size(self):
|
||||
return self.sp_model.get_piece_size()
|
||||
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.get_vocab
|
||||
def get_vocab(self):
|
||||
vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
|
||||
vocab.update(self.added_tokens_encoder)
|
||||
return vocab
|
||||
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.get_special_tokens_mask
|
||||
def get_special_tokens_mask(
|
||||
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
|
||||
) -> List[int]:
|
||||
"""
|
||||
Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
|
||||
special tokens using the tokenizer `prepare_for_model` method.
|
||||
|
||||
Args:
|
||||
token_ids_0 (`List[int]`):
|
||||
List of IDs.
|
||||
token_ids_1 (`List[int]`, *optional*):
|
||||
Optional second list of IDs for sequence pairs.
|
||||
already_has_special_tokens (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not the token list is already formatted with special tokens for the model.
|
||||
|
||||
Returns:
|
||||
`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
|
||||
"""
|
||||
if already_has_special_tokens:
|
||||
return super().get_special_tokens_mask(
|
||||
token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
|
||||
)
|
||||
|
||||
# normal case: some special tokens
|
||||
if token_ids_1 is None:
|
||||
return ([0] * len(token_ids_0)) + [1]
|
||||
return ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
|
||||
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer._add_eos_if_not_present
|
||||
def _add_eos_if_not_present(self, token_ids: List[int]) -> List[int]:
|
||||
"""Do not add eos again if user already added it."""
|
||||
if len(token_ids) > 0 and token_ids[-1] == self.eos_token_id:
|
||||
warnings.warn(
|
||||
f"This sequence already has {self.eos_token}. In future versions this behavior may lead to duplicated"
|
||||
" eos tokens being added."
|
||||
)
|
||||
return token_ids
|
||||
else:
|
||||
return token_ids + [self.eos_token_id]
|
||||
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.create_token_type_ids_from_sequences
|
||||
def create_token_type_ids_from_sequences(
|
||||
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
|
||||
) -> List[int]:
|
||||
"""
|
||||
Create a mask from the two sequences passed to be used in a sequence-pair classification task. T5 does not make
|
||||
use of token type ids, therefore a list of zeros is returned.
|
||||
|
||||
Args:
|
||||
token_ids_0 (`List[int]`):
|
||||
List of IDs.
|
||||
token_ids_1 (`List[int]`, *optional*):
|
||||
Optional second list of IDs for sequence pairs.
|
||||
|
||||
Returns:
|
||||
`List[int]`: List of zeros.
|
||||
"""
|
||||
eos = [self.eos_token_id]
|
||||
|
||||
if token_ids_1 is None:
|
||||
return len(token_ids_0 + eos) * [0]
|
||||
return len(token_ids_0 + eos + token_ids_1 + eos) * [0]
|
||||
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.build_inputs_with_special_tokens
|
||||
def build_inputs_with_special_tokens(
|
||||
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
|
||||
) -> List[int]:
|
||||
"""
|
||||
Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
|
||||
adding special tokens. A sequence has the following format:
|
||||
|
||||
- single sequence: `X </s>`
|
||||
- pair of sequences: `A </s> B </s>`
|
||||
|
||||
Args:
|
||||
token_ids_0 (`List[int]`):
|
||||
List of IDs to which the special tokens will be added.
|
||||
token_ids_1 (`List[int]`, *optional*):
|
||||
Optional second list of IDs for sequence pairs.
|
||||
|
||||
Returns:
|
||||
`List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
|
||||
"""
|
||||
token_ids_0 = self._add_eos_if_not_present(token_ids_0)
|
||||
if token_ids_1 is None:
|
||||
return token_ids_0
|
||||
else:
|
||||
token_ids_1 = self._add_eos_if_not_present(token_ids_1)
|
||||
return token_ids_0 + token_ids_1
|
||||
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.__getstate__
|
||||
def __getstate__(self):
|
||||
state = self.__dict__.copy()
|
||||
state["sp_model"] = None
|
||||
return state
|
||||
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.__setstate__
|
||||
def __setstate__(self, d):
|
||||
self.__dict__ = d
|
||||
|
||||
# for backward compatibility
|
||||
if not hasattr(self, "sp_model_kwargs"):
|
||||
self.sp_model_kwargs = {}
|
||||
|
||||
self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
|
||||
self.sp_model.Load(self.vocab_file)
|
||||
|
||||
def remove_punctuation(self, text: str) -> str:
|
||||
return text.translate(str.maketrans("", "", string.punctuation))
|
||||
|
||||
# source: https://github.com/google-research/big_vision/blob/3b8e5ab6ad4f96e32b32826f9e1b8fd277914f9c/big_vision/evaluators/proj/image_text/prompt_engineering.py#L94
|
||||
def canonicalize_text(self, text, *, keep_punctuation_exact_string=None):
|
||||
"""Returns canonicalized `text` (puncuation removed).
|
||||
|
||||
Args:
|
||||
text (`str`):
|
||||
String to be canonicalized.
|
||||
keep_punctuation_exact_string (`str`, *optional*):
|
||||
If provided, then this exact string is kept. For example providing '{}' will keep any occurrences of '{}'
|
||||
(but will still remove '{' and '}' that appear separately).
|
||||
"""
|
||||
if keep_punctuation_exact_string:
|
||||
text = keep_punctuation_exact_string.join(
|
||||
self.remove_punctuation(part) for part in text.split(keep_punctuation_exact_string)
|
||||
)
|
||||
else:
|
||||
text = self.remove_punctuation(text)
|
||||
text = re.sub(r"\s+", " ", text)
|
||||
text = text.strip()
|
||||
|
||||
return text
|
||||
|
||||
def tokenize(self, text: "TextInput", add_special_tokens=False, **kwargs) -> List[str]:
|
||||
"""
|
||||
Converts a string to a list of tokens.
|
||||
"""
|
||||
tokens = super().tokenize(SPIECE_UNDERLINE + text.replace(SPIECE_UNDERLINE, " "), **kwargs)
|
||||
|
||||
if len(tokens) > 1 and tokens[0] == SPIECE_UNDERLINE and tokens[1] in self.all_special_tokens:
|
||||
tokens = tokens[1:]
|
||||
return tokens
|
||||
|
||||
@property
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.unk_token_length
|
||||
def unk_token_length(self):
|
||||
return len(self.sp_model.encode(str(self.unk_token)))
|
||||
|
||||
def _tokenize(self, text, **kwargs):
|
||||
"""
|
||||
Returns a tokenized string.
|
||||
|
||||
We de-activated the `add_dummy_prefix` option, thus the sentencepiece internals will always strip any
|
||||
SPIECE_UNDERLINE.
|
||||
|
||||
For example: `self.sp_model.encode(f"{SPIECE_UNDERLINE}Hey", out_type = str)` will give `['H', 'e', 'y']` instead of `['▁He', 'y']`.
|
||||
|
||||
Thus we always encode `f"{unk_token}text"` and strip the `unk_token`. Here is an example with `unk_token = "<unk>"` and `unk_token_length = 4`.
|
||||
`self.tokenizer.sp_model.encode("<unk> Hey", out_type = str)[4:]`.
|
||||
"""
|
||||
text = self.canonicalize_text(text, keep_punctuation_exact_string=None)
|
||||
tokens = self.sp_model.encode(text, out_type=str)
|
||||
|
||||
# 1. Encode string + prefix ex: "<unk> Hey"
|
||||
tokens = self.sp_model.encode(self.unk_token + text, out_type=str)
|
||||
# 2. Remove self.unk_token from ['<','unk','>', '▁Hey']
|
||||
return tokens[self.unk_token_length :] if len(tokens) >= self.unk_token_length else tokens
|
||||
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer._convert_token_to_id
|
||||
def _convert_token_to_id(self, token):
|
||||
"""Converts a token (str) in an id using the vocab."""
|
||||
return self.sp_model.piece_to_id(token)
|
||||
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer._convert_id_to_token
|
||||
def _convert_id_to_token(self, index):
|
||||
"""Converts an index (integer) in a token (str) using the vocab."""
|
||||
token = self.sp_model.IdToPiece(index)
|
||||
return token
|
||||
|
||||
def convert_tokens_to_string(self, tokens):
|
||||
"""Converts a sequence of tokens (string) in a single string."""
|
||||
current_sub_tokens = []
|
||||
out_string = ""
|
||||
prev_is_special = False
|
||||
for token in tokens:
|
||||
# make sure that special tokens are not decoded using sentencepiece model
|
||||
if token in self.all_special_tokens:
|
||||
if not prev_is_special:
|
||||
out_string += " "
|
||||
out_string += self.sp_model.decode(current_sub_tokens) + token
|
||||
prev_is_special = True
|
||||
current_sub_tokens = []
|
||||
else:
|
||||
current_sub_tokens.append(token)
|
||||
prev_is_special = False
|
||||
out_string += self.sp_model.decode(current_sub_tokens)
|
||||
return out_string.strip()
|
||||
|
||||
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.save_vocabulary
|
||||
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
|
||||
if not os.path.isdir(save_directory):
|
||||
logger.error(f"Vocabulary path ({save_directory}) should be a directory")
|
||||
return
|
||||
out_vocab_file = os.path.join(
|
||||
save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
|
||||
)
|
||||
|
||||
if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
|
||||
copyfile(self.vocab_file, out_vocab_file)
|
||||
elif not os.path.isfile(self.vocab_file):
|
||||
with open(out_vocab_file, "wb") as fi:
|
||||
content_spiece_model = self.sp_model.serialized_model_proto()
|
||||
fi.write(content_spiece_model)
|
||||
|
||||
return (out_vocab_file,)
|
||||
@@ -0,0 +1,18 @@
|
||||
decord==0.6.0
|
||||
einops==0.8.1
|
||||
huggingface_hub==0.29.1
|
||||
matplotlib==3.7.0
|
||||
numpy==1.24.4
|
||||
opencv_python==4.7.0.72
|
||||
pyarrow==11.0.0
|
||||
PyYAML==6.0.2
|
||||
Requests==2.32.3
|
||||
safetensors==0.4.5
|
||||
scipy==1.10.1
|
||||
sentencepiece==0.1.99
|
||||
torch==2.5.1
|
||||
torchvision==0.20.1
|
||||
transformers==4.49.0
|
||||
flash_attn==2.5.8
|
||||
accelerate>=0.34.0
|
||||
wandb
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 58 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 35 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 85 KiB |
@@ -18,14 +18,16 @@ Origin methods from
|
||||
* [RealCustom](https://github.com/bytedance/RealCustom),
|
||||
* [InstantCharacter](https://github.com/Tencent/InstantCharacter),
|
||||
* [DreamO](https://github.com/bytedance/DreamO)
|
||||
* [Bagel](https://github.com/ByteDance-Seed/Bagel)
|
||||
|
||||
|
||||
## Updates:
|
||||
* 2025/05/23
|
||||
* DreamO的方法ip id style方法实现，双人同框使用ip+ip，默认都是ip模式。带人脸的可以用ip，也可以用id（可以不连入衣服），pos 和neg lora在lora的目录下时默认开启，如果没有就是3 lora模式。开启id和style模式，需要在extra 输入id或 style
|
||||
* 2025/06/02
|
||||
* 新增Bagel模型的支持，支持int8和nf4量化（官方用的十字鱼佬的PR）输入图片则是edit模式，不输入就是文生图，在量化nf4的情况下，显存峰值大约7G，实际跑4G多,edit的编辑能力,在nf4条件下一般；
|
||||
|
||||
|
||||
## previous
|
||||
* DreamO的方法ip id style方法实现，双人同框使用ip+ip，默认都是ip模式。带人脸的可以用ip，也可以用id（可以不连入衣服），pos 和neg lora在lora的目录下时默认开启，如果没有就是3 lora模式。开启id和style模式，需要在extra 输入id或 style
|
||||
* 新增2个ID迁移的方法实现，分别是RealCustom（SDXL）和InstantCharacter（FLUX），基准测试在4070 12G，二个方法的速度都很慢，InstantCharacter支持多种量化，如果使用双截棍量化加速很快，但是没意义，因为IP层没加载进去，具体看示例图和新的工作流文件，RealCustom需要6个单体模型，InstantCharacter需要2个repo形式的clip_vison(暂时没空改)，16G以上显存会好点
|
||||
* 利用uno的功能来实现flux流程的双角色同框，prompt示例见图；
|
||||
* 修复ms-diffusion的双角色提示词错误，使用ms diffusion 角色提示词应该是 [A] a (man)... ,[B] a (woman)...,场景提示词不用改，还是[A] ...[B]...在同一句里时开启；
|
||||
@@ -232,10 +234,18 @@ turbo lora：[alimama-creative/FLUX.1-Turbo-Alpha](https://huggingface.co/alimam
|
||||
├── ComfyUI/models/photomaker/
|
||||
├──FLUX.1-Turbo-Alpha.safetensors #rename 重命名的turbo lora
|
||||
├── anypath/black-forest-labs/FLUX.1-dev
|
||||
├── ComfyUI/models//BEN2_Base.pth #or any path
|
||||
├── ComfyUI/models/BEN2_Base.pth #or any path
|
||||
```
|
||||
|
||||
**3.11 Bagel mode**
|
||||
download [BAGEL-7B-MoT](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT/tree/main)
|
||||
```
|
||||
├── ComfyUI/models/vae/
|
||||
├──ae.safetensors # flux or BAGEL-7B-MoT
|
||||
|
||||
├── Any/path/ByteDance-Seed/BAGEL-7B-MoT/
|
||||
├──all files # 所有文件
|
||||
```
|
||||
|
||||
4 Example
|
||||
----
|
||||
@@ -300,7 +310,14 @@ turbo lora：[alimama-creative/FLUX.1-Turbo-Alpha](https://huggingface.co/alimam
|
||||
* nf4 id
|
||||
<img src="https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/images/dreamo_id.png" width="50%">
|
||||
|
||||
**4.12 comfyUI classic（comfyUI经典模式，可以接任意适配CF的流程，主要是方便使用多角色的clip）**
|
||||
|
||||
**4.11 Bagel**
|
||||
* nf4 image2image
|
||||
<img src="https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/images/bagel_img2img.png" width="50%">
|
||||
* nf4 txt2image
|
||||
<img src="https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/images/bagel_txt2img.png" width="50%">
|
||||
|
||||
**4.13 comfyUI classic（comfyUI经典模式，可以接任意适配CF的流程，主要是方便使用多角色的clip）**
|
||||
* any mode SD1.5 SDXL SD3.5 FLUX...
|
||||
<img src="https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/images/comfyui_classic.png" width="50%">
|
||||
|
||||
@@ -439,3 +456,12 @@ svdquant
|
||||
}
|
||||
```
|
||||
[DreamO](https://github.com/bytedance/DreamO)
|
||||
|
||||
```
|
||||
@article{deng2025bagel,
|
||||
title = {Emerging Properties in Unified Multimodal Pretraining},
|
||||
author = {Deng, Chaorui and Zhu, Deyao and Li, Kunchang and Gou, Chenhui and Li, Feng and Wang, Zeyu and Zhong, Shu and Yu, Weihao and Nie, Xiaonan and Song, Ziang and Shi, Guang and Fan, Haoqi},
|
||||
journal = {arXiv preprint arXiv:2505.14683},
|
||||
year = {2025}
|
||||
}
|
||||
```
|
||||
+100
-20
@@ -137,7 +137,7 @@ class StoryDiffusion_Apply:
|
||||
{
|
||||
"model": ("MODEL",),
|
||||
"vae": ("VAE",),
|
||||
"infer_mode": (["story", "classic","flux_pulid","infiniteyou","uno","realcustom","instant_character","dreamo","story_maker","story_and_maker","consistory","kolor_face","msdiffusion" ],),
|
||||
"infer_mode": (["story", "classic","flux_pulid","infiniteyou","uno","realcustom","instant_character","dreamo","bagel_edit","story_maker","story_and_maker","consistory","kolor_face","msdiffusion" ],),
|
||||
"photomake_ckpt": (["none"] + [i for i in folder_paths.get_filename_list("photomaker") if "v1" in i or "v2" in i],),
|
||||
"ipadapter_ckpt": (["none"] + folder_paths.get_filename_list("photomaker"),),
|
||||
"quantize_mode": ([ "fp8", "nf4","fp16", ],),
|
||||
@@ -268,6 +268,14 @@ class StoryDiffusion_Apply:
|
||||
if dreamo_lora_path is None or cfg_distill_path is None or ipadapter_ckpt_path is None:
|
||||
raise "dreamo need a dreamo lora and cfg distill and turbo lora in ipadapter menu"
|
||||
model = Loader_Dreamo(model,vae,quantize_mode,dreamo_lora_path,cfg_distill_path,ipadapter_ckpt_path,device)
|
||||
elif infer_mode == "bagel_edit":
|
||||
from .Bagel.app import load_bagel_model
|
||||
if not isinstance(model,dict) :
|
||||
raise " must link EasyFunction_Lite node "
|
||||
elif repo is None:
|
||||
raise "EasyFunction_Lite node extra_repo must fill bagel repo"
|
||||
max_mem_per_gpu=str(int(total_vram/1000))+"GIB"
|
||||
model = load_bagel_model(repo,quantize_mode,max_mem_per_gpu)
|
||||
else: # can not choice a mode
|
||||
print("infer use comfyui classic mode")
|
||||
|
||||
@@ -426,6 +434,21 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
|
||||
only_role_emb_ne,x_1_refs_dict,image_emb,x_1_refs_dual=None,None,None,None
|
||||
|
||||
input_id_images_dict={}
|
||||
image_list=[]
|
||||
|
||||
if img2img_mode and image is not None:
|
||||
if len(role_list)==1:
|
||||
image_pil=nomarl_upscale(image, width, height)
|
||||
input_id_images_dict[role_list[0]]=image_pil
|
||||
image_list=[image_pil]
|
||||
else:
|
||||
f1, _, _, _ = image.size()
|
||||
img_list = list(torch.chunk(image, chunks=f1))
|
||||
image_list = [nomarl_upscale(img, width, height) for img in img_list]
|
||||
for index, key in enumerate(role_list):
|
||||
input_id_images_dict[key]=image_list[index]
|
||||
|
||||
if img2img_mode:
|
||||
if infer_mode == "story_and_maker" or infer_mode == "story_maker":
|
||||
k1, _, _, _ = image.size()
|
||||
@@ -466,6 +489,8 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
|
||||
# get emb use insightface
|
||||
pass
|
||||
elif infer_mode == "bagel_edit":
|
||||
image_emb=input_id_images_dict
|
||||
elif infer_mode == "dreamo":
|
||||
from .model_loader_utils import Dreamo_image_encoder
|
||||
from huggingface_hub import hf_hub_download
|
||||
@@ -587,20 +612,7 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
else:
|
||||
input_id_emb_s_dict,input_id_img_s_dict,input_id_emb_un_dict,input_id_cloth_dict={}, {}, {}, {}
|
||||
|
||||
input_id_images_dict={}
|
||||
image_list=[]
|
||||
|
||||
if img2img_mode and image is not None:
|
||||
if len(role_list)==1:
|
||||
image_pil=nomarl_upscale(image, width, height)
|
||||
input_id_images_dict[role_list[0]]=image_pil
|
||||
image_list=[image_pil]
|
||||
else:
|
||||
f1, _, _, _ = image.size()
|
||||
img_list = list(torch.chunk(image, chunks=f1))
|
||||
image_list = [nomarl_upscale(img, width, height) for img in img_list]
|
||||
for index, key in enumerate(role_list):
|
||||
input_id_images_dict[key]=image_list[index]
|
||||
|
||||
|
||||
# pre clip txt emb
|
||||
|
||||
@@ -688,7 +700,6 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
inp = prepare_multi_ip_wrapper(clip,prompt=p, ref_imgs=x_1, pe=uno_pe,device=device,h=height,w=width)
|
||||
ip_emb.append(inp)
|
||||
only_role_emb[key]=ip_emb
|
||||
|
||||
elif infer_mode=="kolor_face":
|
||||
from .model_loader_utils import glm_single_encode
|
||||
from .kolors.models.tokenization_chatglm import ChatGLMTokenizer
|
||||
@@ -727,7 +738,7 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
else:
|
||||
if infer_mode=="classic": #TODO 逆序的角色会出现iD不匹配，受影响的有story文生图
|
||||
only_role_emb= cf_clip(only_role_list, clip, infer_mode,role_list) #story模式需要拆分prompt，所以这里需要传入role_list
|
||||
elif infer_mode=="dreamo":
|
||||
elif infer_mode=="dreamo" or infer_mode=="bagel_edit":
|
||||
pass # TODO 暂时不支持dreamo
|
||||
else:
|
||||
only_role_emb= cf_clip(inf_list_split, clip, infer_mode,role_list) #story,story_maker,story_and_maker,msdiffusion,infinite
|
||||
@@ -746,7 +757,7 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
else:
|
||||
if infer_mode!="kolor_face":
|
||||
nc_emb=cf_clip(nc_txt_list, clip, infer_mode,role_list,input_split=False)
|
||||
elif infer_mode=="dreamo":
|
||||
elif infer_mode=="dreamo" or infer_mode=="bagel_edit":
|
||||
pass # TODO 暂时不支持dreamo
|
||||
else:
|
||||
nc_emb,_= glm_single_encode(chatglm3_model, nc_txt_list,role_list, neg_text, 1,nc_mode=True)
|
||||
@@ -849,6 +860,12 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
only_role_emb[key]=prompts
|
||||
postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb}
|
||||
negative = neg_text # TODO
|
||||
elif infer_mode=="bagel_edit":
|
||||
only_role_emb={}
|
||||
for key ,prompts in zip(role_list,inf_list_split):
|
||||
only_role_emb[key]=prompts
|
||||
postive_dict= {"role": only_role_emb, "nc": None, "daul": None}
|
||||
negative = neg_text # TODO
|
||||
elif infer_mode=="realcustom":
|
||||
postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb}
|
||||
negative=[letent_real]
|
||||
@@ -1594,8 +1611,6 @@ class StoryDiffusion_KSampler:
|
||||
for emb_list in only_role_emb[key]:
|
||||
samples = model(
|
||||
prompt_embeds=emb_list[0],
|
||||
height=height,
|
||||
width=width,
|
||||
pooled_prompt_embeds=emb_list[1],
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=cfg,
|
||||
@@ -1653,6 +1668,71 @@ class StoryDiffusion_KSampler:
|
||||
out = {}
|
||||
out["samples"] = torch.cat(samples_list, dim=0)
|
||||
return (out,)
|
||||
elif infer_mode =="bagel_edit":
|
||||
|
||||
VAE=info.get("VAE")
|
||||
samples_list = []
|
||||
if img2img_mode:
|
||||
from .Bagel.app import edit_image
|
||||
for key in role_list:
|
||||
for prompt in only_role_emb[key]:
|
||||
samples,text_=edit_image(
|
||||
model,
|
||||
image=image_embeds[key],
|
||||
prompt=prompt,
|
||||
show_thinking=False,
|
||||
cfg_text_scale=cfg,
|
||||
cfg_img_scale=2.0,
|
||||
cfg_interval=0.0,
|
||||
timestep_shift=3.0,
|
||||
num_timesteps=steps,
|
||||
cfg_renorm_min=0.0,
|
||||
cfg_renorm_type="text_channel",
|
||||
max_think_token_n=1024,
|
||||
do_sample=False,
|
||||
text_temperature=0.3,
|
||||
seed=seed) # tuple (image,str)
|
||||
print(f'thinking text is {text_}')
|
||||
samples_list.append(samples)
|
||||
else:
|
||||
from .Bagel.app import text_to_image
|
||||
if width==height: #todo
|
||||
image_ratio="1:1"
|
||||
elif width/height==0.75:
|
||||
image_ratio="4:3"
|
||||
elif width/height==3/4:
|
||||
image_ratio="3:4"
|
||||
elif width/height==9/16:
|
||||
image_ratio="9:16"
|
||||
elif width/height==16/9:
|
||||
image_ratio="16:9"
|
||||
else:
|
||||
image_ratio="1:1"
|
||||
|
||||
for key in role_list:
|
||||
for prompt in only_role_emb[key]:
|
||||
samples,text_=text_to_image(
|
||||
model,
|
||||
prompt=prompt,
|
||||
show_thinking=False,
|
||||
cfg_text_scale=cfg,
|
||||
cfg_interval=0.4,
|
||||
timestep_shift=3.0,
|
||||
num_timesteps=steps,
|
||||
cfg_renorm_min=0.0,
|
||||
cfg_renorm_type="global",
|
||||
max_think_token_n=1024,
|
||||
do_sample=False,
|
||||
text_temperature=0.3,
|
||||
seed=seed,
|
||||
image_ratio=image_ratio) # tuple (image,str)
|
||||
print(f'thinking text is {text_}')
|
||||
samples_list.append(samples)
|
||||
|
||||
out = {}
|
||||
samples_list=[VAE.encode(phi2narry(pixels)[:,:,:,:3]) for pixels in samples_list] # TODO pil to tensor
|
||||
out["samples"] = torch.cat(samples_list, dim=0)
|
||||
return (out,)
|
||||
else: #none:
|
||||
return
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Binary file not shown.
|
After Width: | Height: | Size: 362 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 423 KiB |
Reference in New Issue
Block a user