Instructions to use worksimpli/HiDream-O1-Image-endpoint with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use worksimpli/HiDream-O1-Image-endpoint with Transformers:
# pip install -U transformers accelerate # Load model directly from transformers import AutoProcessor, AutoModelForMultimodalLM processor = AutoProcessor.from_pretrained("worksimpli/HiDream-O1-Image-endpoint") model = AutoModelForMultimodalLM.from_pretrained("worksimpli/HiDream-O1-Image-endpoint", device_map="auto") - Notebooks
- Google Colab
- Kaggle
Add files using upload-large-folder tool
Browse files- .gitattributes +5 -0
- .mdl +0 -0
- .msc +0 -0
- .mv +1 -0
- README.md +379 -0
- assets/IP.webp +3 -0
- assets/IP_2.jpg +3 -0
- assets/general.webp +3 -0
- assets/leaderboard.png +3 -0
- assets/text-layout.webp +3 -0
- chat_template.json +3 -0
- config.json +62 -0
- configuration.json +1 -0
- generation_config.json +14 -0
- handler.py +96 -0
- merges.txt +0 -0
- model-00001-of-00008.safetensors +3 -0
- model-00002-of-00008.safetensors +3 -0
- model-00003-of-00008.safetensors +3 -0
- model-00004-of-00008.safetensors +3 -0
- model-00005-of-00008.safetensors +3 -0
- model-00006-of-00008.safetensors +3 -0
- model-00007-of-00008.safetensors +3 -0
- model-00008-of-00008.safetensors +3 -0
- model.safetensors.index.json +767 -0
- models/__init__.py +1 -0
- models/flash_scheduler.py +445 -0
- models/fm_solvers_unipc.py +800 -0
- models/pipeline.py +435 -0
- models/qwen3_vl_transformers.py +0 -0
- models/utils.py +368 -0
- preprocessor_config.json +21 -0
- requirements.txt +5 -0
- tokenizer.json +0 -0
- tokenizer_config.json +244 -0
- video_preprocessor_config.json +21 -0
- vocab.json +0 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
assets/IP_2.jpg filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
assets/IP.webp filter=lfs diff=lfs merge=lfs -text
|
| 38 |
+
assets/leaderboard.png filter=lfs diff=lfs merge=lfs -text
|
| 39 |
+
assets/text-layout.webp filter=lfs diff=lfs merge=lfs -text
|
| 40 |
+
assets/general.webp filter=lfs diff=lfs merge=lfs -text
|
.mdl
ADDED
|
Binary file (48 Bytes). View file
|
|
|
.msc
ADDED
|
Binary file (1.24 kB). View file
|
|
|
.mv
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
Revision:master,CreatedAt:1760445354
|
README.md
ADDED
|
@@ -0,0 +1,379 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: mit
|
| 3 |
+
pipeline_tag: image-text-to-image
|
| 4 |
+
library_name: transformers
|
| 5 |
+
---
|
| 6 |
+
# HiDream-O1-Image
|
| 7 |
+
|
| 8 |
+
HiDream-O1-Image is a natively unified image generative foundation model built on a Pixel-level Unified Transformer (UiT) without external VAEs or disjoint text encoders, which natively encodes raw pixels, text, and task-specific conditions in a single shared token space — supporting text-to-image, image editing, and subject-driven personalization at up to 2,048 × 2,048.
|
| 9 |
+
|
| 10 |
+
## Project Updates
|
| 11 |
+
- 💬 **June 22, 2026:** Join our Discord community for questions, support, and discussions. <a href="https://discord.gg/7ZEnPxdTQ"><img src="https://img.shields.io/badge/Discord-Join%20Community-5865F2.svg?style=flat&logo=discord&logoColor=white" alt="Discord" /></a>
|
| 12 |
+
- 🚀 **May 14, 2026:** We open-sourced [**HiDream-O1-Image-Dev-2604**](https://huggingface.co/HiDream-ai/HiDream-O1-Image-Dev-2604) with its [prompt refiner](https://huggingface.co/HiDream-ai/Prompt-Refine), tailored for text-to-image generation task.
|
| 13 |
+
- 🛠️ **May 13, 2026:** Inference & pipeline updates — accelerated IP inference; the IP pipeline now supports **layout** and **skeleton** conditioning; updated the Dev editing scheduler. For editing tasks we recommend using the **full** model. PyTorch 2.9.x is not recommended due to the [issue](https://github.com/QwenLM/Qwen3-VL/issues/1811).
|
| 14 |
+
- 🤗 **May 10, 2026:** Try **HiDream-O1-Image** online on Hugging Face Spaces — [🤗 HiDream-O1-Image](https://huggingface.co/spaces/HiDream-ai/HiDream-O1-Image) and [🤗 HiDream-O1-Image-Dev](https://huggingface.co/spaces/HiDream-ai/HiDream-O1-Image-Dev).
|
| 15 |
+
- 📕 **May 10, 2026:** Our **technical report** is now available — [📑 HiDream-O1-Image.pdf](https://arxiv.org/pdf/2605.11061v1).
|
| 16 |
+
- 🚀 **May 8, 2026:** We've open-sourced **HiDream-O1-Image (8B)**, including both the undistilled and distilled Dev variants, together with the Reasoning-Driven Prompt Agent.
|
| 17 |
+
|
| 18 |
+
> **HiDream-O1-Image-Dev-2604 debuts at #8 in the Artificial Analysis Text to Image Arena, which is positioned to be the new leading open weights Text to Image model.**
|
| 19 |
+
<p align="center">
|
| 20 |
+
<img src="assets/leaderboard.png" alt="Artificial Analysis Text to Image Arena" width="100%"/>
|
| 21 |
+
<br><sub><b>Artificial Analysis Text to Image Arena</b> at up to 2,048 × 2,048.</sub>
|
| 22 |
+
</p>
|
| 23 |
+
|
| 24 |
+
<p align="center">
|
| 25 |
+
<img src="assets/general.webp" alt="General text-to-image generation" width="100%"/>
|
| 26 |
+
<br><sub><b>General text-to-image generation</b> at up to 2,048 × 2,048.</sub>
|
| 27 |
+
</p>
|
| 28 |
+
|
| 29 |
+
<p align="center">
|
| 30 |
+
<img src="assets/text-layout.webp" alt="Long-text rendering and layout" width="100%"/>
|
| 31 |
+
<br><sub><b>Long-text rendering & layout control</b> — accurate, multi-region, multilingual text.</sub>
|
| 32 |
+
</p>
|
| 33 |
+
|
| 34 |
+
<p align="center">
|
| 35 |
+
<img src="assets/IP_2.jpg" alt="Subject-driven personalization" width="100%"/>
|
| 36 |
+
<br><sub><b>Subject-driven personalization</b> — preserve identity / IP across new scenes.</sub>
|
| 37 |
+
</p>
|
| 38 |
+
|
| 39 |
+
## Key Features
|
| 40 |
+
|
| 41 |
+
- 🧬 **Pixel-Level Unified Transformer** — One end-to-end model on raw pixels, no VAE, no disjoint text encoder.
|
| 42 |
+
- 🎨 **One Model, Many Tasks** — Text-to-image, long-text rendering, instruction editing, subject-driven personalization, and storyboard generation in a single architecture.
|
| 43 |
+
- 🧠 **Reasoning-Driven Prompt Agent** — Built-in "thinking" agent that resolves implicit knowledge, layout, and text rendering before generation.
|
| 44 |
+
- 🖼️ **Native High Resolution** — Direct synthesis up to 2,048 × 2,048 with sharp fine-grained detail.
|
| 45 |
+
- ⚡ **Exceptional Efficiency and Versatility at 8B Scale** — With only 8B parameters, achieves performance parity with or even surpasses larger open-source DiTs and leading closed-source models.
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
## Models
|
| 50 |
+
|
| 51 |
+
| Name | Script | Inference Steps | HuggingFace Repo |
|
| 52 |
+
| :--- | :--- | :---: | :--- |
|
| 53 |
+
| HiDream-O1-Image | [`inference.py`](https://github.com/HiDream-ai/HiDream-O1-Image/blob/main/inference.py) | 50 | [🤗 HiDream-O1-Image](https://huggingface.co/HiDream-ai/HiDream-O1-Image) |
|
| 54 |
+
| HiDream-O1-Image-Dev | [`inference.py`](https://github.com/HiDream-ai/HiDream-O1-Image/blob/main/inference.py) | 28 | [🤗 HiDream-O1-Image-Dev](https://huggingface.co/HiDream-ai/HiDream-O1-Image-Dev) |
|
| 55 |
+
| Prompt Agent | [`prompt_agent.py`](https://github.com/HiDream-ai/HiDream-O1-Image/blob/main/prompt_agent.py) | — | [🤗 google/gemma-4-31B-it](https://huggingface.co/google/gemma-4-31B-it) |
|
| 56 |
+
| Web Demo | [`app.py`](https://github.com/HiDream-ai/HiDream-O1-Image/blob/main/app.py) | — | — |
|
| 57 |
+
| HiDream-O1-Image-Dev-2604 | [`inference.py` (dev branch)](https://github.com/HiDream-ai/HiDream-O1-Image/blob/dev/inference.py) | 28 | [🤗 HiDream-O1-Image-Dev-2604](https://huggingface.co/HiDream-ai/HiDream-O1-Image-Dev-2604) |
|
| 58 |
+
| Prompt Agent 2604 | [`prompt_agent_v2.py` (dev branch)](https://github.com/HiDream-ai/HiDream-O1-Image/blob/dev/prompt_agent_v2.py) | — | [🤗 HiDream-ai/Prompt-Refine](https://huggingface.co/HiDream-ai/Prompt-Refine) |
|
| 59 |
+
|
| 60 |
+
## Evaluation
|
| 61 |
+
|
| 62 |
+
We benchmark HiDream-O1-Image against state-of-the-art open-source and proprietary models on five widely-used evaluation suites covering compositional generation, dense prompt alignment, human preference, complex visual text generation, and long-text rendering. In each table, the **best** result is highlighted in bold and the <u>second-best</u> is underlined. Click any benchmark below to expand or collapse.
|
| 63 |
+
|
| 64 |
+
<details open>
|
| 65 |
+
<summary><b>GenEval</b> — compositional generation</summary>
|
| 66 |
+
|
| 67 |
+
| Model | #Params | Single-Obj | Two-Obj | Count | Color | Position | Attr | Overall |
|
| 68 |
+
| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
|
| 69 |
+
| Nano Banana 2.0 | – | 1.00 | 0.96 | 0.71 | 0.84 | 0.86 | 0.65 | 0.83 |
|
| 70 |
+
| Seedream-4.0 | – | 1.00 | 0.92 | 0.71 | 0.93 | 0.78 | 0.68 | 0.84 |
|
| 71 |
+
| GPT Image 1 [High] | – | 0.99 | 0.92 | 0.85 | 0.92 | 0.75 | 0.61 | 0.84 |
|
| 72 |
+
| GPT Image 2 | – | 0.99 | 0.98 | 0.85 | 0.93 | 0.85 | 0.77 | 0.89 |
|
| 73 |
+
| PixArt | 4.3B + 0.6B | 0.98 | 0.50 | 0.44 | 0.80 | 0.08 | 0.07 | 0.48 |
|
| 74 |
+
| Show-o | 1.3B | 0.95 | 0.52 | 0.49 | 0.82 | 0.11 | 0.28 | 0.53 |
|
| 75 |
+
| Emu3-Gen | 8B | 0.98 | 0.71 | 0.34 | 0.81 | 0.17 | 0.21 | 0.54 |
|
| 76 |
+
| SD3-Medium | 5.5B + 2B | 0.98 | 0.74 | 0.63 | 0.67 | 0.34 | 0.36 | 0.62 |
|
| 77 |
+
| JanusFlow | 1.3B | 0.97 | 0.59 | 0.45 | 0.83 | 0.53 | 0.42 | 0.63 |
|
| 78 |
+
| FLUX.1 [Dev] | 4.8B + 12B | 0.98 | 0.81 | 0.74 | 0.79 | 0.22 | 0.45 | 0.66 |
|
| 79 |
+
| SD3.5 Large | 5.5B + 8.1B | 0.98 | 0.89 | 0.73 | 0.83 | 0.34 | 0.47 | 0.71 |
|
| 80 |
+
| Janus-Pro-7B | 7B | 0.99 | 0.89 | 0.59 | 0.90 | 0.79 | 0.66 | 0.80 |
|
| 81 |
+
| Z-Image-Turbo | 4B + 6B | 1.00 | 0.95 | 0.77 | 0.89 | 0.65 | 0.68 | 0.82 |
|
| 82 |
+
| FLUX.2 [Dev] | 24B + 32B | 1.00 | 0.99 | 0.79 | 0.93 | 0.73 | 0.78 | 0.87 |
|
| 83 |
+
| Qwen-Image | 7B + 20B | 0.99 | 0.92 | 0.89 | 0.88 | 0.76 | 0.77 | 0.87 |
|
| 84 |
+
| **HiDream-O1-Image** | 8B | 1.00 | 0.99 | 0.79 | 0.89 | 0.93 | 0.78 | <u>0.90</u> |
|
| 85 |
+
| **HiDream-O1-Image-Pro** | 200B+ | 1.00 | 0.99 | 0.85 | 0.94 | 0.94 | 0.79 | **0.92** |
|
| 86 |
+
|
| 87 |
+
</details>
|
| 88 |
+
|
| 89 |
+
<details open>
|
| 90 |
+
<summary><b>DPG-Bench</b> — dense prompt alignment</summary>
|
| 91 |
+
|
| 92 |
+
| Model | #Params | Global | Entity | Attribute | Relation | Other | Overall |
|
| 93 |
+
| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
|
| 94 |
+
| GPT Image 1 [High] | – | 88.89 | 88.94 | 89.84 | 92.63 | 90.96 | 85.15 |
|
| 95 |
+
| GPT Image 2 | – | 87.27 | 91.91 | 90.85 | 91.59 | 91.58 | 85.98 |
|
| 96 |
+
| Nano Banana 2.0 | – | 85.17 | 92.55 | 91.16 | 90.45 | 91.08 | 86.90 |
|
| 97 |
+
| Seedream-4.0 | – | 87.17 | 92.41 | 92.29 | 93.33 | 95.48 | 88.63 |
|
| 98 |
+
| SD v1.5 | 0.12B + 0.86B | 74.63 | 74.23 | 75.39 | 73.49 | 67.81 | 63.18 |
|
| 99 |
+
| PixArt | 4.3B + 0.6B | 74.97 | 79.32 | 78.60 | 82.57 | 76.96 | 71.11 |
|
| 100 |
+
| Lumina-Next | 2B + 2B | 82.82 | 88.65 | 86.44 | 80.53 | 81.82 | 74.63 |
|
| 101 |
+
| SDXL | 0.81B + 2.6B | 83.27 | 82.43 | 80.91 | 86.76 | 80.41 | 74.65 |
|
| 102 |
+
| Hunyuan-DiT | 4.8B + 1.5B | 84.59 | 80.59 | 88.01 | 74.36 | 86.41 | 78.87 |
|
| 103 |
+
| Emu3-Gen | 8B | 85.21 | 86.68 | 86.84 | 90.22 | 83.15 | 80.60 |
|
| 104 |
+
| DALL-E 3 | – | 90.97 | 89.61 | 88.39 | 90.58 | 89.83 | 83.50 |
|
| 105 |
+
| FLUX.1 [Dev] | 4.8B + 12B | 74.35 | 90.00 | 88.96 | 90.87 | 88.33 | 83.84 |
|
| 106 |
+
| SD3 Medium | 5.5B + 2B | 87.90 | 91.01 | 88.83 | 80.70 | 88.68 | 84.08 |
|
| 107 |
+
| Janus-Pro-7B | 7B | 86.90 | 88.90 | 89.40 | 89.32 | 89.48 | 84.19 |
|
| 108 |
+
| Z-Image-Turbo | 4B + 6B | 91.29 | 89.59 | 90.14 | 92.16 | 88.68 | 84.86 |
|
| 109 |
+
| HiDream-I1-Full | 13.5B + 17B | 76.44 | 90.22 | 89.48 | 93.74 | 91.83 | 85.89 |
|
| 110 |
+
| FLUX.2 [Dev] | 24B + 32B | 92.20 | 91.36 | 93.28 | 93.52 | 89.72 | 87.57 |
|
| 111 |
+
| Qwen-Image | 7B + 20B | 91.32 | 91.56 | 92.02 | 94.31 | 92.73 | 88.32 |
|
| 112 |
+
| **HiDream-O1-Image** | 8B | 95.15 | 92.32 | 93.74 | 92.88 | 90.25 | <u>89.83</u> |
|
| 113 |
+
| **HiDream-O1-Image-Pro** | 200B+ | 94.97 | 95.42 | 92.59 | 90.82 | 89.50 | **90.30** |
|
| 114 |
+
|
| 115 |
+
</details>
|
| 116 |
+
|
| 117 |
+
<details open>
|
| 118 |
+
<summary><b>HPSv3</b> — human preference across 12 categories</summary>
|
| 119 |
+
|
| 120 |
+
| Model | #Params | All | Characters | Arts | Design | Architecture | Animals | Natural Scenery | Transportation | Products | Plants | Food | Science | Others |
|
| 121 |
+
| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
|
| 122 |
+
| Seedream-4.0 | – | 9.32 | 9.83 | 9.20 | 8.83 | 9.95 | 8.99 | 9.40 | 9.58 | 9.12 | 9.26 | 9.75 | 9.11 | 9.51 |
|
| 123 |
+
| Nano Banana 2.0 | – | 10.01 | 10.18 | 9.18 | 9.58 | 10.96 | 9.71 | 10.04 | 10.38 | 10.36 | 10.14 | 10.61 | 9.14 | 9.89 |
|
| 124 |
+
| GPT Image 2 | – | 10.21 | 10.75 | 9.91 | 10.15 | 10.59 | 10.05 | 10.29 | 10.17 | 10.26 | 10.07 | 10.75 | 10.05 | 10.00 |
|
| 125 |
+
| Z-Image-Turbo | 4B + 6B | 8.35 | 8.98 | 8.29 | 7.65 | 9.26 | 8.51 | 8.33 | 8.81 | 7.83 | 8.46 | 8.64 | 7.93 | 8.57 |
|
| 126 |
+
| FLUX.2 [Dev] | 24B + 32B | 9.28 | 10.23 | 9.56 | 8.80 | 9.73 | 9.43 | 9.21 | 9.44 | 8.93 | 9.23 | 9.82 | 8.67 | 9.11 |
|
| 127 |
+
| Qwen-Image | 7B + 20B | 9.94 | 10.91 | 10.47 | 9.56 | 10.22 | 10.61 | 9.87 | 10.10 | 9.15 | 9.99 | 10.08 | 9.19 | 9.83 |
|
| 128 |
+
| **HiDream-O1-Image** | 8B | <u>10.37</u> | 10.59 | 10.44 | 10.29 | 11.02 | 10.34 | 10.37 | 10.54 | 10.50 | 10.38 | 10.85 | 9.68 | 10.09 |
|
| 129 |
+
| **HiDream-O1-Image-Pro** | 200B+ | **10.47** | 10.63 | 10.51 | 10.33 | 11.11 | 10.08 | 10.45 | 10.37 | 10.75 | 10.29 | 11.13 | 10.09 | 10.39 |
|
| 130 |
+
|
| 131 |
+
</details>
|
| 132 |
+
|
| 133 |
+
<details>
|
| 134 |
+
<summary><b>CVTG-2K</b> — complex visual text generation (click to expand)</summary>
|
| 135 |
+
|
| 136 |
+
| Model | #Params | 2 regions | 3 regions | 4 regions | 5 regions | Average | NED | CLIP Score |
|
| 137 |
+
| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
|
| 138 |
+
| Nano Banana 2.0 | – | 0.7465 | 0.7720 | 0.8067 | 0.7980 | 0.7875 | 0.8945 | 0.7212 |
|
| 139 |
+
| GPT Image 1 [High] | – | 0.8779 | 0.8659 | 0.8731 | 0.8218 | 0.8569 | 0.9478 | 0.7982 |
|
| 140 |
+
| Seedream-4.0 | – | 0.8980 | 0.8949 | 0.9044 | 0.9015 | 0.9003 | 0.9511 | 0.8033 |
|
| 141 |
+
| GPT Image 2 | – | 0.8904 | 0.8887 | 0.9101 | 0.9044 | 0.9003 | 0.9515 | 0.7798 |
|
| 142 |
+
| TextDiffuser-2 | 0.12B + 0.9B | 0.5322 | 0.3255 | 0.1787 | 0.0809 | 0.2326 | 0.4353 | 0.6765 |
|
| 143 |
+
| RAG-Diffusion | 4.8B + 12B | 0.4388 | 0.3316 | 0.2116 | 0.1910 | 0.2648 | 0.4498 | 0.7797 |
|
| 144 |
+
| AnyText | 0.123B + 1.2B | 0.0513 | 0.1739 | 0.1948 | 0.2249 | 0.1804 | 0.4675 | 0.7432 |
|
| 145 |
+
| 3DIS | 0.81B + 2.6B | 0.4495 | 0.3959 | 0.3880 | 0.3303 | 0.3813 | 0.6505 | 0.7767 |
|
| 146 |
+
| FLUX.1 [Dev] | 4.8B + 12B | 0.6089 | 0.5531 | 0.4661 | 0.4316 | 0.4965 | 0.6879 | 0.7401 |
|
| 147 |
+
| SD3.5 Large | 5.5B + 8.1B | 0.7293 | 0.6825 | 0.6574 | 0.5940 | 0.6548 | 0.8470 | 0.7797 |
|
| 148 |
+
| TextCrafter | 7B + 20B | 0.7628 | 0.7628 | 0.7406 | 0.6977 | 0.7370 | 0.8679 | 0.7868 |
|
| 149 |
+
| Qwen-Image | 7B + 20B | 0.8370 | 0.8364 | 0.8313 | 0.8158 | 0.8288 | 0.9116 | 0.8017 |
|
| 150 |
+
| Z-Image-Turbo | 4B + 6B | 0.8872 | 0.8662 | 0.8628 | 0.8347 | 0.8585 | 0.9281 | 0.8048 |
|
| 151 |
+
| FLUX.2 [Dev] | 24B + 32B | 0.9261 | 0.8897 | 0.8995 | 0.8732 | 0.8926 | 0.9475 | <u>0.8104</u> |
|
| 152 |
+
| **HiDream-O1-Image** | 8B | 0.9085 | 0.9159 | 0.9216 | 0.9015 | <u>0.9128</u> | <u>0.9561</u> | 0.8076 |
|
| 153 |
+
| **HiDream-O1-Image-Pro** | 200B+ | 0.9133 | 0.9221 | 0.9365 | 0.9175 | **0.9222** | **0.9628** | **0.8349** |
|
| 154 |
+
|
| 155 |
+
</details>
|
| 156 |
+
|
| 157 |
+
<details>
|
| 158 |
+
<summary><b>LongText-Bench</b> — long-text rendering, EN & ZH (click to expand)</summary>
|
| 159 |
+
|
| 160 |
+
| Model | #Params | LongText-Bench-EN | LongText-Bench-ZH |
|
| 161 |
+
| :--- | :---: | :---: | :---: |
|
| 162 |
+
| Seedream-4.0 | – | 0.936 | 0.946 |
|
| 163 |
+
| GPT Image 1 [High] | – | 0.956 | 0.619 |
|
| 164 |
+
| GPT Image 2 | – | 0.960 | 0.961 |
|
| 165 |
+
| Nano Banana 2.0 | – | <u>0.980</u> | 0.965 |
|
| 166 |
+
| Janus-Pro-7B | 7B | 0.019 | 0.006 |
|
| 167 |
+
| BLIP3-o | 7B + 1.4B | 0.021 | 0.018 |
|
| 168 |
+
| Kolors 2.0 | – | 0.258 | 0.329 |
|
| 169 |
+
| BAGEL | 7B + 7B | 0.373 | 0.310 |
|
| 170 |
+
| OmniGen2 | 3B + 4B | 0.561 | 0.059 |
|
| 171 |
+
| X-Omni | 7B | 0.900 | 0.814 |
|
| 172 |
+
| HiDream-I1-Full | 13.5B + 17B | 0.543 | 0.024 |
|
| 173 |
+
| FLUX.1 [Dev] | 4.8B + 12B | 0.607 | 0.005 |
|
| 174 |
+
| Z-Image-Turbo | 4B + 6B | 0.917 | 0.926 |
|
| 175 |
+
| FLUX.2 [Dev] | 24B + 32B | 0.963 | 0.757 |
|
| 176 |
+
| Qwen-Image | 7B + 20B | 0.943 | 0.946 |
|
| 177 |
+
| **HiDream-O1-Image** | 8B | 0.979 | <u>0.978</u> |
|
| 178 |
+
| **HiDream-O1-Image-Pro** | 200B+ | **0.982** | **0.980** |
|
| 179 |
+
|
| 180 |
+
</details>
|
| 181 |
+
|
| 182 |
+
## Installation
|
| 183 |
+
|
| 184 |
+
1. Clone this repository:
|
| 185 |
+
```bash
|
| 186 |
+
git clone https://github.com/HiDream-ai/HiDream-O1-Image.git
|
| 187 |
+
cd HiDream-O1-Image
|
| 188 |
+
```
|
| 189 |
+
|
| 190 |
+
2. Install the required dependencies:
|
| 191 |
+
```bash
|
| 192 |
+
pip install -r requirements.txt
|
| 193 |
+
```
|
| 194 |
+
|
| 195 |
+
> **Note on `flash-attn`.** We highly recommend installing [`flash-attn`](https://github.com/Dao-AILab/flash-attention) for optimized attention computation. **If you do not (or cannot) install `flash-attn`, you must edit `models/pipeline.py` line 341 and change `"use_flash_attn": True` to `"use_flash_attn": False`** — otherwise inference will fail to import the kernel.
|
| 196 |
+
|
| 197 |
+
## Reasoning-Driven Prompt Agent
|
| 198 |
+
|
| 199 |
+
HiDream-O1-Image ships with a Reasoning-Driven Prompt Agent (`prompt_agent.py`) that explicitly reasons through layout, subject attributes, physical logic, and text-rendering details, then rewrites a raw user instruction into a self-contained English prompt. It supports two backends — pick one with `--backend`.
|
| 200 |
+
|
| 201 |
+
The agent prints a JSON object with three fields: `prompt` (rewritten English prompt), `reasoning`, and `resolved_knowledge`. Feed the `prompt` field into `inference.py` for best results on intricate, reasoning-heavy requests.
|
| 202 |
+
|
| 203 |
+
### Option A — Local Backend (Gemma-4-31B-it)
|
| 204 |
+
|
| 205 |
+
1. Download the Gemma weights (requires accepting the Gemma license on HuggingFace):
|
| 206 |
+
```bash
|
| 207 |
+
huggingface-cli download google/gemma-4-31B-it --local-dir /path/to/gemma-4-31B-it
|
| 208 |
+
```
|
| 209 |
+
|
| 210 |
+
2. Run the refiner locally:
|
| 211 |
+
```bash
|
| 212 |
+
python prompt_agent.py \
|
| 213 |
+
--backend local \
|
| 214 |
+
--model_id /path/to/gemma-4-31B-it \
|
| 215 |
+
--prompt "李白的静夜思写在古墙上"
|
| 216 |
+
```
|
| 217 |
+
|
| 218 |
+
### Option B — External OpenAI-Compatible API
|
| 219 |
+
|
| 220 |
+
Use any OpenAI-compatible endpoint (OpenAI, Azure, vLLM, SGLang, DeepSeek, etc.) by providing `--base_url`, `--api_key`, and `--model_name`:
|
| 221 |
+
|
| 222 |
+
```bash
|
| 223 |
+
python prompt_agent.py \
|
| 224 |
+
--backend api \
|
| 225 |
+
--base_url https://api.openai.com/v1 \
|
| 226 |
+
--api_key $OPENAI_API_KEY \
|
| 227 |
+
--model_name deepseek-v4-pro \
|
| 228 |
+
--prompt "李白的静夜思写在古墙上"
|
| 229 |
+
```
|
| 230 |
+
|
| 231 |
+
## Usage
|
| 232 |
+
|
| 233 |
+
A CUDA-capable GPU is required for inference. The examples below use the **undistilled** model (`--model_type full`); see the last subsection for running the same tasks with the **distilled** model (`--model_type dev`).
|
| 234 |
+
|
| 235 |
+
### 1. Text-to-Image Generation
|
| 236 |
+
Generate an image from a text prompt:
|
| 237 |
+
|
| 238 |
+
```bash
|
| 239 |
+
python inference.py \
|
| 240 |
+
--model_path /path/to/HiDream-O1-Image \
|
| 241 |
+
--prompt "medium shot, eye-level, front view. A woman is seated in an ornate bedroom, illuminated by candlelight, with a calm and composed expression. The subject is a young woman with fair skin, light brown hair styled in an updo with loose tendrils framing her face, and blue eyes. She wears a cream-colored satin robe with delicate floral embroidery and lace trim along the neckline. Her ears are adorned with pearl drop earrings. She is seated on a bed with a dark, intricately carved wooden headboard. To her left, a wooden nightstand holds three lit white candles and a candelabra with multiple lit candles in the background. The bed is covered with patterned pillows and a dark, textured blanket. The walls are paneled with dark wood and feature a large, ornate tapestry with muted earth tones. The lighting creates soft highlights on her face and robe, with warm shadows cast across the room." \
|
| 242 |
+
--output_image results/t2i.png \
|
| 243 |
+
--height 2048 \
|
| 244 |
+
--width 2048
|
| 245 |
+
```
|
| 246 |
+
|
| 247 |
+
### 2. Instruction-Based Image Editing
|
| 248 |
+
Provide a single reference image and an editing instruction:
|
| 249 |
+
|
| 250 |
+
```bash
|
| 251 |
+
python inference.py \
|
| 252 |
+
--model_path /path/to/HiDream-O1-Image \
|
| 253 |
+
--prompt "remove the earphones" \
|
| 254 |
+
--ref_images assets/edit/test.jpg \
|
| 255 |
+
--output_image results/edit.png \
|
| 256 |
+
--keep_original_aspect
|
| 257 |
+
```
|
| 258 |
+
|
| 259 |
+
### 3. Multi-Reference Subject-Driven Personalization
|
| 260 |
+
Provide two or more reference images that define the subject(s), and a prompt that places them in a new scene:
|
| 261 |
+
|
| 262 |
+
```bash
|
| 263 |
+
python inference.py \
|
| 264 |
+
--model_path /path/to/HiDream-O1-Image \
|
| 265 |
+
--shift 1 \
|
| 266 |
+
--prompt "A young boy with blonde hair stands on steps wearing light blue jeans, a white t-shirt with logo, and blue and white sneakers. He wears a brown cord necklace with beads, a black wristwatch with digital display, and carries a yellow fanny pack with white zipper. In his hand is a red boxing glove with white top, a teal plastic toy car, and a plastic toy figure of Captain America. He wears a straw hat with cream band. Natural light illuminates the scene." \
|
| 267 |
+
--ref_images assets/IP/1.jpg assets/IP/2.jpg assets/IP/3.jpg assets/IP/4.jpg assets/IP/5.jpg assets/IP/6.jpg assets/IP/7.jpg assets/IP/8.jpg assets/IP/9.jpg assets/IP/10.jpg \
|
| 268 |
+
--output_image results/subject.png
|
| 269 |
+
```
|
| 270 |
+
|
| 271 |
+
### 4. Multi-Reference Subject-Driven Personalization with Skeleton
|
| 272 |
+
```bash
|
| 273 |
+
python inference.py \
|
| 274 |
+
--model_path /path/to/HiDream-O1-Image \
|
| 275 |
+
--shift 1 \
|
| 276 |
+
--seed 42 \
|
| 277 |
+
--prompt "Create a realistic try-on image of the person wearing the provided clothing." \
|
| 278 |
+
--ref_images assets/IP_skeleton/0.face.jpg assets/IP_skeleton/0.bg.jpg assets/IP_skeleton/0.openpose.jpg assets/IP_skeleton/0.part_1.jpg assets/IP_skeleton/0.part_2.jpg assets/IP_skeleton/0.part_3.jpg \
|
| 279 |
+
--output_image results/subject.png
|
| 280 |
+
```
|
| 281 |
+
|
| 282 |
+
### 5. Multi-Reference Subject-Driven Personalization with Layout
|
| 283 |
+
```bash
|
| 284 |
+
python inference.py \
|
| 285 |
+
--model_path /path/to/HiDream-O1-Image \
|
| 286 |
+
--shift 1 \
|
| 287 |
+
--seed 42 \
|
| 288 |
+
--prompt "City council members pose with relaxed smiles on a sunlit terrace, warm approachable mood, golden hour, cinematic soft glow." \
|
| 289 |
+
--ref_images assets/IP_layout/0.jpg assets/IP_layout/1.jpg \
|
| 290 |
+
--layout_bboxes "[[0.20507812, 0.43945312, 0.48828125, 0.7421875 ], [0.57617188, 0.80078125, 0.08789062, 0.34179688]]" \
|
| 291 |
+
--output_image results/ip_layout.png
|
| 292 |
+
```
|
| 293 |
+
|
| 294 |
+
### 6. Running with the Dev Model
|
| 295 |
+
All three tasks above can be run with the **Dev** model by switching `--model_path` to the Dev checkpoint and setting `--model_type dev`. For example:
|
| 296 |
+
|
| 297 |
+
```bash
|
| 298 |
+
python inference.py \
|
| 299 |
+
--model_path /path/to/HiDream-O1-Image-Dev \
|
| 300 |
+
--prompt "A dog holds a sign that says \"HiDream-O1-Image release.\"" \
|
| 301 |
+
--output_image results/t2i_dev.png \
|
| 302 |
+
--model_type dev
|
| 303 |
+
```
|
| 304 |
+
|
| 305 |
+
For **editing** tasks (exactly one reference image), the Dev model defaults to the `flow_match` scheduler. `flow_match` is recommended for editing tasks. Pass `--editing_scheduler flash` to use the flash scheduler instead. This flag has no effect on the `full` model or on non-editing tasks.
|
| 306 |
+
|
| 307 |
+
### Command Line Arguments
|
| 308 |
+
|
| 309 |
+
- `--model_path`: Path to the complete HuggingFace model directory (undistilled or distilled).
|
| 310 |
+
- `--prompt`: Text prompt for the generation or editing task.
|
| 311 |
+
- `--ref_images`: Paths to one or more reference images (optional; space-separated).
|
| 312 |
+
- `--output_image`: Path to save the generated image (default: `output.png`).
|
| 313 |
+
- `--height` / `--width`: Output image dimensions (default: `2048` × `2048`; values snap to valid resolutions internally).
|
| 314 |
+
- `--model_type`: `full` or `dev` (default: `full`). Selects the inference recipe:
|
| 315 |
+
- `full`: 50 steps, guidance scale `5.0`, shift `3.0`, default scheduler.
|
| 316 |
+
- `dev`: 28 steps, guidance scale `0.0`, shift `1.0`, flash scheduler with predefined timesteps. For editing tasks (exactly one reference image), the default scheduler is `flow_match` instead — see `--editing_scheduler`.
|
| 317 |
+
- `--seed`: Random seed (default: `32`).
|
| 318 |
+
- `--guidance_scale`: Guidance scale (default: `5.0`). Only effective when `--model_type` is `full`.
|
| 319 |
+
- `--noise_scale_start`, `--noise_scale_end`: Control the scale of the noise injected by the scheduler at each denoising step; the per-step scale linearly interpolates from `noise_scale_start` (first step) to `noise_scale_end` (last step). See `models/pipeline.py:313` (initial noise) and `models/pipeline.py:323-326` (per-step linear interpolation). Defaults: `7.5`, `7.5`.
|
| 320 |
+
- `--noise_clip_std`: Per-step clipping threshold (in units of the injected noise's standard deviation) applied to the noise added during scheduler stepping. See `models/flash_scheduler.py:350-354`. Default: `2.5`.
|
| 321 |
+
- `--editing_scheduler`: Scheduler to use for editing tasks (exactly one reference image) when `--model_type dev`. Choices: `flow_match` (default) or `flash`. Ignored for the `full` model and for non-editing tasks.
|
| 322 |
+
- `--keep_original_aspect`: When exactly one reference image is provided, resize it with `max_size=2048` and use its dimensions for the target image (preserves the reference's aspect ratio) if `True`.
|
| 323 |
+
|
| 324 |
+
## Web Demo
|
| 325 |
+
|
| 326 |
+
`app.py` is a single-file Flask web UI (with HTML / CSS / JS embedded inline) that exposes all generation modes. It also integrates the Reasoning-Driven Prompt Agent.
|
| 327 |
+
|
| 328 |
+
### Starting the server
|
| 329 |
+
|
| 330 |
+
```bash
|
| 331 |
+
python app.py \
|
| 332 |
+
--model_path /path/to/HiDream-O1-Image \
|
| 333 |
+
--host 0.0.0.0 \
|
| 334 |
+
--port 7860
|
| 335 |
+
```
|
| 336 |
+
|
| 337 |
+
Then open `http://localhost:7860` in your browser.
|
| 338 |
+
|
| 339 |
+
### Command-line arguments
|
| 340 |
+
|
| 341 |
+
| Argument | Default | Description |
|
| 342 |
+
| :--- | :--- | :--- |
|
| 343 |
+
| `--model_path` | `$HIDREAM_MODEL_PATH` | Path to the checkpoint directory (`HiDream-O1-Image` or `HiDream-O1-Image-Dev`). |
|
| 344 |
+
| `--model_type` | `full` | `full` (50-step) or `dev` (28-step). |
|
| 345 |
+
| `--host` | `0.0.0.0` | Bind address for the Flask server. |
|
| 346 |
+
| `--port` | `7860` | Port for the Flask server. |
|
| 347 |
+
|
| 348 |
+
All four CLI arguments above can also be set via environment variables (see `.env.example`): `HIDREAM_MODEL_PATH`, `HIDREAM_MODEL_TYPE`, `HIDREAM_HOST`, and `HIDREAM_PORT`.
|
| 349 |
+
|
| 350 |
+
The Prompt Agent panel in the Web Demo reads additional environment variables from `.env`:
|
| 351 |
+
|
| 352 |
+
| Env Var | Used by | Description |
|
| 353 |
+
| :--- | :--- | :--- |
|
| 354 |
+
| `HIDREAM_AGENT_MODEL` | Local · Gemma backend | Path or HF repo id of the local Gemma weights. |
|
| 355 |
+
| `OPENAI_BASE_URL` | OpenAI-compatible API backend | Default base URL pre-filled in the UI. |
|
| 356 |
+
| `OPENAI_API_KEY` | OpenAI-compatible API backend | Default API key pre-filled in the UI. |
|
| 357 |
+
| `OPENAI_MODEL` | OpenAI-compatible API backend | Default model name pre-filled in the UI. |
|
| 358 |
+
|
| 359 |
+
### Prompt Agent in the UI
|
| 360 |
+
|
| 361 |
+
The sidebar contains a Prompt Agent panel that calls the same Reasoning-Driven Prompt Agent used by `prompt_agent.py`. Select either the *OpenAI-compatible API* backend (any endpoint, key, and model name) or the *Local · Gemma* backend (set `HIDREAM_AGENT_MODEL` in `.env` or the environment to point to your local Gemma-4-31B-it weights).
|
| 362 |
+
|
| 363 |
+
### Editing Scheduler (Dev model only)
|
| 364 |
+
|
| 365 |
+
When the server is launched with `--model_type dev`, the **Edit** tab exposes a *Scheduler* dropdown with two options: `flow_match` (default) and `flash`. The selector is hidden for the `full` model and for the Text → Image / Subject tabs, where the scheduler is fixed.
|
| 366 |
+
|
| 367 |
+
## License
|
| 368 |
+
The code in this repository and the HiDream-O1-Image models are licensed under MIT License.
|
| 369 |
+
|
| 370 |
+
## Citation
|
| 371 |
+
|
| 372 |
+
```bibtex
|
| 373 |
+
@article{hidreamolimage,
|
| 374 |
+
title={HiDream-O1-Image: A Natively Unified Image Generative Foundation Model with Pixel-level Unified Transformer},
|
| 375 |
+
author={Cai, Qi and Chen, Jingwen and Gao, Chengmin and Gong, Zijian and Li, Yehao and Mei, Tao and Pan, Yingwei and Peng, Yi and Qiu, Zhaofan and Yao, Ting and Yu, Kai and Zhang, Yiheng and others},
|
| 376 |
+
journal={arXiv preprint arXiv:2605.11061},
|
| 377 |
+
year={2026}
|
| 378 |
+
}
|
| 379 |
+
```
|
assets/IP.webp
ADDED
|
Git LFS Details
|
assets/IP_2.jpg
ADDED
|
Git LFS Details
|
assets/general.webp
ADDED
|
Git LFS Details
|
assets/leaderboard.png
ADDED
|
Git LFS Details
|
assets/text-layout.webp
ADDED
|
Git LFS Details
|
chat_template.json
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"chat_template": "{%- if tools %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0].role == 'system' %}\n {%- if messages[0].content is string %}\n {{- messages[0].content }}\n {%- else %}\n {%- for content in messages[0].content %}\n {%- if 'text' in content %}\n {{- content.text }}\n {%- endif %}\n {%- endfor %}\n {%- endif %}\n {{- '\\n\\n' }}\n {%- endif %}\n {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0].role == 'system' %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0].content is string %}\n {{- messages[0].content }}\n {%- else %}\n {%- for content in messages[0].content %}\n {%- if 'text' in content %}\n {{- content.text }}\n {%- endif %}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- for message in messages %}\n {%- if message.role == \"user\" %}\n {{- '<|im_start|>' + message.role + '\\n' }}\n {%- if message.content is string %}\n {{- message.content }}\n {%- else %}\n {%- for content in message.content %}\n {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n {%- set image_count.value = image_count.value + 1 %}\n {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n <|vision_start|><|image_pad|><|vision_end|>\n {%- elif content.type == 'video' or 'video' in content %}\n {%- set video_count.value = video_count.value + 1 %}\n {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n <|vision_start|><|video_pad|><|vision_end|>\n {%- elif 'text' in content %}\n {{- content.text }}\n {%- endif %}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"assistant\" %}\n {{- '<|im_start|>' + message.role + '\\n' }}\n {%- if message.content is string %}\n {{- message.content }}\n {%- else %}\n {%- for content_item in message.content %}\n {%- if 'text' in content_item %}\n {{- content_item.text }}\n {%- endif %}\n {%- endfor %}\n {%- endif %}\n {%- if message.tool_calls %}\n {%- for tool_call in message.tool_calls %}\n {%- if (loop.first and message.content) or (not loop.first) %}\n {{- '\\n' }}\n {%- endif %}\n {%- if tool_call.function %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {{- '<tool_call>\\n{\"name\": \"' }}\n {{- tool_call.name }}\n {{- '\", \"arguments\": ' }}\n {%- if tool_call.arguments is string %}\n {{- tool_call.arguments }}\n {%- else %}\n {{- tool_call.arguments | tojson }}\n {%- endif %}\n {{- '}\\n</tool_call>' }}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n {{- '<|im_start|>user' }}\n {%- endif %}\n {{- '\\n<tool_response>\\n' }}\n {%- if message.content is string %}\n {{- message.content }}\n {%- else %}\n {%- for content in message.content %}\n {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n {%- set image_count.value = image_count.value + 1 %}\n {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n <|vision_start|><|image_pad|><|vision_end|>\n {%- elif content.type == 'video' or 'video' in content %}\n {%- set video_count.value = video_count.value + 1 %}\n {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n <|vision_start|><|video_pad|><|vision_end|>\n {%- elif 'text' in content %}\n {{- content.text }}\n {%- endif %}\n {%- endfor %}\n {%- endif %}\n {{- '\\n</tool_response>' }}\n {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
|
| 3 |
+
}
|
config.json
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"Qwen3VLForConditionalGeneration"
|
| 4 |
+
],
|
| 5 |
+
"image_token_id": 151655,
|
| 6 |
+
"model_type": "qwen3_vl",
|
| 7 |
+
"text_config": {
|
| 8 |
+
"attention_bias": false,
|
| 9 |
+
"attention_dropout": 0.0,
|
| 10 |
+
"bos_token_id": 151643,
|
| 11 |
+
"dtype": "bfloat16",
|
| 12 |
+
"eos_token_id": 151645,
|
| 13 |
+
"head_dim": 128,
|
| 14 |
+
"hidden_act": "silu",
|
| 15 |
+
"hidden_size": 4096,
|
| 16 |
+
"initializer_range": 0.02,
|
| 17 |
+
"intermediate_size": 12288,
|
| 18 |
+
"max_position_embeddings": 262144,
|
| 19 |
+
"model_type": "qwen3_vl_text",
|
| 20 |
+
"num_attention_heads": 32,
|
| 21 |
+
"num_hidden_layers": 36,
|
| 22 |
+
"num_key_value_heads": 8,
|
| 23 |
+
"rms_norm_eps": 1e-06,
|
| 24 |
+
"rope_scaling": {
|
| 25 |
+
"mrope_interleaved": true,
|
| 26 |
+
"mrope_section": [
|
| 27 |
+
24,
|
| 28 |
+
20,
|
| 29 |
+
20
|
| 30 |
+
],
|
| 31 |
+
"rope_type": "default"
|
| 32 |
+
},
|
| 33 |
+
"rope_theta": 5000000,
|
| 34 |
+
"use_cache": true,
|
| 35 |
+
"vocab_size": 151936
|
| 36 |
+
},
|
| 37 |
+
"tie_word_embeddings": false,
|
| 38 |
+
"transformers_version": "4.57.0.dev0",
|
| 39 |
+
"video_token_id": 151656,
|
| 40 |
+
"vision_config": {
|
| 41 |
+
"deepstack_visual_indexes": [
|
| 42 |
+
8,
|
| 43 |
+
16,
|
| 44 |
+
24
|
| 45 |
+
],
|
| 46 |
+
"depth": 27,
|
| 47 |
+
"hidden_act": "gelu_pytorch_tanh",
|
| 48 |
+
"hidden_size": 1152,
|
| 49 |
+
"in_channels": 3,
|
| 50 |
+
"initializer_range": 0.02,
|
| 51 |
+
"intermediate_size": 4304,
|
| 52 |
+
"model_type": "qwen3_vl",
|
| 53 |
+
"num_heads": 16,
|
| 54 |
+
"num_position_embeddings": 2304,
|
| 55 |
+
"out_hidden_size": 4096,
|
| 56 |
+
"patch_size": 16,
|
| 57 |
+
"spatial_merge_size": 2,
|
| 58 |
+
"temporal_patch_size": 2
|
| 59 |
+
},
|
| 60 |
+
"vision_end_token_id": 151653,
|
| 61 |
+
"vision_start_token_id": 151652
|
| 62 |
+
}
|
configuration.json
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
{"framework": "pytorch", "task": "text-generation", "allow_remote": true}
|
generation_config.json
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token_id": 151643,
|
| 3 |
+
"pad_token_id": 151643,
|
| 4 |
+
"do_sample": true,
|
| 5 |
+
"eos_token_id": [
|
| 6 |
+
151645,
|
| 7 |
+
151643
|
| 8 |
+
],
|
| 9 |
+
"top_k": 20,
|
| 10 |
+
"top_p": 0.8,
|
| 11 |
+
"repetition_penalty": 1.0,
|
| 12 |
+
"temperature": 0.7,
|
| 13 |
+
"transformers_version": "4.56.0"
|
| 14 |
+
}
|
handler.py
ADDED
|
@@ -0,0 +1,96 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
|
| 2 |
+
import os, base64, tempfile
|
| 3 |
+
from typing import Dict, List, Any
|
| 4 |
+
from io import BytesIO
|
| 5 |
+
import torch
|
| 6 |
+
from transformers import AutoProcessor
|
| 7 |
+
from models.qwen3_vl_transformers import Qwen3VLForConditionalGeneration
|
| 8 |
+
from models.pipeline import generate_image, DEFAULT_TIMESTEPS
|
| 9 |
+
|
| 10 |
+
def add_special_tokens(tokenizer):
|
| 11 |
+
tokenizer.boi_token = "<|boi_token|>"
|
| 12 |
+
tokenizer.bor_token = "<|bor_token|>"
|
| 13 |
+
tokenizer.eor_token = "<|eor_token|>"
|
| 14 |
+
tokenizer.bot_token = "<|bot_token|>"
|
| 15 |
+
tokenizer.tms_token = "<|tms_token|>"
|
| 16 |
+
|
| 17 |
+
def get_tokenizer(processor):
|
| 18 |
+
from transformers import PreTrainedTokenizerBase
|
| 19 |
+
if isinstance(processor, PreTrainedTokenizerBase):
|
| 20 |
+
return processor
|
| 21 |
+
return processor.tokenizer
|
| 22 |
+
|
| 23 |
+
class EndpointHandler:
|
| 24 |
+
def __init__(self, path=""):
|
| 25 |
+
self.processor = AutoProcessor.from_pretrained(path)
|
| 26 |
+
self.model = Qwen3VLForConditionalGeneration.from_pretrained(
|
| 27 |
+
path, torch_dtype=torch.float32, device_map="cuda"
|
| 28 |
+
).eval()
|
| 29 |
+
self.tokenizer = get_tokenizer(self.processor)
|
| 30 |
+
add_special_tokens(self.tokenizer)
|
| 31 |
+
|
| 32 |
+
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
|
| 33 |
+
inputs = data.get("inputs")
|
| 34 |
+
params = data.get("parameters", {})
|
| 35 |
+
if not inputs:
|
| 36 |
+
return [{"error": "Missing \'inputs\' (text prompt)"}]
|
| 37 |
+
|
| 38 |
+
model_type = params.get("model_type", "full")
|
| 39 |
+
ref_images_b64 = params.get("ref_images", [])
|
| 40 |
+
|
| 41 |
+
ref_image_paths = []
|
| 42 |
+
tmp_files = []
|
| 43 |
+
try:
|
| 44 |
+
for i, b64_str in enumerate(ref_images_b64):
|
| 45 |
+
img_bytes = base64.b64decode(b64_str)
|
| 46 |
+
tmp = tempfile.NamedTemporaryFile(suffix=f"_{i}.png", delete=False)
|
| 47 |
+
tmp.write(img_bytes)
|
| 48 |
+
tmp.close()
|
| 49 |
+
ref_image_paths.append(tmp.name)
|
| 50 |
+
tmp_files.append(tmp.name)
|
| 51 |
+
|
| 52 |
+
if model_type == "full":
|
| 53 |
+
num_inference_steps = 50
|
| 54 |
+
guidance_scale = params.get("guidance_scale", 5.0)
|
| 55 |
+
shift = 3.0
|
| 56 |
+
timesteps_list = None
|
| 57 |
+
scheduler_name = "default"
|
| 58 |
+
extra_kwargs = {}
|
| 59 |
+
else:
|
| 60 |
+
num_inference_steps = 28
|
| 61 |
+
guidance_scale = 0.0
|
| 62 |
+
shift = 1.0
|
| 63 |
+
timesteps_list = DEFAULT_TIMESTEPS
|
| 64 |
+
scheduler_name = "flash"
|
| 65 |
+
extra_kwargs = {
|
| 66 |
+
"noise_scale_start": params.get("noise_scale_start", 7.5),
|
| 67 |
+
"noise_scale_end": params.get("noise_scale_end", 7.5),
|
| 68 |
+
"noise_clip_std": params.get("noise_clip_std", 2.5),
|
| 69 |
+
}
|
| 70 |
+
|
| 71 |
+
image = generate_image(
|
| 72 |
+
model=self.model,
|
| 73 |
+
processor=self.processor,
|
| 74 |
+
prompt=inputs,
|
| 75 |
+
ref_image_paths=ref_image_paths,
|
| 76 |
+
height=params.get("height", 2048),
|
| 77 |
+
width=params.get("width", 2048),
|
| 78 |
+
num_inference_steps=num_inference_steps,
|
| 79 |
+
guidance_scale=guidance_scale,
|
| 80 |
+
shift=shift,
|
| 81 |
+
timesteps_list=timesteps_list,
|
| 82 |
+
scheduler_name=scheduler_name,
|
| 83 |
+
seed=params.get("seed", 32),
|
| 84 |
+
keep_original_aspect=params.get("keep_original_aspect", False),
|
| 85 |
+
**extra_kwargs,
|
| 86 |
+
)
|
| 87 |
+
except Exception as e:
|
| 88 |
+
return [{"error": str(e)}]
|
| 89 |
+
finally:
|
| 90 |
+
for f in tmp_files:
|
| 91 |
+
os.remove(f)
|
| 92 |
+
|
| 93 |
+
buffered = BytesIO()
|
| 94 |
+
image.save(buffered, format="PNG")
|
| 95 |
+
img_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
| 96 |
+
return [{"image_base64": img_b64}]
|
merges.txt
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
model-00001-of-00008.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:db5d56d92c145eac59d3f69e744244523344c674312d6302e0bac370d72db0b0
|
| 3 |
+
size 4962686568
|
model-00002-of-00008.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:2a06772156fe7f8d7b36b8d6e5ec4c96ade10a7791183ace3b2558a781157333
|
| 3 |
+
size 4832049608
|
model-00003-of-00008.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:270b64408925c6ac760b3b457e66719ef1342a8b5c38f9782a469dd068a7dbf3
|
| 3 |
+
size 4832049624
|
model-00004-of-00008.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:303aa9d7a1ca1598f94e391e5686feb714d9cec1800169ba9777495c7e38777b
|
| 3 |
+
size 4999856656
|
model-00005-of-00008.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:483babbd0f1c63fa69a390f740ce354dac8dea57ecb2e01e70b7e4b1ab3d3e8e
|
| 3 |
+
size 4832049680
|
model-00006-of-00008.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:3d225758a2a35a95aa95d42843b084e226cb7bd83b2fc0be5586619b2bc00ff3
|
| 3 |
+
size 4832049672
|
model-00007-of-00008.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:817c6ac9a332d344d63b82a1481414205f2bf36da5021345c6ed739558b8e3d9
|
| 3 |
+
size 3439582192
|
model-00008-of-00008.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:19c7495b4a90129d52b02c4d4fd1f8b0dcba8a6f9c99fb7f56098cd8a32675c6
|
| 3 |
+
size 2489319552
|
model.safetensors.index.json
ADDED
|
@@ -0,0 +1,767 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"metadata": {
|
| 3 |
+
"total_parameters": 8804887792,
|
| 4 |
+
"total_size": 35219551168
|
| 5 |
+
},
|
| 6 |
+
"weight_map": {
|
| 7 |
+
"lm_head.weight": "model-00008-of-00008.safetensors",
|
| 8 |
+
"model.final_layer2.linear.bias": "model-00007-of-00008.safetensors",
|
| 9 |
+
"model.final_layer2.linear.weight": "model-00007-of-00008.safetensors",
|
| 10 |
+
"model.language_model.embed_tokens.weight": "model-00001-of-00008.safetensors",
|
| 11 |
+
"model.language_model.layers.0.input_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 12 |
+
"model.language_model.layers.0.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
|
| 13 |
+
"model.language_model.layers.0.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
|
| 14 |
+
"model.language_model.layers.0.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
|
| 15 |
+
"model.language_model.layers.0.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 16 |
+
"model.language_model.layers.0.self_attn.k_norm.weight": "model-00001-of-00008.safetensors",
|
| 17 |
+
"model.language_model.layers.0.self_attn.k_proj.weight": "model-00001-of-00008.safetensors",
|
| 18 |
+
"model.language_model.layers.0.self_attn.o_proj.weight": "model-00001-of-00008.safetensors",
|
| 19 |
+
"model.language_model.layers.0.self_attn.q_norm.weight": "model-00001-of-00008.safetensors",
|
| 20 |
+
"model.language_model.layers.0.self_attn.q_proj.weight": "model-00001-of-00008.safetensors",
|
| 21 |
+
"model.language_model.layers.0.self_attn.v_proj.weight": "model-00001-of-00008.safetensors",
|
| 22 |
+
"model.language_model.layers.1.input_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 23 |
+
"model.language_model.layers.1.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
|
| 24 |
+
"model.language_model.layers.1.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
|
| 25 |
+
"model.language_model.layers.1.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
|
| 26 |
+
"model.language_model.layers.1.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 27 |
+
"model.language_model.layers.1.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
|
| 28 |
+
"model.language_model.layers.1.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
|
| 29 |
+
"model.language_model.layers.1.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
|
| 30 |
+
"model.language_model.layers.1.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
|
| 31 |
+
"model.language_model.layers.1.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
|
| 32 |
+
"model.language_model.layers.1.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
|
| 33 |
+
"model.language_model.layers.10.input_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 34 |
+
"model.language_model.layers.10.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
|
| 35 |
+
"model.language_model.layers.10.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
|
| 36 |
+
"model.language_model.layers.10.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
|
| 37 |
+
"model.language_model.layers.10.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 38 |
+
"model.language_model.layers.10.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
|
| 39 |
+
"model.language_model.layers.10.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
|
| 40 |
+
"model.language_model.layers.10.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
|
| 41 |
+
"model.language_model.layers.10.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
|
| 42 |
+
"model.language_model.layers.10.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
|
| 43 |
+
"model.language_model.layers.10.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
|
| 44 |
+
"model.language_model.layers.11.input_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 45 |
+
"model.language_model.layers.11.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
|
| 46 |
+
"model.language_model.layers.11.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
|
| 47 |
+
"model.language_model.layers.11.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
|
| 48 |
+
"model.language_model.layers.11.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 49 |
+
"model.language_model.layers.11.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
|
| 50 |
+
"model.language_model.layers.11.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
|
| 51 |
+
"model.language_model.layers.11.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
|
| 52 |
+
"model.language_model.layers.11.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
|
| 53 |
+
"model.language_model.layers.11.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
|
| 54 |
+
"model.language_model.layers.11.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
|
| 55 |
+
"model.language_model.layers.12.input_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 56 |
+
"model.language_model.layers.12.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
|
| 57 |
+
"model.language_model.layers.12.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
|
| 58 |
+
"model.language_model.layers.12.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
|
| 59 |
+
"model.language_model.layers.12.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 60 |
+
"model.language_model.layers.12.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
|
| 61 |
+
"model.language_model.layers.12.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
|
| 62 |
+
"model.language_model.layers.12.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
|
| 63 |
+
"model.language_model.layers.12.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
|
| 64 |
+
"model.language_model.layers.12.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
|
| 65 |
+
"model.language_model.layers.12.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
|
| 66 |
+
"model.language_model.layers.13.input_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 67 |
+
"model.language_model.layers.13.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
|
| 68 |
+
"model.language_model.layers.13.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
|
| 69 |
+
"model.language_model.layers.13.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
|
| 70 |
+
"model.language_model.layers.13.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 71 |
+
"model.language_model.layers.13.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
|
| 72 |
+
"model.language_model.layers.13.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
|
| 73 |
+
"model.language_model.layers.13.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
|
| 74 |
+
"model.language_model.layers.13.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
|
| 75 |
+
"model.language_model.layers.13.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
|
| 76 |
+
"model.language_model.layers.13.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
|
| 77 |
+
"model.language_model.layers.14.input_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 78 |
+
"model.language_model.layers.14.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
|
| 79 |
+
"model.language_model.layers.14.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
|
| 80 |
+
"model.language_model.layers.14.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
|
| 81 |
+
"model.language_model.layers.14.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 82 |
+
"model.language_model.layers.14.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
|
| 83 |
+
"model.language_model.layers.14.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
|
| 84 |
+
"model.language_model.layers.14.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
|
| 85 |
+
"model.language_model.layers.14.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
|
| 86 |
+
"model.language_model.layers.14.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
|
| 87 |
+
"model.language_model.layers.14.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
|
| 88 |
+
"model.language_model.layers.15.input_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 89 |
+
"model.language_model.layers.15.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
|
| 90 |
+
"model.language_model.layers.15.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
|
| 91 |
+
"model.language_model.layers.15.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
|
| 92 |
+
"model.language_model.layers.15.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 93 |
+
"model.language_model.layers.15.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
|
| 94 |
+
"model.language_model.layers.15.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
|
| 95 |
+
"model.language_model.layers.15.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
|
| 96 |
+
"model.language_model.layers.15.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
|
| 97 |
+
"model.language_model.layers.15.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
|
| 98 |
+
"model.language_model.layers.15.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
|
| 99 |
+
"model.language_model.layers.16.input_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 100 |
+
"model.language_model.layers.16.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
|
| 101 |
+
"model.language_model.layers.16.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
|
| 102 |
+
"model.language_model.layers.16.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
|
| 103 |
+
"model.language_model.layers.16.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 104 |
+
"model.language_model.layers.16.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
|
| 105 |
+
"model.language_model.layers.16.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
|
| 106 |
+
"model.language_model.layers.16.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
|
| 107 |
+
"model.language_model.layers.16.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
|
| 108 |
+
"model.language_model.layers.16.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
|
| 109 |
+
"model.language_model.layers.16.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
|
| 110 |
+
"model.language_model.layers.17.input_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 111 |
+
"model.language_model.layers.17.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
|
| 112 |
+
"model.language_model.layers.17.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
|
| 113 |
+
"model.language_model.layers.17.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
|
| 114 |
+
"model.language_model.layers.17.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 115 |
+
"model.language_model.layers.17.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
|
| 116 |
+
"model.language_model.layers.17.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
|
| 117 |
+
"model.language_model.layers.17.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
|
| 118 |
+
"model.language_model.layers.17.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
|
| 119 |
+
"model.language_model.layers.17.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
|
| 120 |
+
"model.language_model.layers.17.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
|
| 121 |
+
"model.language_model.layers.18.input_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 122 |
+
"model.language_model.layers.18.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
|
| 123 |
+
"model.language_model.layers.18.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
|
| 124 |
+
"model.language_model.layers.18.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
|
| 125 |
+
"model.language_model.layers.18.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
|
| 126 |
+
"model.language_model.layers.18.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
|
| 127 |
+
"model.language_model.layers.18.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
|
| 128 |
+
"model.language_model.layers.18.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
|
| 129 |
+
"model.language_model.layers.18.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
|
| 130 |
+
"model.language_model.layers.18.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
|
| 131 |
+
"model.language_model.layers.18.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
|
| 132 |
+
"model.language_model.layers.19.input_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 133 |
+
"model.language_model.layers.19.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
|
| 134 |
+
"model.language_model.layers.19.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
|
| 135 |
+
"model.language_model.layers.19.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
|
| 136 |
+
"model.language_model.layers.19.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 137 |
+
"model.language_model.layers.19.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
|
| 138 |
+
"model.language_model.layers.19.self_attn.k_proj.weight": "model-00004-of-00008.safetensors",
|
| 139 |
+
"model.language_model.layers.19.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
|
| 140 |
+
"model.language_model.layers.19.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
|
| 141 |
+
"model.language_model.layers.19.self_attn.q_proj.weight": "model-00004-of-00008.safetensors",
|
| 142 |
+
"model.language_model.layers.19.self_attn.v_proj.weight": "model-00004-of-00008.safetensors",
|
| 143 |
+
"model.language_model.layers.2.input_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 144 |
+
"model.language_model.layers.2.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
|
| 145 |
+
"model.language_model.layers.2.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
|
| 146 |
+
"model.language_model.layers.2.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
|
| 147 |
+
"model.language_model.layers.2.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 148 |
+
"model.language_model.layers.2.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
|
| 149 |
+
"model.language_model.layers.2.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
|
| 150 |
+
"model.language_model.layers.2.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
|
| 151 |
+
"model.language_model.layers.2.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
|
| 152 |
+
"model.language_model.layers.2.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
|
| 153 |
+
"model.language_model.layers.2.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
|
| 154 |
+
"model.language_model.layers.20.input_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 155 |
+
"model.language_model.layers.20.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
|
| 156 |
+
"model.language_model.layers.20.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
|
| 157 |
+
"model.language_model.layers.20.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
|
| 158 |
+
"model.language_model.layers.20.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 159 |
+
"model.language_model.layers.20.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
|
| 160 |
+
"model.language_model.layers.20.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
|
| 161 |
+
"model.language_model.layers.20.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
|
| 162 |
+
"model.language_model.layers.20.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
|
| 163 |
+
"model.language_model.layers.20.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
|
| 164 |
+
"model.language_model.layers.20.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
|
| 165 |
+
"model.language_model.layers.21.input_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 166 |
+
"model.language_model.layers.21.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
|
| 167 |
+
"model.language_model.layers.21.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
|
| 168 |
+
"model.language_model.layers.21.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
|
| 169 |
+
"model.language_model.layers.21.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 170 |
+
"model.language_model.layers.21.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
|
| 171 |
+
"model.language_model.layers.21.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
|
| 172 |
+
"model.language_model.layers.21.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
|
| 173 |
+
"model.language_model.layers.21.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
|
| 174 |
+
"model.language_model.layers.21.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
|
| 175 |
+
"model.language_model.layers.21.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
|
| 176 |
+
"model.language_model.layers.22.input_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 177 |
+
"model.language_model.layers.22.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
|
| 178 |
+
"model.language_model.layers.22.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
|
| 179 |
+
"model.language_model.layers.22.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
|
| 180 |
+
"model.language_model.layers.22.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 181 |
+
"model.language_model.layers.22.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
|
| 182 |
+
"model.language_model.layers.22.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
|
| 183 |
+
"model.language_model.layers.22.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
|
| 184 |
+
"model.language_model.layers.22.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
|
| 185 |
+
"model.language_model.layers.22.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
|
| 186 |
+
"model.language_model.layers.22.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
|
| 187 |
+
"model.language_model.layers.23.input_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 188 |
+
"model.language_model.layers.23.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
|
| 189 |
+
"model.language_model.layers.23.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
|
| 190 |
+
"model.language_model.layers.23.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
|
| 191 |
+
"model.language_model.layers.23.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 192 |
+
"model.language_model.layers.23.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
|
| 193 |
+
"model.language_model.layers.23.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
|
| 194 |
+
"model.language_model.layers.23.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
|
| 195 |
+
"model.language_model.layers.23.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
|
| 196 |
+
"model.language_model.layers.23.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
|
| 197 |
+
"model.language_model.layers.23.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
|
| 198 |
+
"model.language_model.layers.24.input_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 199 |
+
"model.language_model.layers.24.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
|
| 200 |
+
"model.language_model.layers.24.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
|
| 201 |
+
"model.language_model.layers.24.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
|
| 202 |
+
"model.language_model.layers.24.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
|
| 203 |
+
"model.language_model.layers.24.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
|
| 204 |
+
"model.language_model.layers.24.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
|
| 205 |
+
"model.language_model.layers.24.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
|
| 206 |
+
"model.language_model.layers.24.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
|
| 207 |
+
"model.language_model.layers.24.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
|
| 208 |
+
"model.language_model.layers.24.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
|
| 209 |
+
"model.language_model.layers.25.input_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 210 |
+
"model.language_model.layers.25.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
|
| 211 |
+
"model.language_model.layers.25.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
|
| 212 |
+
"model.language_model.layers.25.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
|
| 213 |
+
"model.language_model.layers.25.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 214 |
+
"model.language_model.layers.25.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
|
| 215 |
+
"model.language_model.layers.25.self_attn.k_proj.weight": "model-00005-of-00008.safetensors",
|
| 216 |
+
"model.language_model.layers.25.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
|
| 217 |
+
"model.language_model.layers.25.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
|
| 218 |
+
"model.language_model.layers.25.self_attn.q_proj.weight": "model-00005-of-00008.safetensors",
|
| 219 |
+
"model.language_model.layers.25.self_attn.v_proj.weight": "model-00005-of-00008.safetensors",
|
| 220 |
+
"model.language_model.layers.26.input_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 221 |
+
"model.language_model.layers.26.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
|
| 222 |
+
"model.language_model.layers.26.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
|
| 223 |
+
"model.language_model.layers.26.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
|
| 224 |
+
"model.language_model.layers.26.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 225 |
+
"model.language_model.layers.26.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
|
| 226 |
+
"model.language_model.layers.26.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
|
| 227 |
+
"model.language_model.layers.26.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
|
| 228 |
+
"model.language_model.layers.26.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
|
| 229 |
+
"model.language_model.layers.26.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
|
| 230 |
+
"model.language_model.layers.26.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
|
| 231 |
+
"model.language_model.layers.27.input_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 232 |
+
"model.language_model.layers.27.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
|
| 233 |
+
"model.language_model.layers.27.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
|
| 234 |
+
"model.language_model.layers.27.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
|
| 235 |
+
"model.language_model.layers.27.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 236 |
+
"model.language_model.layers.27.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
|
| 237 |
+
"model.language_model.layers.27.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
|
| 238 |
+
"model.language_model.layers.27.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
|
| 239 |
+
"model.language_model.layers.27.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
|
| 240 |
+
"model.language_model.layers.27.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
|
| 241 |
+
"model.language_model.layers.27.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
|
| 242 |
+
"model.language_model.layers.28.input_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 243 |
+
"model.language_model.layers.28.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
|
| 244 |
+
"model.language_model.layers.28.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
|
| 245 |
+
"model.language_model.layers.28.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
|
| 246 |
+
"model.language_model.layers.28.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 247 |
+
"model.language_model.layers.28.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
|
| 248 |
+
"model.language_model.layers.28.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
|
| 249 |
+
"model.language_model.layers.28.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
|
| 250 |
+
"model.language_model.layers.28.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
|
| 251 |
+
"model.language_model.layers.28.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
|
| 252 |
+
"model.language_model.layers.28.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
|
| 253 |
+
"model.language_model.layers.29.input_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 254 |
+
"model.language_model.layers.29.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
|
| 255 |
+
"model.language_model.layers.29.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
|
| 256 |
+
"model.language_model.layers.29.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
|
| 257 |
+
"model.language_model.layers.29.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 258 |
+
"model.language_model.layers.29.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
|
| 259 |
+
"model.language_model.layers.29.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
|
| 260 |
+
"model.language_model.layers.29.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
|
| 261 |
+
"model.language_model.layers.29.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
|
| 262 |
+
"model.language_model.layers.29.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
|
| 263 |
+
"model.language_model.layers.29.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
|
| 264 |
+
"model.language_model.layers.3.input_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 265 |
+
"model.language_model.layers.3.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
|
| 266 |
+
"model.language_model.layers.3.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
|
| 267 |
+
"model.language_model.layers.3.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
|
| 268 |
+
"model.language_model.layers.3.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 269 |
+
"model.language_model.layers.3.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
|
| 270 |
+
"model.language_model.layers.3.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
|
| 271 |
+
"model.language_model.layers.3.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
|
| 272 |
+
"model.language_model.layers.3.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
|
| 273 |
+
"model.language_model.layers.3.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
|
| 274 |
+
"model.language_model.layers.3.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
|
| 275 |
+
"model.language_model.layers.30.input_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 276 |
+
"model.language_model.layers.30.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
|
| 277 |
+
"model.language_model.layers.30.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
|
| 278 |
+
"model.language_model.layers.30.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
|
| 279 |
+
"model.language_model.layers.30.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
|
| 280 |
+
"model.language_model.layers.30.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
|
| 281 |
+
"model.language_model.layers.30.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
|
| 282 |
+
"model.language_model.layers.30.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
|
| 283 |
+
"model.language_model.layers.30.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
|
| 284 |
+
"model.language_model.layers.30.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
|
| 285 |
+
"model.language_model.layers.30.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
|
| 286 |
+
"model.language_model.layers.31.input_layernorm.weight": "model-00007-of-00008.safetensors",
|
| 287 |
+
"model.language_model.layers.31.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
|
| 288 |
+
"model.language_model.layers.31.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
|
| 289 |
+
"model.language_model.layers.31.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
|
| 290 |
+
"model.language_model.layers.31.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
|
| 291 |
+
"model.language_model.layers.31.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
|
| 292 |
+
"model.language_model.layers.31.self_attn.k_proj.weight": "model-00006-of-00008.safetensors",
|
| 293 |
+
"model.language_model.layers.31.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
|
| 294 |
+
"model.language_model.layers.31.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
|
| 295 |
+
"model.language_model.layers.31.self_attn.q_proj.weight": "model-00006-of-00008.safetensors",
|
| 296 |
+
"model.language_model.layers.31.self_attn.v_proj.weight": "model-00006-of-00008.safetensors",
|
| 297 |
+
"model.language_model.layers.32.input_layernorm.weight": "model-00007-of-00008.safetensors",
|
| 298 |
+
"model.language_model.layers.32.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
|
| 299 |
+
"model.language_model.layers.32.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
|
| 300 |
+
"model.language_model.layers.32.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
|
| 301 |
+
"model.language_model.layers.32.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
|
| 302 |
+
"model.language_model.layers.32.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
|
| 303 |
+
"model.language_model.layers.32.self_attn.k_proj.weight": "model-00007-of-00008.safetensors",
|
| 304 |
+
"model.language_model.layers.32.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
|
| 305 |
+
"model.language_model.layers.32.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
|
| 306 |
+
"model.language_model.layers.32.self_attn.q_proj.weight": "model-00007-of-00008.safetensors",
|
| 307 |
+
"model.language_model.layers.32.self_attn.v_proj.weight": "model-00007-of-00008.safetensors",
|
| 308 |
+
"model.language_model.layers.33.input_layernorm.weight": "model-00007-of-00008.safetensors",
|
| 309 |
+
"model.language_model.layers.33.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
|
| 310 |
+
"model.language_model.layers.33.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
|
| 311 |
+
"model.language_model.layers.33.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
|
| 312 |
+
"model.language_model.layers.33.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
|
| 313 |
+
"model.language_model.layers.33.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
|
| 314 |
+
"model.language_model.layers.33.self_attn.k_proj.weight": "model-00007-of-00008.safetensors",
|
| 315 |
+
"model.language_model.layers.33.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
|
| 316 |
+
"model.language_model.layers.33.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
|
| 317 |
+
"model.language_model.layers.33.self_attn.q_proj.weight": "model-00007-of-00008.safetensors",
|
| 318 |
+
"model.language_model.layers.33.self_attn.v_proj.weight": "model-00007-of-00008.safetensors",
|
| 319 |
+
"model.language_model.layers.34.input_layernorm.weight": "model-00007-of-00008.safetensors",
|
| 320 |
+
"model.language_model.layers.34.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
|
| 321 |
+
"model.language_model.layers.34.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
|
| 322 |
+
"model.language_model.layers.34.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
|
| 323 |
+
"model.language_model.layers.34.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
|
| 324 |
+
"model.language_model.layers.34.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
|
| 325 |
+
"model.language_model.layers.34.self_attn.k_proj.weight": "model-00007-of-00008.safetensors",
|
| 326 |
+
"model.language_model.layers.34.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
|
| 327 |
+
"model.language_model.layers.34.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
|
| 328 |
+
"model.language_model.layers.34.self_attn.q_proj.weight": "model-00007-of-00008.safetensors",
|
| 329 |
+
"model.language_model.layers.34.self_attn.v_proj.weight": "model-00007-of-00008.safetensors",
|
| 330 |
+
"model.language_model.layers.35.input_layernorm.weight": "model-00007-of-00008.safetensors",
|
| 331 |
+
"model.language_model.layers.35.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
|
| 332 |
+
"model.language_model.layers.35.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
|
| 333 |
+
"model.language_model.layers.35.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
|
| 334 |
+
"model.language_model.layers.35.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
|
| 335 |
+
"model.language_model.layers.35.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
|
| 336 |
+
"model.language_model.layers.35.self_attn.k_proj.weight": "model-00007-of-00008.safetensors",
|
| 337 |
+
"model.language_model.layers.35.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
|
| 338 |
+
"model.language_model.layers.35.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
|
| 339 |
+
"model.language_model.layers.35.self_attn.q_proj.weight": "model-00007-of-00008.safetensors",
|
| 340 |
+
"model.language_model.layers.35.self_attn.v_proj.weight": "model-00007-of-00008.safetensors",
|
| 341 |
+
"model.language_model.layers.4.input_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 342 |
+
"model.language_model.layers.4.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
|
| 343 |
+
"model.language_model.layers.4.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
|
| 344 |
+
"model.language_model.layers.4.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
|
| 345 |
+
"model.language_model.layers.4.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 346 |
+
"model.language_model.layers.4.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
|
| 347 |
+
"model.language_model.layers.4.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
|
| 348 |
+
"model.language_model.layers.4.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
|
| 349 |
+
"model.language_model.layers.4.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
|
| 350 |
+
"model.language_model.layers.4.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
|
| 351 |
+
"model.language_model.layers.4.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
|
| 352 |
+
"model.language_model.layers.5.input_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 353 |
+
"model.language_model.layers.5.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
|
| 354 |
+
"model.language_model.layers.5.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
|
| 355 |
+
"model.language_model.layers.5.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
|
| 356 |
+
"model.language_model.layers.5.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
|
| 357 |
+
"model.language_model.layers.5.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
|
| 358 |
+
"model.language_model.layers.5.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
|
| 359 |
+
"model.language_model.layers.5.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
|
| 360 |
+
"model.language_model.layers.5.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
|
| 361 |
+
"model.language_model.layers.5.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
|
| 362 |
+
"model.language_model.layers.5.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
|
| 363 |
+
"model.language_model.layers.6.input_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 364 |
+
"model.language_model.layers.6.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
|
| 365 |
+
"model.language_model.layers.6.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
|
| 366 |
+
"model.language_model.layers.6.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
|
| 367 |
+
"model.language_model.layers.6.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 368 |
+
"model.language_model.layers.6.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
|
| 369 |
+
"model.language_model.layers.6.self_attn.k_proj.weight": "model-00002-of-00008.safetensors",
|
| 370 |
+
"model.language_model.layers.6.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
|
| 371 |
+
"model.language_model.layers.6.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
|
| 372 |
+
"model.language_model.layers.6.self_attn.q_proj.weight": "model-00002-of-00008.safetensors",
|
| 373 |
+
"model.language_model.layers.6.self_attn.v_proj.weight": "model-00002-of-00008.safetensors",
|
| 374 |
+
"model.language_model.layers.7.input_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 375 |
+
"model.language_model.layers.7.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
|
| 376 |
+
"model.language_model.layers.7.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
|
| 377 |
+
"model.language_model.layers.7.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
|
| 378 |
+
"model.language_model.layers.7.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 379 |
+
"model.language_model.layers.7.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
|
| 380 |
+
"model.language_model.layers.7.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
|
| 381 |
+
"model.language_model.layers.7.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
|
| 382 |
+
"model.language_model.layers.7.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
|
| 383 |
+
"model.language_model.layers.7.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
|
| 384 |
+
"model.language_model.layers.7.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
|
| 385 |
+
"model.language_model.layers.8.input_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 386 |
+
"model.language_model.layers.8.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
|
| 387 |
+
"model.language_model.layers.8.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
|
| 388 |
+
"model.language_model.layers.8.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
|
| 389 |
+
"model.language_model.layers.8.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 390 |
+
"model.language_model.layers.8.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
|
| 391 |
+
"model.language_model.layers.8.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
|
| 392 |
+
"model.language_model.layers.8.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
|
| 393 |
+
"model.language_model.layers.8.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
|
| 394 |
+
"model.language_model.layers.8.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
|
| 395 |
+
"model.language_model.layers.8.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
|
| 396 |
+
"model.language_model.layers.9.input_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 397 |
+
"model.language_model.layers.9.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
|
| 398 |
+
"model.language_model.layers.9.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
|
| 399 |
+
"model.language_model.layers.9.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
|
| 400 |
+
"model.language_model.layers.9.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
|
| 401 |
+
"model.language_model.layers.9.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
|
| 402 |
+
"model.language_model.layers.9.self_attn.k_proj.weight": "model-00003-of-00008.safetensors",
|
| 403 |
+
"model.language_model.layers.9.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
|
| 404 |
+
"model.language_model.layers.9.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
|
| 405 |
+
"model.language_model.layers.9.self_attn.q_proj.weight": "model-00003-of-00008.safetensors",
|
| 406 |
+
"model.language_model.layers.9.self_attn.v_proj.weight": "model-00003-of-00008.safetensors",
|
| 407 |
+
"model.language_model.norm.weight": "model-00007-of-00008.safetensors",
|
| 408 |
+
"model.t_embedder1.mlp.0.bias": "model-00007-of-00008.safetensors",
|
| 409 |
+
"model.t_embedder1.mlp.0.weight": "model-00007-of-00008.safetensors",
|
| 410 |
+
"model.t_embedder1.mlp.2.bias": "model-00007-of-00008.safetensors",
|
| 411 |
+
"model.t_embedder1.mlp.2.weight": "model-00007-of-00008.safetensors",
|
| 412 |
+
"model.visual.blocks.0.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 413 |
+
"model.visual.blocks.0.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 414 |
+
"model.visual.blocks.0.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 415 |
+
"model.visual.blocks.0.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 416 |
+
"model.visual.blocks.0.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 417 |
+
"model.visual.blocks.0.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 418 |
+
"model.visual.blocks.0.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 419 |
+
"model.visual.blocks.0.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 420 |
+
"model.visual.blocks.0.norm1.bias": "model-00001-of-00008.safetensors",
|
| 421 |
+
"model.visual.blocks.0.norm1.weight": "model-00001-of-00008.safetensors",
|
| 422 |
+
"model.visual.blocks.0.norm2.bias": "model-00001-of-00008.safetensors",
|
| 423 |
+
"model.visual.blocks.0.norm2.weight": "model-00001-of-00008.safetensors",
|
| 424 |
+
"model.visual.blocks.1.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 425 |
+
"model.visual.blocks.1.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 426 |
+
"model.visual.blocks.1.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 427 |
+
"model.visual.blocks.1.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 428 |
+
"model.visual.blocks.1.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 429 |
+
"model.visual.blocks.1.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 430 |
+
"model.visual.blocks.1.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 431 |
+
"model.visual.blocks.1.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 432 |
+
"model.visual.blocks.1.norm1.bias": "model-00001-of-00008.safetensors",
|
| 433 |
+
"model.visual.blocks.1.norm1.weight": "model-00001-of-00008.safetensors",
|
| 434 |
+
"model.visual.blocks.1.norm2.bias": "model-00001-of-00008.safetensors",
|
| 435 |
+
"model.visual.blocks.1.norm2.weight": "model-00001-of-00008.safetensors",
|
| 436 |
+
"model.visual.blocks.10.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 437 |
+
"model.visual.blocks.10.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 438 |
+
"model.visual.blocks.10.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 439 |
+
"model.visual.blocks.10.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 440 |
+
"model.visual.blocks.10.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 441 |
+
"model.visual.blocks.10.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 442 |
+
"model.visual.blocks.10.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 443 |
+
"model.visual.blocks.10.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 444 |
+
"model.visual.blocks.10.norm1.bias": "model-00001-of-00008.safetensors",
|
| 445 |
+
"model.visual.blocks.10.norm1.weight": "model-00001-of-00008.safetensors",
|
| 446 |
+
"model.visual.blocks.10.norm2.bias": "model-00001-of-00008.safetensors",
|
| 447 |
+
"model.visual.blocks.10.norm2.weight": "model-00001-of-00008.safetensors",
|
| 448 |
+
"model.visual.blocks.11.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 449 |
+
"model.visual.blocks.11.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 450 |
+
"model.visual.blocks.11.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 451 |
+
"model.visual.blocks.11.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 452 |
+
"model.visual.blocks.11.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 453 |
+
"model.visual.blocks.11.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 454 |
+
"model.visual.blocks.11.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 455 |
+
"model.visual.blocks.11.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 456 |
+
"model.visual.blocks.11.norm1.bias": "model-00001-of-00008.safetensors",
|
| 457 |
+
"model.visual.blocks.11.norm1.weight": "model-00001-of-00008.safetensors",
|
| 458 |
+
"model.visual.blocks.11.norm2.bias": "model-00001-of-00008.safetensors",
|
| 459 |
+
"model.visual.blocks.11.norm2.weight": "model-00001-of-00008.safetensors",
|
| 460 |
+
"model.visual.blocks.12.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 461 |
+
"model.visual.blocks.12.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 462 |
+
"model.visual.blocks.12.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 463 |
+
"model.visual.blocks.12.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 464 |
+
"model.visual.blocks.12.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 465 |
+
"model.visual.blocks.12.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 466 |
+
"model.visual.blocks.12.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 467 |
+
"model.visual.blocks.12.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 468 |
+
"model.visual.blocks.12.norm1.bias": "model-00001-of-00008.safetensors",
|
| 469 |
+
"model.visual.blocks.12.norm1.weight": "model-00001-of-00008.safetensors",
|
| 470 |
+
"model.visual.blocks.12.norm2.bias": "model-00001-of-00008.safetensors",
|
| 471 |
+
"model.visual.blocks.12.norm2.weight": "model-00001-of-00008.safetensors",
|
| 472 |
+
"model.visual.blocks.13.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 473 |
+
"model.visual.blocks.13.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 474 |
+
"model.visual.blocks.13.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 475 |
+
"model.visual.blocks.13.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 476 |
+
"model.visual.blocks.13.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 477 |
+
"model.visual.blocks.13.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 478 |
+
"model.visual.blocks.13.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 479 |
+
"model.visual.blocks.13.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 480 |
+
"model.visual.blocks.13.norm1.bias": "model-00001-of-00008.safetensors",
|
| 481 |
+
"model.visual.blocks.13.norm1.weight": "model-00001-of-00008.safetensors",
|
| 482 |
+
"model.visual.blocks.13.norm2.bias": "model-00001-of-00008.safetensors",
|
| 483 |
+
"model.visual.blocks.13.norm2.weight": "model-00001-of-00008.safetensors",
|
| 484 |
+
"model.visual.blocks.14.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 485 |
+
"model.visual.blocks.14.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 486 |
+
"model.visual.blocks.14.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 487 |
+
"model.visual.blocks.14.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 488 |
+
"model.visual.blocks.14.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 489 |
+
"model.visual.blocks.14.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 490 |
+
"model.visual.blocks.14.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 491 |
+
"model.visual.blocks.14.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 492 |
+
"model.visual.blocks.14.norm1.bias": "model-00001-of-00008.safetensors",
|
| 493 |
+
"model.visual.blocks.14.norm1.weight": "model-00001-of-00008.safetensors",
|
| 494 |
+
"model.visual.blocks.14.norm2.bias": "model-00001-of-00008.safetensors",
|
| 495 |
+
"model.visual.blocks.14.norm2.weight": "model-00001-of-00008.safetensors",
|
| 496 |
+
"model.visual.blocks.15.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 497 |
+
"model.visual.blocks.15.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 498 |
+
"model.visual.blocks.15.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 499 |
+
"model.visual.blocks.15.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 500 |
+
"model.visual.blocks.15.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 501 |
+
"model.visual.blocks.15.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 502 |
+
"model.visual.blocks.15.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 503 |
+
"model.visual.blocks.15.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 504 |
+
"model.visual.blocks.15.norm1.bias": "model-00001-of-00008.safetensors",
|
| 505 |
+
"model.visual.blocks.15.norm1.weight": "model-00001-of-00008.safetensors",
|
| 506 |
+
"model.visual.blocks.15.norm2.bias": "model-00001-of-00008.safetensors",
|
| 507 |
+
"model.visual.blocks.15.norm2.weight": "model-00001-of-00008.safetensors",
|
| 508 |
+
"model.visual.blocks.16.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 509 |
+
"model.visual.blocks.16.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 510 |
+
"model.visual.blocks.16.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 511 |
+
"model.visual.blocks.16.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 512 |
+
"model.visual.blocks.16.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 513 |
+
"model.visual.blocks.16.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 514 |
+
"model.visual.blocks.16.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 515 |
+
"model.visual.blocks.16.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 516 |
+
"model.visual.blocks.16.norm1.bias": "model-00001-of-00008.safetensors",
|
| 517 |
+
"model.visual.blocks.16.norm1.weight": "model-00001-of-00008.safetensors",
|
| 518 |
+
"model.visual.blocks.16.norm2.bias": "model-00001-of-00008.safetensors",
|
| 519 |
+
"model.visual.blocks.16.norm2.weight": "model-00001-of-00008.safetensors",
|
| 520 |
+
"model.visual.blocks.17.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 521 |
+
"model.visual.blocks.17.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 522 |
+
"model.visual.blocks.17.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 523 |
+
"model.visual.blocks.17.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 524 |
+
"model.visual.blocks.17.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 525 |
+
"model.visual.blocks.17.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 526 |
+
"model.visual.blocks.17.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 527 |
+
"model.visual.blocks.17.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 528 |
+
"model.visual.blocks.17.norm1.bias": "model-00001-of-00008.safetensors",
|
| 529 |
+
"model.visual.blocks.17.norm1.weight": "model-00001-of-00008.safetensors",
|
| 530 |
+
"model.visual.blocks.17.norm2.bias": "model-00001-of-00008.safetensors",
|
| 531 |
+
"model.visual.blocks.17.norm2.weight": "model-00001-of-00008.safetensors",
|
| 532 |
+
"model.visual.blocks.18.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 533 |
+
"model.visual.blocks.18.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 534 |
+
"model.visual.blocks.18.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 535 |
+
"model.visual.blocks.18.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 536 |
+
"model.visual.blocks.18.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 537 |
+
"model.visual.blocks.18.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 538 |
+
"model.visual.blocks.18.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 539 |
+
"model.visual.blocks.18.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 540 |
+
"model.visual.blocks.18.norm1.bias": "model-00001-of-00008.safetensors",
|
| 541 |
+
"model.visual.blocks.18.norm1.weight": "model-00001-of-00008.safetensors",
|
| 542 |
+
"model.visual.blocks.18.norm2.bias": "model-00001-of-00008.safetensors",
|
| 543 |
+
"model.visual.blocks.18.norm2.weight": "model-00001-of-00008.safetensors",
|
| 544 |
+
"model.visual.blocks.19.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 545 |
+
"model.visual.blocks.19.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 546 |
+
"model.visual.blocks.19.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 547 |
+
"model.visual.blocks.19.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 548 |
+
"model.visual.blocks.19.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 549 |
+
"model.visual.blocks.19.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 550 |
+
"model.visual.blocks.19.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 551 |
+
"model.visual.blocks.19.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 552 |
+
"model.visual.blocks.19.norm1.bias": "model-00001-of-00008.safetensors",
|
| 553 |
+
"model.visual.blocks.19.norm1.weight": "model-00001-of-00008.safetensors",
|
| 554 |
+
"model.visual.blocks.19.norm2.bias": "model-00001-of-00008.safetensors",
|
| 555 |
+
"model.visual.blocks.19.norm2.weight": "model-00001-of-00008.safetensors",
|
| 556 |
+
"model.visual.blocks.2.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 557 |
+
"model.visual.blocks.2.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 558 |
+
"model.visual.blocks.2.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 559 |
+
"model.visual.blocks.2.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 560 |
+
"model.visual.blocks.2.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 561 |
+
"model.visual.blocks.2.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 562 |
+
"model.visual.blocks.2.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 563 |
+
"model.visual.blocks.2.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 564 |
+
"model.visual.blocks.2.norm1.bias": "model-00001-of-00008.safetensors",
|
| 565 |
+
"model.visual.blocks.2.norm1.weight": "model-00001-of-00008.safetensors",
|
| 566 |
+
"model.visual.blocks.2.norm2.bias": "model-00001-of-00008.safetensors",
|
| 567 |
+
"model.visual.blocks.2.norm2.weight": "model-00001-of-00008.safetensors",
|
| 568 |
+
"model.visual.blocks.20.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 569 |
+
"model.visual.blocks.20.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 570 |
+
"model.visual.blocks.20.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 571 |
+
"model.visual.blocks.20.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 572 |
+
"model.visual.blocks.20.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 573 |
+
"model.visual.blocks.20.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 574 |
+
"model.visual.blocks.20.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 575 |
+
"model.visual.blocks.20.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 576 |
+
"model.visual.blocks.20.norm1.bias": "model-00001-of-00008.safetensors",
|
| 577 |
+
"model.visual.blocks.20.norm1.weight": "model-00001-of-00008.safetensors",
|
| 578 |
+
"model.visual.blocks.20.norm2.bias": "model-00001-of-00008.safetensors",
|
| 579 |
+
"model.visual.blocks.20.norm2.weight": "model-00001-of-00008.safetensors",
|
| 580 |
+
"model.visual.blocks.21.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 581 |
+
"model.visual.blocks.21.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 582 |
+
"model.visual.blocks.21.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 583 |
+
"model.visual.blocks.21.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 584 |
+
"model.visual.blocks.21.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 585 |
+
"model.visual.blocks.21.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 586 |
+
"model.visual.blocks.21.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 587 |
+
"model.visual.blocks.21.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 588 |
+
"model.visual.blocks.21.norm1.bias": "model-00001-of-00008.safetensors",
|
| 589 |
+
"model.visual.blocks.21.norm1.weight": "model-00001-of-00008.safetensors",
|
| 590 |
+
"model.visual.blocks.21.norm2.bias": "model-00001-of-00008.safetensors",
|
| 591 |
+
"model.visual.blocks.21.norm2.weight": "model-00001-of-00008.safetensors",
|
| 592 |
+
"model.visual.blocks.22.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 593 |
+
"model.visual.blocks.22.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 594 |
+
"model.visual.blocks.22.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 595 |
+
"model.visual.blocks.22.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 596 |
+
"model.visual.blocks.22.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 597 |
+
"model.visual.blocks.22.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 598 |
+
"model.visual.blocks.22.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 599 |
+
"model.visual.blocks.22.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 600 |
+
"model.visual.blocks.22.norm1.bias": "model-00001-of-00008.safetensors",
|
| 601 |
+
"model.visual.blocks.22.norm1.weight": "model-00001-of-00008.safetensors",
|
| 602 |
+
"model.visual.blocks.22.norm2.bias": "model-00001-of-00008.safetensors",
|
| 603 |
+
"model.visual.blocks.22.norm2.weight": "model-00001-of-00008.safetensors",
|
| 604 |
+
"model.visual.blocks.23.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 605 |
+
"model.visual.blocks.23.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 606 |
+
"model.visual.blocks.23.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 607 |
+
"model.visual.blocks.23.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 608 |
+
"model.visual.blocks.23.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 609 |
+
"model.visual.blocks.23.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 610 |
+
"model.visual.blocks.23.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 611 |
+
"model.visual.blocks.23.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 612 |
+
"model.visual.blocks.23.norm1.bias": "model-00001-of-00008.safetensors",
|
| 613 |
+
"model.visual.blocks.23.norm1.weight": "model-00001-of-00008.safetensors",
|
| 614 |
+
"model.visual.blocks.23.norm2.bias": "model-00001-of-00008.safetensors",
|
| 615 |
+
"model.visual.blocks.23.norm2.weight": "model-00001-of-00008.safetensors",
|
| 616 |
+
"model.visual.blocks.24.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 617 |
+
"model.visual.blocks.24.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 618 |
+
"model.visual.blocks.24.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 619 |
+
"model.visual.blocks.24.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 620 |
+
"model.visual.blocks.24.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 621 |
+
"model.visual.blocks.24.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 622 |
+
"model.visual.blocks.24.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 623 |
+
"model.visual.blocks.24.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 624 |
+
"model.visual.blocks.24.norm1.bias": "model-00001-of-00008.safetensors",
|
| 625 |
+
"model.visual.blocks.24.norm1.weight": "model-00001-of-00008.safetensors",
|
| 626 |
+
"model.visual.blocks.24.norm2.bias": "model-00001-of-00008.safetensors",
|
| 627 |
+
"model.visual.blocks.24.norm2.weight": "model-00001-of-00008.safetensors",
|
| 628 |
+
"model.visual.blocks.25.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 629 |
+
"model.visual.blocks.25.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 630 |
+
"model.visual.blocks.25.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 631 |
+
"model.visual.blocks.25.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 632 |
+
"model.visual.blocks.25.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 633 |
+
"model.visual.blocks.25.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 634 |
+
"model.visual.blocks.25.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 635 |
+
"model.visual.blocks.25.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 636 |
+
"model.visual.blocks.25.norm1.bias": "model-00001-of-00008.safetensors",
|
| 637 |
+
"model.visual.blocks.25.norm1.weight": "model-00001-of-00008.safetensors",
|
| 638 |
+
"model.visual.blocks.25.norm2.bias": "model-00001-of-00008.safetensors",
|
| 639 |
+
"model.visual.blocks.25.norm2.weight": "model-00001-of-00008.safetensors",
|
| 640 |
+
"model.visual.blocks.26.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 641 |
+
"model.visual.blocks.26.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 642 |
+
"model.visual.blocks.26.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 643 |
+
"model.visual.blocks.26.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 644 |
+
"model.visual.blocks.26.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 645 |
+
"model.visual.blocks.26.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 646 |
+
"model.visual.blocks.26.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 647 |
+
"model.visual.blocks.26.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 648 |
+
"model.visual.blocks.26.norm1.bias": "model-00001-of-00008.safetensors",
|
| 649 |
+
"model.visual.blocks.26.norm1.weight": "model-00001-of-00008.safetensors",
|
| 650 |
+
"model.visual.blocks.26.norm2.bias": "model-00001-of-00008.safetensors",
|
| 651 |
+
"model.visual.blocks.26.norm2.weight": "model-00001-of-00008.safetensors",
|
| 652 |
+
"model.visual.blocks.3.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 653 |
+
"model.visual.blocks.3.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 654 |
+
"model.visual.blocks.3.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 655 |
+
"model.visual.blocks.3.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 656 |
+
"model.visual.blocks.3.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 657 |
+
"model.visual.blocks.3.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 658 |
+
"model.visual.blocks.3.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 659 |
+
"model.visual.blocks.3.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 660 |
+
"model.visual.blocks.3.norm1.bias": "model-00001-of-00008.safetensors",
|
| 661 |
+
"model.visual.blocks.3.norm1.weight": "model-00001-of-00008.safetensors",
|
| 662 |
+
"model.visual.blocks.3.norm2.bias": "model-00001-of-00008.safetensors",
|
| 663 |
+
"model.visual.blocks.3.norm2.weight": "model-00001-of-00008.safetensors",
|
| 664 |
+
"model.visual.blocks.4.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 665 |
+
"model.visual.blocks.4.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 666 |
+
"model.visual.blocks.4.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 667 |
+
"model.visual.blocks.4.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 668 |
+
"model.visual.blocks.4.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 669 |
+
"model.visual.blocks.4.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 670 |
+
"model.visual.blocks.4.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 671 |
+
"model.visual.blocks.4.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 672 |
+
"model.visual.blocks.4.norm1.bias": "model-00001-of-00008.safetensors",
|
| 673 |
+
"model.visual.blocks.4.norm1.weight": "model-00001-of-00008.safetensors",
|
| 674 |
+
"model.visual.blocks.4.norm2.bias": "model-00001-of-00008.safetensors",
|
| 675 |
+
"model.visual.blocks.4.norm2.weight": "model-00001-of-00008.safetensors",
|
| 676 |
+
"model.visual.blocks.5.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 677 |
+
"model.visual.blocks.5.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 678 |
+
"model.visual.blocks.5.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 679 |
+
"model.visual.blocks.5.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 680 |
+
"model.visual.blocks.5.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 681 |
+
"model.visual.blocks.5.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 682 |
+
"model.visual.blocks.5.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 683 |
+
"model.visual.blocks.5.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 684 |
+
"model.visual.blocks.5.norm1.bias": "model-00001-of-00008.safetensors",
|
| 685 |
+
"model.visual.blocks.5.norm1.weight": "model-00001-of-00008.safetensors",
|
| 686 |
+
"model.visual.blocks.5.norm2.bias": "model-00001-of-00008.safetensors",
|
| 687 |
+
"model.visual.blocks.5.norm2.weight": "model-00001-of-00008.safetensors",
|
| 688 |
+
"model.visual.blocks.6.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 689 |
+
"model.visual.blocks.6.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 690 |
+
"model.visual.blocks.6.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 691 |
+
"model.visual.blocks.6.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 692 |
+
"model.visual.blocks.6.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 693 |
+
"model.visual.blocks.6.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 694 |
+
"model.visual.blocks.6.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 695 |
+
"model.visual.blocks.6.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 696 |
+
"model.visual.blocks.6.norm1.bias": "model-00001-of-00008.safetensors",
|
| 697 |
+
"model.visual.blocks.6.norm1.weight": "model-00001-of-00008.safetensors",
|
| 698 |
+
"model.visual.blocks.6.norm2.bias": "model-00001-of-00008.safetensors",
|
| 699 |
+
"model.visual.blocks.6.norm2.weight": "model-00001-of-00008.safetensors",
|
| 700 |
+
"model.visual.blocks.7.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 701 |
+
"model.visual.blocks.7.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 702 |
+
"model.visual.blocks.7.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 703 |
+
"model.visual.blocks.7.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 704 |
+
"model.visual.blocks.7.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 705 |
+
"model.visual.blocks.7.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 706 |
+
"model.visual.blocks.7.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 707 |
+
"model.visual.blocks.7.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 708 |
+
"model.visual.blocks.7.norm1.bias": "model-00001-of-00008.safetensors",
|
| 709 |
+
"model.visual.blocks.7.norm1.weight": "model-00001-of-00008.safetensors",
|
| 710 |
+
"model.visual.blocks.7.norm2.bias": "model-00001-of-00008.safetensors",
|
| 711 |
+
"model.visual.blocks.7.norm2.weight": "model-00001-of-00008.safetensors",
|
| 712 |
+
"model.visual.blocks.8.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 713 |
+
"model.visual.blocks.8.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 714 |
+
"model.visual.blocks.8.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 715 |
+
"model.visual.blocks.8.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 716 |
+
"model.visual.blocks.8.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 717 |
+
"model.visual.blocks.8.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 718 |
+
"model.visual.blocks.8.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 719 |
+
"model.visual.blocks.8.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 720 |
+
"model.visual.blocks.8.norm1.bias": "model-00001-of-00008.safetensors",
|
| 721 |
+
"model.visual.blocks.8.norm1.weight": "model-00001-of-00008.safetensors",
|
| 722 |
+
"model.visual.blocks.8.norm2.bias": "model-00001-of-00008.safetensors",
|
| 723 |
+
"model.visual.blocks.8.norm2.weight": "model-00001-of-00008.safetensors",
|
| 724 |
+
"model.visual.blocks.9.attn.proj.bias": "model-00001-of-00008.safetensors",
|
| 725 |
+
"model.visual.blocks.9.attn.proj.weight": "model-00001-of-00008.safetensors",
|
| 726 |
+
"model.visual.blocks.9.attn.qkv.bias": "model-00001-of-00008.safetensors",
|
| 727 |
+
"model.visual.blocks.9.attn.qkv.weight": "model-00001-of-00008.safetensors",
|
| 728 |
+
"model.visual.blocks.9.mlp.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 729 |
+
"model.visual.blocks.9.mlp.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 730 |
+
"model.visual.blocks.9.mlp.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 731 |
+
"model.visual.blocks.9.mlp.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 732 |
+
"model.visual.blocks.9.norm1.bias": "model-00001-of-00008.safetensors",
|
| 733 |
+
"model.visual.blocks.9.norm1.weight": "model-00001-of-00008.safetensors",
|
| 734 |
+
"model.visual.blocks.9.norm2.bias": "model-00001-of-00008.safetensors",
|
| 735 |
+
"model.visual.blocks.9.norm2.weight": "model-00001-of-00008.safetensors",
|
| 736 |
+
"model.visual.deepstack_merger_list.0.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 737 |
+
"model.visual.deepstack_merger_list.0.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 738 |
+
"model.visual.deepstack_merger_list.0.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 739 |
+
"model.visual.deepstack_merger_list.0.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 740 |
+
"model.visual.deepstack_merger_list.0.norm.bias": "model-00001-of-00008.safetensors",
|
| 741 |
+
"model.visual.deepstack_merger_list.0.norm.weight": "model-00001-of-00008.safetensors",
|
| 742 |
+
"model.visual.deepstack_merger_list.1.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 743 |
+
"model.visual.deepstack_merger_list.1.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 744 |
+
"model.visual.deepstack_merger_list.1.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 745 |
+
"model.visual.deepstack_merger_list.1.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 746 |
+
"model.visual.deepstack_merger_list.1.norm.bias": "model-00001-of-00008.safetensors",
|
| 747 |
+
"model.visual.deepstack_merger_list.1.norm.weight": "model-00001-of-00008.safetensors",
|
| 748 |
+
"model.visual.deepstack_merger_list.2.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 749 |
+
"model.visual.deepstack_merger_list.2.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 750 |
+
"model.visual.deepstack_merger_list.2.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 751 |
+
"model.visual.deepstack_merger_list.2.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 752 |
+
"model.visual.deepstack_merger_list.2.norm.bias": "model-00001-of-00008.safetensors",
|
| 753 |
+
"model.visual.deepstack_merger_list.2.norm.weight": "model-00001-of-00008.safetensors",
|
| 754 |
+
"model.visual.merger.linear_fc1.bias": "model-00001-of-00008.safetensors",
|
| 755 |
+
"model.visual.merger.linear_fc1.weight": "model-00001-of-00008.safetensors",
|
| 756 |
+
"model.visual.merger.linear_fc2.bias": "model-00001-of-00008.safetensors",
|
| 757 |
+
"model.visual.merger.linear_fc2.weight": "model-00001-of-00008.safetensors",
|
| 758 |
+
"model.visual.merger.norm.bias": "model-00001-of-00008.safetensors",
|
| 759 |
+
"model.visual.merger.norm.weight": "model-00001-of-00008.safetensors",
|
| 760 |
+
"model.visual.patch_embed.proj.bias": "model-00001-of-00008.safetensors",
|
| 761 |
+
"model.visual.patch_embed.proj.weight": "model-00001-of-00008.safetensors",
|
| 762 |
+
"model.visual.pos_embed.weight": "model-00001-of-00008.safetensors",
|
| 763 |
+
"model.x_embedder.proj1.weight": "model-00007-of-00008.safetensors",
|
| 764 |
+
"model.x_embedder.proj2.bias": "model-00007-of-00008.safetensors",
|
| 765 |
+
"model.x_embedder.proj2.weight": "model-00007-of-00008.safetensors"
|
| 766 |
+
}
|
| 767 |
+
}
|
models/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
"""Init."""
|
models/flash_scheduler.py
ADDED
|
@@ -0,0 +1,445 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
import math
|
| 16 |
+
from dataclasses import dataclass
|
| 17 |
+
from typing import List, Optional, Tuple, Union
|
| 18 |
+
|
| 19 |
+
import numpy as np
|
| 20 |
+
import torch
|
| 21 |
+
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
| 22 |
+
from diffusers.schedulers.scheduling_utils import SchedulerMixin
|
| 23 |
+
from diffusers.utils import BaseOutput, is_scipy_available, logging
|
| 24 |
+
from diffusers.utils.torch_utils import randn_tensor
|
| 25 |
+
|
| 26 |
+
if is_scipy_available():
|
| 27 |
+
import scipy.stats
|
| 28 |
+
|
| 29 |
+
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def hack_randn_like(*args, **kwargs):
|
| 33 |
+
# for dist broadcast
|
| 34 |
+
value = torch.randn_like(*args, **kwargs)
|
| 35 |
+
return value
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def hack_randn_tensor(*args, **kwargs):
|
| 39 |
+
# for dist broadcast
|
| 40 |
+
value = randn_tensor(*args, **kwargs)
|
| 41 |
+
return value
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
@dataclass
|
| 45 |
+
class FlashFlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
|
| 46 |
+
"""
|
| 47 |
+
Output class for the scheduler's `step` function output.
|
| 48 |
+
|
| 49 |
+
Args:
|
| 50 |
+
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
|
| 51 |
+
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
|
| 52 |
+
denoising loop.
|
| 53 |
+
"""
|
| 54 |
+
|
| 55 |
+
prev_sample: torch.FloatTensor
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
|
| 59 |
+
"""
|
| 60 |
+
Euler scheduler.
|
| 61 |
+
|
| 62 |
+
This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
|
| 63 |
+
methods the library implements for all schedulers such as loading and saving.
|
| 64 |
+
|
| 65 |
+
Args:
|
| 66 |
+
num_train_timesteps (`int`, defaults to 1000):
|
| 67 |
+
The number of diffusion steps to train the model.
|
| 68 |
+
timestep_spacing (`str`, defaults to `"linspace"`):
|
| 69 |
+
The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
|
| 70 |
+
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
|
| 71 |
+
shift (`float`, defaults to 1.0):
|
| 72 |
+
The shift value for the timestep schedule.
|
| 73 |
+
"""
|
| 74 |
+
|
| 75 |
+
_compatibles = []
|
| 76 |
+
order = 1
|
| 77 |
+
|
| 78 |
+
@register_to_config
|
| 79 |
+
def __init__(
|
| 80 |
+
self,
|
| 81 |
+
num_train_timesteps: int = 1000,
|
| 82 |
+
shift: float = 1.0,
|
| 83 |
+
use_dynamic_shifting=True,
|
| 84 |
+
base_shift: Optional[float] = 0.5,
|
| 85 |
+
max_shift: Optional[float] = 1.15,
|
| 86 |
+
base_image_seq_len: Optional[int] = 256,
|
| 87 |
+
max_image_seq_len: Optional[int] = 4096,
|
| 88 |
+
invert_sigmas: bool = False,
|
| 89 |
+
use_karras_sigmas: Optional[bool] = False,
|
| 90 |
+
use_exponential_sigmas: Optional[bool] = False,
|
| 91 |
+
use_beta_sigmas: Optional[bool] = False,
|
| 92 |
+
):
|
| 93 |
+
if self.config.use_beta_sigmas and not is_scipy_available():
|
| 94 |
+
raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
|
| 95 |
+
if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1:
|
| 96 |
+
raise ValueError(
|
| 97 |
+
"Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
|
| 98 |
+
)
|
| 99 |
+
timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
|
| 100 |
+
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
|
| 101 |
+
|
| 102 |
+
sigmas = timesteps / num_train_timesteps
|
| 103 |
+
if not use_dynamic_shifting:
|
| 104 |
+
# when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
|
| 105 |
+
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
|
| 106 |
+
|
| 107 |
+
self.timesteps = sigmas * num_train_timesteps
|
| 108 |
+
|
| 109 |
+
self._step_index = None
|
| 110 |
+
self._begin_index = None
|
| 111 |
+
|
| 112 |
+
self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication
|
| 113 |
+
self.sigma_min = self.sigmas[-1].item()
|
| 114 |
+
self.sigma_max = self.sigmas[0].item()
|
| 115 |
+
|
| 116 |
+
@property
|
| 117 |
+
def step_index(self):
|
| 118 |
+
"""
|
| 119 |
+
The index counter for current timestep. It will increase 1 after each scheduler step.
|
| 120 |
+
"""
|
| 121 |
+
return self._step_index
|
| 122 |
+
|
| 123 |
+
@property
|
| 124 |
+
def begin_index(self):
|
| 125 |
+
"""
|
| 126 |
+
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
|
| 127 |
+
"""
|
| 128 |
+
return self._begin_index
|
| 129 |
+
|
| 130 |
+
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
|
| 131 |
+
def set_begin_index(self, begin_index: int = 0):
|
| 132 |
+
"""
|
| 133 |
+
Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
|
| 134 |
+
|
| 135 |
+
Args:
|
| 136 |
+
begin_index (`int`):
|
| 137 |
+
The begin index for the scheduler.
|
| 138 |
+
"""
|
| 139 |
+
self._begin_index = begin_index
|
| 140 |
+
|
| 141 |
+
def scale_noise(
|
| 142 |
+
self,
|
| 143 |
+
sample: torch.FloatTensor,
|
| 144 |
+
timestep: Union[float, torch.FloatTensor],
|
| 145 |
+
noise: Optional[torch.FloatTensor] = None,
|
| 146 |
+
) -> torch.FloatTensor:
|
| 147 |
+
"""
|
| 148 |
+
Forward process in flow-matching
|
| 149 |
+
|
| 150 |
+
Args:
|
| 151 |
+
sample (`torch.FloatTensor`):
|
| 152 |
+
The input sample.
|
| 153 |
+
timestep (`int`, *optional*):
|
| 154 |
+
The current timestep in the diffusion chain.
|
| 155 |
+
|
| 156 |
+
Returns:
|
| 157 |
+
`torch.FloatTensor`:
|
| 158 |
+
A scaled input sample.
|
| 159 |
+
"""
|
| 160 |
+
# Make sure sigmas and timesteps have the same device and dtype as original_samples
|
| 161 |
+
sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
|
| 162 |
+
|
| 163 |
+
if sample.device.type == "mps" and torch.is_floating_point(timestep):
|
| 164 |
+
# mps does not support float64
|
| 165 |
+
schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
|
| 166 |
+
timestep = timestep.to(sample.device, dtype=torch.float32)
|
| 167 |
+
else:
|
| 168 |
+
schedule_timesteps = self.timesteps.to(sample.device)
|
| 169 |
+
timestep = timestep.to(sample.device)
|
| 170 |
+
|
| 171 |
+
# self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
|
| 172 |
+
if self.begin_index is None:
|
| 173 |
+
step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
|
| 174 |
+
elif self.step_index is not None:
|
| 175 |
+
# add_noise is called after first denoising step (for inpainting)
|
| 176 |
+
step_indices = [self.step_index] * timestep.shape[0]
|
| 177 |
+
else:
|
| 178 |
+
# add noise is called before first denoising step to create initial latent(img2img)
|
| 179 |
+
step_indices = [self.begin_index] * timestep.shape[0]
|
| 180 |
+
|
| 181 |
+
sigma = sigmas[step_indices].flatten()
|
| 182 |
+
while len(sigma.shape) < len(sample.shape):
|
| 183 |
+
sigma = sigma.unsqueeze(-1)
|
| 184 |
+
|
| 185 |
+
sample = sigma * noise + (1.0 - sigma) * sample
|
| 186 |
+
|
| 187 |
+
return sample
|
| 188 |
+
|
| 189 |
+
def _sigma_to_t(self, sigma):
|
| 190 |
+
return sigma * self.config.num_train_timesteps
|
| 191 |
+
|
| 192 |
+
def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
|
| 193 |
+
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
| 194 |
+
|
| 195 |
+
def set_timesteps(
|
| 196 |
+
self,
|
| 197 |
+
num_inference_steps: int = None,
|
| 198 |
+
device: Union[str, torch.device] = None,
|
| 199 |
+
sigmas: Optional[List[float]] = None,
|
| 200 |
+
mu: Optional[float] = None,
|
| 201 |
+
):
|
| 202 |
+
"""
|
| 203 |
+
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
|
| 204 |
+
|
| 205 |
+
Args:
|
| 206 |
+
num_inference_steps (`int`):
|
| 207 |
+
The number of diffusion steps used when generating samples with a pre-trained model.
|
| 208 |
+
device (`str` or `torch.device`, *optional*):
|
| 209 |
+
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
|
| 210 |
+
"""
|
| 211 |
+
if self.config.use_dynamic_shifting and mu is None:
|
| 212 |
+
raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`")
|
| 213 |
+
|
| 214 |
+
if sigmas is None:
|
| 215 |
+
timesteps = np.linspace(
|
| 216 |
+
self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
|
| 217 |
+
)
|
| 218 |
+
|
| 219 |
+
sigmas = timesteps / self.config.num_train_timesteps
|
| 220 |
+
else:
|
| 221 |
+
sigmas = np.array(sigmas).astype(np.float32)
|
| 222 |
+
num_inference_steps = len(sigmas)
|
| 223 |
+
self.num_inference_steps = num_inference_steps
|
| 224 |
+
|
| 225 |
+
if self.config.use_dynamic_shifting:
|
| 226 |
+
sigmas = self.time_shift(mu, 1.0, sigmas)
|
| 227 |
+
else:
|
| 228 |
+
sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
|
| 229 |
+
|
| 230 |
+
if self.config.use_karras_sigmas:
|
| 231 |
+
sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
|
| 232 |
+
|
| 233 |
+
elif self.config.use_exponential_sigmas:
|
| 234 |
+
sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
|
| 235 |
+
|
| 236 |
+
elif self.config.use_beta_sigmas:
|
| 237 |
+
sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
|
| 238 |
+
|
| 239 |
+
sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
|
| 240 |
+
timesteps = sigmas * self.config.num_train_timesteps
|
| 241 |
+
|
| 242 |
+
if self.config.invert_sigmas:
|
| 243 |
+
sigmas = 1.0 - sigmas
|
| 244 |
+
timesteps = sigmas * self.config.num_train_timesteps
|
| 245 |
+
sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)])
|
| 246 |
+
else:
|
| 247 |
+
sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
|
| 248 |
+
|
| 249 |
+
self.timesteps = timesteps.to(device=device)
|
| 250 |
+
self.sigmas = sigmas
|
| 251 |
+
self._step_index = None
|
| 252 |
+
self._begin_index = None
|
| 253 |
+
|
| 254 |
+
def index_for_timestep(self, timestep, schedule_timesteps=None):
|
| 255 |
+
if schedule_timesteps is None:
|
| 256 |
+
schedule_timesteps = self.timesteps
|
| 257 |
+
|
| 258 |
+
indices = (schedule_timesteps == timestep).nonzero()
|
| 259 |
+
|
| 260 |
+
# The sigma index that is taken for the **very** first `step`
|
| 261 |
+
# is always the second index (or the last index if there is only 1)
|
| 262 |
+
# This way we can ensure we don't accidentally skip a sigma in
|
| 263 |
+
# case we start in the middle of the denoising schedule (e.g. for image-to-image)
|
| 264 |
+
pos = 1 if len(indices) > 1 else 0
|
| 265 |
+
|
| 266 |
+
return indices[pos].item()
|
| 267 |
+
|
| 268 |
+
def _init_step_index(self, timestep):
|
| 269 |
+
if self.begin_index is None:
|
| 270 |
+
if isinstance(timestep, torch.Tensor):
|
| 271 |
+
timestep = timestep.to(self.timesteps.device)
|
| 272 |
+
self._step_index = self.index_for_timestep(timestep)
|
| 273 |
+
else:
|
| 274 |
+
self._step_index = self._begin_index
|
| 275 |
+
|
| 276 |
+
def step(
|
| 277 |
+
self,
|
| 278 |
+
model_output: torch.FloatTensor,
|
| 279 |
+
timestep: Union[float, torch.FloatTensor],
|
| 280 |
+
sample: torch.FloatTensor,
|
| 281 |
+
s_churn: float = 0.0,
|
| 282 |
+
s_tmin: float = 0.0,
|
| 283 |
+
s_tmax: float = float("inf"),
|
| 284 |
+
s_noise: float = 1.0,
|
| 285 |
+
noise_clip_std: float = 0.0,
|
| 286 |
+
generator: Optional[torch.Generator] = None,
|
| 287 |
+
return_dict: bool = True,
|
| 288 |
+
) -> Union[FlashFlowMatchEulerDiscreteSchedulerOutput, Tuple]:
|
| 289 |
+
"""
|
| 290 |
+
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
|
| 291 |
+
process from the learned model outputs (most often the predicted noise).
|
| 292 |
+
|
| 293 |
+
Args:
|
| 294 |
+
model_output (`torch.FloatTensor`):
|
| 295 |
+
The direct output from learned diffusion model.
|
| 296 |
+
timestep (`float`):
|
| 297 |
+
The current discrete timestep in the diffusion chain.
|
| 298 |
+
sample (`torch.FloatTensor`):
|
| 299 |
+
A current instance of a sample created by the diffusion process.
|
| 300 |
+
s_churn (`float`):
|
| 301 |
+
s_tmin (`float`):
|
| 302 |
+
s_tmax (`float`):
|
| 303 |
+
s_noise (`float`, defaults to 1.0):
|
| 304 |
+
Scaling factor for noise added to the sample.
|
| 305 |
+
generator (`torch.Generator`, *optional*):
|
| 306 |
+
A random number generator.
|
| 307 |
+
return_dict (`bool`):
|
| 308 |
+
Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
|
| 309 |
+
tuple.
|
| 310 |
+
|
| 311 |
+
Returns:
|
| 312 |
+
[`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
|
| 313 |
+
If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
|
| 314 |
+
returned, otherwise a tuple is returned where the first element is the sample tensor.
|
| 315 |
+
"""
|
| 316 |
+
|
| 317 |
+
if (
|
| 318 |
+
isinstance(timestep, int)
|
| 319 |
+
or isinstance(timestep, torch.IntTensor)
|
| 320 |
+
or isinstance(timestep, torch.LongTensor)
|
| 321 |
+
):
|
| 322 |
+
raise ValueError(
|
| 323 |
+
(
|
| 324 |
+
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
|
| 325 |
+
" `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
|
| 326 |
+
" one of the `scheduler.timesteps` as a timestep."
|
| 327 |
+
),
|
| 328 |
+
)
|
| 329 |
+
|
| 330 |
+
if self.step_index is None:
|
| 331 |
+
self._init_step_index(timestep)
|
| 332 |
+
|
| 333 |
+
# Upcast to avoid precision issues when computing prev_sample
|
| 334 |
+
|
| 335 |
+
sigma = self.sigmas[self.step_index]
|
| 336 |
+
|
| 337 |
+
# Upcast to avoid precision issues when computing prev_sample
|
| 338 |
+
sample = sample.to(torch.float32)
|
| 339 |
+
|
| 340 |
+
denoised = sample - model_output * sigma
|
| 341 |
+
|
| 342 |
+
if self.step_index < self.num_inference_steps:
|
| 343 |
+
sigma_next = self.sigmas[self.step_index + 1]
|
| 344 |
+
noise = hack_randn_tensor(
|
| 345 |
+
model_output.shape,
|
| 346 |
+
generator=generator,
|
| 347 |
+
device=model_output.device,
|
| 348 |
+
dtype=denoised.dtype,
|
| 349 |
+
)
|
| 350 |
+
if noise_clip_std > 0:
|
| 351 |
+
noise_std = noise.std().item()
|
| 352 |
+
clip_val = noise_clip_std * noise_std
|
| 353 |
+
noise = noise.clamp(min=-clip_val, max=clip_val)
|
| 354 |
+
sample = sigma_next * noise * s_noise + (1.0 - sigma_next) * denoised
|
| 355 |
+
|
| 356 |
+
self._step_index += 1
|
| 357 |
+
sample = sample.to(model_output.dtype)
|
| 358 |
+
|
| 359 |
+
if not return_dict:
|
| 360 |
+
return (sample,)
|
| 361 |
+
|
| 362 |
+
return FlashFlowMatchEulerDiscreteSchedulerOutput(prev_sample=sample)
|
| 363 |
+
|
| 364 |
+
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
|
| 365 |
+
def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
|
| 366 |
+
"""Constructs the noise schedule of Karras et al. (2022)."""
|
| 367 |
+
|
| 368 |
+
# Hack to make sure that other schedulers which copy this function don't break
|
| 369 |
+
# TODO: Add this logic to the other schedulers
|
| 370 |
+
if hasattr(self.config, "sigma_min"):
|
| 371 |
+
sigma_min = self.config.sigma_min
|
| 372 |
+
else:
|
| 373 |
+
sigma_min = None
|
| 374 |
+
|
| 375 |
+
if hasattr(self.config, "sigma_max"):
|
| 376 |
+
sigma_max = self.config.sigma_max
|
| 377 |
+
else:
|
| 378 |
+
sigma_max = None
|
| 379 |
+
|
| 380 |
+
sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
|
| 381 |
+
sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
|
| 382 |
+
|
| 383 |
+
rho = 7.0 # 7.0 is the value used in the paper
|
| 384 |
+
ramp = np.linspace(0, 1, num_inference_steps)
|
| 385 |
+
min_inv_rho = sigma_min ** (1 / rho)
|
| 386 |
+
max_inv_rho = sigma_max ** (1 / rho)
|
| 387 |
+
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
|
| 388 |
+
return sigmas
|
| 389 |
+
|
| 390 |
+
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential
|
| 391 |
+
def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
|
| 392 |
+
"""Constructs an exponential noise schedule."""
|
| 393 |
+
|
| 394 |
+
# Hack to make sure that other schedulers which copy this function don't break
|
| 395 |
+
# TODO: Add this logic to the other schedulers
|
| 396 |
+
if hasattr(self.config, "sigma_min"):
|
| 397 |
+
sigma_min = self.config.sigma_min
|
| 398 |
+
else:
|
| 399 |
+
sigma_min = None
|
| 400 |
+
|
| 401 |
+
if hasattr(self.config, "sigma_max"):
|
| 402 |
+
sigma_max = self.config.sigma_max
|
| 403 |
+
else:
|
| 404 |
+
sigma_max = None
|
| 405 |
+
|
| 406 |
+
sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
|
| 407 |
+
sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
|
| 408 |
+
|
| 409 |
+
sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps))
|
| 410 |
+
return sigmas
|
| 411 |
+
|
| 412 |
+
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta
|
| 413 |
+
def _convert_to_beta(
|
| 414 |
+
self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6
|
| 415 |
+
) -> torch.Tensor:
|
| 416 |
+
"""From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)"""
|
| 417 |
+
|
| 418 |
+
# Hack to make sure that other schedulers which copy this function don't break
|
| 419 |
+
# TODO: Add this logic to the other schedulers
|
| 420 |
+
if hasattr(self.config, "sigma_min"):
|
| 421 |
+
sigma_min = self.config.sigma_min
|
| 422 |
+
else:
|
| 423 |
+
sigma_min = None
|
| 424 |
+
|
| 425 |
+
if hasattr(self.config, "sigma_max"):
|
| 426 |
+
sigma_max = self.config.sigma_max
|
| 427 |
+
else:
|
| 428 |
+
sigma_max = None
|
| 429 |
+
|
| 430 |
+
sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
|
| 431 |
+
sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
|
| 432 |
+
|
| 433 |
+
sigmas = np.array(
|
| 434 |
+
[
|
| 435 |
+
sigma_min + (ppf * (sigma_max - sigma_min))
|
| 436 |
+
for ppf in [
|
| 437 |
+
scipy.stats.beta.ppf(timestep, alpha, beta)
|
| 438 |
+
for timestep in 1 - np.linspace(0, 1, num_inference_steps)
|
| 439 |
+
]
|
| 440 |
+
]
|
| 441 |
+
)
|
| 442 |
+
return sigmas
|
| 443 |
+
|
| 444 |
+
def __len__(self):
|
| 445 |
+
return self.config.num_train_timesteps
|
models/fm_solvers_unipc.py
ADDED
|
@@ -0,0 +1,800 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py
|
| 2 |
+
# Convert unipc for flow matching
|
| 3 |
+
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
| 4 |
+
|
| 5 |
+
import math
|
| 6 |
+
from typing import List, Optional, Tuple, Union
|
| 7 |
+
|
| 8 |
+
import numpy as np
|
| 9 |
+
import torch
|
| 10 |
+
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
| 11 |
+
from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
|
| 12 |
+
SchedulerMixin,
|
| 13 |
+
SchedulerOutput)
|
| 14 |
+
from diffusers.utils import deprecate, is_scipy_available
|
| 15 |
+
|
| 16 |
+
if is_scipy_available():
|
| 17 |
+
import scipy.stats
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
|
| 21 |
+
"""
|
| 22 |
+
`UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models.
|
| 23 |
+
|
| 24 |
+
This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
|
| 25 |
+
methods the library implements for all schedulers such as loading and saving.
|
| 26 |
+
|
| 27 |
+
Args:
|
| 28 |
+
num_train_timesteps (`int`, defaults to 1000):
|
| 29 |
+
The number of diffusion steps to train the model.
|
| 30 |
+
solver_order (`int`, default `2`):
|
| 31 |
+
The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1`
|
| 32 |
+
due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for
|
| 33 |
+
unconditional sampling.
|
| 34 |
+
prediction_type (`str`, defaults to "flow_prediction"):
|
| 35 |
+
Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
|
| 36 |
+
the flow of the diffusion process.
|
| 37 |
+
thresholding (`bool`, defaults to `False`):
|
| 38 |
+
Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
|
| 39 |
+
as Stable Diffusion.
|
| 40 |
+
dynamic_thresholding_ratio (`float`, defaults to 0.995):
|
| 41 |
+
The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
|
| 42 |
+
sample_max_value (`float`, defaults to 1.0):
|
| 43 |
+
The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`.
|
| 44 |
+
predict_x0 (`bool`, defaults to `True`):
|
| 45 |
+
Whether to use the updating algorithm on the predicted x0.
|
| 46 |
+
solver_type (`str`, default `bh2`):
|
| 47 |
+
Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2`
|
| 48 |
+
otherwise.
|
| 49 |
+
lower_order_final (`bool`, default `True`):
|
| 50 |
+
Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
|
| 51 |
+
stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
|
| 52 |
+
disable_corrector (`list`, default `[]`):
|
| 53 |
+
Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)`
|
| 54 |
+
and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is
|
| 55 |
+
usually disabled during the first few steps.
|
| 56 |
+
solver_p (`SchedulerMixin`, default `None`):
|
| 57 |
+
Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
|
| 58 |
+
use_karras_sigmas (`bool`, *optional*, defaults to `False`):
|
| 59 |
+
Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
|
| 60 |
+
the sigmas are determined according to a sequence of noise levels {σi}.
|
| 61 |
+
use_exponential_sigmas (`bool`, *optional*, defaults to `False`):
|
| 62 |
+
Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
|
| 63 |
+
timestep_spacing (`str`, defaults to `"linspace"`):
|
| 64 |
+
The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
|
| 65 |
+
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
|
| 66 |
+
steps_offset (`int`, defaults to 0):
|
| 67 |
+
An offset added to the inference steps, as required by some model families.
|
| 68 |
+
final_sigmas_type (`str`, defaults to `"zero"`):
|
| 69 |
+
The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
|
| 70 |
+
sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
|
| 71 |
+
"""
|
| 72 |
+
|
| 73 |
+
_compatibles = [e.name for e in KarrasDiffusionSchedulers]
|
| 74 |
+
order = 1
|
| 75 |
+
|
| 76 |
+
@register_to_config
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
num_train_timesteps: int = 1000,
|
| 80 |
+
solver_order: int = 2,
|
| 81 |
+
prediction_type: str = "flow_prediction",
|
| 82 |
+
shift: Optional[float] = 1.0,
|
| 83 |
+
use_dynamic_shifting=False,
|
| 84 |
+
thresholding: bool = False,
|
| 85 |
+
dynamic_thresholding_ratio: float = 0.995,
|
| 86 |
+
sample_max_value: float = 1.0,
|
| 87 |
+
predict_x0: bool = True,
|
| 88 |
+
solver_type: str = "bh2",
|
| 89 |
+
lower_order_final: bool = True,
|
| 90 |
+
disable_corrector: List[int] = [],
|
| 91 |
+
solver_p: SchedulerMixin = None,
|
| 92 |
+
timestep_spacing: str = "linspace",
|
| 93 |
+
steps_offset: int = 0,
|
| 94 |
+
final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min"
|
| 95 |
+
):
|
| 96 |
+
|
| 97 |
+
if solver_type not in ["bh1", "bh2"]:
|
| 98 |
+
if solver_type in ["midpoint", "heun", "logrho"]:
|
| 99 |
+
self.register_to_config(solver_type="bh2")
|
| 100 |
+
else:
|
| 101 |
+
raise NotImplementedError(
|
| 102 |
+
f"{solver_type} is not implemented for {self.__class__}")
|
| 103 |
+
|
| 104 |
+
self.predict_x0 = predict_x0
|
| 105 |
+
# setable values
|
| 106 |
+
self.num_inference_steps = None
|
| 107 |
+
alphas = np.linspace(1, 1 / num_train_timesteps,
|
| 108 |
+
num_train_timesteps)[::-1].copy()
|
| 109 |
+
sigmas = 1.0 - alphas
|
| 110 |
+
sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
|
| 111 |
+
|
| 112 |
+
if not use_dynamic_shifting:
|
| 113 |
+
# when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
|
| 114 |
+
sigmas = shift * sigmas / (1 +
|
| 115 |
+
(shift - 1) * sigmas) # pyright: ignore
|
| 116 |
+
|
| 117 |
+
self.sigmas = sigmas
|
| 118 |
+
self.timesteps = sigmas * num_train_timesteps
|
| 119 |
+
|
| 120 |
+
self.model_outputs = [None] * solver_order
|
| 121 |
+
self.timestep_list = [None] * solver_order
|
| 122 |
+
self.lower_order_nums = 0
|
| 123 |
+
self.disable_corrector = disable_corrector
|
| 124 |
+
self.solver_p = solver_p
|
| 125 |
+
self.last_sample = None
|
| 126 |
+
self._step_index = None
|
| 127 |
+
self._begin_index = None
|
| 128 |
+
|
| 129 |
+
self.sigmas = self.sigmas.to(
|
| 130 |
+
"cpu") # to avoid too much CPU/GPU communication
|
| 131 |
+
self.sigma_min = self.sigmas[-1].item()
|
| 132 |
+
self.sigma_max = self.sigmas[0].item()
|
| 133 |
+
|
| 134 |
+
@property
|
| 135 |
+
def step_index(self):
|
| 136 |
+
"""
|
| 137 |
+
The index counter for current timestep. It will increase 1 after each scheduler step.
|
| 138 |
+
"""
|
| 139 |
+
return self._step_index
|
| 140 |
+
|
| 141 |
+
@property
|
| 142 |
+
def begin_index(self):
|
| 143 |
+
"""
|
| 144 |
+
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
|
| 145 |
+
"""
|
| 146 |
+
return self._begin_index
|
| 147 |
+
|
| 148 |
+
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
|
| 149 |
+
def set_begin_index(self, begin_index: int = 0):
|
| 150 |
+
"""
|
| 151 |
+
Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
|
| 152 |
+
|
| 153 |
+
Args:
|
| 154 |
+
begin_index (`int`):
|
| 155 |
+
The begin index for the scheduler.
|
| 156 |
+
"""
|
| 157 |
+
self._begin_index = begin_index
|
| 158 |
+
|
| 159 |
+
# Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
|
| 160 |
+
def set_timesteps(
|
| 161 |
+
self,
|
| 162 |
+
num_inference_steps: Union[int, None] = None,
|
| 163 |
+
device: Union[str, torch.device] = None,
|
| 164 |
+
sigmas: Optional[List[float]] = None,
|
| 165 |
+
mu: Optional[Union[float, None]] = None,
|
| 166 |
+
shift: Optional[Union[float, None]] = None,
|
| 167 |
+
):
|
| 168 |
+
"""
|
| 169 |
+
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
|
| 170 |
+
Args:
|
| 171 |
+
num_inference_steps (`int`):
|
| 172 |
+
Total number of the spacing of the time steps.
|
| 173 |
+
device (`str` or `torch.device`, *optional*):
|
| 174 |
+
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
|
| 175 |
+
"""
|
| 176 |
+
|
| 177 |
+
if self.config.use_dynamic_shifting and mu is None:
|
| 178 |
+
raise ValueError(
|
| 179 |
+
" you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
|
| 180 |
+
)
|
| 181 |
+
|
| 182 |
+
if sigmas is None:
|
| 183 |
+
sigmas = np.linspace(self.sigma_max, self.sigma_min,
|
| 184 |
+
num_inference_steps +
|
| 185 |
+
1).copy()[:-1] # pyright: ignore
|
| 186 |
+
|
| 187 |
+
if self.config.use_dynamic_shifting:
|
| 188 |
+
sigmas = self.time_shift(mu, 1.0, sigmas) # pyright: ignore
|
| 189 |
+
else:
|
| 190 |
+
if shift is None:
|
| 191 |
+
shift = self.config.shift
|
| 192 |
+
sigmas = shift * sigmas / (1 +
|
| 193 |
+
(shift - 1) * sigmas) # pyright: ignore
|
| 194 |
+
|
| 195 |
+
if self.config.final_sigmas_type == "sigma_min":
|
| 196 |
+
sigma_last = ((1 - self.alphas_cumprod[0]) /
|
| 197 |
+
self.alphas_cumprod[0])**0.5
|
| 198 |
+
elif self.config.final_sigmas_type == "zero":
|
| 199 |
+
sigma_last = 0
|
| 200 |
+
else:
|
| 201 |
+
raise ValueError(
|
| 202 |
+
f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
|
| 203 |
+
)
|
| 204 |
+
|
| 205 |
+
timesteps = sigmas * self.config.num_train_timesteps
|
| 206 |
+
sigmas = np.concatenate([sigmas, [sigma_last]
|
| 207 |
+
]).astype(np.float32) # pyright: ignore
|
| 208 |
+
|
| 209 |
+
self.sigmas = torch.from_numpy(sigmas)
|
| 210 |
+
self.timesteps = torch.from_numpy(timesteps).to(
|
| 211 |
+
device=device, dtype=torch.int64)
|
| 212 |
+
|
| 213 |
+
self.num_inference_steps = len(timesteps)
|
| 214 |
+
|
| 215 |
+
self.model_outputs = [
|
| 216 |
+
None,
|
| 217 |
+
] * self.config.solver_order
|
| 218 |
+
self.lower_order_nums = 0
|
| 219 |
+
self.last_sample = None
|
| 220 |
+
if self.solver_p:
|
| 221 |
+
self.solver_p.set_timesteps(self.num_inference_steps, device=device)
|
| 222 |
+
|
| 223 |
+
# add an index counter for schedulers that allow duplicated timesteps
|
| 224 |
+
self._step_index = None
|
| 225 |
+
self._begin_index = None
|
| 226 |
+
self.sigmas = self.sigmas.to(
|
| 227 |
+
"cpu") # to avoid too much CPU/GPU communication
|
| 228 |
+
|
| 229 |
+
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
|
| 230 |
+
def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
|
| 231 |
+
"""
|
| 232 |
+
"Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
|
| 233 |
+
prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
|
| 234 |
+
s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
|
| 235 |
+
pixels from saturation at each step. We find that dynamic thresholding results in significantly better
|
| 236 |
+
photorealism as well as better image-text alignment, especially when using very large guidance weights."
|
| 237 |
+
|
| 238 |
+
https://arxiv.org/abs/2205.11487
|
| 239 |
+
"""
|
| 240 |
+
dtype = sample.dtype
|
| 241 |
+
batch_size, channels, *remaining_dims = sample.shape
|
| 242 |
+
|
| 243 |
+
if dtype not in (torch.float32, torch.float64):
|
| 244 |
+
sample = sample.float(
|
| 245 |
+
) # upcast for quantile calculation, and clamp not implemented for cpu half
|
| 246 |
+
|
| 247 |
+
# Flatten sample for doing quantile calculation along each image
|
| 248 |
+
sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
|
| 249 |
+
|
| 250 |
+
abs_sample = sample.abs() # "a certain percentile absolute pixel value"
|
| 251 |
+
|
| 252 |
+
s = torch.quantile(
|
| 253 |
+
abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
|
| 254 |
+
s = torch.clamp(
|
| 255 |
+
s, min=1, max=self.config.sample_max_value
|
| 256 |
+
) # When clamped to min=1, equivalent to standard clipping to [-1, 1]
|
| 257 |
+
s = s.unsqueeze(
|
| 258 |
+
1) # (batch_size, 1) because clamp will broadcast along dim=0
|
| 259 |
+
sample = torch.clamp(
|
| 260 |
+
sample, -s, s
|
| 261 |
+
) / s # "we threshold xt0 to the range [-s, s] and then divide by s"
|
| 262 |
+
|
| 263 |
+
sample = sample.reshape(batch_size, channels, *remaining_dims)
|
| 264 |
+
sample = sample.to(dtype)
|
| 265 |
+
|
| 266 |
+
return sample
|
| 267 |
+
|
| 268 |
+
# Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t
|
| 269 |
+
def _sigma_to_t(self, sigma):
|
| 270 |
+
return sigma * self.config.num_train_timesteps
|
| 271 |
+
|
| 272 |
+
def _sigma_to_alpha_sigma_t(self, sigma):
|
| 273 |
+
return 1 - sigma, sigma
|
| 274 |
+
|
| 275 |
+
# Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
|
| 276 |
+
def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
|
| 277 |
+
return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)
|
| 278 |
+
|
| 279 |
+
def convert_model_output(
|
| 280 |
+
self,
|
| 281 |
+
model_output: torch.Tensor,
|
| 282 |
+
*args,
|
| 283 |
+
sample: torch.Tensor = None,
|
| 284 |
+
**kwargs,
|
| 285 |
+
) -> torch.Tensor:
|
| 286 |
+
r"""
|
| 287 |
+
Convert the model output to the corresponding type the UniPC algorithm needs.
|
| 288 |
+
|
| 289 |
+
Args:
|
| 290 |
+
model_output (`torch.Tensor`):
|
| 291 |
+
The direct output from the learned diffusion model.
|
| 292 |
+
timestep (`int`):
|
| 293 |
+
The current discrete timestep in the diffusion chain.
|
| 294 |
+
sample (`torch.Tensor`):
|
| 295 |
+
A current instance of a sample created by the diffusion process.
|
| 296 |
+
|
| 297 |
+
Returns:
|
| 298 |
+
`torch.Tensor`:
|
| 299 |
+
The converted model output.
|
| 300 |
+
"""
|
| 301 |
+
timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
|
| 302 |
+
if sample is None:
|
| 303 |
+
if len(args) > 1:
|
| 304 |
+
sample = args[1]
|
| 305 |
+
else:
|
| 306 |
+
raise ValueError(
|
| 307 |
+
"missing `sample` as a required keyward argument")
|
| 308 |
+
if timestep is not None:
|
| 309 |
+
deprecate(
|
| 310 |
+
"timesteps",
|
| 311 |
+
"1.0.0",
|
| 312 |
+
"Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
|
| 313 |
+
)
|
| 314 |
+
|
| 315 |
+
sigma = self.sigmas[self.step_index]
|
| 316 |
+
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
|
| 317 |
+
|
| 318 |
+
if self.predict_x0:
|
| 319 |
+
if self.config.prediction_type == "flow_prediction":
|
| 320 |
+
sigma_t = self.sigmas[self.step_index]
|
| 321 |
+
x0_pred = sample - sigma_t * model_output
|
| 322 |
+
else:
|
| 323 |
+
raise ValueError(
|
| 324 |
+
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
|
| 325 |
+
" `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
|
| 326 |
+
)
|
| 327 |
+
|
| 328 |
+
if self.config.thresholding:
|
| 329 |
+
x0_pred = self._threshold_sample(x0_pred)
|
| 330 |
+
|
| 331 |
+
return x0_pred
|
| 332 |
+
else:
|
| 333 |
+
if self.config.prediction_type == "flow_prediction":
|
| 334 |
+
sigma_t = self.sigmas[self.step_index]
|
| 335 |
+
epsilon = sample - (1 - sigma_t) * model_output
|
| 336 |
+
else:
|
| 337 |
+
raise ValueError(
|
| 338 |
+
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
|
| 339 |
+
" `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
|
| 340 |
+
)
|
| 341 |
+
|
| 342 |
+
if self.config.thresholding:
|
| 343 |
+
sigma_t = self.sigmas[self.step_index]
|
| 344 |
+
x0_pred = sample - sigma_t * model_output
|
| 345 |
+
x0_pred = self._threshold_sample(x0_pred)
|
| 346 |
+
epsilon = model_output + x0_pred
|
| 347 |
+
|
| 348 |
+
return epsilon
|
| 349 |
+
|
| 350 |
+
def multistep_uni_p_bh_update(
|
| 351 |
+
self,
|
| 352 |
+
model_output: torch.Tensor,
|
| 353 |
+
*args,
|
| 354 |
+
sample: torch.Tensor = None,
|
| 355 |
+
order: int = None, # pyright: ignore
|
| 356 |
+
**kwargs,
|
| 357 |
+
) -> torch.Tensor:
|
| 358 |
+
"""
|
| 359 |
+
One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified.
|
| 360 |
+
|
| 361 |
+
Args:
|
| 362 |
+
model_output (`torch.Tensor`):
|
| 363 |
+
The direct output from the learned diffusion model at the current timestep.
|
| 364 |
+
prev_timestep (`int`):
|
| 365 |
+
The previous discrete timestep in the diffusion chain.
|
| 366 |
+
sample (`torch.Tensor`):
|
| 367 |
+
A current instance of a sample created by the diffusion process.
|
| 368 |
+
order (`int`):
|
| 369 |
+
The order of UniP at this timestep (corresponds to the *p* in UniPC-p).
|
| 370 |
+
|
| 371 |
+
Returns:
|
| 372 |
+
`torch.Tensor`:
|
| 373 |
+
The sample tensor at the previous timestep.
|
| 374 |
+
"""
|
| 375 |
+
prev_timestep = args[0] if len(args) > 0 else kwargs.pop(
|
| 376 |
+
"prev_timestep", None)
|
| 377 |
+
if sample is None:
|
| 378 |
+
if len(args) > 1:
|
| 379 |
+
sample = args[1]
|
| 380 |
+
else:
|
| 381 |
+
raise ValueError(
|
| 382 |
+
" missing `sample` as a required keyward argument")
|
| 383 |
+
if order is None:
|
| 384 |
+
if len(args) > 2:
|
| 385 |
+
order = args[2]
|
| 386 |
+
else:
|
| 387 |
+
raise ValueError(
|
| 388 |
+
" missing `order` as a required keyward argument")
|
| 389 |
+
if prev_timestep is not None:
|
| 390 |
+
deprecate(
|
| 391 |
+
"prev_timestep",
|
| 392 |
+
"1.0.0",
|
| 393 |
+
"Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
|
| 394 |
+
)
|
| 395 |
+
model_output_list = self.model_outputs
|
| 396 |
+
|
| 397 |
+
s0 = self.timestep_list[-1]
|
| 398 |
+
m0 = model_output_list[-1]
|
| 399 |
+
x = sample
|
| 400 |
+
|
| 401 |
+
if self.solver_p:
|
| 402 |
+
x_t = self.solver_p.step(model_output, s0, x).prev_sample
|
| 403 |
+
return x_t
|
| 404 |
+
|
| 405 |
+
sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[
|
| 406 |
+
self.step_index] # pyright: ignore
|
| 407 |
+
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
|
| 408 |
+
alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
|
| 409 |
+
|
| 410 |
+
lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
|
| 411 |
+
lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
|
| 412 |
+
|
| 413 |
+
h = lambda_t - lambda_s0
|
| 414 |
+
device = sample.device
|
| 415 |
+
|
| 416 |
+
rks = []
|
| 417 |
+
D1s = []
|
| 418 |
+
for i in range(1, order):
|
| 419 |
+
si = self.step_index - i # pyright: ignore
|
| 420 |
+
mi = model_output_list[-(i + 1)]
|
| 421 |
+
alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
|
| 422 |
+
lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
|
| 423 |
+
rk = (lambda_si - lambda_s0) / h
|
| 424 |
+
rks.append(rk)
|
| 425 |
+
D1s.append((mi - m0) / rk) # pyright: ignore
|
| 426 |
+
|
| 427 |
+
rks.append(1.0)
|
| 428 |
+
rks = torch.tensor(rks, device=device)
|
| 429 |
+
|
| 430 |
+
R = []
|
| 431 |
+
b = []
|
| 432 |
+
|
| 433 |
+
hh = -h if self.predict_x0 else h
|
| 434 |
+
h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
|
| 435 |
+
h_phi_k = h_phi_1 / hh - 1
|
| 436 |
+
|
| 437 |
+
factorial_i = 1
|
| 438 |
+
|
| 439 |
+
if self.config.solver_type == "bh1":
|
| 440 |
+
B_h = hh
|
| 441 |
+
elif self.config.solver_type == "bh2":
|
| 442 |
+
B_h = torch.expm1(hh)
|
| 443 |
+
else:
|
| 444 |
+
raise NotImplementedError()
|
| 445 |
+
|
| 446 |
+
for i in range(1, order + 1):
|
| 447 |
+
R.append(torch.pow(rks, i - 1))
|
| 448 |
+
b.append(h_phi_k * factorial_i / B_h)
|
| 449 |
+
factorial_i *= i + 1
|
| 450 |
+
h_phi_k = h_phi_k / hh - 1 / factorial_i
|
| 451 |
+
|
| 452 |
+
R = torch.stack(R)
|
| 453 |
+
b = torch.tensor(b, device=device)
|
| 454 |
+
|
| 455 |
+
if len(D1s) > 0:
|
| 456 |
+
D1s = torch.stack(D1s, dim=1) # (B, K)
|
| 457 |
+
# for order 2, we use a simplified version
|
| 458 |
+
if order == 2:
|
| 459 |
+
rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
|
| 460 |
+
else:
|
| 461 |
+
rhos_p = torch.linalg.solve(R[:-1, :-1],
|
| 462 |
+
b[:-1]).to(device).to(x.dtype)
|
| 463 |
+
else:
|
| 464 |
+
D1s = None
|
| 465 |
+
|
| 466 |
+
if self.predict_x0:
|
| 467 |
+
x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
|
| 468 |
+
if D1s is not None:
|
| 469 |
+
pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
|
| 470 |
+
D1s) # pyright: ignore
|
| 471 |
+
else:
|
| 472 |
+
pred_res = 0
|
| 473 |
+
x_t = x_t_ - alpha_t * B_h * pred_res
|
| 474 |
+
else:
|
| 475 |
+
x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
|
| 476 |
+
if D1s is not None:
|
| 477 |
+
pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
|
| 478 |
+
D1s) # pyright: ignore
|
| 479 |
+
else:
|
| 480 |
+
pred_res = 0
|
| 481 |
+
x_t = x_t_ - sigma_t * B_h * pred_res
|
| 482 |
+
|
| 483 |
+
x_t = x_t.to(x.dtype)
|
| 484 |
+
return x_t
|
| 485 |
+
|
| 486 |
+
def multistep_uni_c_bh_update(
|
| 487 |
+
self,
|
| 488 |
+
this_model_output: torch.Tensor,
|
| 489 |
+
*args,
|
| 490 |
+
last_sample: torch.Tensor = None,
|
| 491 |
+
this_sample: torch.Tensor = None,
|
| 492 |
+
order: int = None, # pyright: ignore
|
| 493 |
+
**kwargs,
|
| 494 |
+
) -> torch.Tensor:
|
| 495 |
+
"""
|
| 496 |
+
One step for the UniC (B(h) version).
|
| 497 |
+
|
| 498 |
+
Args:
|
| 499 |
+
this_model_output (`torch.Tensor`):
|
| 500 |
+
The model outputs at `x_t`.
|
| 501 |
+
this_timestep (`int`):
|
| 502 |
+
The current timestep `t`.
|
| 503 |
+
last_sample (`torch.Tensor`):
|
| 504 |
+
The generated sample before the last predictor `x_{t-1}`.
|
| 505 |
+
this_sample (`torch.Tensor`):
|
| 506 |
+
The generated sample after the last predictor `x_{t}`.
|
| 507 |
+
order (`int`):
|
| 508 |
+
The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.
|
| 509 |
+
|
| 510 |
+
Returns:
|
| 511 |
+
`torch.Tensor`:
|
| 512 |
+
The corrected sample tensor at the current timestep.
|
| 513 |
+
"""
|
| 514 |
+
this_timestep = args[0] if len(args) > 0 else kwargs.pop(
|
| 515 |
+
"this_timestep", None)
|
| 516 |
+
if last_sample is None:
|
| 517 |
+
if len(args) > 1:
|
| 518 |
+
last_sample = args[1]
|
| 519 |
+
else:
|
| 520 |
+
raise ValueError(
|
| 521 |
+
" missing`last_sample` as a required keyward argument")
|
| 522 |
+
if this_sample is None:
|
| 523 |
+
if len(args) > 2:
|
| 524 |
+
this_sample = args[2]
|
| 525 |
+
else:
|
| 526 |
+
raise ValueError(
|
| 527 |
+
" missing`this_sample` as a required keyward argument")
|
| 528 |
+
if order is None:
|
| 529 |
+
if len(args) > 3:
|
| 530 |
+
order = args[3]
|
| 531 |
+
else:
|
| 532 |
+
raise ValueError(
|
| 533 |
+
" missing`order` as a required keyward argument")
|
| 534 |
+
if this_timestep is not None:
|
| 535 |
+
deprecate(
|
| 536 |
+
"this_timestep",
|
| 537 |
+
"1.0.0",
|
| 538 |
+
"Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
|
| 539 |
+
)
|
| 540 |
+
|
| 541 |
+
model_output_list = self.model_outputs
|
| 542 |
+
|
| 543 |
+
m0 = model_output_list[-1]
|
| 544 |
+
x = last_sample
|
| 545 |
+
x_t = this_sample
|
| 546 |
+
model_t = this_model_output
|
| 547 |
+
|
| 548 |
+
sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[
|
| 549 |
+
self.step_index - 1] # pyright: ignore
|
| 550 |
+
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
|
| 551 |
+
alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
|
| 552 |
+
|
| 553 |
+
lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
|
| 554 |
+
lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
|
| 555 |
+
|
| 556 |
+
h = lambda_t - lambda_s0
|
| 557 |
+
device = this_sample.device
|
| 558 |
+
|
| 559 |
+
rks = []
|
| 560 |
+
D1s = []
|
| 561 |
+
for i in range(1, order):
|
| 562 |
+
si = self.step_index - (i + 1) # pyright: ignore
|
| 563 |
+
mi = model_output_list[-(i + 1)]
|
| 564 |
+
alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
|
| 565 |
+
lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
|
| 566 |
+
rk = (lambda_si - lambda_s0) / h
|
| 567 |
+
rks.append(rk)
|
| 568 |
+
D1s.append((mi - m0) / rk) # pyright: ignore
|
| 569 |
+
|
| 570 |
+
rks.append(1.0)
|
| 571 |
+
rks = torch.tensor(rks, device=device)
|
| 572 |
+
|
| 573 |
+
R = []
|
| 574 |
+
b = []
|
| 575 |
+
|
| 576 |
+
hh = -h if self.predict_x0 else h
|
| 577 |
+
h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
|
| 578 |
+
h_phi_k = h_phi_1 / hh - 1
|
| 579 |
+
|
| 580 |
+
factorial_i = 1
|
| 581 |
+
|
| 582 |
+
if self.config.solver_type == "bh1":
|
| 583 |
+
B_h = hh
|
| 584 |
+
elif self.config.solver_type == "bh2":
|
| 585 |
+
B_h = torch.expm1(hh)
|
| 586 |
+
else:
|
| 587 |
+
raise NotImplementedError()
|
| 588 |
+
|
| 589 |
+
for i in range(1, order + 1):
|
| 590 |
+
R.append(torch.pow(rks, i - 1))
|
| 591 |
+
b.append(h_phi_k * factorial_i / B_h)
|
| 592 |
+
factorial_i *= i + 1
|
| 593 |
+
h_phi_k = h_phi_k / hh - 1 / factorial_i
|
| 594 |
+
|
| 595 |
+
R = torch.stack(R)
|
| 596 |
+
b = torch.tensor(b, device=device)
|
| 597 |
+
|
| 598 |
+
if len(D1s) > 0:
|
| 599 |
+
D1s = torch.stack(D1s, dim=1)
|
| 600 |
+
else:
|
| 601 |
+
D1s = None
|
| 602 |
+
|
| 603 |
+
# for order 1, we use a simplified version
|
| 604 |
+
if order == 1:
|
| 605 |
+
rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
|
| 606 |
+
else:
|
| 607 |
+
rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)
|
| 608 |
+
|
| 609 |
+
if self.predict_x0:
|
| 610 |
+
x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
|
| 611 |
+
if D1s is not None:
|
| 612 |
+
corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
|
| 613 |
+
else:
|
| 614 |
+
corr_res = 0
|
| 615 |
+
D1_t = model_t - m0
|
| 616 |
+
x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
|
| 617 |
+
else:
|
| 618 |
+
x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
|
| 619 |
+
if D1s is not None:
|
| 620 |
+
corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
|
| 621 |
+
else:
|
| 622 |
+
corr_res = 0
|
| 623 |
+
D1_t = model_t - m0
|
| 624 |
+
x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
|
| 625 |
+
x_t = x_t.to(x.dtype)
|
| 626 |
+
return x_t
|
| 627 |
+
|
| 628 |
+
def index_for_timestep(self, timestep, schedule_timesteps=None):
|
| 629 |
+
if schedule_timesteps is None:
|
| 630 |
+
schedule_timesteps = self.timesteps
|
| 631 |
+
|
| 632 |
+
indices = (schedule_timesteps == timestep).nonzero()
|
| 633 |
+
|
| 634 |
+
# The sigma index that is taken for the **very** first `step`
|
| 635 |
+
# is always the second index (or the last index if there is only 1)
|
| 636 |
+
# This way we can ensure we don't accidentally skip a sigma in
|
| 637 |
+
# case we start in the middle of the denoising schedule (e.g. for image-to-image)
|
| 638 |
+
pos = 1 if len(indices) > 1 else 0
|
| 639 |
+
|
| 640 |
+
return indices[pos].item()
|
| 641 |
+
|
| 642 |
+
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
|
| 643 |
+
def _init_step_index(self, timestep):
|
| 644 |
+
"""
|
| 645 |
+
Initialize the step_index counter for the scheduler.
|
| 646 |
+
"""
|
| 647 |
+
|
| 648 |
+
if self.begin_index is None:
|
| 649 |
+
if isinstance(timestep, torch.Tensor):
|
| 650 |
+
timestep = timestep.to(self.timesteps.device)
|
| 651 |
+
self._step_index = self.index_for_timestep(timestep)
|
| 652 |
+
else:
|
| 653 |
+
self._step_index = self._begin_index
|
| 654 |
+
|
| 655 |
+
def step(self,
|
| 656 |
+
model_output: torch.Tensor,
|
| 657 |
+
timestep: Union[int, torch.Tensor],
|
| 658 |
+
sample: torch.Tensor,
|
| 659 |
+
return_dict: bool = True,
|
| 660 |
+
generator=None) -> Union[SchedulerOutput, Tuple]:
|
| 661 |
+
"""
|
| 662 |
+
Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
|
| 663 |
+
the multistep UniPC.
|
| 664 |
+
|
| 665 |
+
Args:
|
| 666 |
+
model_output (`torch.Tensor`):
|
| 667 |
+
The direct output from learned diffusion model.
|
| 668 |
+
timestep (`int`):
|
| 669 |
+
The current discrete timestep in the diffusion chain.
|
| 670 |
+
sample (`torch.Tensor`):
|
| 671 |
+
A current instance of a sample created by the diffusion process.
|
| 672 |
+
return_dict (`bool`):
|
| 673 |
+
Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
|
| 674 |
+
|
| 675 |
+
Returns:
|
| 676 |
+
[`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
|
| 677 |
+
If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
|
| 678 |
+
tuple is returned where the first element is the sample tensor.
|
| 679 |
+
|
| 680 |
+
"""
|
| 681 |
+
if self.num_inference_steps is None:
|
| 682 |
+
raise ValueError(
|
| 683 |
+
"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
|
| 684 |
+
)
|
| 685 |
+
|
| 686 |
+
if self.step_index is None:
|
| 687 |
+
self._init_step_index(timestep)
|
| 688 |
+
|
| 689 |
+
use_corrector = (
|
| 690 |
+
self.step_index > 0 and
|
| 691 |
+
self.step_index - 1 not in self.disable_corrector and
|
| 692 |
+
self.last_sample is not None # pyright: ignore
|
| 693 |
+
)
|
| 694 |
+
|
| 695 |
+
model_output_convert = self.convert_model_output(
|
| 696 |
+
model_output, sample=sample)
|
| 697 |
+
if use_corrector:
|
| 698 |
+
sample = self.multistep_uni_c_bh_update(
|
| 699 |
+
this_model_output=model_output_convert,
|
| 700 |
+
last_sample=self.last_sample,
|
| 701 |
+
this_sample=sample,
|
| 702 |
+
order=self.this_order,
|
| 703 |
+
)
|
| 704 |
+
|
| 705 |
+
for i in range(self.config.solver_order - 1):
|
| 706 |
+
self.model_outputs[i] = self.model_outputs[i + 1]
|
| 707 |
+
self.timestep_list[i] = self.timestep_list[i + 1]
|
| 708 |
+
|
| 709 |
+
self.model_outputs[-1] = model_output_convert
|
| 710 |
+
self.timestep_list[-1] = timestep # pyright: ignore
|
| 711 |
+
|
| 712 |
+
if self.config.lower_order_final:
|
| 713 |
+
this_order = min(self.config.solver_order,
|
| 714 |
+
len(self.timesteps) -
|
| 715 |
+
self.step_index) # pyright: ignore
|
| 716 |
+
else:
|
| 717 |
+
this_order = self.config.solver_order
|
| 718 |
+
|
| 719 |
+
self.this_order = min(this_order,
|
| 720 |
+
self.lower_order_nums + 1) # warmup for multistep
|
| 721 |
+
assert self.this_order > 0
|
| 722 |
+
|
| 723 |
+
self.last_sample = sample
|
| 724 |
+
prev_sample = self.multistep_uni_p_bh_update(
|
| 725 |
+
model_output=model_output, # pass the original non-converted model output, in case solver-p is used
|
| 726 |
+
sample=sample,
|
| 727 |
+
order=self.this_order,
|
| 728 |
+
)
|
| 729 |
+
|
| 730 |
+
if self.lower_order_nums < self.config.solver_order:
|
| 731 |
+
self.lower_order_nums += 1
|
| 732 |
+
|
| 733 |
+
# upon completion increase step index by one
|
| 734 |
+
self._step_index += 1 # pyright: ignore
|
| 735 |
+
|
| 736 |
+
if not return_dict:
|
| 737 |
+
return (prev_sample,)
|
| 738 |
+
|
| 739 |
+
return SchedulerOutput(prev_sample=prev_sample)
|
| 740 |
+
|
| 741 |
+
def scale_model_input(self, sample: torch.Tensor, *args,
|
| 742 |
+
**kwargs) -> torch.Tensor:
|
| 743 |
+
"""
|
| 744 |
+
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
|
| 745 |
+
current timestep.
|
| 746 |
+
|
| 747 |
+
Args:
|
| 748 |
+
sample (`torch.Tensor`):
|
| 749 |
+
The input sample.
|
| 750 |
+
|
| 751 |
+
Returns:
|
| 752 |
+
`torch.Tensor`:
|
| 753 |
+
A scaled input sample.
|
| 754 |
+
"""
|
| 755 |
+
return sample
|
| 756 |
+
|
| 757 |
+
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise
|
| 758 |
+
def add_noise(
|
| 759 |
+
self,
|
| 760 |
+
original_samples: torch.Tensor,
|
| 761 |
+
noise: torch.Tensor,
|
| 762 |
+
timesteps: torch.IntTensor,
|
| 763 |
+
) -> torch.Tensor:
|
| 764 |
+
# Make sure sigmas and timesteps have the same device and dtype as original_samples
|
| 765 |
+
sigmas = self.sigmas.to(
|
| 766 |
+
device=original_samples.device, dtype=original_samples.dtype)
|
| 767 |
+
if original_samples.device.type == "mps" and torch.is_floating_point(
|
| 768 |
+
timesteps):
|
| 769 |
+
# mps does not support float64
|
| 770 |
+
schedule_timesteps = self.timesteps.to(
|
| 771 |
+
original_samples.device, dtype=torch.float32)
|
| 772 |
+
timesteps = timesteps.to(
|
| 773 |
+
original_samples.device, dtype=torch.float32)
|
| 774 |
+
else:
|
| 775 |
+
schedule_timesteps = self.timesteps.to(original_samples.device)
|
| 776 |
+
timesteps = timesteps.to(original_samples.device)
|
| 777 |
+
|
| 778 |
+
# begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
|
| 779 |
+
if self.begin_index is None:
|
| 780 |
+
step_indices = [
|
| 781 |
+
self.index_for_timestep(t, schedule_timesteps)
|
| 782 |
+
for t in timesteps
|
| 783 |
+
]
|
| 784 |
+
elif self.step_index is not None:
|
| 785 |
+
# add_noise is called after first denoising step (for inpainting)
|
| 786 |
+
step_indices = [self.step_index] * timesteps.shape[0]
|
| 787 |
+
else:
|
| 788 |
+
# add noise is called before first denoising step to create initial latent(img2img)
|
| 789 |
+
step_indices = [self.begin_index] * timesteps.shape[0]
|
| 790 |
+
|
| 791 |
+
sigma = sigmas[step_indices].flatten()
|
| 792 |
+
while len(sigma.shape) < len(original_samples.shape):
|
| 793 |
+
sigma = sigma.unsqueeze(-1)
|
| 794 |
+
|
| 795 |
+
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
|
| 796 |
+
noisy_samples = alpha_t * original_samples + sigma_t * noise
|
| 797 |
+
return noisy_samples
|
| 798 |
+
|
| 799 |
+
def __len__(self):
|
| 800 |
+
return self.config.num_train_timesteps
|
models/pipeline.py
ADDED
|
@@ -0,0 +1,435 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
import einops
|
| 3 |
+
import numpy as np
|
| 4 |
+
import tqdm
|
| 5 |
+
from PIL import Image
|
| 6 |
+
import torchvision.transforms.v2 as transforms
|
| 7 |
+
from diffusers import FlowMatchEulerDiscreteScheduler
|
| 8 |
+
# FlowUniPCMultistepScheduler generates more details than FlowMatchEulerDiscreteScheduler
|
| 9 |
+
from .fm_solvers_unipc import FlowUniPCMultistepScheduler # noqa: E402
|
| 10 |
+
from .flash_scheduler import FlashFlowMatchEulerDiscreteScheduler
|
| 11 |
+
from .utils import resize_pilimage, calculate_dimensions, get_rope_index_fix_point, find_closest_resolution
|
| 12 |
+
from .utils import create_layout_reference_images, load_layout_bboxes
|
| 13 |
+
|
| 14 |
+
TIMESTEP_TOKEN_NUM = 1
|
| 15 |
+
NOISE_SCALE = 8.0
|
| 16 |
+
T_EPS = 0.001
|
| 17 |
+
CONDITION_IMAGE_SIZE = 384
|
| 18 |
+
PATCH_SIZE = 32
|
| 19 |
+
|
| 20 |
+
TENSOR_TRANSFORM = transforms.Compose([
|
| 21 |
+
transforms.ToImage(),
|
| 22 |
+
transforms.ToDtype(torch.float32, scale=True),
|
| 23 |
+
transforms.Normalize([0.5], [0.5]),
|
| 24 |
+
])
|
| 25 |
+
|
| 26 |
+
DEFAULT_TIMESTEPS = [
|
| 27 |
+
999, 987, 974, 960, 945, 929, 913, 895, 877, 857, 836, 814, 790, 764, 737,
|
| 28 |
+
707, 675, 640, 602, 560, 515, 464, 409, 347, 278, 199, 110, 8,
|
| 29 |
+
]
|
| 30 |
+
|
| 31 |
+
def build_t2i_text_sample(prompt, height, width, tokenizer, processor, model_config):
|
| 32 |
+
image_token_id = model_config.image_token_id
|
| 33 |
+
video_token_id = model_config.video_token_id
|
| 34 |
+
vision_start_token_id = model_config.vision_start_token_id
|
| 35 |
+
image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE)
|
| 36 |
+
|
| 37 |
+
boi_token = getattr(tokenizer, "boi_token", "<|boi_token|>")
|
| 38 |
+
tms_token = getattr(tokenizer, "tms_token", "<|tms_token|>")
|
| 39 |
+
|
| 40 |
+
messages = [{"role": "user", "content": prompt}]
|
| 41 |
+
template_caption = (
|
| 42 |
+
processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
| 43 |
+
+ boi_token
|
| 44 |
+
+ tms_token * TIMESTEP_TOKEN_NUM
|
| 45 |
+
)
|
| 46 |
+
input_ids = tokenizer.encode(template_caption, return_tensors="pt", add_special_tokens=False)
|
| 47 |
+
|
| 48 |
+
image_grid_thw = torch.tensor(
|
| 49 |
+
[1, height // PATCH_SIZE, width // PATCH_SIZE], dtype=torch.int64
|
| 50 |
+
).unsqueeze(0)
|
| 51 |
+
|
| 52 |
+
vision_tokens = torch.zeros((1, image_len), dtype=input_ids.dtype) + image_token_id
|
| 53 |
+
vision_tokens[0, 0] = vision_start_token_id
|
| 54 |
+
input_ids_pad = torch.cat([input_ids, vision_tokens], dim=-1)
|
| 55 |
+
|
| 56 |
+
position_ids, _ = get_rope_index_fix_point(
|
| 57 |
+
1, image_token_id, video_token_id, vision_start_token_id,
|
| 58 |
+
input_ids=input_ids_pad, image_grid_thw=image_grid_thw,
|
| 59 |
+
video_grid_thw=None, attention_mask=None, skip_vision_start_token=[1],
|
| 60 |
+
)
|
| 61 |
+
|
| 62 |
+
txt_seq_len = input_ids.shape[-1]
|
| 63 |
+
all_seq_len = position_ids.shape[-1]
|
| 64 |
+
|
| 65 |
+
token_types = torch.zeros((1, all_seq_len), dtype=input_ids.dtype)
|
| 66 |
+
bgn = txt_seq_len - TIMESTEP_TOKEN_NUM
|
| 67 |
+
token_types[0, bgn: bgn + image_len + TIMESTEP_TOKEN_NUM] = 1
|
| 68 |
+
token_types[0, txt_seq_len - TIMESTEP_TOKEN_NUM: txt_seq_len] = 3
|
| 69 |
+
|
| 70 |
+
vinput_mask = (token_types == 1)
|
| 71 |
+
token_types_bin = (token_types > 0).to(token_types.dtype)
|
| 72 |
+
|
| 73 |
+
return {
|
| 74 |
+
'input_ids': input_ids,
|
| 75 |
+
'position_ids': position_ids,
|
| 76 |
+
'token_types': token_types_bin,
|
| 77 |
+
'vinput_mask': vinput_mask,
|
| 78 |
+
}
|
| 79 |
+
|
| 80 |
+
def build_scheduler(num_inference_steps, timesteps_list, shift, device, scheduler_name="default"):
|
| 81 |
+
if scheduler_name == "flash":
|
| 82 |
+
sched = FlashFlowMatchEulerDiscreteScheduler(
|
| 83 |
+
num_train_timesteps=1000, shift=shift, use_dynamic_shifting=False)
|
| 84 |
+
elif scheduler_name == "flow_match":
|
| 85 |
+
sched = FlowMatchEulerDiscreteScheduler(
|
| 86 |
+
num_train_timesteps=1000, shift=shift)
|
| 87 |
+
elif scheduler_name == "default":
|
| 88 |
+
sched = FlowUniPCMultistepScheduler(use_dynamic_shifting=False, shift=shift)
|
| 89 |
+
else:
|
| 90 |
+
raise ValueError(f"Unknown scheduler_name={scheduler_name!r}")
|
| 91 |
+
sched.set_timesteps(num_inference_steps, device=device)
|
| 92 |
+
if timesteps_list is not None:
|
| 93 |
+
sched.timesteps = torch.tensor(timesteps_list, device=device, dtype=torch.long)
|
| 94 |
+
sigmas = [t.item() / 1000.0 for t in sched.timesteps]
|
| 95 |
+
sigmas.append(0.0)
|
| 96 |
+
sched.sigmas = torch.tensor(sigmas, device=device)
|
| 97 |
+
return sched
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
def clamp_tensor(tensor, percentage = 0.1):
|
| 101 |
+
lower_bound = torch.quantile(tensor.float(), percentage)
|
| 102 |
+
upper_bound = torch.quantile(tensor.float(), 1 - percentage)
|
| 103 |
+
src_dtype = tensor.dtype
|
| 104 |
+
return torch.clamp(tensor.float(), min=lower_bound, max=upper_bound).to(src_dtype)
|
| 105 |
+
|
| 106 |
+
@torch.no_grad()
|
| 107 |
+
def generate_image(
|
| 108 |
+
model,
|
| 109 |
+
processor,
|
| 110 |
+
prompt: str,
|
| 111 |
+
ref_image_paths: list = None,
|
| 112 |
+
height: int = 1440,
|
| 113 |
+
width: int = 2560,
|
| 114 |
+
num_inference_steps: int = 50,
|
| 115 |
+
guidance_scale: float = 5.0,
|
| 116 |
+
shift: float = 3.0,
|
| 117 |
+
timesteps_list=None,
|
| 118 |
+
scheduler_name: str = "default",
|
| 119 |
+
seed: int = 42,
|
| 120 |
+
noise_scale_start: float = NOISE_SCALE,
|
| 121 |
+
noise_scale_end: float = NOISE_SCALE,
|
| 122 |
+
noise_clip_std: float = 0.0,
|
| 123 |
+
keep_original_aspect: bool = False,
|
| 124 |
+
layout_bboxes: str = None,
|
| 125 |
+
callback=None,
|
| 126 |
+
) -> Image.Image:
|
| 127 |
+
device = model.device
|
| 128 |
+
dtype = torch.bfloat16
|
| 129 |
+
model_config = model.config
|
| 130 |
+
tokenizer = processor.tokenizer if hasattr(processor, 'tokenizer') else processor
|
| 131 |
+
|
| 132 |
+
# When `keep_original_aspect` is enabled and exactly one reference image is
|
| 133 |
+
# provided, resize the reference to max_size=2048 (patch-aligned) and derive
|
| 134 |
+
# the target image dimensions from the resized reference. This bypasses the
|
| 135 |
+
# predefined-resolution snapping so the output preserves the reference's
|
| 136 |
+
# original aspect ratio.
|
| 137 |
+
preresized_ref_pil = None
|
| 138 |
+
if keep_original_aspect and ref_image_paths and len(ref_image_paths) == 1:
|
| 139 |
+
pil_orig = Image.open(ref_image_paths[0]).convert("RGB")
|
| 140 |
+
preresized_ref_pil = resize_pilimage(pil_orig, 2048, PATCH_SIZE)
|
| 141 |
+
width, height = preresized_ref_pil.size
|
| 142 |
+
print(
|
| 143 |
+
f"[info] keep_original_aspect: target size set to {width}x{height} "
|
| 144 |
+
f"from reference image"
|
| 145 |
+
)
|
| 146 |
+
else:
|
| 147 |
+
if keep_original_aspect:
|
| 148 |
+
print(
|
| 149 |
+
"[warning] keep_original_aspect requires exactly one reference "
|
| 150 |
+
"image; falling back to default resolution snapping."
|
| 151 |
+
)
|
| 152 |
+
w, h = find_closest_resolution(width, height)
|
| 153 |
+
if w != width or h != height:
|
| 154 |
+
print(f"[warning] Resolution snapped from {width}x{height} to {w}x{h}")
|
| 155 |
+
width, height = w, h
|
| 156 |
+
|
| 157 |
+
h_patches = height // PATCH_SIZE
|
| 158 |
+
w_patches = width // PATCH_SIZE
|
| 159 |
+
|
| 160 |
+
if not ref_image_paths:
|
| 161 |
+
cond_sample = build_t2i_text_sample(prompt, height, width, tokenizer, processor, model_config)
|
| 162 |
+
uncond_sample = None
|
| 163 |
+
if guidance_scale > 1.0:
|
| 164 |
+
uncond_sample = build_t2i_text_sample(" ", height, width, tokenizer, processor, model_config)
|
| 165 |
+
|
| 166 |
+
def to_device(s):
|
| 167 |
+
return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in s.items()}
|
| 168 |
+
|
| 169 |
+
cond_sample = to_device(cond_sample)
|
| 170 |
+
if uncond_sample is not None:
|
| 171 |
+
uncond_sample = to_device(uncond_sample)
|
| 172 |
+
|
| 173 |
+
ref_patches = None
|
| 174 |
+
tgt_image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE)
|
| 175 |
+
samples = [cond_sample]
|
| 176 |
+
if uncond_sample:
|
| 177 |
+
samples.append(uncond_sample)
|
| 178 |
+
else:
|
| 179 |
+
image_token_id = model_config.image_token_id
|
| 180 |
+
video_token_id = model_config.video_token_id
|
| 181 |
+
vision_start_token_id = model_config.vision_start_token_id
|
| 182 |
+
spatial_merge_size = model_config.vision_config.spatial_merge_size
|
| 183 |
+
|
| 184 |
+
if preresized_ref_pil is not None:
|
| 185 |
+
ref_pils = [preresized_ref_pil]
|
| 186 |
+
else:
|
| 187 |
+
ref_pils = [Image.open(p).convert("RGB") for p in ref_image_paths]
|
| 188 |
+
|
| 189 |
+
K = len(ref_pils)
|
| 190 |
+
layout_data = None
|
| 191 |
+
if layout_bboxes is not None and len(layout_bboxes) > 0 and preresized_ref_pil is None:
|
| 192 |
+
try:
|
| 193 |
+
layout_data = load_layout_bboxes(layout_bboxes)
|
| 194 |
+
K += 1
|
| 195 |
+
except Exception as e:
|
| 196 |
+
print(f"Incorrect layout_bboxes: {layout_bboxes}, {e}")
|
| 197 |
+
|
| 198 |
+
if K == 1: max_size = max(height, width)
|
| 199 |
+
elif K == 2: max_size = max(height, width) * 48 // 64
|
| 200 |
+
elif K <= 4: max_size = max(height, width) // 2
|
| 201 |
+
elif K <= 8: max_size = max(height, width) * 24 // 64
|
| 202 |
+
else: max_size = max(height, width) // 4
|
| 203 |
+
|
| 204 |
+
if layout_data is not None:
|
| 205 |
+
ref_pils = create_layout_reference_images(
|
| 206 |
+
ref_pils = ref_pils,
|
| 207 |
+
layout_bboxes = layout_data,
|
| 208 |
+
image_width = width,
|
| 209 |
+
image_height = height,
|
| 210 |
+
ref_max_size = max_size,
|
| 211 |
+
patch_size = PATCH_SIZE,
|
| 212 |
+
)
|
| 213 |
+
|
| 214 |
+
ref_pils_resized, ref_images = [], []
|
| 215 |
+
for pil in ref_pils:
|
| 216 |
+
# Skip resizing when caller already produced a patch-aligned ref via
|
| 217 |
+
# `keep_original_aspect` — re-running resize_pilimage on it would
|
| 218 |
+
# upscale (since max_size == max(width, height) of the resized ref).
|
| 219 |
+
if preresized_ref_pil is not None and pil is preresized_ref_pil:
|
| 220 |
+
pil_r = pil
|
| 221 |
+
else:
|
| 222 |
+
pil_r = resize_pilimage(pil, max_size, PATCH_SIZE)
|
| 223 |
+
ref_pils_resized.append(pil_r)
|
| 224 |
+
x = TENSOR_TRANSFORM(pil_r)
|
| 225 |
+
x = einops.rearrange(x, "C (H p1) (W p2) -> (H W) (C p1 p2)", p1=PATCH_SIZE, p2=PATCH_SIZE)
|
| 226 |
+
ref_images.append(x)
|
| 227 |
+
|
| 228 |
+
ref_image_lens = [img.shape[0] for img in ref_images]
|
| 229 |
+
total_ref_len = sum(ref_image_lens)
|
| 230 |
+
ref_patches = torch.cat(ref_images, dim=0).unsqueeze(0).to(device, dtype)
|
| 231 |
+
|
| 232 |
+
tgt_image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE)
|
| 233 |
+
h_patches = height // PATCH_SIZE
|
| 234 |
+
w_patches = width // PATCH_SIZE
|
| 235 |
+
|
| 236 |
+
if K <= 4: cond_img_size = CONDITION_IMAGE_SIZE
|
| 237 |
+
elif K <= 8: cond_img_size = CONDITION_IMAGE_SIZE * 48 // 64
|
| 238 |
+
else: cond_img_size = CONDITION_IMAGE_SIZE // 2
|
| 239 |
+
|
| 240 |
+
ref_pils_vlm = []
|
| 241 |
+
for pil_r in ref_pils_resized:
|
| 242 |
+
cond_w, cond_h = calculate_dimensions(cond_img_size, pil_r.width / pil_r.height)
|
| 243 |
+
ref_pils_vlm.append(pil_r.resize((cond_w, cond_h), resample=Image.LANCZOS))
|
| 244 |
+
|
| 245 |
+
image_grid_thw_tgt = torch.tensor([1, height // PATCH_SIZE, width // PATCH_SIZE], dtype=torch.int64).unsqueeze(0)
|
| 246 |
+
image_grid_thw_ref = torch.zeros((K, 3), dtype=torch.int64)
|
| 247 |
+
for i, pil_r in enumerate(ref_pils_resized):
|
| 248 |
+
rw, rh = pil_r.size
|
| 249 |
+
image_grid_thw_ref[i] = torch.tensor([1, rh // PATCH_SIZE, rw // PATCH_SIZE], dtype=torch.int64)
|
| 250 |
+
|
| 251 |
+
samples = []
|
| 252 |
+
captions = [prompt]
|
| 253 |
+
if guidance_scale > 1.0:
|
| 254 |
+
captions.append(" ")
|
| 255 |
+
|
| 256 |
+
for caption in captions:
|
| 257 |
+
boi_token = getattr(tokenizer, "boi_token", "<|boi_token|>")
|
| 258 |
+
tms_token = getattr(tokenizer, "tms_token", "<|tms_token|>")
|
| 259 |
+
|
| 260 |
+
content = [{"type": "image"} for _ in range(K)]
|
| 261 |
+
content.append({"type": "text", "text": caption})
|
| 262 |
+
messages = [{"role": "user", "content": content}]
|
| 263 |
+
template_caption = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
| 264 |
+
proc = processor(text=[template_caption], images=ref_pils_vlm, padding="longest", return_tensors="pt")
|
| 265 |
+
input_ids_2 = tokenizer.encode(boi_token + tms_token * TIMESTEP_TOKEN_NUM, return_tensors="pt", add_special_tokens=False)
|
| 266 |
+
input_ids = torch.cat([proc.input_ids, input_ids_2], dim=-1)
|
| 267 |
+
|
| 268 |
+
igthw_cond = proc.image_grid_thw.clone()
|
| 269 |
+
for i in range(K):
|
| 270 |
+
igthw_cond[i, 1] //= spatial_merge_size
|
| 271 |
+
igthw_cond[i, 2] //= spatial_merge_size
|
| 272 |
+
igthw_all = torch.cat([igthw_cond, image_grid_thw_tgt, image_grid_thw_ref], dim=0)
|
| 273 |
+
|
| 274 |
+
vision_tokens_list = []
|
| 275 |
+
vt_tgt = torch.full((1, tgt_image_len), image_token_id, dtype=input_ids.dtype)
|
| 276 |
+
vt_tgt[0, 0] = vision_start_token_id
|
| 277 |
+
vision_tokens_list.append(vt_tgt)
|
| 278 |
+
for rl in ref_image_lens:
|
| 279 |
+
vt_ref = torch.full((1, rl), image_token_id, dtype=input_ids.dtype)
|
| 280 |
+
vt_ref[0, 0] = vision_start_token_id
|
| 281 |
+
vision_tokens_list.append(vt_ref)
|
| 282 |
+
vision_tokens = torch.cat(vision_tokens_list, dim=1)
|
| 283 |
+
input_ids_pad = torch.cat([input_ids, vision_tokens], dim=-1)
|
| 284 |
+
|
| 285 |
+
position_ids, _ = get_rope_index_fix_point(
|
| 286 |
+
1, image_token_id, video_token_id, vision_start_token_id,
|
| 287 |
+
input_ids=input_ids_pad, image_grid_thw=igthw_all,
|
| 288 |
+
video_grid_thw=None, attention_mask=None,
|
| 289 |
+
skip_vision_start_token=[0] * K + [1] + [1] * K,
|
| 290 |
+
)
|
| 291 |
+
txt_seq_len = input_ids.shape[-1]
|
| 292 |
+
all_seq_len = position_ids.shape[-1]
|
| 293 |
+
|
| 294 |
+
token_types_raw = torch.zeros((1, all_seq_len), dtype=input_ids.dtype)
|
| 295 |
+
bgn = txt_seq_len - TIMESTEP_TOKEN_NUM
|
| 296 |
+
end = bgn + tgt_image_len + TIMESTEP_TOKEN_NUM
|
| 297 |
+
token_types_raw[0, bgn:end] = 1
|
| 298 |
+
token_types_raw[0, end: end + total_ref_len] = 2
|
| 299 |
+
token_types_raw[0, txt_seq_len - TIMESTEP_TOKEN_NUM: txt_seq_len] = 3
|
| 300 |
+
|
| 301 |
+
vinput_mask = torch.logical_or(token_types_raw == 1, token_types_raw == 2)
|
| 302 |
+
token_types_bin = (token_types_raw > 0).to(token_types_raw.dtype)
|
| 303 |
+
|
| 304 |
+
samples.append({
|
| 305 |
+
"input_ids": input_ids.to(device),
|
| 306 |
+
"position_ids": position_ids.to(device),
|
| 307 |
+
"token_types": token_types_bin.to(device),
|
| 308 |
+
"vinput_mask": vinput_mask.to(device),
|
| 309 |
+
"pixel_values": proc.pixel_values.to(device, dtype),
|
| 310 |
+
"image_grid_thw": proc.image_grid_thw.to(device),
|
| 311 |
+
})
|
| 312 |
+
|
| 313 |
+
noise = noise_scale_start * torch.randn(
|
| 314 |
+
(1, 3, height, width),
|
| 315 |
+
generator=torch.Generator('cpu').manual_seed(seed + 1),
|
| 316 |
+
).to(device, dtype)
|
| 317 |
+
z = einops.rearrange(noise, 'B C (H p1) (W p2) -> B (H W) (C p1 p2)', p1=PATCH_SIZE, p2=PATCH_SIZE)
|
| 318 |
+
|
| 319 |
+
sched = build_scheduler(num_inference_steps, timesteps_list, shift, device, scheduler_name)
|
| 320 |
+
|
| 321 |
+
num_steps = len(sched.timesteps)
|
| 322 |
+
if num_steps > 1:
|
| 323 |
+
noise_scale_schedule = [
|
| 324 |
+
noise_scale_start + (noise_scale_end - noise_scale_start) * i / (num_steps - 1)
|
| 325 |
+
for i in range(num_steps)
|
| 326 |
+
]
|
| 327 |
+
else:
|
| 328 |
+
noise_scale_schedule = [noise_scale_start]
|
| 329 |
+
|
| 330 |
+
torch.manual_seed(seed + 1)
|
| 331 |
+
if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed + 1)
|
| 332 |
+
|
| 333 |
+
def forward_once(sample, z_in, t_pixeldit, precomputed_image_embeds=None, precomputed_deepstack_image_embeds=None):
|
| 334 |
+
with torch.autocast(device.type, dtype=dtype, cache_enabled=False):
|
| 335 |
+
kwargs = {
|
| 336 |
+
"input_ids": sample['input_ids'],
|
| 337 |
+
"position_ids": sample['position_ids'],
|
| 338 |
+
"vinputs": z_in,
|
| 339 |
+
"timestep": t_pixeldit.reshape(-1).to(device),
|
| 340 |
+
"token_types": sample['token_types'],
|
| 341 |
+
"use_flash_attn": False,
|
| 342 |
+
"precomputed_image_embeds": precomputed_image_embeds,
|
| 343 |
+
"precomputed_deepstack_image_embeds": precomputed_deepstack_image_embeds
|
| 344 |
+
}
|
| 345 |
+
if "pixel_values" in sample: kwargs["pixel_values"] = sample["pixel_values"]
|
| 346 |
+
if "image_grid_thw" in sample: kwargs["image_grid_thw"] = sample["image_grid_thw"]
|
| 347 |
+
|
| 348 |
+
outputs = model(**kwargs)
|
| 349 |
+
|
| 350 |
+
x_pred = outputs.x_pred
|
| 351 |
+
emb = getattr(outputs, "cond_image_embeds", None)
|
| 352 |
+
ds_emb = getattr(outputs, "cond_deepstack_image_embeds", None)
|
| 353 |
+
# x_pred = clamp_tensor(x_pred, percentage = 0.01)
|
| 354 |
+
if ref_patches is None:
|
| 355 |
+
return x_pred[0, sample['vinput_mask'][0]].unsqueeze(0)
|
| 356 |
+
else:
|
| 357 |
+
return x_pred[0, sample['vinput_mask'][0]][:tgt_image_len].unsqueeze(0), emb, ds_emb
|
| 358 |
+
|
| 359 |
+
def _decode_x0_preview(x0_pred):
|
| 360 |
+
"""Convert a model-predicted x_0 (patch layout, [-1,1]) to a PIL image."""
|
| 361 |
+
img_t = (x0_pred.float() + 1) / 2
|
| 362 |
+
img_t = einops.rearrange(
|
| 363 |
+
img_t.cpu(), 'B (H W) (C p1 p2) -> B C (H p1) (W p2)',
|
| 364 |
+
H=h_patches, W=w_patches, p1=PATCH_SIZE, p2=PATCH_SIZE,
|
| 365 |
+
)
|
| 366 |
+
arr_p = np.round(np.clip(img_t[0].numpy().transpose(1, 2, 0) * 255, 0, 255)).astype(np.uint8)
|
| 367 |
+
return Image.fromarray(arr_p).convert("RGB")
|
| 368 |
+
|
| 369 |
+
cond_image_embeds = None
|
| 370 |
+
cond_deepstack_image_embeds = None
|
| 371 |
+
|
| 372 |
+
for step_idx, step_t in enumerate(tqdm.tqdm(sched.timesteps, desc="Generating")):
|
| 373 |
+
t_pixeldit = 1.0 - step_t.float() / 1000.0
|
| 374 |
+
sigma = (step_t.float() / 1000.0).to(dtype=torch.float32).clamp_min(T_EPS)
|
| 375 |
+
|
| 376 |
+
if ref_patches is None:
|
| 377 |
+
x_pred_cond = forward_once(samples[0], z.clone(), t_pixeldit)
|
| 378 |
+
v_cond = (x_pred_cond.to(dtype=torch.float32) - z.to(dtype=torch.float32)) / sigma
|
| 379 |
+
|
| 380 |
+
if len(samples) > 1:
|
| 381 |
+
x_pred_uncond = forward_once(samples[1], z.clone(), t_pixeldit)
|
| 382 |
+
v_uncond = (x_pred_uncond.to(dtype=torch.float32) - z.to(dtype=torch.float32)) / sigma
|
| 383 |
+
v_guided = v_uncond + guidance_scale * (v_cond - v_uncond)
|
| 384 |
+
else:
|
| 385 |
+
v_guided = v_cond
|
| 386 |
+
preview_x0 = x_pred_cond
|
| 387 |
+
else:
|
| 388 |
+
vinputs = torch.cat([z, ref_patches], dim=1)
|
| 389 |
+
x_vis_list = []
|
| 390 |
+
for sample in samples:
|
| 391 |
+
xp, emb, ds_emb = forward_once(
|
| 392 |
+
sample,
|
| 393 |
+
vinputs,
|
| 394 |
+
t_pixeldit,
|
| 395 |
+
precomputed_image_embeds=cond_image_embeds,
|
| 396 |
+
precomputed_deepstack_image_embeds=cond_deepstack_image_embeds
|
| 397 |
+
)
|
| 398 |
+
if emb is not None and ds_emb is not None:
|
| 399 |
+
cond_image_embeds = emb.detach()
|
| 400 |
+
cond_deepstack_image_embeds = [d.detach() for d in ds_emb]
|
| 401 |
+
x_vis_list.append(xp)
|
| 402 |
+
#x_vis_list = [forward_once(sample, vinputs, t_pixeldit) for sample in samples]
|
| 403 |
+
x_vis_stacked = torch.cat(x_vis_list, dim=0)
|
| 404 |
+
|
| 405 |
+
z_rep = z.expand(len(samples), -1, -1)
|
| 406 |
+
v_pred = (x_vis_stacked.to(dtype=torch.float32) - z_rep.to(dtype=torch.float32)) / sigma
|
| 407 |
+
|
| 408 |
+
v_cond = v_pred[0:1]
|
| 409 |
+
if len(samples) > 1:
|
| 410 |
+
v_uncond = v_pred[1:]
|
| 411 |
+
v_guided = v_uncond + guidance_scale * (v_cond - v_uncond)
|
| 412 |
+
else:
|
| 413 |
+
v_guided = v_cond
|
| 414 |
+
preview_x0 = x_vis_list[0]
|
| 415 |
+
|
| 416 |
+
model_output = -v_guided
|
| 417 |
+
# model_output = clamp_tensor(model_output, percentage = 0.05)
|
| 418 |
+
if scheduler_name == "flash":
|
| 419 |
+
z = sched.step(model_output.float(), step_t.to(dtype=torch.float32), z.float(), s_noise=noise_scale_schedule[step_idx], noise_clip_std=noise_clip_std, return_dict=False)[0].to(dtype)
|
| 420 |
+
else:
|
| 421 |
+
z = sched.step(model_output.float(), step_t.to(dtype=torch.float32), z.float(), return_dict=False)[0].to(dtype)
|
| 422 |
+
|
| 423 |
+
if callback is not None:
|
| 424 |
+
try:
|
| 425 |
+
# Pass a closure that captures the current step's x0 prediction.
|
| 426 |
+
# Use a default-arg binding to avoid late-binding issues.
|
| 427 |
+
callback(step_idx, len(sched.timesteps),
|
| 428 |
+
lambda x0=preview_x0: _decode_x0_preview(x0))
|
| 429 |
+
except Exception:
|
| 430 |
+
pass
|
| 431 |
+
|
| 432 |
+
img = (z + 1) / 2
|
| 433 |
+
img = einops.rearrange(img.cpu().float(), 'B (H W) (C p1 p2) -> B C (H p1) (W p2)', H=h_patches, W=w_patches, p1=PATCH_SIZE, p2=PATCH_SIZE)
|
| 434 |
+
arr = np.round(np.clip(img[0].numpy().transpose(1, 2, 0) * 255, 0, 255)).astype(np.uint8)
|
| 435 |
+
return Image.fromarray(arr).convert("RGB")
|
models/qwen3_vl_transformers.py
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
models/utils.py
ADDED
|
@@ -0,0 +1,368 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import math
|
| 3 |
+
import torch
|
| 4 |
+
from typing import Optional
|
| 5 |
+
from PIL import Image, ImageDraw
|
| 6 |
+
import json
|
| 7 |
+
from typing import Any, Dict, Iterable, List, Sequence, Tuple
|
| 8 |
+
|
| 9 |
+
MAX_BOX = 5
|
| 10 |
+
PREDEFINED_RESOLUTIONS = [
|
| 11 |
+
(2048, 2048),
|
| 12 |
+
(2304, 1728),
|
| 13 |
+
(1728, 2304),
|
| 14 |
+
(2560, 1440),
|
| 15 |
+
(1440, 2560),
|
| 16 |
+
(2496, 1664),
|
| 17 |
+
(1664, 2496),
|
| 18 |
+
(3104, 1312),
|
| 19 |
+
(1312, 3104),
|
| 20 |
+
(2304, 1792),
|
| 21 |
+
(1792, 2304),
|
| 22 |
+
]
|
| 23 |
+
DEFAULT_COLORS = [
|
| 24 |
+
(255, 0, 0),
|
| 25 |
+
(0, 180, 0),
|
| 26 |
+
(0, 0, 255),
|
| 27 |
+
(204, 180, 0),
|
| 28 |
+
(255, 0, 255),
|
| 29 |
+
(0, 255, 255),
|
| 30 |
+
(128, 0, 0),
|
| 31 |
+
(0, 128, 0),
|
| 32 |
+
(0, 0, 128),
|
| 33 |
+
(128, 128, 0),
|
| 34 |
+
]
|
| 35 |
+
|
| 36 |
+
def load_layout_bboxes(layout_bboxes: str) -> Any:
|
| 37 |
+
"""Load layout boxes from either a JSON string or a JSON file path."""
|
| 38 |
+
if os.path.exists(layout_bboxes):
|
| 39 |
+
with open(layout_bboxes, "r", encoding="utf-8") as f:
|
| 40 |
+
return json.load(f)
|
| 41 |
+
return json.loads(layout_bboxes)
|
| 42 |
+
|
| 43 |
+
def _unwrap_boxes(data: Any) -> Any:
|
| 44 |
+
if isinstance(data, dict):
|
| 45 |
+
for key in ("layout_bboxes", "bboxes", "boxes", "bbox_list"):
|
| 46 |
+
if key in data:
|
| 47 |
+
return data[key]
|
| 48 |
+
return data
|
| 49 |
+
|
| 50 |
+
def _as_bbox_and_text(item: Any) -> Tuple[Sequence[float], str]:
|
| 51 |
+
if isinstance(item, dict):
|
| 52 |
+
bbox = item.get("bbox") or item.get("box")
|
| 53 |
+
text = str(item.get("text") or item.get("label") or "")
|
| 54 |
+
if bbox is None:
|
| 55 |
+
raise ValueError(f"Missing bbox in layout item: {item!r}")
|
| 56 |
+
return bbox, text
|
| 57 |
+
if isinstance(item, (list, tuple)) and len(item) == 4:
|
| 58 |
+
return item, ""
|
| 59 |
+
raise ValueError(f"Unsupported layout bbox item: {item!r}")
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def _xxyy_relative_to_absolute_bbox(bbox: Sequence[float], width: int, height: int) -> List[int]:
|
| 63 |
+
if len(bbox) != 4:
|
| 64 |
+
raise ValueError(f"Expected bbox with 4 values, got: {bbox!r}")
|
| 65 |
+
x1, x2, y1, y2 = [float(v) for v in bbox]
|
| 66 |
+
|
| 67 |
+
# Inference layout input is xxyy relative coordinates: [x1, x2, y1, y2].
|
| 68 |
+
# Values in [0, 1] are the intended format. Keep 0-100 support for convenience.
|
| 69 |
+
max_abs = max(abs(x1), abs(y1), abs(x2), abs(y2))
|
| 70 |
+
if max_abs <= 1.0:
|
| 71 |
+
x1, x2 = x1 * width, x2 * width
|
| 72 |
+
y1, y2 = y1 * height, y2 * height
|
| 73 |
+
elif max_abs <= 100.0:
|
| 74 |
+
x1, x2 = x1 / 100.0 * width, x2 / 100.0 * width
|
| 75 |
+
y1, y2 = y1 / 100.0 * height, y2 / 100.0 * height
|
| 76 |
+
|
| 77 |
+
x1, x2 = sorted((x1, x2))
|
| 78 |
+
y1, y2 = sorted((y1, y2))
|
| 79 |
+
x1 = max(0, min(width - 1, int(round(x1))))
|
| 80 |
+
y1 = max(0, min(height - 1, int(round(y1))))
|
| 81 |
+
x2 = max(0, min(width - 1, int(round(x2))))
|
| 82 |
+
y2 = max(0, min(height - 1, int(round(y2))))
|
| 83 |
+
if x2 <= x1 or y2 <= y1:
|
| 84 |
+
raise ValueError(f"Invalid bbox after scaling/clamping: {[x1, y1, x2, y2]!r}")
|
| 85 |
+
return [x1, y1, x2, y2]
|
| 86 |
+
|
| 87 |
+
def parse_layout_bboxes(layout_bboxes: Any, width: int, height: int) -> List[Dict[str, Any]]:
|
| 88 |
+
"""Convert xxyy relative layout boxes into the training-side bbox layout format."""
|
| 89 |
+
raw_boxes = _unwrap_boxes(layout_bboxes)
|
| 90 |
+
if not isinstance(raw_boxes, list):
|
| 91 |
+
raise ValueError("layout_bboxes must be a list, or a dict containing one of: layout_bboxes/bboxes/boxes")
|
| 92 |
+
|
| 93 |
+
parsed = []
|
| 94 |
+
for idx, item in enumerate(raw_boxes):
|
| 95 |
+
bbox, text = _as_bbox_and_text(item)
|
| 96 |
+
parsed.append({
|
| 97 |
+
"bbox": _xxyy_relative_to_absolute_bbox(bbox, width, height),
|
| 98 |
+
"color": "",
|
| 99 |
+
"text": text,
|
| 100 |
+
"image": None,
|
| 101 |
+
"_orig_idx": idx,
|
| 102 |
+
})
|
| 103 |
+
return parsed
|
| 104 |
+
|
| 105 |
+
def _bbox_area(item: Dict[str, Any]) -> int:
|
| 106 |
+
x1, y1, x2, y2 = item["bbox"]
|
| 107 |
+
return max(0, x2 - x1) * max(0, y2 - y1)
|
| 108 |
+
|
| 109 |
+
def get_render_params(image_width: int, image_height: int) -> Tuple[int, int]:
|
| 110 |
+
edge = math.sqrt(image_width * image_height)
|
| 111 |
+
max_font_size = int(edge * 0.07)
|
| 112 |
+
max_bbox_line_width = int(edge * 0.05)
|
| 113 |
+
return max_font_size, max_bbox_line_width
|
| 114 |
+
|
| 115 |
+
def draw_bbox_layout(
|
| 116 |
+
bbox_list: List[Dict[str, Any]],
|
| 117 |
+
image_width: int,
|
| 118 |
+
image_height: int,
|
| 119 |
+
max_bbox: int = MAX_BOX,
|
| 120 |
+
max_bbox_line_width: int | None = None,
|
| 121 |
+
bbox_line_gap: int | None = None,
|
| 122 |
+
return_color: bool = False,
|
| 123 |
+
):
|
| 124 |
+
"""Draw a black layout image with colored boxes, matching the training-side layout style."""
|
| 125 |
+
if max_bbox_line_width is None:
|
| 126 |
+
_, max_bbox_line_width = get_render_params(image_width, image_height)
|
| 127 |
+
if bbox_line_gap is None:
|
| 128 |
+
bbox_line_gap = max(1, max_bbox_line_width // max_bbox)
|
| 129 |
+
|
| 130 |
+
image = Image.new("RGB", (image_width, image_height), (0, 0, 0))
|
| 131 |
+
draw = ImageDraw.Draw(image)
|
| 132 |
+
color_list = [None] * len(bbox_list)
|
| 133 |
+
sorted_bboxes = sorted(bbox_list, key=_bbox_area, reverse=True)[:max_bbox]
|
| 134 |
+
|
| 135 |
+
for sorted_idx, item in enumerate(sorted_bboxes):
|
| 136 |
+
color = DEFAULT_COLORS[sorted_idx % len(DEFAULT_COLORS)]
|
| 137 |
+
orig_idx = int(item.get("_orig_idx", sorted_idx))
|
| 138 |
+
if 0 <= orig_idx < len(color_list):
|
| 139 |
+
color_list[orig_idx] = color
|
| 140 |
+
line_width = max(max_bbox_line_width - sorted_idx * bbox_line_gap, 5)
|
| 141 |
+
draw.rectangle([int(v) for v in item["bbox"]], outline=color, width=line_width)
|
| 142 |
+
|
| 143 |
+
if return_color:
|
| 144 |
+
return image, color_list
|
| 145 |
+
return image
|
| 146 |
+
|
| 147 |
+
def add_outer_border_keep_size(pil: Image.Image, color: Iterable[int], width: int) -> Image.Image:
|
| 148 |
+
"""Draw a border inside the image without changing its size."""
|
| 149 |
+
img = pil.convert("RGB").copy()
|
| 150 |
+
color_tuple = tuple(int(c) for c in color)
|
| 151 |
+
width = max(0, int(width))
|
| 152 |
+
if width == 0:
|
| 153 |
+
return img
|
| 154 |
+
|
| 155 |
+
draw = ImageDraw.Draw(img)
|
| 156 |
+
w, h = img.size
|
| 157 |
+
for t in range(width):
|
| 158 |
+
draw.rectangle([t, t, w - 1 - t, h - 1 - t], outline=color_tuple)
|
| 159 |
+
return img
|
| 160 |
+
|
| 161 |
+
def create_layout_reference_images(
|
| 162 |
+
ref_pils: Sequence[str],
|
| 163 |
+
layout_bboxes: Any,
|
| 164 |
+
image_width: int,
|
| 165 |
+
image_height: int,
|
| 166 |
+
ref_max_size: int | None = None,
|
| 167 |
+
patch_size: int = 32,
|
| 168 |
+
) -> Tuple[List[str], str]:
|
| 169 |
+
"""Create bordered ref images plus one layout image; returns paths to pass as ref_images."""
|
| 170 |
+
parsed_boxes = parse_layout_bboxes(layout_bboxes, image_width, image_height)
|
| 171 |
+
layout_image, color_list = draw_bbox_layout(
|
| 172 |
+
parsed_boxes,
|
| 173 |
+
image_width=image_width,
|
| 174 |
+
image_height=image_height,
|
| 175 |
+
return_color=True,
|
| 176 |
+
)
|
| 177 |
+
|
| 178 |
+
output_refs: List[str] = []
|
| 179 |
+
for idx, ref in enumerate(ref_pils):
|
| 180 |
+
if ref_max_size is not None:
|
| 181 |
+
ref = resize_pilimage(ref, ref_max_size, patch_size)
|
| 182 |
+
color = color_list[idx] if idx < len(color_list) and color_list[idx] is not None else DEFAULT_COLORS[idx % len(DEFAULT_COLORS)]
|
| 183 |
+
line_width = int(math.sqrt(ref.width * ref.height) * 0.04)
|
| 184 |
+
bordered = add_outer_border_keep_size(ref, color, line_width)
|
| 185 |
+
output_refs.append(bordered)
|
| 186 |
+
output_refs.append(layout_image)
|
| 187 |
+
return output_refs
|
| 188 |
+
|
| 189 |
+
|
| 190 |
+
def find_closest_resolution(width, height):
|
| 191 |
+
img_ratio = width / height
|
| 192 |
+
best_res = None
|
| 193 |
+
min_diff = float("inf")
|
| 194 |
+
for w, h in PREDEFINED_RESOLUTIONS:
|
| 195 |
+
ratio = w / h
|
| 196 |
+
diff = abs(ratio - img_ratio)
|
| 197 |
+
if diff < min_diff:
|
| 198 |
+
min_diff = diff
|
| 199 |
+
best_res = (w, h)
|
| 200 |
+
return best_res
|
| 201 |
+
|
| 202 |
+
def resize_pilimage(pil_image, image_size, patch_size=16, resampler=Image.BICUBIC):
|
| 203 |
+
while min(*pil_image.size) >= 2 * image_size:
|
| 204 |
+
pil_image = pil_image.resize(
|
| 205 |
+
tuple(x // 2 for x in pil_image.size), resample=Image.BOX
|
| 206 |
+
)
|
| 207 |
+
|
| 208 |
+
m = patch_size
|
| 209 |
+
width, height = pil_image.width, pil_image.height
|
| 210 |
+
S_max = image_size * image_size
|
| 211 |
+
scale = S_max / (width * height)
|
| 212 |
+
scale = math.sqrt(scale)
|
| 213 |
+
|
| 214 |
+
new_sizes = [
|
| 215 |
+
(round(width * scale) // m * m, round(height * scale) // m * m),
|
| 216 |
+
(round(width * scale) // m * m, math.floor(height * scale) // m * m),
|
| 217 |
+
(math.floor(width * scale) // m * m, round(height * scale) // m * m),
|
| 218 |
+
(math.floor(width * scale) // m * m, math.floor(height * scale) // m * m),
|
| 219 |
+
]
|
| 220 |
+
new_sizes = sorted(new_sizes, key=lambda x: x[0] * x[1], reverse=True)
|
| 221 |
+
|
| 222 |
+
for new_size in new_sizes:
|
| 223 |
+
if new_size[0] * new_size[1] <= S_max:
|
| 224 |
+
break
|
| 225 |
+
|
| 226 |
+
s1 = width / new_size[0]
|
| 227 |
+
s2 = height / new_size[1]
|
| 228 |
+
if s1 < s2:
|
| 229 |
+
pil_image = pil_image.resize([new_size[0], round(height / s1)], resample=resampler)
|
| 230 |
+
top = (round(height / s1) - new_size[1]) // 2
|
| 231 |
+
pil_image = pil_image.crop((0, top, new_size[0], top + new_size[1]))
|
| 232 |
+
else:
|
| 233 |
+
pil_image = pil_image.resize([round(width / s2), new_size[1]], resample=resampler)
|
| 234 |
+
left = (round(width / s2) - new_size[0]) // 2
|
| 235 |
+
pil_image = pil_image.crop((left, 0, left + new_size[0], new_size[1]))
|
| 236 |
+
|
| 237 |
+
return pil_image
|
| 238 |
+
|
| 239 |
+
def calculate_dimensions(max_size, ratio):
|
| 240 |
+
width = math.sqrt(max_size * max_size * ratio)
|
| 241 |
+
height = width / ratio
|
| 242 |
+
width = int(width / 32) * 32
|
| 243 |
+
height = int(height / 32) * 32
|
| 244 |
+
return width, height
|
| 245 |
+
|
| 246 |
+
def get_rope_index_fix_point(
|
| 247 |
+
spatial_merge_size,
|
| 248 |
+
image_token_id,
|
| 249 |
+
video_token_id,
|
| 250 |
+
vision_start_token_id,
|
| 251 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 252 |
+
image_grid_thw: Optional[torch.LongTensor] = None,
|
| 253 |
+
video_grid_thw: Optional[torch.LongTensor] = None,
|
| 254 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 255 |
+
skip_vision_start_token=None,
|
| 256 |
+
fix_point=4096,
|
| 257 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 258 |
+
if video_grid_thw is not None:
|
| 259 |
+
video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
|
| 260 |
+
video_grid_thw[:, 0] = 1
|
| 261 |
+
|
| 262 |
+
mrope_position_deltas = []
|
| 263 |
+
if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
|
| 264 |
+
total_input_ids = input_ids
|
| 265 |
+
if attention_mask is None:
|
| 266 |
+
attention_mask = torch.ones_like(total_input_ids)
|
| 267 |
+
position_ids = torch.ones(
|
| 268 |
+
3,
|
| 269 |
+
input_ids.shape[0],
|
| 270 |
+
input_ids.shape[1],
|
| 271 |
+
dtype=input_ids.dtype,
|
| 272 |
+
device=input_ids.device,
|
| 273 |
+
)
|
| 274 |
+
image_index, video_index = 0, 0
|
| 275 |
+
attention_mask = attention_mask.to(total_input_ids.device)
|
| 276 |
+
for i, input_ids in enumerate(total_input_ids):
|
| 277 |
+
input_ids = input_ids[attention_mask[i] == 1]
|
| 278 |
+
image_nums, video_nums = 0, 0
|
| 279 |
+
vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
|
| 280 |
+
vision_tokens = input_ids[vision_start_indices + 1]
|
| 281 |
+
image_nums = (vision_tokens == image_token_id).sum()
|
| 282 |
+
video_nums = (vision_tokens == video_token_id).sum()
|
| 283 |
+
input_tokens = input_ids.tolist()
|
| 284 |
+
llm_pos_ids_list: list = []
|
| 285 |
+
st = 0
|
| 286 |
+
remain_images, remain_videos = image_nums, video_nums
|
| 287 |
+
for _ in range(image_nums + video_nums):
|
| 288 |
+
if image_token_id in input_tokens and remain_images > 0:
|
| 289 |
+
ed_image = input_tokens.index(image_token_id, st)
|
| 290 |
+
else:
|
| 291 |
+
ed_image = len(input_tokens) + 1
|
| 292 |
+
if video_token_id in input_tokens and remain_videos > 0:
|
| 293 |
+
ed_video = input_tokens.index(video_token_id, st)
|
| 294 |
+
else:
|
| 295 |
+
ed_video = len(input_tokens) + 1
|
| 296 |
+
if ed_image < ed_video:
|
| 297 |
+
t, h, w = (
|
| 298 |
+
image_grid_thw[image_index][0],
|
| 299 |
+
image_grid_thw[image_index][1],
|
| 300 |
+
image_grid_thw[image_index][2],
|
| 301 |
+
)
|
| 302 |
+
image_index += 1
|
| 303 |
+
remain_images -= 1
|
| 304 |
+
ed = ed_image
|
| 305 |
+
else:
|
| 306 |
+
t, h, w = (
|
| 307 |
+
video_grid_thw[video_index][0],
|
| 308 |
+
video_grid_thw[video_index][1],
|
| 309 |
+
video_grid_thw[video_index][2],
|
| 310 |
+
)
|
| 311 |
+
video_index += 1
|
| 312 |
+
remain_videos -= 1
|
| 313 |
+
ed = ed_video
|
| 314 |
+
llm_grid_t, llm_grid_h, llm_grid_w = (
|
| 315 |
+
t.item(),
|
| 316 |
+
h.item() // spatial_merge_size,
|
| 317 |
+
w.item() // spatial_merge_size,
|
| 318 |
+
)
|
| 319 |
+
text_len = ed - st
|
| 320 |
+
|
| 321 |
+
text_len -= skip_vision_start_token[image_index - 1]
|
| 322 |
+
text_len = max(0, text_len)
|
| 323 |
+
|
| 324 |
+
st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
|
| 325 |
+
llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
|
| 326 |
+
|
| 327 |
+
t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
|
| 328 |
+
h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
|
| 329 |
+
w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
|
| 330 |
+
|
| 331 |
+
if skip_vision_start_token[image_index - 1]:
|
| 332 |
+
if fix_point > 0:
|
| 333 |
+
fix_point = fix_point - st_idx
|
| 334 |
+
llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + fix_point + st_idx)
|
| 335 |
+
fix_point = 0
|
| 336 |
+
else:
|
| 337 |
+
llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
|
| 338 |
+
st = ed + llm_grid_t * llm_grid_h * llm_grid_w
|
| 339 |
+
|
| 340 |
+
if st < len(input_tokens):
|
| 341 |
+
st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
|
| 342 |
+
text_len = len(input_tokens) - st
|
| 343 |
+
llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
|
| 344 |
+
|
| 345 |
+
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
|
| 346 |
+
position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
|
| 347 |
+
mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
|
| 348 |
+
mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
|
| 349 |
+
return position_ids, mrope_position_deltas
|
| 350 |
+
else:
|
| 351 |
+
if attention_mask is not None:
|
| 352 |
+
position_ids = attention_mask.long().cumsum(-1) - 1
|
| 353 |
+
position_ids.masked_fill_(attention_mask == 0, 1)
|
| 354 |
+
position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
|
| 355 |
+
max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
|
| 356 |
+
mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
|
| 357 |
+
else:
|
| 358 |
+
position_ids = (
|
| 359 |
+
torch.arange(input_ids.shape[1], device=input_ids.device)
|
| 360 |
+
.view(1, 1, -1)
|
| 361 |
+
.expand(3, input_ids.shape[0], -1)
|
| 362 |
+
)
|
| 363 |
+
mrope_position_deltas = torch.zeros(
|
| 364 |
+
[input_ids.shape[0], 1],
|
| 365 |
+
device=input_ids.device,
|
| 366 |
+
dtype=input_ids.dtype,
|
| 367 |
+
)
|
| 368 |
+
return position_ids, mrope_position_deltas
|
preprocessor_config.json
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"size": {
|
| 3 |
+
"longest_edge": 16777216,
|
| 4 |
+
"shortest_edge": 65536
|
| 5 |
+
},
|
| 6 |
+
"patch_size": 16,
|
| 7 |
+
"temporal_patch_size": 2,
|
| 8 |
+
"merge_size": 2,
|
| 9 |
+
"image_mean": [
|
| 10 |
+
0.5,
|
| 11 |
+
0.5,
|
| 12 |
+
0.5
|
| 13 |
+
],
|
| 14 |
+
"image_std": [
|
| 15 |
+
0.5,
|
| 16 |
+
0.5,
|
| 17 |
+
0.5
|
| 18 |
+
],
|
| 19 |
+
"processor_class": "Qwen3VLProcessor",
|
| 20 |
+
"image_processor_type": "Qwen2VLImageProcessorFast"
|
| 21 |
+
}
|
requirements.txt
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
transformers
|
| 2 |
+
accelerate
|
| 3 |
+
torch
|
| 4 |
+
safetensors
|
| 5 |
+
pillow
|
tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
tokenizer_config.json
ADDED
|
@@ -0,0 +1,244 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"add_bos_token": false,
|
| 3 |
+
"add_prefix_space": false,
|
| 4 |
+
"added_tokens_decoder": {
|
| 5 |
+
"151643": {
|
| 6 |
+
"content": "<|endoftext|>",
|
| 7 |
+
"lstrip": false,
|
| 8 |
+
"normalized": false,
|
| 9 |
+
"rstrip": false,
|
| 10 |
+
"single_word": false,
|
| 11 |
+
"special": true
|
| 12 |
+
},
|
| 13 |
+
"151644": {
|
| 14 |
+
"content": "<|im_start|>",
|
| 15 |
+
"lstrip": false,
|
| 16 |
+
"normalized": false,
|
| 17 |
+
"rstrip": false,
|
| 18 |
+
"single_word": false,
|
| 19 |
+
"special": true
|
| 20 |
+
},
|
| 21 |
+
"151645": {
|
| 22 |
+
"content": "<|im_end|>",
|
| 23 |
+
"lstrip": false,
|
| 24 |
+
"normalized": false,
|
| 25 |
+
"rstrip": false,
|
| 26 |
+
"single_word": false,
|
| 27 |
+
"special": true
|
| 28 |
+
},
|
| 29 |
+
"151646": {
|
| 30 |
+
"content": "<|object_ref_start|>",
|
| 31 |
+
"lstrip": false,
|
| 32 |
+
"normalized": false,
|
| 33 |
+
"rstrip": false,
|
| 34 |
+
"single_word": false,
|
| 35 |
+
"special": true
|
| 36 |
+
},
|
| 37 |
+
"151647": {
|
| 38 |
+
"content": "<|object_ref_end|>",
|
| 39 |
+
"lstrip": false,
|
| 40 |
+
"normalized": false,
|
| 41 |
+
"rstrip": false,
|
| 42 |
+
"single_word": false,
|
| 43 |
+
"special": true
|
| 44 |
+
},
|
| 45 |
+
"151648": {
|
| 46 |
+
"content": "<|box_start|>",
|
| 47 |
+
"lstrip": false,
|
| 48 |
+
"normalized": false,
|
| 49 |
+
"rstrip": false,
|
| 50 |
+
"single_word": false,
|
| 51 |
+
"special": true
|
| 52 |
+
},
|
| 53 |
+
"151649": {
|
| 54 |
+
"content": "<|box_end|>",
|
| 55 |
+
"lstrip": false,
|
| 56 |
+
"normalized": false,
|
| 57 |
+
"rstrip": false,
|
| 58 |
+
"single_word": false,
|
| 59 |
+
"special": true
|
| 60 |
+
},
|
| 61 |
+
"151650": {
|
| 62 |
+
"content": "<|quad_start|>",
|
| 63 |
+
"lstrip": false,
|
| 64 |
+
"normalized": false,
|
| 65 |
+
"rstrip": false,
|
| 66 |
+
"single_word": false,
|
| 67 |
+
"special": true
|
| 68 |
+
},
|
| 69 |
+
"151651": {
|
| 70 |
+
"content": "<|quad_end|>",
|
| 71 |
+
"lstrip": false,
|
| 72 |
+
"normalized": false,
|
| 73 |
+
"rstrip": false,
|
| 74 |
+
"single_word": false,
|
| 75 |
+
"special": true
|
| 76 |
+
},
|
| 77 |
+
"151652": {
|
| 78 |
+
"content": "<|vision_start|>",
|
| 79 |
+
"lstrip": false,
|
| 80 |
+
"normalized": false,
|
| 81 |
+
"rstrip": false,
|
| 82 |
+
"single_word": false,
|
| 83 |
+
"special": true
|
| 84 |
+
},
|
| 85 |
+
"151653": {
|
| 86 |
+
"content": "<|vision_end|>",
|
| 87 |
+
"lstrip": false,
|
| 88 |
+
"normalized": false,
|
| 89 |
+
"rstrip": false,
|
| 90 |
+
"single_word": false,
|
| 91 |
+
"special": true
|
| 92 |
+
},
|
| 93 |
+
"151654": {
|
| 94 |
+
"content": "<|vision_pad|>",
|
| 95 |
+
"lstrip": false,
|
| 96 |
+
"normalized": false,
|
| 97 |
+
"rstrip": false,
|
| 98 |
+
"single_word": false,
|
| 99 |
+
"special": true
|
| 100 |
+
},
|
| 101 |
+
"151655": {
|
| 102 |
+
"content": "<|image_pad|>",
|
| 103 |
+
"lstrip": false,
|
| 104 |
+
"normalized": false,
|
| 105 |
+
"rstrip": false,
|
| 106 |
+
"single_word": false,
|
| 107 |
+
"special": true
|
| 108 |
+
},
|
| 109 |
+
"151656": {
|
| 110 |
+
"content": "<|video_pad|>",
|
| 111 |
+
"lstrip": false,
|
| 112 |
+
"normalized": false,
|
| 113 |
+
"rstrip": false,
|
| 114 |
+
"single_word": false,
|
| 115 |
+
"special": true
|
| 116 |
+
},
|
| 117 |
+
"151657": {
|
| 118 |
+
"content": "<tool_call>",
|
| 119 |
+
"lstrip": false,
|
| 120 |
+
"normalized": false,
|
| 121 |
+
"rstrip": false,
|
| 122 |
+
"single_word": false,
|
| 123 |
+
"special": false
|
| 124 |
+
},
|
| 125 |
+
"151658": {
|
| 126 |
+
"content": "</tool_call>",
|
| 127 |
+
"lstrip": false,
|
| 128 |
+
"normalized": false,
|
| 129 |
+
"rstrip": false,
|
| 130 |
+
"single_word": false,
|
| 131 |
+
"special": false
|
| 132 |
+
},
|
| 133 |
+
"151659": {
|
| 134 |
+
"content": "<|fim_prefix|>",
|
| 135 |
+
"lstrip": false,
|
| 136 |
+
"normalized": false,
|
| 137 |
+
"rstrip": false,
|
| 138 |
+
"single_word": false,
|
| 139 |
+
"special": false
|
| 140 |
+
},
|
| 141 |
+
"151660": {
|
| 142 |
+
"content": "<|fim_middle|>",
|
| 143 |
+
"lstrip": false,
|
| 144 |
+
"normalized": false,
|
| 145 |
+
"rstrip": false,
|
| 146 |
+
"single_word": false,
|
| 147 |
+
"special": false
|
| 148 |
+
},
|
| 149 |
+
"151661": {
|
| 150 |
+
"content": "<|fim_suffix|>",
|
| 151 |
+
"lstrip": false,
|
| 152 |
+
"normalized": false,
|
| 153 |
+
"rstrip": false,
|
| 154 |
+
"single_word": false,
|
| 155 |
+
"special": false
|
| 156 |
+
},
|
| 157 |
+
"151662": {
|
| 158 |
+
"content": "<|fim_pad|>",
|
| 159 |
+
"lstrip": false,
|
| 160 |
+
"normalized": false,
|
| 161 |
+
"rstrip": false,
|
| 162 |
+
"single_word": false,
|
| 163 |
+
"special": false
|
| 164 |
+
},
|
| 165 |
+
"151663": {
|
| 166 |
+
"content": "<|repo_name|>",
|
| 167 |
+
"lstrip": false,
|
| 168 |
+
"normalized": false,
|
| 169 |
+
"rstrip": false,
|
| 170 |
+
"single_word": false,
|
| 171 |
+
"special": false
|
| 172 |
+
},
|
| 173 |
+
"151664": {
|
| 174 |
+
"content": "<|file_sep|>",
|
| 175 |
+
"lstrip": false,
|
| 176 |
+
"normalized": false,
|
| 177 |
+
"rstrip": false,
|
| 178 |
+
"single_word": false,
|
| 179 |
+
"special": false
|
| 180 |
+
},
|
| 181 |
+
"151665": {
|
| 182 |
+
"content": "<tool_response>",
|
| 183 |
+
"lstrip": false,
|
| 184 |
+
"normalized": false,
|
| 185 |
+
"rstrip": false,
|
| 186 |
+
"single_word": false,
|
| 187 |
+
"special": false
|
| 188 |
+
},
|
| 189 |
+
"151666": {
|
| 190 |
+
"content": "</tool_response>",
|
| 191 |
+
"lstrip": false,
|
| 192 |
+
"normalized": false,
|
| 193 |
+
"rstrip": false,
|
| 194 |
+
"single_word": false,
|
| 195 |
+
"special": false
|
| 196 |
+
},
|
| 197 |
+
"151667": {
|
| 198 |
+
"content": "<think>",
|
| 199 |
+
"lstrip": false,
|
| 200 |
+
"normalized": false,
|
| 201 |
+
"rstrip": false,
|
| 202 |
+
"single_word": false,
|
| 203 |
+
"special": false
|
| 204 |
+
},
|
| 205 |
+
"151668": {
|
| 206 |
+
"content": "</think>",
|
| 207 |
+
"lstrip": false,
|
| 208 |
+
"normalized": false,
|
| 209 |
+
"rstrip": false,
|
| 210 |
+
"single_word": false,
|
| 211 |
+
"special": false
|
| 212 |
+
}
|
| 213 |
+
},
|
| 214 |
+
"additional_special_tokens": [
|
| 215 |
+
"<|im_start|>",
|
| 216 |
+
"<|im_end|>",
|
| 217 |
+
"<|object_ref_start|>",
|
| 218 |
+
"<|object_ref_end|>",
|
| 219 |
+
"<|box_start|>",
|
| 220 |
+
"<|box_end|>",
|
| 221 |
+
"<|quad_start|>",
|
| 222 |
+
"<|quad_end|>",
|
| 223 |
+
"<|vision_start|>",
|
| 224 |
+
"<|vision_end|>",
|
| 225 |
+
"<|vision_pad|>",
|
| 226 |
+
"<|image_pad|>",
|
| 227 |
+
"<|video_pad|>",
|
| 228 |
+
"<|boi_token|>",
|
| 229 |
+
"<|bor_token|>",
|
| 230 |
+
"<|eor_token|>",
|
| 231 |
+
"<|bot_token|>",
|
| 232 |
+
"<|tms_token|>"
|
| 233 |
+
],
|
| 234 |
+
"bos_token": null,
|
| 235 |
+
"chat_template": "{%- if tools %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0].role == 'system' %}\n {%- if messages[0].content is string %}\n {{- messages[0].content }}\n {%- else %}\n {%- for content in messages[0].content %}\n {%- if 'text' in content %}\n {{- content.text }}\n {%- endif %}\n {%- endfor %}\n {%- endif %}\n {{- '\\n\\n' }}\n {%- endif %}\n {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0].role == 'system' %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0].content is string %}\n {{- messages[0].content }}\n {%- else %}\n {%- for content in messages[0].content %}\n {%- if 'text' in content %}\n {{- content.text }}\n {%- endif %}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- for message in messages %}\n {%- if message.role == \"user\" %}\n {{- '<|im_start|>' + message.role + '\\n' }}\n {%- if message.content is string %}\n {{- message.content }}\n {%- else %}\n {%- for content in message.content %}\n {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n {%- set image_count.value = image_count.value + 1 %}\n {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n <|vision_start|><|image_pad|><|vision_end|>\n {%- elif content.type == 'video' or 'video' in content %}\n {%- set video_count.value = video_count.value + 1 %}\n {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n <|vision_start|><|video_pad|><|vision_end|>\n {%- elif 'text' in content %}\n {{- content.text }}\n {%- endif %}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"assistant\" %}\n {{- '<|im_start|>' + message.role + '\\n' }}\n {%- if message.content is string %}\n {{- message.content }}\n {%- else %}\n {%- for content_item in message.content %}\n {%- if 'text' in content_item %}\n {{- content_item.text }}\n {%- endif %}\n {%- endfor %}\n {%- endif %}\n {%- if message.tool_calls %}\n {%- for tool_call in message.tool_calls %}\n {%- if (loop.first and message.content) or (not loop.first) %}\n {{- '\\n' }}\n {%- endif %}\n {%- if tool_call.function %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {{- '<tool_call>\\n{\"name\": \"' }}\n {{- tool_call.name }}\n {{- '\", \"arguments\": ' }}\n {%- if tool_call.arguments is string %}\n {{- tool_call.arguments }}\n {%- else %}\n {{- tool_call.arguments | tojson }}\n {%- endif %}\n {{- '}\\n</tool_call>' }}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n {{- '<|im_start|>user' }}\n {%- endif %}\n {{- '\\n<tool_response>\\n' }}\n {%- if message.content is string %}\n {{- message.content }}\n {%- else %}\n {%- for content in message.content %}\n {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}\n {%- set image_count.value = image_count.value + 1 %}\n {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}\n <|vision_start|><|image_pad|><|vision_end|>\n {%- elif content.type == 'video' or 'video' in content %}\n {%- set video_count.value = video_count.value + 1 %}\n {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}\n <|vision_start|><|video_pad|><|vision_end|>\n {%- elif 'text' in content %}\n {{- content.text }}\n {%- endif %}\n {%- endfor %}\n {%- endif %}\n {{- '\\n</tool_response>' }}\n {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
|
| 236 |
+
"clean_up_tokenization_spaces": false,
|
| 237 |
+
"eos_token": "<|im_end|>",
|
| 238 |
+
"errors": "replace",
|
| 239 |
+
"model_max_length": 262144,
|
| 240 |
+
"pad_token": "<|endoftext|>",
|
| 241 |
+
"split_special_tokens": false,
|
| 242 |
+
"tokenizer_class": "Qwen2Tokenizer",
|
| 243 |
+
"unk_token": null
|
| 244 |
+
}
|
video_preprocessor_config.json
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"size": {
|
| 3 |
+
"longest_edge": 25165824,
|
| 4 |
+
"shortest_edge": 4096
|
| 5 |
+
},
|
| 6 |
+
"patch_size": 16,
|
| 7 |
+
"temporal_patch_size": 2,
|
| 8 |
+
"merge_size": 2,
|
| 9 |
+
"image_mean": [
|
| 10 |
+
0.5,
|
| 11 |
+
0.5,
|
| 12 |
+
0.5
|
| 13 |
+
],
|
| 14 |
+
"image_std": [
|
| 15 |
+
0.5,
|
| 16 |
+
0.5,
|
| 17 |
+
0.5
|
| 18 |
+
],
|
| 19 |
+
"processor_class": "Qwen3VLProcessor",
|
| 20 |
+
"video_processor_type": "Qwen3VLVideoProcessor"
|
| 21 |
+
}
|
vocab.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|