worksimpli commited on
Commit
dfdba22
·
verified ·
1 Parent(s): 4aad089

Add files using upload-large-folder tool

Browse files
.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&amp;logo=discord&amp;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

  • SHA256: 9dc8f61b5e7cb68c2c82a82f930234adec0d9214f45d827f67cd073d4665d3e2
  • Pointer size: 132 Bytes
  • Size of remote file: 2.58 MB
assets/IP_2.jpg ADDED

Git LFS Details

  • SHA256: d62ab82d69046ecd423151096628a58e961bc44a264040aa1aa893fb5aab114a
  • Pointer size: 132 Bytes
  • Size of remote file: 1.27 MB
assets/general.webp ADDED

Git LFS Details

  • SHA256: 32a93bb7f0aa2792216bb4353176ee2e3fd574771448fcbae097a4f82db41fdd
  • Pointer size: 132 Bytes
  • Size of remote file: 5.57 MB
assets/leaderboard.png ADDED

Git LFS Details

  • SHA256: 4a5c034b0956cbbbec1f254f66bc7a71cdc9b37b1514042964cd5df4e5d14b44
  • Pointer size: 131 Bytes
  • Size of remote file: 773 kB
assets/text-layout.webp ADDED

Git LFS Details

  • SHA256: f12b57ab352811f259134a75aec3befb38ec956684474184d5f3b1aac2959184
  • Pointer size: 132 Bytes
  • Size of remote file: 5.54 MB
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