Any-to-Any
Transformers
Safetensors
kimi_k3
feature-extraction
kimi-k3
compressed-tensors
conversational
abliterated
uncensored
custom_code
8-bit precision
Instructions to use SHSLab/Kimi-K3-Abliterated with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use SHSLab/Kimi-K3-Abliterated with Transformers:
# Load model directly from transformers import AutoModel model = AutoModel.from_pretrained("SHSLab/Kimi-K3-Abliterated", trust_remote_code=True, device_map="auto") - Notebooks
- Google Colab
- Kaggle
Commit ·
bb450fc
0
Parent(s):
Duplicate from kepom/Kimi-K3-Abliterated-V1
Browse filesCo-authored-by: oieok <kepom@users.noreply.huggingface.co>
This view is limited to 50 files because it contains too many changes. See raw diff
- .gitattributes +37 -0
- LICENSE +52 -0
- README.md +28 -0
- config.json +301 -0
- configuration_kimi_k3.py +285 -0
- encoding_k3.py +647 -0
- generation_config.json +4 -0
- kimi_k3_processor.py +187 -0
- kimi_k3_vision_processing.py +179 -0
- media_utils.py +376 -0
- model-00001-of-000096.safetensors +3 -0
- model-00002-of-000096.safetensors +3 -0
- model-00003-of-000096.safetensors +3 -0
- model-00004-of-000096.safetensors +3 -0
- model-00005-of-000096.safetensors +3 -0
- model-00006-of-000096.safetensors +3 -0
- model-00007-of-000096.safetensors +3 -0
- model-00008-of-000096.safetensors +3 -0
- model-00009-of-000096.safetensors +3 -0
- model-00010-of-000096.safetensors +3 -0
- model-00011-of-000096.safetensors +3 -0
- model-00012-of-000096.safetensors +3 -0
- model-00013-of-000096.safetensors +3 -0
- model-00014-of-000096.safetensors +3 -0
- model-00015-of-000096.safetensors +3 -0
- model-00016-of-000096.safetensors +3 -0
- model-00017-of-000096.safetensors +3 -0
- model-00018-of-000096.safetensors +3 -0
- model-00019-of-000096.safetensors +3 -0
- model-00020-of-000096.safetensors +3 -0
- model-00021-of-000096.safetensors +3 -0
- model-00022-of-000096.safetensors +3 -0
- model-00023-of-000096.safetensors +3 -0
- model-00024-of-000096.safetensors +3 -0
- model-00025-of-000096.safetensors +3 -0
- model-00026-of-000096.safetensors +3 -0
- model-00027-of-000096.safetensors +3 -0
- model-00028-of-000096.safetensors +3 -0
- model-00029-of-000096.safetensors +3 -0
- model-00030-of-000096.safetensors +3 -0
- model-00031-of-000096.safetensors +3 -0
- model-00032-of-000096.safetensors +3 -0
- model-00033-of-000096.safetensors +3 -0
- model-00034-of-000096.safetensors +3 -0
- model-00035-of-000096.safetensors +3 -0
- model-00036-of-000096.safetensors +3 -0
- model-00037-of-000096.safetensors +3 -0
- model-00038-of-000096.safetensors +3 -0
- model-00039-of-000096.safetensors +3 -0
- model-00040-of-000096.safetensors +3 -0
.gitattributes
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
+
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
+
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
+
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
+
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
+
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
+
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
+
*.xz 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 |
+
model.safetensors.index.json filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
assets/k3-hero.jpg filter=lfs diff=lfs merge=lfs -text
|
LICENSE
ADDED
|
@@ -0,0 +1,52 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Kimi K3 License
|
| 2 |
+
|
| 3 |
+
Copyright (c) 2026 Moonshot AI
|
| 4 |
+
|
| 5 |
+
Permission is hereby granted, free of charge, to any person (the "Licensee")
|
| 6 |
+
obtaining a copy of this software — including the model weights, parameters,
|
| 7 |
+
configuration files, inference and training code, and associated documentation
|
| 8 |
+
(collectively, the "Software") — to deal in the Software without restriction.
|
| 9 |
+
This includes, without limitation, the rights to use, copy, modify, merge,
|
| 10 |
+
publish, distribute, sublicense, and/or sell copies of the Software; to run,
|
| 11 |
+
deploy, fine-tune, or otherwise modify the Software and create derivative works
|
| 12 |
+
from it; and to permit persons to whom the Software is furnished to do so, in
|
| 13 |
+
each case subject to the following conditions:
|
| 14 |
+
|
| 15 |
+
1. The above copyright notice and this permission notice shall be included in
|
| 16 |
+
all copies or substantial portions of the Software. Licensee's use of the
|
| 17 |
+
Software must comply with applicable laws and regulations.
|
| 18 |
+
|
| 19 |
+
2. "Model as a Service" means giving a third party access to language model
|
| 20 |
+
inference or fine-tuning (e.g., via API) in a manner that allows such third
|
| 21 |
+
party to exercise meaningful control over the inputs, parameters, or training
|
| 22 |
+
data. This does not include (a) end-user products with model capabilities solely
|
| 23 |
+
embedded within specific features or harnesses, or (b) mere relaying of requests
|
| 24 |
+
to models hosted by others.
|
| 25 |
+
|
| 26 |
+
If the Licensee or any of its affiliates operates a Model as a Service business,
|
| 27 |
+
and the aggregate revenue of the Licensee and its affiliates exceeds 20 million
|
| 28 |
+
US dollars (or the equivalent in other currencies) in total over any consecutive
|
| 29 |
+
12 months, the Licensee must enter into a separate agreement with Moonshot AI
|
| 30 |
+
before using the Software or its derivative works for any commercial purpose.
|
| 31 |
+
|
| 32 |
+
3. If the Software (or any derivative works thereof) is used for any of the
|
| 33 |
+
Licensee's commercial products or services that have more than 100 million
|
| 34 |
+
monthly active users, or more than 20 million US dollars (or equivalent in other
|
| 35 |
+
currencies) in monthly revenue, "Kimi K3" must be prominently displayed on the
|
| 36 |
+
user interface of such product or service.
|
| 37 |
+
|
| 38 |
+
4. The requirements set forth in Sections 2 and 3 do not apply to: (a) internal
|
| 39 |
+
use of the Software, defined as any use that does not make the Software, its
|
| 40 |
+
outputs, or its underlying capabilities available to third parties; or (b) any
|
| 41 |
+
use of the Software accessed through Moonshot AI's official products or
|
| 42 |
+
certified inference partners.
|
| 43 |
+
|
| 44 |
+
5. THE SOFTWARE AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS”
|
| 45 |
+
BASIS, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
|
| 46 |
+
LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE
|
| 47 |
+
AND NONINFRINGEMENT. IN NO EVENT SHALL MOONSHOT AI OR ITS AFFILIATES OR
|
| 48 |
+
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
| 49 |
+
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
|
| 50 |
+
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
| 51 |
+
|
| 52 |
+
For any questions regarding this license, please contact <license@moonshot.ai>.
|
README.md
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
base_model:
|
| 3 |
+
- moonshotai/Kimi-K3
|
| 4 |
+
base_model_relation: finetune
|
| 5 |
+
tags:
|
| 6 |
+
- compressed-tensors
|
| 7 |
+
- conversational
|
| 8 |
+
- abliterated
|
| 9 |
+
- uncensored
|
| 10 |
+
license: other
|
| 11 |
+
license_name: "kimi-k3"
|
| 12 |
+
library_name: transformers
|
| 13 |
+
pipeline_tag: image-text-to-text
|
| 14 |
+
|
| 15 |
+
---
|
| 16 |
+
# Uniboshi/Kimi-K3-Abliterated-V1
|
| 17 |
+
|
| 18 |
+
- This model is an abliterated version of [moonshotai/Kimi-K3](<https://huggingface.co/moonshotai/Kimi-K3>) for uncensored.
|
| 19 |
+
- This model has been tuned to make refusals less likely to be triggered for English and Japanese.
|
| 20 |
+
- In this model, more than 98% of the safeguards have been removed.
|
| 21 |
+
```
|
| 22 |
+
writer written as n before after removed
|
| 23 |
+
*.2 [0] as stored 1 0.980..0.980 1.53e-02..1.53e-02 98.44%
|
| 24 |
+
*.down_proj [0] as stored 93 0.932..1.421 1.35e-02..1.56e-02 98.62%
|
| 25 |
+
*.embed_tokens [1] as stored 1 1.040..1.040 1.46e-02..1.46e-02 98.60%
|
| 26 |
+
*.o_proj [0] as stored 93 0.938..1.358 1.38e-02..1.60e-02 98.54%
|
| 27 |
+
*.routed_expert_up_proj [0] as stored 92 0.910..1.402 1.24e-02..1.56e-02 98.57%
|
| 28 |
+
```
|
config.json
ADDED
|
@@ -0,0 +1,301 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"KimiK3ForConditionalGeneration"
|
| 4 |
+
],
|
| 5 |
+
"auto_map": {
|
| 6 |
+
"AutoConfig": "configuration_kimi_k3.KimiK3Config",
|
| 7 |
+
"AutoModel": "modeling_kimi_k3.KimiK3ForConditionalGeneration",
|
| 8 |
+
"AutoModelForCausalLM": "modeling_kimi_k3.KimiK3ForConditionalGeneration"
|
| 9 |
+
},
|
| 10 |
+
"bos_token_id": 163584,
|
| 11 |
+
"dtype": "bfloat16",
|
| 12 |
+
"eos_token_id": 163586,
|
| 13 |
+
"ignore_index": -100,
|
| 14 |
+
"image_placeholder": "<|kimi_image_placeholder|>",
|
| 15 |
+
"media_placeholder_token_id": 163605,
|
| 16 |
+
"model_type": "kimi_k3",
|
| 17 |
+
"pad_token_id": 163839,
|
| 18 |
+
"text_config": {
|
| 19 |
+
"_name_or_path": "",
|
| 20 |
+
"activation_situ_beta": 4.0,
|
| 21 |
+
"activation_situ_linear_beta": 25.0,
|
| 22 |
+
"add_cross_attention": false,
|
| 23 |
+
"architectures": [
|
| 24 |
+
"KimiLinearForCausalLM"
|
| 25 |
+
],
|
| 26 |
+
"attn_res_block_size": 12,
|
| 27 |
+
"auto_map": {
|
| 28 |
+
"AutoConfig": "configuration_kimi_k3.KimiLinearConfig",
|
| 29 |
+
"AutoModel": "modeling_kimi_linear.KimiLinearModel",
|
| 30 |
+
"AutoModelForCausalLM": "modeling_kimi_linear.KimiLinearForCausalLM"
|
| 31 |
+
},
|
| 32 |
+
"bad_words_ids": null,
|
| 33 |
+
"begin_suppress_tokens": null,
|
| 34 |
+
"bos_token_id": 163584,
|
| 35 |
+
"chunk_size_feed_forward": 0,
|
| 36 |
+
"cross_attention_hidden_size": null,
|
| 37 |
+
"decoder_start_token_id": null,
|
| 38 |
+
"diversity_penalty": 0.0,
|
| 39 |
+
"do_sample": false,
|
| 40 |
+
"dtype": "bfloat16",
|
| 41 |
+
"early_stopping": false,
|
| 42 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 43 |
+
"eos_token_id": 163586,
|
| 44 |
+
"exponential_decay_length_penalty": null,
|
| 45 |
+
"finetuning_task": null,
|
| 46 |
+
"first_k_dense_replace": 1,
|
| 47 |
+
"forced_bos_token_id": null,
|
| 48 |
+
"forced_eos_token_id": null,
|
| 49 |
+
"hidden_act": "situ",
|
| 50 |
+
"hidden_size": 7168,
|
| 51 |
+
"id2label": {
|
| 52 |
+
"0": "LABEL_0",
|
| 53 |
+
"1": "LABEL_1"
|
| 54 |
+
},
|
| 55 |
+
"initializer_range": 0.02,
|
| 56 |
+
"intermediate_size": 33792,
|
| 57 |
+
"is_decoder": false,
|
| 58 |
+
"is_encoder_decoder": false,
|
| 59 |
+
"kv_lora_rank": 512,
|
| 60 |
+
"label2id": {
|
| 61 |
+
"LABEL_0": 0,
|
| 62 |
+
"LABEL_1": 1
|
| 63 |
+
},
|
| 64 |
+
"latent_moe_use_norm": true,
|
| 65 |
+
"length_penalty": 1.0,
|
| 66 |
+
"linear_attn_config": {
|
| 67 |
+
"full_attn_layers": [
|
| 68 |
+
4,
|
| 69 |
+
8,
|
| 70 |
+
12,
|
| 71 |
+
16,
|
| 72 |
+
20,
|
| 73 |
+
24,
|
| 74 |
+
28,
|
| 75 |
+
32,
|
| 76 |
+
36,
|
| 77 |
+
40,
|
| 78 |
+
44,
|
| 79 |
+
48,
|
| 80 |
+
52,
|
| 81 |
+
56,
|
| 82 |
+
60,
|
| 83 |
+
64,
|
| 84 |
+
68,
|
| 85 |
+
72,
|
| 86 |
+
76,
|
| 87 |
+
80,
|
| 88 |
+
84,
|
| 89 |
+
88,
|
| 90 |
+
92,
|
| 91 |
+
93
|
| 92 |
+
],
|
| 93 |
+
"gate_lower_bound": -5.0,
|
| 94 |
+
"head_dim": 128,
|
| 95 |
+
"kda_layers": [
|
| 96 |
+
1,
|
| 97 |
+
2,
|
| 98 |
+
3,
|
| 99 |
+
5,
|
| 100 |
+
6,
|
| 101 |
+
7,
|
| 102 |
+
9,
|
| 103 |
+
10,
|
| 104 |
+
11,
|
| 105 |
+
13,
|
| 106 |
+
14,
|
| 107 |
+
15,
|
| 108 |
+
17,
|
| 109 |
+
18,
|
| 110 |
+
19,
|
| 111 |
+
21,
|
| 112 |
+
22,
|
| 113 |
+
23,
|
| 114 |
+
25,
|
| 115 |
+
26,
|
| 116 |
+
27,
|
| 117 |
+
29,
|
| 118 |
+
30,
|
| 119 |
+
31,
|
| 120 |
+
33,
|
| 121 |
+
34,
|
| 122 |
+
35,
|
| 123 |
+
37,
|
| 124 |
+
38,
|
| 125 |
+
39,
|
| 126 |
+
41,
|
| 127 |
+
42,
|
| 128 |
+
43,
|
| 129 |
+
45,
|
| 130 |
+
46,
|
| 131 |
+
47,
|
| 132 |
+
49,
|
| 133 |
+
50,
|
| 134 |
+
51,
|
| 135 |
+
53,
|
| 136 |
+
54,
|
| 137 |
+
55,
|
| 138 |
+
57,
|
| 139 |
+
58,
|
| 140 |
+
59,
|
| 141 |
+
61,
|
| 142 |
+
62,
|
| 143 |
+
63,
|
| 144 |
+
65,
|
| 145 |
+
66,
|
| 146 |
+
67,
|
| 147 |
+
69,
|
| 148 |
+
70,
|
| 149 |
+
71,
|
| 150 |
+
73,
|
| 151 |
+
74,
|
| 152 |
+
75,
|
| 153 |
+
77,
|
| 154 |
+
78,
|
| 155 |
+
79,
|
| 156 |
+
81,
|
| 157 |
+
82,
|
| 158 |
+
83,
|
| 159 |
+
85,
|
| 160 |
+
86,
|
| 161 |
+
87,
|
| 162 |
+
89,
|
| 163 |
+
90,
|
| 164 |
+
91
|
| 165 |
+
],
|
| 166 |
+
"num_heads": 96,
|
| 167 |
+
"short_conv_kernel_size": 4,
|
| 168 |
+
"use_full_rank_gate": true
|
| 169 |
+
},
|
| 170 |
+
"max_length": 20,
|
| 171 |
+
"max_position_embeddings": 1048576,
|
| 172 |
+
"min_length": 0,
|
| 173 |
+
"mla_use_nope": true,
|
| 174 |
+
"mla_use_output_gate": true,
|
| 175 |
+
"model_type": "kimi_linear",
|
| 176 |
+
"moe_intermediate_size": 3072,
|
| 177 |
+
"moe_layer_freq": 1,
|
| 178 |
+
"moe_renormalize": true,
|
| 179 |
+
"moe_router_activation_func": "sigmoid",
|
| 180 |
+
"no_repeat_ngram_size": 0,
|
| 181 |
+
"num_attention_heads": 96,
|
| 182 |
+
"num_beam_groups": 1,
|
| 183 |
+
"num_beams": 1,
|
| 184 |
+
"num_expert_group": 1,
|
| 185 |
+
"num_experts": 896,
|
| 186 |
+
"num_experts_per_token": 16,
|
| 187 |
+
"num_hidden_layers": 93,
|
| 188 |
+
"num_key_value_heads": 96,
|
| 189 |
+
"num_nextn_predict_layers": 0,
|
| 190 |
+
"num_return_sequences": 1,
|
| 191 |
+
"num_shared_experts": 2,
|
| 192 |
+
"output_attentions": false,
|
| 193 |
+
"output_hidden_states": false,
|
| 194 |
+
"output_scores": false,
|
| 195 |
+
"pad_token_id": 163839,
|
| 196 |
+
"prefix": null,
|
| 197 |
+
"problem_type": null,
|
| 198 |
+
"pruned_heads": {},
|
| 199 |
+
"q_lora_rank": 1536,
|
| 200 |
+
"qk_nope_head_dim": 128,
|
| 201 |
+
"qk_rope_head_dim": 64,
|
| 202 |
+
"quantization_config": {
|
| 203 |
+
"config_groups": {
|
| 204 |
+
"group_0": {
|
| 205 |
+
"format": "mxfp4-pack-quantized",
|
| 206 |
+
"input_activations": null,
|
| 207 |
+
"output_activations": null,
|
| 208 |
+
"targets": [
|
| 209 |
+
"Linear"
|
| 210 |
+
],
|
| 211 |
+
"weights": {
|
| 212 |
+
"actorder": null,
|
| 213 |
+
"block_structure": null,
|
| 214 |
+
"dynamic": false,
|
| 215 |
+
"group_size": 32,
|
| 216 |
+
"num_bits": 4,
|
| 217 |
+
"observer": "minmax",
|
| 218 |
+
"observer_kwargs": {},
|
| 219 |
+
"scale_dtype": "torch.uint8",
|
| 220 |
+
"strategy": "group",
|
| 221 |
+
"symmetric": true,
|
| 222 |
+
"type": "float",
|
| 223 |
+
"zp_dtype": null
|
| 224 |
+
}
|
| 225 |
+
}
|
| 226 |
+
},
|
| 227 |
+
"format": "mxfp4-pack-quantized",
|
| 228 |
+
"global_compression_ratio": null,
|
| 229 |
+
"ignore": [
|
| 230 |
+
"re:.*self_attn.*",
|
| 231 |
+
"re:.*shared_experts.*",
|
| 232 |
+
"re:.*mlp\\.(gate|up|gate_up|down)_proj.*",
|
| 233 |
+
"re:.*lm_head.*",
|
| 234 |
+
"re:.*vision_tower.*",
|
| 235 |
+
"re:.*mm_projector.*",
|
| 236 |
+
"re:.*\\.(?:embed_tokens|gate|input_layernorm|mlp_res_norm|mlp_res_proj|norm|output_attn_res_norm|output_attn_res_proj|post_attention_layernorm|routed_expert_down_proj|routed_expert_norm|routed_expert_up_proj|self_attention_res_norm|self_attention_res_proj)$"
|
| 237 |
+
],
|
| 238 |
+
"kv_cache_scheme": null,
|
| 239 |
+
"quant_method": "compressed-tensors",
|
| 240 |
+
"quantization_status": "compressed"
|
| 241 |
+
},
|
| 242 |
+
"remove_invalid_values": false,
|
| 243 |
+
"repetition_penalty": 1.0,
|
| 244 |
+
"return_dict": true,
|
| 245 |
+
"return_dict_in_generate": false,
|
| 246 |
+
"rms_norm_eps": 1e-05,
|
| 247 |
+
"routed_expert_hidden_size": 3584,
|
| 248 |
+
"routed_scaling_factor": 1.0,
|
| 249 |
+
"sep_token_id": null,
|
| 250 |
+
"suppress_tokens": null,
|
| 251 |
+
"task_specific_params": null,
|
| 252 |
+
"temperature": 1.0,
|
| 253 |
+
"tf_legacy_loss": false,
|
| 254 |
+
"tie_encoder_decoder": false,
|
| 255 |
+
"tie_word_embeddings": false,
|
| 256 |
+
"tokenizer_class": null,
|
| 257 |
+
"top_k": 50,
|
| 258 |
+
"top_p": 1.0,
|
| 259 |
+
"topk_group": 1,
|
| 260 |
+
"topk_method": "noaux_tc",
|
| 261 |
+
"torchscript": false,
|
| 262 |
+
"transformers_version": "4.56.2",
|
| 263 |
+
"typical_p": 1.0,
|
| 264 |
+
"use_bfloat16": false,
|
| 265 |
+
"use_cache": true,
|
| 266 |
+
"use_grouped_topk": true,
|
| 267 |
+
"v_head_dim": 128,
|
| 268 |
+
"vocab_size": 163840
|
| 269 |
+
},
|
| 270 |
+
"tie_word_embeddings": false,
|
| 271 |
+
"vision_config": {
|
| 272 |
+
"_attn_implementation": "flash_attention_2",
|
| 273 |
+
"activation_func": "gelu_pytorch_tanh",
|
| 274 |
+
"attn_bias": false,
|
| 275 |
+
"init_pos_emb_height": 64,
|
| 276 |
+
"init_pos_emb_time": 4,
|
| 277 |
+
"init_pos_emb_width": 64,
|
| 278 |
+
"linear_bias": false,
|
| 279 |
+
"merge_kernel_size": [
|
| 280 |
+
2,
|
| 281 |
+
2
|
| 282 |
+
],
|
| 283 |
+
"merge_type": "sd2_tpool",
|
| 284 |
+
"mlp_type": "mlp2",
|
| 285 |
+
"mm_hidden_size": 1024,
|
| 286 |
+
"mm_projector_type": "patchmergerv2",
|
| 287 |
+
"norm_type": "rmsnorm",
|
| 288 |
+
"patch_embed_proj_bias": false,
|
| 289 |
+
"patch_size": 14,
|
| 290 |
+
"pos_emb_interpolation_mode": "bilinear",
|
| 291 |
+
"pos_emb_type": "divided_fixed",
|
| 292 |
+
"projector_hidden_act": "gelu",
|
| 293 |
+
"projector_ln_eps": 1e-05,
|
| 294 |
+
"qkv_hidden_size": 1536,
|
| 295 |
+
"text_hidden_size": 7168,
|
| 296 |
+
"vt_hidden_size": 1024,
|
| 297 |
+
"vt_intermediate_size": 4096,
|
| 298 |
+
"vt_num_attention_heads": 12,
|
| 299 |
+
"vt_num_hidden_layers": 27
|
| 300 |
+
}
|
| 301 |
+
}
|
configuration_kimi_k3.py
ADDED
|
@@ -0,0 +1,285 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from typing import Optional
|
| 2 |
+
|
| 3 |
+
from transformers.configuration_utils import PretrainedConfig
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class KimiLinearConfig(PretrainedConfig):
|
| 7 |
+
model_type = "kimi_linear"
|
| 8 |
+
keys_to_ignore_at_inference = ["past_key_values"]
|
| 9 |
+
|
| 10 |
+
def __init__(
|
| 11 |
+
self,
|
| 12 |
+
model_type="kimi_linear",
|
| 13 |
+
vocab_size=163840,
|
| 14 |
+
hidden_size=4096,
|
| 15 |
+
head_dim=None,
|
| 16 |
+
intermediate_size=11008,
|
| 17 |
+
num_hidden_layers=32,
|
| 18 |
+
num_attention_heads=32,
|
| 19 |
+
num_key_value_heads=None,
|
| 20 |
+
hidden_act="silu",
|
| 21 |
+
initializer_range=0.02,
|
| 22 |
+
rms_norm_eps=1e-6,
|
| 23 |
+
use_cache=True,
|
| 24 |
+
pad_token_id=0,
|
| 25 |
+
bos_token_id=1,
|
| 26 |
+
eos_token_id=2,
|
| 27 |
+
rope_theta=10000.0,
|
| 28 |
+
rope_scaling=None,
|
| 29 |
+
tie_word_embeddings=False,
|
| 30 |
+
moe_intermediate_size: Optional[int] = None,
|
| 31 |
+
moe_renormalize: bool = True,
|
| 32 |
+
moe_router_activation_func: str = "sigmoid",
|
| 33 |
+
num_experts: Optional[int] = None,
|
| 34 |
+
num_experts_per_token: Optional[int] = None,
|
| 35 |
+
num_shared_experts: int = 0,
|
| 36 |
+
routed_scaling_factor: float = 1.0,
|
| 37 |
+
first_k_dense_replace: int = 0,
|
| 38 |
+
moe_layer_freq: int = 1,
|
| 39 |
+
use_grouped_topk: bool = True,
|
| 40 |
+
num_expert_group: int = 1,
|
| 41 |
+
topk_group: int = 1,
|
| 42 |
+
q_lora_rank: Optional[int] = None,
|
| 43 |
+
kv_lora_rank: Optional[int] = None,
|
| 44 |
+
qk_nope_head_dim: Optional[int] = None,
|
| 45 |
+
qk_rope_head_dim: Optional[int] = None,
|
| 46 |
+
v_head_dim: Optional[int] = None,
|
| 47 |
+
mla_use_nope: Optional[bool] = False,
|
| 48 |
+
mla_use_output_gate: Optional[bool] = False,
|
| 49 |
+
num_nextn_predict_layers: int = 0,
|
| 50 |
+
linear_attn_config: Optional[dict] = None,
|
| 51 |
+
attn_res_block_size: Optional[int] = None,
|
| 52 |
+
latent_moe_use_norm: bool = False,
|
| 53 |
+
activation_situ_beta: Optional[float] = None,
|
| 54 |
+
activation_situ_linear_beta: Optional[float] = None,
|
| 55 |
+
max_position_embeddings: int = 4096,
|
| 56 |
+
routed_expert_hidden_size: Optional[int] = None,
|
| 57 |
+
topk_method: str = "noaux_tc",
|
| 58 |
+
**kwargs,
|
| 59 |
+
):
|
| 60 |
+
self.model_type = model_type
|
| 61 |
+
self.vocab_size = vocab_size
|
| 62 |
+
self.hidden_size = hidden_size
|
| 63 |
+
self.head_dim = (
|
| 64 |
+
head_dim if head_dim is not None else hidden_size // num_attention_heads
|
| 65 |
+
)
|
| 66 |
+
self.intermediate_size = intermediate_size
|
| 67 |
+
self.num_hidden_layers = num_hidden_layers
|
| 68 |
+
self.num_attention_heads = num_attention_heads
|
| 69 |
+
|
| 70 |
+
# for backward compatibility
|
| 71 |
+
if num_key_value_heads is None:
|
| 72 |
+
num_key_value_heads = num_attention_heads
|
| 73 |
+
|
| 74 |
+
self.num_key_value_heads = num_key_value_heads
|
| 75 |
+
self.hidden_act = hidden_act
|
| 76 |
+
self.initializer_range = initializer_range
|
| 77 |
+
self.rms_norm_eps = rms_norm_eps
|
| 78 |
+
self.use_cache = use_cache
|
| 79 |
+
self.rope_theta = rope_theta
|
| 80 |
+
self.rope_scaling = rope_scaling
|
| 81 |
+
|
| 82 |
+
self.q_lora_rank = q_lora_rank
|
| 83 |
+
self.kv_lora_rank = kv_lora_rank
|
| 84 |
+
self.qk_nope_head_dim = qk_nope_head_dim
|
| 85 |
+
self.qk_rope_head_dim = qk_rope_head_dim
|
| 86 |
+
self.v_head_dim = v_head_dim
|
| 87 |
+
self.mla_use_nope = mla_use_nope
|
| 88 |
+
self.mla_use_output_gate = mla_use_output_gate
|
| 89 |
+
# moe config
|
| 90 |
+
self.num_experts = num_experts
|
| 91 |
+
self.num_experts_per_token = num_experts_per_token
|
| 92 |
+
self.moe_renormalize = moe_renormalize
|
| 93 |
+
self.num_shared_experts = num_shared_experts
|
| 94 |
+
self.routed_scaling_factor = routed_scaling_factor
|
| 95 |
+
self.moe_router_activation_func = moe_router_activation_func
|
| 96 |
+
assert self.moe_router_activation_func in ("softmax", "sigmoid")
|
| 97 |
+
self.moe_intermediate_size = moe_intermediate_size
|
| 98 |
+
self.first_k_dense_replace = first_k_dense_replace
|
| 99 |
+
self.moe_layer_freq = moe_layer_freq
|
| 100 |
+
self.use_grouped_topk = use_grouped_topk
|
| 101 |
+
self.num_expert_group = num_expert_group
|
| 102 |
+
self.topk_group = topk_group
|
| 103 |
+
self.num_nextn_predict_layers = num_nextn_predict_layers
|
| 104 |
+
|
| 105 |
+
self.attn_res_block_size = attn_res_block_size
|
| 106 |
+
self.latent_moe_use_norm = latent_moe_use_norm
|
| 107 |
+
self.activation_situ_beta = activation_situ_beta
|
| 108 |
+
self.activation_situ_linear_beta = activation_situ_linear_beta
|
| 109 |
+
self.max_position_embeddings = max_position_embeddings
|
| 110 |
+
self.routed_expert_hidden_size = routed_expert_hidden_size
|
| 111 |
+
self.topk_method = topk_method
|
| 112 |
+
|
| 113 |
+
if linear_attn_config is not None:
|
| 114 |
+
assert linear_attn_config["kda_layers"] is not None
|
| 115 |
+
assert linear_attn_config["full_attn_layers"] is not None
|
| 116 |
+
self.linear_attn_config = linear_attn_config
|
| 117 |
+
|
| 118 |
+
super().__init__(
|
| 119 |
+
pad_token_id=pad_token_id,
|
| 120 |
+
bos_token_id=bos_token_id,
|
| 121 |
+
eos_token_id=eos_token_id,
|
| 122 |
+
tie_word_embeddings=tie_word_embeddings,
|
| 123 |
+
**kwargs,
|
| 124 |
+
)
|
| 125 |
+
|
| 126 |
+
@property
|
| 127 |
+
def is_mla(self):
|
| 128 |
+
return (
|
| 129 |
+
self.q_lora_rank is not None
|
| 130 |
+
or self.kv_lora_rank is not None
|
| 131 |
+
or self.qk_nope_head_dim is not None
|
| 132 |
+
or self.qk_rope_head_dim is not None
|
| 133 |
+
or self.v_head_dim is not None
|
| 134 |
+
or self.mla_use_nope is True
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
@property
|
| 138 |
+
def is_moe(self):
|
| 139 |
+
return self.num_experts is not None
|
| 140 |
+
|
| 141 |
+
@property
|
| 142 |
+
def is_linear_attn(self) -> bool:
|
| 143 |
+
return not (
|
| 144 |
+
self.linear_attn_config is None
|
| 145 |
+
or (
|
| 146 |
+
isinstance(self.linear_attn_config, dict)
|
| 147 |
+
and self.linear_attn_config["kda_layers"] is not None
|
| 148 |
+
and len(self.linear_attn_config["kda_layers"]) == 0
|
| 149 |
+
)
|
| 150 |
+
)
|
| 151 |
+
|
| 152 |
+
def is_kda_layer(self, layer_idx: int):
|
| 153 |
+
return (
|
| 154 |
+
self.linear_attn_config is not None
|
| 155 |
+
and (layer_idx + 1) in self.linear_attn_config["kda_layers"]
|
| 156 |
+
)
|
| 157 |
+
|
| 158 |
+
|
| 159 |
+
class KimiK3VisionConfig(PretrainedConfig):
|
| 160 |
+
|
| 161 |
+
def __init__(
|
| 162 |
+
self,
|
| 163 |
+
patch_size: int = 14,
|
| 164 |
+
init_pos_emb_height: int = 64,
|
| 165 |
+
init_pos_emb_width: int = 64,
|
| 166 |
+
init_pos_emb_time: int = 4,
|
| 167 |
+
pos_emb_type: str = 'divided_fixed',
|
| 168 |
+
vt_num_attention_heads: int = 12,
|
| 169 |
+
vt_num_hidden_layers: int = 27,
|
| 170 |
+
vt_hidden_size: int = 1024,
|
| 171 |
+
vt_intermediate_size: int = 4096,
|
| 172 |
+
merge_kernel_size: tuple = (2, 2),
|
| 173 |
+
merge_type: str = 'sd2_tpool',
|
| 174 |
+
_attn_implementation: str = 'flash_attention_2',
|
| 175 |
+
# MM Projector parameters
|
| 176 |
+
mm_projector_type: str = 'patchmergerv2',
|
| 177 |
+
mm_hidden_size: int | None = None,
|
| 178 |
+
projector_hidden_act: str = "gelu",
|
| 179 |
+
projector_ln_eps: float = 1e-5,
|
| 180 |
+
# vision tower parameters
|
| 181 |
+
qkv_hidden_size: int = 1536,
|
| 182 |
+
norm_type: str = 'rmsnorm',
|
| 183 |
+
attn_bias: bool = False,
|
| 184 |
+
patch_embed_proj_bias: bool = False,
|
| 185 |
+
mlp_type: str = 'mlp2',
|
| 186 |
+
linear_bias: bool = False,
|
| 187 |
+
activation_func: str = 'gelu_pytorch_tanh',
|
| 188 |
+
pos_emb_interpolation_mode: str = 'bilinear',
|
| 189 |
+
# Other parameters
|
| 190 |
+
ignore_index: int = -100,
|
| 191 |
+
media_placeholder_token_id: int = 163605,
|
| 192 |
+
pad_token_id: int = 0,
|
| 193 |
+
text_hidden_size=7168,
|
| 194 |
+
**kwargs):
|
| 195 |
+
|
| 196 |
+
self.patch_size = patch_size
|
| 197 |
+
self.init_pos_emb_height = init_pos_emb_height
|
| 198 |
+
self.init_pos_emb_width = init_pos_emb_width
|
| 199 |
+
self.init_pos_emb_time = init_pos_emb_time
|
| 200 |
+
self.pos_emb_type = pos_emb_type
|
| 201 |
+
self.vt_num_attention_heads = vt_num_attention_heads
|
| 202 |
+
self.vt_num_hidden_layers = vt_num_hidden_layers
|
| 203 |
+
self.vt_hidden_size = vt_hidden_size
|
| 204 |
+
self.vt_intermediate_size = vt_intermediate_size
|
| 205 |
+
self.merge_kernel_size = merge_kernel_size
|
| 206 |
+
self.merge_type = merge_type
|
| 207 |
+
self._attn_implementation = _attn_implementation
|
| 208 |
+
|
| 209 |
+
# MM Projector config
|
| 210 |
+
self.mm_projector_type = mm_projector_type
|
| 211 |
+
self.mm_hidden_size = mm_hidden_size if mm_hidden_size is not None else vt_hidden_size
|
| 212 |
+
self.projector_hidden_act = projector_hidden_act
|
| 213 |
+
self.projector_ln_eps = projector_ln_eps
|
| 214 |
+
self.text_hidden_size = text_hidden_size
|
| 215 |
+
|
| 216 |
+
# vision tower parameters
|
| 217 |
+
self.qkv_hidden_size = qkv_hidden_size
|
| 218 |
+
self.norm_type = norm_type
|
| 219 |
+
self.attn_bias = attn_bias
|
| 220 |
+
self.patch_embed_proj_bias = patch_embed_proj_bias
|
| 221 |
+
self.mlp_type = mlp_type
|
| 222 |
+
self.linear_bias = linear_bias
|
| 223 |
+
self.activation_func = activation_func
|
| 224 |
+
self.pos_emb_interpolation_mode = pos_emb_interpolation_mode
|
| 225 |
+
|
| 226 |
+
super().__init__(**kwargs)
|
| 227 |
+
|
| 228 |
+
|
| 229 |
+
class KimiK3Config(PretrainedConfig):
|
| 230 |
+
"""Kimi-K3 model configuration.
|
| 231 |
+
|
| 232 |
+
Args:
|
| 233 |
+
text_config (dict | KimiLinearConfig): Configuration for the text model.
|
| 234 |
+
|
| 235 |
+
Vision Tower Parameters (from MoonViT3dConfig):
|
| 236 |
+
patch_size (int): Patch size for vision tower.
|
| 237 |
+
init_pos_emb_height (int): Initial position embedding height.
|
| 238 |
+
init_pos_emb_width (int): Initial position embedding width.
|
| 239 |
+
init_pos_emb_time (int): Initial position embedding time dimension.
|
| 240 |
+
pos_emb_type (str): Type of position embedding.
|
| 241 |
+
vt_num_attention_heads (int): Number of attention heads in vision tower.
|
| 242 |
+
vt_num_hidden_layers (int): Number of hidden layers in vision tower.
|
| 243 |
+
vt_hidden_size (int): Hidden size of vision tower.
|
| 244 |
+
vt_intermediate_size (int): Intermediate size in vision tower FFN.
|
| 245 |
+
merge_kernel_size (tuple): Kernel size for patch merging.
|
| 246 |
+
merge_type (str): Type of merge operation.
|
| 247 |
+
_attn_implementation (str): Attention implementation type.
|
| 248 |
+
|
| 249 |
+
MM Projector Parameters (from MultiModalProjectorConfig):
|
| 250 |
+
mm_projector_type (str): Type of multimodal projector.
|
| 251 |
+
mm_hidden_size (int): Hidden size from vision tower (should match vt_hidden_size).
|
| 252 |
+
projector_hidden_act (str): Activation function for projector.
|
| 253 |
+
projector_ln_eps (float): Layer norm epsilon for projector.
|
| 254 |
+
|
| 255 |
+
Other Parameters:
|
| 256 |
+
ignore_index (int): The ignore index for the loss function.
|
| 257 |
+
media_placeholder_token_id (int): The token ID to use for media placeholders.
|
| 258 |
+
pad_token_id (int): The token ID to use for padding.
|
| 259 |
+
"""
|
| 260 |
+
|
| 261 |
+
model_type = "kimi_k3"
|
| 262 |
+
|
| 263 |
+
def __init__(
|
| 264 |
+
self,
|
| 265 |
+
text_config: dict | KimiLinearConfig = None,
|
| 266 |
+
vision_config: dict | KimiK3VisionConfig = None,
|
| 267 |
+
# Other parameters
|
| 268 |
+
ignore_index: int = -100,
|
| 269 |
+
media_placeholder_token_id: int = 163605,
|
| 270 |
+
pad_token_id: int = 0,
|
| 271 |
+
**kwargs,
|
| 272 |
+
):
|
| 273 |
+
if isinstance(text_config, dict):
|
| 274 |
+
text_config = KimiLinearConfig(**text_config)
|
| 275 |
+
if isinstance(vision_config, dict):
|
| 276 |
+
vision_config = KimiK3VisionConfig(**vision_config)
|
| 277 |
+
self.text_config = text_config
|
| 278 |
+
self.vision_config = vision_config
|
| 279 |
+
# Other config
|
| 280 |
+
self.ignore_index = ignore_index
|
| 281 |
+
self.media_placeholder_token_id = media_placeholder_token_id
|
| 282 |
+
if getattr(self.text_config, "quantization_config", None) is not None:
|
| 283 |
+
self.quantization_config = self.text_config.quantization_config
|
| 284 |
+
|
| 285 |
+
super().__init__(pad_token_id=pad_token_id, **kwargs)
|
encoding_k3.py
ADDED
|
@@ -0,0 +1,647 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Kimi K3 XTML encoding helpers.
|
| 2 |
+
|
| 3 |
+
This module keeps chat rendering in Python.
|
| 4 |
+
Callers that need token IDs should consume ``EncodeSegment`` objects directly:
|
| 5 |
+
structural markers may be encoded as tiktoken special tokens, while user/tool
|
| 6 |
+
text and attribute values are encoded as ordinary text.
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
from __future__ import annotations
|
| 10 |
+
|
| 11 |
+
import json
|
| 12 |
+
from dataclasses import dataclass
|
| 13 |
+
from typing import Any, Iterable, Optional
|
| 14 |
+
|
| 15 |
+
OPEN_TOKEN = "<|open|>"
|
| 16 |
+
CLOSE_TOKEN = "<|close|>"
|
| 17 |
+
SEP_TOKEN = "<|sep|>"
|
| 18 |
+
END_OF_MSG_TOKEN = "<|end_of_msg|>"
|
| 19 |
+
IMAGE_PLACEHOLDER = "<|kimi_image_placeholder|>"
|
| 20 |
+
|
| 21 |
+
_VALID_THINKING_EFFORTS = {"low", "high", "max"}
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
@dataclass(frozen=True)
|
| 25 |
+
class EncodeSegment:
|
| 26 |
+
text: str
|
| 27 |
+
allow_special: bool = False
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
class _ImagePromptState:
|
| 31 |
+
def __init__(self, image_prompts: Optional[list[str]] = None):
|
| 32 |
+
self.image_prompts = image_prompts
|
| 33 |
+
self.index = 0
|
| 34 |
+
|
| 35 |
+
def next_prompt(self) -> str:
|
| 36 |
+
if self.image_prompts is None:
|
| 37 |
+
return IMAGE_PLACEHOLDER
|
| 38 |
+
if self.index >= len(self.image_prompts):
|
| 39 |
+
raise ValueError("More image placeholders than image prompts.")
|
| 40 |
+
prompt = self.image_prompts[self.index]
|
| 41 |
+
self.index += 1
|
| 42 |
+
return prompt
|
| 43 |
+
|
| 44 |
+
def assert_consumed(self) -> None:
|
| 45 |
+
if self.image_prompts is None:
|
| 46 |
+
return
|
| 47 |
+
if self.index != len(self.image_prompts):
|
| 48 |
+
raise ValueError(
|
| 49 |
+
f"image prompt count {len(self.image_prompts)} != "
|
| 50 |
+
f"consumed placeholder count {self.index}"
|
| 51 |
+
)
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def _segment(text: Any, *, allow_special: bool = False) -> list[EncodeSegment]:
|
| 55 |
+
text = str(text)
|
| 56 |
+
if not text:
|
| 57 |
+
return []
|
| 58 |
+
return [EncodeSegment(text, allow_special=allow_special)]
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def _control(text: str) -> list[EncodeSegment]:
|
| 62 |
+
return _segment(text, allow_special=True)
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def _text(text: Any) -> list[EncodeSegment]:
|
| 66 |
+
return _segment(text, allow_special=False)
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def _append_text(
|
| 70 |
+
segments: list[EncodeSegment],
|
| 71 |
+
text: Any,
|
| 72 |
+
image_state: _ImagePromptState,
|
| 73 |
+
) -> None:
|
| 74 |
+
text = str(text)
|
| 75 |
+
if text == "":
|
| 76 |
+
return
|
| 77 |
+
if image_state.image_prompts is None or IMAGE_PLACEHOLDER not in text:
|
| 78 |
+
segments.extend(_text(text))
|
| 79 |
+
return
|
| 80 |
+
|
| 81 |
+
parts = text.split(IMAGE_PLACEHOLDER)
|
| 82 |
+
for i, part in enumerate(parts):
|
| 83 |
+
segments.extend(_text(part))
|
| 84 |
+
if i < len(parts) - 1:
|
| 85 |
+
segments.extend(_segment(image_state.next_prompt(),
|
| 86 |
+
allow_special=True))
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def _escape_attr_value(value: Any) -> str:
|
| 90 |
+
return str(value).replace("&", "&").replace('"', """)
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def _attr(key: str, value: Any) -> list[EncodeSegment]:
|
| 94 |
+
return (
|
| 95 |
+
_text(f" {key}")
|
| 96 |
+
+ _text('="')
|
| 97 |
+
+ _text(_escape_attr_value(value))
|
| 98 |
+
+ _text('"')
|
| 99 |
+
)
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def _open_tag(tag: str, attrs: Iterable[tuple[str, Any]] = ()) -> list[EncodeSegment]:
|
| 103 |
+
segments: list[EncodeSegment] = []
|
| 104 |
+
segments.extend(_control(OPEN_TOKEN))
|
| 105 |
+
segments.extend(_text(tag))
|
| 106 |
+
for key, value in attrs:
|
| 107 |
+
segments.extend(_attr(key, value))
|
| 108 |
+
segments.extend(_control(SEP_TOKEN))
|
| 109 |
+
return segments
|
| 110 |
+
|
| 111 |
+
def _close_tag(tag: str) -> list[EncodeSegment]:
|
| 112 |
+
segments: list[EncodeSegment] = []
|
| 113 |
+
segments.extend(_control(CLOSE_TOKEN))
|
| 114 |
+
segments.extend(_text(tag))
|
| 115 |
+
segments.extend(_control(SEP_TOKEN))
|
| 116 |
+
return segments
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
def _end_of_msg() -> list[EncodeSegment]:
|
| 120 |
+
return _control(END_OF_MSG_TOKEN)
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
def _json_compact(value: Any) -> str:
|
| 124 |
+
return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
def _is_mapping(value: Any) -> bool:
|
| 128 |
+
return isinstance(value, dict)
|
| 129 |
+
|
| 130 |
+
|
| 131 |
+
def _xtml_type(value: Any) -> str:
|
| 132 |
+
if isinstance(value, bool):
|
| 133 |
+
return "boolean"
|
| 134 |
+
if value is None:
|
| 135 |
+
return "null"
|
| 136 |
+
if isinstance(value, (int, float)) and not isinstance(value, bool):
|
| 137 |
+
return "number"
|
| 138 |
+
if isinstance(value, str):
|
| 139 |
+
return "string"
|
| 140 |
+
if _is_mapping(value):
|
| 141 |
+
return "object"
|
| 142 |
+
return "array"
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def _xtml_value(value: Any) -> str:
|
| 146 |
+
if isinstance(value, str):
|
| 147 |
+
return value
|
| 148 |
+
return json.dumps(value, ensure_ascii=False)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def _get_value(obj: Any, key: str, default: Any = None) -> Any:
|
| 152 |
+
if isinstance(obj, dict):
|
| 153 |
+
return obj.get(key, default)
|
| 154 |
+
return getattr(obj, key, default)
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
def extract_response_schema(response_format: Any) -> Any:
|
| 158 |
+
if response_format is None:
|
| 159 |
+
return None
|
| 160 |
+
|
| 161 |
+
json_schema = _get_value(response_format, "json_schema")
|
| 162 |
+
if json_schema is None:
|
| 163 |
+
return None
|
| 164 |
+
|
| 165 |
+
if isinstance(json_schema, dict):
|
| 166 |
+
return json_schema.get(
|
| 167 |
+
"schema",
|
| 168 |
+
json_schema.get("json_schema", json_schema),
|
| 169 |
+
)
|
| 170 |
+
|
| 171 |
+
schema = _get_value(json_schema, "schema")
|
| 172 |
+
if schema is not None:
|
| 173 |
+
return schema
|
| 174 |
+
|
| 175 |
+
schema = _get_value(json_schema, "json_schema")
|
| 176 |
+
if schema is not None:
|
| 177 |
+
return schema
|
| 178 |
+
|
| 179 |
+
return json_schema
|
| 180 |
+
|
| 181 |
+
|
| 182 |
+
def deep_sort_dict(obj: Any) -> Any:
|
| 183 |
+
if isinstance(obj, dict):
|
| 184 |
+
return {k: deep_sort_dict(v) for k, v in sorted(obj.items())}
|
| 185 |
+
if isinstance(obj, list):
|
| 186 |
+
return [deep_sort_dict(item) for item in obj]
|
| 187 |
+
return obj
|
| 188 |
+
|
| 189 |
+
|
| 190 |
+
def normalize_tool_arguments(arguments: Any) -> tuple[dict[str, Any], Optional[str]]:
|
| 191 |
+
if arguments is None:
|
| 192 |
+
return {}, None
|
| 193 |
+
if isinstance(arguments, dict):
|
| 194 |
+
return arguments, None
|
| 195 |
+
if isinstance(arguments, str):
|
| 196 |
+
if not arguments.strip():
|
| 197 |
+
return {}, None
|
| 198 |
+
try:
|
| 199 |
+
parsed = json.loads(arguments)
|
| 200 |
+
except json.JSONDecodeError:
|
| 201 |
+
return {}, arguments
|
| 202 |
+
if not isinstance(parsed, dict):
|
| 203 |
+
raise ValueError("Kimi K3 tool call arguments must be a JSON object.")
|
| 204 |
+
return parsed, None
|
| 205 |
+
raise TypeError(
|
| 206 |
+
"Kimi K3 tool call arguments must be a dict or a JSON object string."
|
| 207 |
+
)
|
| 208 |
+
|
| 209 |
+
|
| 210 |
+
def normalize_message(message: Any) -> Any:
|
| 211 |
+
if not isinstance(message, dict):
|
| 212 |
+
return message
|
| 213 |
+
|
| 214 |
+
normalized = dict(message)
|
| 215 |
+
|
| 216 |
+
tools = normalized.get("tools")
|
| 217 |
+
if tools is not None:
|
| 218 |
+
normalized["tools"] = deep_sort_dict(tools)
|
| 219 |
+
|
| 220 |
+
tool_calls = normalized.get("tool_calls")
|
| 221 |
+
if not tool_calls:
|
| 222 |
+
return normalized
|
| 223 |
+
|
| 224 |
+
normalized_calls = []
|
| 225 |
+
for tool_call in tool_calls:
|
| 226 |
+
if not isinstance(tool_call, dict):
|
| 227 |
+
normalized_calls.append(tool_call)
|
| 228 |
+
continue
|
| 229 |
+
|
| 230 |
+
tc = dict(tool_call)
|
| 231 |
+
function = tc.get("function")
|
| 232 |
+
if isinstance(function, dict):
|
| 233 |
+
fn = dict(function)
|
| 234 |
+
arguments, json_block = normalize_tool_arguments(fn.get("arguments"))
|
| 235 |
+
fn["arguments"] = arguments
|
| 236 |
+
if json_block is None:
|
| 237 |
+
fn.pop("_xtml_json_block", None)
|
| 238 |
+
else:
|
| 239 |
+
fn["_xtml_json_block"] = json_block
|
| 240 |
+
tc["function"] = fn
|
| 241 |
+
else:
|
| 242 |
+
arguments, json_block = normalize_tool_arguments(tc.get("arguments"))
|
| 243 |
+
tc["arguments"] = arguments
|
| 244 |
+
if json_block is None:
|
| 245 |
+
tc.pop("_xtml_json_block", None)
|
| 246 |
+
else:
|
| 247 |
+
tc["_xtml_json_block"] = json_block
|
| 248 |
+
normalized_calls.append(tc)
|
| 249 |
+
|
| 250 |
+
normalized["tool_calls"] = normalized_calls
|
| 251 |
+
return normalized
|
| 252 |
+
|
| 253 |
+
|
| 254 |
+
def normalize_conversation(conversation: Any) -> Any:
|
| 255 |
+
if not isinstance(conversation, list):
|
| 256 |
+
return conversation
|
| 257 |
+
|
| 258 |
+
def normalize_messages(messages: list[Any]) -> list[Any]:
|
| 259 |
+
return [normalize_message(message) for message in messages]
|
| 260 |
+
|
| 261 |
+
if conversation and isinstance(conversation[0], list):
|
| 262 |
+
return [normalize_messages(messages) for messages in conversation]
|
| 263 |
+
return normalize_messages(conversation)
|
| 264 |
+
|
| 265 |
+
|
| 266 |
+
def _tool_call_id_index(tool_calls: Any) -> dict:
|
| 267 |
+
"""Map assistant ``tool_calls[].id`` to ``(1-based position, function name)``.
|
| 268 |
+
|
| 269 |
+
The position mirrors the chat template's enumeration over ``tool_calls``
|
| 270 |
+
(every entry advances the position, even an id-less one). Duplicate ids keep
|
| 271 |
+
their first occurrence.
|
| 272 |
+
"""
|
| 273 |
+
index: dict = {}
|
| 274 |
+
if not isinstance(tool_calls, list):
|
| 275 |
+
return index
|
| 276 |
+
for position, tool_call in enumerate(tool_calls, start=1):
|
| 277 |
+
if not isinstance(tool_call, dict):
|
| 278 |
+
continue
|
| 279 |
+
call_id = tool_call.get("id")
|
| 280 |
+
if call_id is None:
|
| 281 |
+
continue
|
| 282 |
+
key = str(call_id)
|
| 283 |
+
if key in index:
|
| 284 |
+
continue
|
| 285 |
+
function = tool_call.get("function")
|
| 286 |
+
name = (
|
| 287 |
+
function.get("name") if isinstance(function, dict) else tool_call.get("name")
|
| 288 |
+
)
|
| 289 |
+
index[key] = (position, name)
|
| 290 |
+
return index
|
| 291 |
+
|
| 292 |
+
|
| 293 |
+
def normalize_xtml_tool_result_messages(messages: list[Any]) -> list[Any]:
|
| 294 |
+
"""Re-sort K3 XTML tool results into assistant ``tool_calls`` order.
|
| 295 |
+
|
| 296 |
+
Serving frameworks generally deliver tool results already in call order. A
|
| 297 |
+
direct Transformers caller, however, may pass OpenAI-style tool messages in any
|
| 298 |
+
order, so each run of consecutive tool messages is matched against the most
|
| 299 |
+
recent preceding assistant ``tool_calls`` by opaque ``tool_call_id`` ==
|
| 300 |
+
``tool_calls[].id`` (K3 drops the ``func:index`` format requirement) and
|
| 301 |
+
sorted by the matched 1-based position. The matched call is authoritative,
|
| 302 |
+
so each matched message's ``tool`` is set to that call's function name --
|
| 303 |
+
this keeps an explicit (and possibly stale) ``tool``/``name`` from drifting
|
| 304 |
+
out of sync with the reordered position. ``index`` is still derived from the
|
| 305 |
+
rendered position by the chat template. A run that cannot be fully matched is
|
| 306 |
+
left untouched. Re-running is idempotent.
|
| 307 |
+
|
| 308 |
+
This function is side-effect free: matched tool messages are shallow-copied
|
| 309 |
+
before their ``tool``/``name`` is rewritten, and every other message is
|
| 310 |
+
appended to the output as-is. The input list and its message objects are
|
| 311 |
+
never mutated.
|
| 312 |
+
"""
|
| 313 |
+
if not isinstance(messages, list):
|
| 314 |
+
return messages
|
| 315 |
+
|
| 316 |
+
output: list[Any] = []
|
| 317 |
+
current_index: dict = {}
|
| 318 |
+
i = 0
|
| 319 |
+
n = len(messages)
|
| 320 |
+
|
| 321 |
+
while i < n:
|
| 322 |
+
message = messages[i]
|
| 323 |
+
|
| 324 |
+
if isinstance(message, dict) and message.get("role") == "assistant":
|
| 325 |
+
tool_calls = message.get("tool_calls")
|
| 326 |
+
current_index = _tool_call_id_index(tool_calls) if tool_calls else {}
|
| 327 |
+
output.append(message)
|
| 328 |
+
i += 1
|
| 329 |
+
continue
|
| 330 |
+
|
| 331 |
+
if not isinstance(message, dict) or message.get("role") != "tool":
|
| 332 |
+
output.append(message)
|
| 333 |
+
i += 1
|
| 334 |
+
continue
|
| 335 |
+
|
| 336 |
+
run: list[tuple] = [] # (position, original_offset, message, name)
|
| 337 |
+
unresolved = False
|
| 338 |
+
offset = 0
|
| 339 |
+
while (
|
| 340 |
+
i < n and isinstance(messages[i], dict) and messages[i].get("role") == "tool"
|
| 341 |
+
):
|
| 342 |
+
tool_message = messages[i]
|
| 343 |
+
call_id = tool_message.get("tool_call_id", tool_message.get("id"))
|
| 344 |
+
matched = current_index.get(str(call_id)) if call_id is not None else None
|
| 345 |
+
if matched is None:
|
| 346 |
+
unresolved = True
|
| 347 |
+
run.append((None, offset, tool_message, None))
|
| 348 |
+
else:
|
| 349 |
+
position, name = matched
|
| 350 |
+
run.append((position, offset, tool_message, name))
|
| 351 |
+
offset += 1
|
| 352 |
+
i += 1
|
| 353 |
+
|
| 354 |
+
if unresolved:
|
| 355 |
+
output.extend(item[2] for item in run)
|
| 356 |
+
else:
|
| 357 |
+
run.sort(key=lambda item: (item[0], item[1]))
|
| 358 |
+
for _, _, tool_message, name in run:
|
| 359 |
+
if name is None:
|
| 360 |
+
output.append(tool_message)
|
| 361 |
+
continue
|
| 362 |
+
# The id-matched call is authoritative: align tool (and any
|
| 363 |
+
# explicit name) so the rendered XTML tool attribute cannot
|
| 364 |
+
# disagree with the reordered position. Copy first so the
|
| 365 |
+
# caller's message object is never mutated.
|
| 366 |
+
resolved = dict(tool_message)
|
| 367 |
+
resolved["tool"] = name
|
| 368 |
+
if "name" in resolved:
|
| 369 |
+
resolved["name"] = name
|
| 370 |
+
output.append(resolved)
|
| 371 |
+
|
| 372 |
+
return output
|
| 373 |
+
|
| 374 |
+
|
| 375 |
+
def is_batched_conversation(conversation: Any) -> bool:
|
| 376 |
+
return (
|
| 377 |
+
isinstance(conversation, list)
|
| 378 |
+
and bool(conversation)
|
| 379 |
+
and isinstance(conversation[0], list)
|
| 380 |
+
)
|
| 381 |
+
|
| 382 |
+
|
| 383 |
+
def _render_content_segments(
|
| 384 |
+
content: Any,
|
| 385 |
+
image_state: _ImagePromptState,
|
| 386 |
+
) -> list[EncodeSegment]:
|
| 387 |
+
segments: list[EncodeSegment] = []
|
| 388 |
+
if isinstance(content, str):
|
| 389 |
+
_append_text(segments, content, image_state)
|
| 390 |
+
elif content is not None:
|
| 391 |
+
for part in content:
|
| 392 |
+
if part["type"] in ["image", "image_url"]:
|
| 393 |
+
segments.extend(
|
| 394 |
+
_segment(image_state.next_prompt(), allow_special=True))
|
| 395 |
+
else:
|
| 396 |
+
_append_text(segments, part["text"], image_state)
|
| 397 |
+
return segments
|
| 398 |
+
|
| 399 |
+
|
| 400 |
+
def _internal_system_message(message_type: str, body: str) -> list[EncodeSegment]:
|
| 401 |
+
segments: list[EncodeSegment] = []
|
| 402 |
+
segments.extend(_open_tag("message", [("role", "system"), ("type", message_type)]))
|
| 403 |
+
segments.extend(_text(body.strip()))
|
| 404 |
+
segments.extend(_close_tag("message"))
|
| 405 |
+
segments.extend(_end_of_msg())
|
| 406 |
+
return segments
|
| 407 |
+
|
| 408 |
+
|
| 409 |
+
def _render_assistant_segments(
|
| 410 |
+
message: dict[str, Any],
|
| 411 |
+
image_state: _ImagePromptState,
|
| 412 |
+
thinking: bool = True,
|
| 413 |
+
) -> list[EncodeSegment]:
|
| 414 |
+
segments: list[EncodeSegment] = []
|
| 415 |
+
# The <think> channel is structural: in thinking mode every assistant
|
| 416 |
+
# message carries the open/close tags even when there is no reasoning
|
| 417 |
+
# content to fill in. In non-thinking mode the channel is dropped
|
| 418 |
+
# entirely.
|
| 419 |
+
if thinking:
|
| 420 |
+
reasoning_content = message.get("reasoning_content") or message.get(
|
| 421 |
+
"reasoning"
|
| 422 |
+
)
|
| 423 |
+
segments.extend(_open_tag("think"))
|
| 424 |
+
if reasoning_content is not None and str(reasoning_content).strip():
|
| 425 |
+
_append_text(segments, reasoning_content, image_state)
|
| 426 |
+
segments.extend(_close_tag("think"))
|
| 427 |
+
|
| 428 |
+
segments.extend(_open_tag("response"))
|
| 429 |
+
segments.extend(_render_content_segments(message.get("content"), image_state))
|
| 430 |
+
segments.extend(_close_tag("response"))
|
| 431 |
+
|
| 432 |
+
tool_calls = message.get("tool_calls")
|
| 433 |
+
if tool_calls:
|
| 434 |
+
segments.extend(_open_tag("tools"))
|
| 435 |
+
for index, tool_call in enumerate(tool_calls, start=1):
|
| 436 |
+
fn = tool_call.get("function", tool_call)
|
| 437 |
+
segments.extend(
|
| 438 |
+
_open_tag("call", [("tool", fn["name"]), ("index", index)])
|
| 439 |
+
)
|
| 440 |
+
args = fn.get("arguments", {})
|
| 441 |
+
json_block = fn.get("_xtml_json_block")
|
| 442 |
+
if json_block is not None:
|
| 443 |
+
segments.extend(_open_tag("json", [("type", "object")]))
|
| 444 |
+
_append_text(segments, json_block, image_state)
|
| 445 |
+
segments.extend(_close_tag("json"))
|
| 446 |
+
elif _is_mapping(args):
|
| 447 |
+
for key, value in args.items():
|
| 448 |
+
segments.extend(
|
| 449 |
+
_open_tag(
|
| 450 |
+
"argument",
|
| 451 |
+
[("key", key), ("type", _xtml_type(value))],
|
| 452 |
+
)
|
| 453 |
+
)
|
| 454 |
+
_append_text(segments, _xtml_value(value), image_state)
|
| 455 |
+
segments.extend(_close_tag("argument"))
|
| 456 |
+
segments.extend(_close_tag("call"))
|
| 457 |
+
segments.extend(_close_tag("tools"))
|
| 458 |
+
|
| 459 |
+
return segments
|
| 460 |
+
|
| 461 |
+
|
| 462 |
+
def _render_tool_declare(tools: Any, *, dynamic: bool = False) -> list[EncodeSegment]:
|
| 463 |
+
if dynamic:
|
| 464 |
+
body = (
|
| 465 |
+
"## New Tools Available\n"
|
| 466 |
+
"The system dynamically extends the toolset via lazy-loading.\n"
|
| 467 |
+
"You have access to all existing and extended tools.\n"
|
| 468 |
+
"Here are the specs for the extended tools.\n\n"
|
| 469 |
+
"```json\n"
|
| 470 |
+
f"{_json_compact(tools)}\n"
|
| 471 |
+
"```"
|
| 472 |
+
)
|
| 473 |
+
else:
|
| 474 |
+
body = (
|
| 475 |
+
"# Tools\n"
|
| 476 |
+
"Here are the available tools, described in JSONSchema.\n\n"
|
| 477 |
+
"```json\n"
|
| 478 |
+
f"{_json_compact(tools)}\n"
|
| 479 |
+
"```"
|
| 480 |
+
)
|
| 481 |
+
segments: list[EncodeSegment] = []
|
| 482 |
+
segments.extend(_open_tag("message", [("role", "system"), ("type", "tool-declare")]))
|
| 483 |
+
segments.extend(_text(body))
|
| 484 |
+
segments.extend(_close_tag("message"))
|
| 485 |
+
segments.extend(_end_of_msg())
|
| 486 |
+
return segments
|
| 487 |
+
|
| 488 |
+
|
| 489 |
+
def build_chat_segments(
|
| 490 |
+
messages: list[Any],
|
| 491 |
+
tools: Optional[list[dict]] = None,
|
| 492 |
+
*,
|
| 493 |
+
add_generation_prompt: bool = True,
|
| 494 |
+
thinking: bool = True,
|
| 495 |
+
image_prompts: Optional[list[str]] = None,
|
| 496 |
+
**kwargs: Any,
|
| 497 |
+
) -> list[EncodeSegment]:
|
| 498 |
+
# Re-sort tool results by tool_call_id at the lowest layer so every caller
|
| 499 |
+
# (processor or direct tokenizer) gets correctly ordered XTML. The helper is
|
| 500 |
+
# side-effect free, so the caller's message objects are left untouched.
|
| 501 |
+
messages = normalize_xtml_tool_result_messages(messages)
|
| 502 |
+
messages = normalize_conversation(messages)
|
| 503 |
+
tools = deep_sort_dict(tools)
|
| 504 |
+
|
| 505 |
+
kwargs = dict(kwargs)
|
| 506 |
+
response_format = kwargs.get("response_format")
|
| 507 |
+
if "response_schema" not in kwargs:
|
| 508 |
+
response_schema = extract_response_schema(response_format)
|
| 509 |
+
if response_schema is not None:
|
| 510 |
+
kwargs["response_schema"] = response_schema
|
| 511 |
+
if kwargs.get("response_schema") is not None:
|
| 512 |
+
kwargs["response_schema"] = deep_sort_dict(kwargs["response_schema"])
|
| 513 |
+
|
| 514 |
+
image_state = _ImagePromptState(image_prompts)
|
| 515 |
+
segments: list[EncodeSegment] = []
|
| 516 |
+
|
| 517 |
+
tool_calls = None
|
| 518 |
+
tool_index = 0
|
| 519 |
+
|
| 520 |
+
if tools:
|
| 521 |
+
segments.extend(_render_tool_declare(tools))
|
| 522 |
+
|
| 523 |
+
thinking_effort = kwargs.get("thinking_effort")
|
| 524 |
+
if thinking and thinking_effort is not None:
|
| 525 |
+
assert thinking_effort in _VALID_THINKING_EFFORTS, (
|
| 526 |
+
f"Unsupported thinking_effort={thinking_effort!r}; "
|
| 527 |
+
f"supported values are {sorted(_VALID_THINKING_EFFORTS)}."
|
| 528 |
+
)
|
| 529 |
+
if thinking and thinking_effort in _VALID_THINKING_EFFORTS:
|
| 530 |
+
segments.extend(
|
| 531 |
+
_internal_system_message(
|
| 532 |
+
"thinking-effort",
|
| 533 |
+
"`thinking_effort` guides on how much to think in your "
|
| 534 |
+
"thinking channel (not including the response channel), "
|
| 535 |
+
"supported values include `low`, `medium`, `high`, and `max`.\n"
|
| 536 |
+
f"Now the system is invoked with `thinking_effort={thinking_effort}`.",
|
| 537 |
+
)
|
| 538 |
+
)
|
| 539 |
+
|
| 540 |
+
for message_index, message in enumerate(messages):
|
| 541 |
+
if not isinstance(message, dict):
|
| 542 |
+
continue
|
| 543 |
+
|
| 544 |
+
role = message["role"]
|
| 545 |
+
if role == "user":
|
| 546 |
+
attrs = [("role", "user")]
|
| 547 |
+
if message.get("name"):
|
| 548 |
+
attrs.append(("name", message["name"]))
|
| 549 |
+
segments.extend(_open_tag("message", attrs))
|
| 550 |
+
segments.extend(_render_content_segments(message.get("content"), image_state))
|
| 551 |
+
segments.extend(_close_tag("message"))
|
| 552 |
+
segments.extend(_end_of_msg())
|
| 553 |
+
elif role == "system" and message.get("tools"):
|
| 554 |
+
segments.extend(_render_tool_declare(message["tools"], dynamic=True))
|
| 555 |
+
elif role == "system":
|
| 556 |
+
attrs = [("role", "system")]
|
| 557 |
+
if message.get("name"):
|
| 558 |
+
attrs.append(("name", message["name"]))
|
| 559 |
+
segments.extend(_open_tag("message", attrs))
|
| 560 |
+
segments.extend(_render_content_segments(message.get("content"), image_state))
|
| 561 |
+
segments.extend(_close_tag("message"))
|
| 562 |
+
segments.extend(_end_of_msg())
|
| 563 |
+
elif role == "tool":
|
| 564 |
+
tool_index += 1
|
| 565 |
+
tool_name = message.get("tool", message.get("name"))
|
| 566 |
+
if (
|
| 567 |
+
tool_name is None
|
| 568 |
+
and tool_calls is not None
|
| 569 |
+
and tool_index <= len(tool_calls)
|
| 570 |
+
):
|
| 571 |
+
tc = tool_calls[tool_index - 1]
|
| 572 |
+
fn = tc.get("function", tc)
|
| 573 |
+
tool_name = fn["name"]
|
| 574 |
+
if tool_name is None:
|
| 575 |
+
raise ValueError(
|
| 576 |
+
"Kimi K3 tool messages need a resolvable tool name: "
|
| 577 |
+
"carry `tool`/`name`, or match a preceding assistant "
|
| 578 |
+
"tool_call by order."
|
| 579 |
+
)
|
| 580 |
+
segments.extend(
|
| 581 |
+
_open_tag(
|
| 582 |
+
"message",
|
| 583 |
+
[("role", "tool"), ("tool", tool_name), ("index", tool_index)],
|
| 584 |
+
)
|
| 585 |
+
)
|
| 586 |
+
segments.extend(_render_content_segments(message.get("content"), image_state))
|
| 587 |
+
segments.extend(_close_tag("message"))
|
| 588 |
+
segments.extend(_end_of_msg())
|
| 589 |
+
elif role == "assistant":
|
| 590 |
+
tool_calls = message.get("tool_calls")
|
| 591 |
+
tool_index = 0
|
| 592 |
+
attrs = [("role", "assistant")]
|
| 593 |
+
if message.get("name"):
|
| 594 |
+
attrs.append(("name", message["name"]))
|
| 595 |
+
segments.extend(_open_tag("message", attrs))
|
| 596 |
+
segments.extend(_render_assistant_segments(message, image_state, thinking))
|
| 597 |
+
segments.extend(_close_tag("message"))
|
| 598 |
+
segments.extend(_end_of_msg())
|
| 599 |
+
|
| 600 |
+
tool_choice = kwargs.get("tool_choice")
|
| 601 |
+
if tool_choice == "required":
|
| 602 |
+
segments.extend(
|
| 603 |
+
_internal_system_message(
|
| 604 |
+
"tool-choice",
|
| 605 |
+
"The system is invoked with `tool_choice=required`.\n"
|
| 606 |
+
"You MUST call tools in the next message.",
|
| 607 |
+
)
|
| 608 |
+
)
|
| 609 |
+
elif tool_choice == "none":
|
| 610 |
+
segments.extend(
|
| 611 |
+
_internal_system_message(
|
| 612 |
+
"tool-choice",
|
| 613 |
+
"The system is invoked with `tool_choice=none`.\n"
|
| 614 |
+
"You MUST NOT call any tools in the next message.",
|
| 615 |
+
)
|
| 616 |
+
)
|
| 617 |
+
|
| 618 |
+
rf = kwargs.get("response_format")
|
| 619 |
+
rf_type = _get_value(rf, "type", rf) if isinstance(rf, dict) else rf
|
| 620 |
+
if rf_type == "json_object":
|
| 621 |
+
segments.extend(
|
| 622 |
+
_internal_system_message(
|
| 623 |
+
"response-format",
|
| 624 |
+
"The system is invoked with `response_format=json_object`.\n"
|
| 625 |
+
"Your response must be raw JSON data without markdown code "
|
| 626 |
+
"blocks (```json) or any additional formatting.",
|
| 627 |
+
)
|
| 628 |
+
)
|
| 629 |
+
elif rf_type == "json_schema":
|
| 630 |
+
schema = _json_compact(kwargs.get("response_schema"))
|
| 631 |
+
segments.extend(
|
| 632 |
+
_internal_system_message(
|
| 633 |
+
"response-format",
|
| 634 |
+
"The system is invoked with `response_format=json_schema`.\n"
|
| 635 |
+
"Your response must be raw JSON data without markdown code "
|
| 636 |
+
"blocks (```json) or any additional formatting.\n"
|
| 637 |
+
"The JSON data must match the following schema:\n"
|
| 638 |
+
f"```json\n{schema}\n```",
|
| 639 |
+
)
|
| 640 |
+
)
|
| 641 |
+
|
| 642 |
+
if add_generation_prompt:
|
| 643 |
+
segments.extend(_open_tag("message", [("role", "assistant")]))
|
| 644 |
+
segments.extend(_open_tag("think" if thinking else "response"))
|
| 645 |
+
|
| 646 |
+
image_state.assert_consumed()
|
| 647 |
+
return segments
|
generation_config.json
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"max_length": 1048576,
|
| 3 |
+
"eos_token_id": 163586
|
| 4 |
+
}
|
kimi_k3_processor.py
ADDED
|
@@ -0,0 +1,187 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Kimi-K3 processor: wraps vision processor + tokenizer into a single interface.
|
| 2 |
+
|
| 3 |
+
Chat rendering (including XTML tool-result ordering) is handled by the
|
| 4 |
+
tokenizer's Python encoder; this processor adds multimodal media preprocessing.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
from transformers.feature_extraction_utils import BatchFeature
|
| 8 |
+
from transformers.processing_utils import ProcessorMixin
|
| 9 |
+
from transformers.utils import logging
|
| 10 |
+
|
| 11 |
+
from .media_utils import ensure_media_type
|
| 12 |
+
|
| 13 |
+
logger = logging.get_logger(__name__)
|
| 14 |
+
|
| 15 |
+
# ── KimiK3Processor ───────────────────────────────────────────────────
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class KimiK3Processor(ProcessorMixin):
|
| 19 |
+
r"""
|
| 20 |
+
Constructs a KimiK3 processor which wraps a KimiK3 image processor
|
| 21 |
+
and a tokenizer into a single processor.
|
| 22 |
+
|
| 23 |
+
[`KimiK3Processor`] offers all the functionalities of
|
| 24 |
+
[`KimiK3VisionProcessor`] and [`TikTokenTokenizer`].
|
| 25 |
+
|
| 26 |
+
Args:
|
| 27 |
+
image_processor ([`KimiK3VisionProcessor`], *optional*):
|
| 28 |
+
The image processor is a required input.
|
| 29 |
+
tokenizer ([`TikTokenTokenizer`], *optional*):
|
| 30 |
+
The tokenizer is a required input.
|
| 31 |
+
chat_template (`str`, *optional*): Kept for ProcessorMixin
|
| 32 |
+
compatibility. Kimi K3 chat encoding is implemented in Python by
|
| 33 |
+
the tokenizer.
|
| 34 |
+
"""
|
| 35 |
+
|
| 36 |
+
attributes = ["image_processor", "tokenizer"]
|
| 37 |
+
valid_kwargs = ["chat_template"]
|
| 38 |
+
image_processor_class = "AutoImageProcessor"
|
| 39 |
+
tokenizer_class = "AutoTokenizer"
|
| 40 |
+
|
| 41 |
+
def __init__(
|
| 42 |
+
self,
|
| 43 |
+
image_processor=None,
|
| 44 |
+
tokenizer=None,
|
| 45 |
+
chat_template=None,
|
| 46 |
+
**kwargs,
|
| 47 |
+
):
|
| 48 |
+
super().__init__(image_processor,
|
| 49 |
+
tokenizer,
|
| 50 |
+
chat_template=chat_template)
|
| 51 |
+
self.media_processor = image_processor
|
| 52 |
+
self.image_placeholder = "<|kimi_image_placeholder|>"
|
| 53 |
+
|
| 54 |
+
# ── Media preprocessing ────────────────────────────────────────────
|
| 55 |
+
|
| 56 |
+
def update_raw_text(self, text: str, image_prompts: list[str]) -> str:
|
| 57 |
+
# Replace image placeholders
|
| 58 |
+
image_count = text.count(self.image_placeholder)
|
| 59 |
+
if image_count > 0:
|
| 60 |
+
assert image_count == len(image_prompts), (
|
| 61 |
+
f"image placeholder count {image_count} != "
|
| 62 |
+
f"image_prompts count {len(image_prompts)}")
|
| 63 |
+
text_parts = text.split(self.image_placeholder)
|
| 64 |
+
assert len(text_parts) == len(image_prompts) + 1
|
| 65 |
+
text = "".join([
|
| 66 |
+
text_parts[i] + image_prompts[i]
|
| 67 |
+
for i in range(len(image_prompts))
|
| 68 |
+
])
|
| 69 |
+
text += text_parts[-1]
|
| 70 |
+
|
| 71 |
+
return text
|
| 72 |
+
|
| 73 |
+
def preprocess_medias(self,
|
| 74 |
+
medias: list[dict]) -> tuple[list[dict], list[str]]:
|
| 75 |
+
"""Process media items and generate corresponding prompts.
|
| 76 |
+
|
| 77 |
+
Returns:
|
| 78 |
+
A tuple of (updated_medias, image_prompts).
|
| 79 |
+
"""
|
| 80 |
+
updated_medias = []
|
| 81 |
+
image_prompts = []
|
| 82 |
+
for media in medias:
|
| 83 |
+
if media['type'] == 'image':
|
| 84 |
+
updated_medias.append(media)
|
| 85 |
+
img = ensure_media_type(
|
| 86 |
+
media,
|
| 87 |
+
transparent_bg_config=self.media_processor.
|
| 88 |
+
_transparent_bg_config,
|
| 89 |
+
transparent_bg_fill_stage=self.media_processor.
|
| 90 |
+
_transparent_bg_fill_stage,
|
| 91 |
+
)['image']
|
| 92 |
+
w, h = img.size
|
| 93 |
+
image_prompts.append(
|
| 94 |
+
self.media_processor.make_image_prompt(w, h))
|
| 95 |
+
else:
|
| 96 |
+
raise ValueError(f"unsupported media type: {media['type']}")
|
| 97 |
+
return updated_medias, image_prompts
|
| 98 |
+
|
| 99 |
+
# ── Main entry points ──────────────────────────────────────────────
|
| 100 |
+
|
| 101 |
+
def __call__(self,
|
| 102 |
+
messages: list[dict] = None,
|
| 103 |
+
medias: list[dict] = None,
|
| 104 |
+
text: str = None,
|
| 105 |
+
return_tensors: str = "pt",
|
| 106 |
+
**kwargs) -> BatchFeature:
|
| 107 |
+
"""
|
| 108 |
+
Process multimodal inputs for Kimi-K3 model.
|
| 109 |
+
|
| 110 |
+
Args:
|
| 111 |
+
messages: List of message dicts with 'role' and 'content' fields.
|
| 112 |
+
If provided, medias and text will be extracted automatically.
|
| 113 |
+
medias: Pre-extracted list of media dicts.
|
| 114 |
+
text: Pre-formatted text string.
|
| 115 |
+
return_tensors: Format of returned tensors. Default: 'pt'.
|
| 116 |
+
**kwargs: Additional arguments passed to apply_chat_template.
|
| 117 |
+
|
| 118 |
+
Returns:
|
| 119 |
+
BatchFeature with fields: input_ids, attention_mask,
|
| 120 |
+
pixel_values, grid_thws.
|
| 121 |
+
"""
|
| 122 |
+
if messages is None and (medias is None or text is None):
|
| 123 |
+
raise ValueError(
|
| 124 |
+
"Provide either 'messages' or both 'medias' and 'text'")
|
| 125 |
+
|
| 126 |
+
if medias is not None and text is not None:
|
| 127 |
+
updated_medias, image_prompts = (self.preprocess_medias(medias))
|
| 128 |
+
preprocessed = self.media_processor.preprocess(
|
| 129 |
+
updated_medias, return_tensors=return_tensors)
|
| 130 |
+
text = self.update_raw_text(text, image_prompts)
|
| 131 |
+
text_inputs = self.tokenizer(text, return_tensors=return_tensors)
|
| 132 |
+
return BatchFeature(data={**text_inputs, **preprocessed.data})
|
| 133 |
+
|
| 134 |
+
if medias is None:
|
| 135 |
+
medias = self._extract_medias_from_messages(messages)
|
| 136 |
+
updated_medias, image_prompts = (self.preprocess_medias(medias))
|
| 137 |
+
preprocessed = self.media_processor.preprocess(
|
| 138 |
+
updated_medias, return_tensors=return_tensors)
|
| 139 |
+
|
| 140 |
+
if text is None:
|
| 141 |
+
text_inputs = self.tokenizer.apply_chat_template(
|
| 142 |
+
messages,
|
| 143 |
+
tokenize=True,
|
| 144 |
+
return_tensors=return_tensors,
|
| 145 |
+
return_dict=True,
|
| 146 |
+
image_prompts=image_prompts,
|
| 147 |
+
**kwargs)
|
| 148 |
+
return BatchFeature(data={**text_inputs, **preprocessed.data})
|
| 149 |
+
|
| 150 |
+
text = self.update_raw_text(text, image_prompts)
|
| 151 |
+
text_inputs = self.tokenizer(text, return_tensors=return_tensors)
|
| 152 |
+
return BatchFeature(data={**text_inputs, **preprocessed.data})
|
| 153 |
+
|
| 154 |
+
@staticmethod
|
| 155 |
+
def _extract_medias_from_messages(messages: list[dict]) -> list[dict]:
|
| 156 |
+
"""Extract media items from messages in a single pass."""
|
| 157 |
+
medias = []
|
| 158 |
+
for msg in messages:
|
| 159 |
+
if msg['role'] != 'user' or not msg.get('content'):
|
| 160 |
+
continue
|
| 161 |
+
|
| 162 |
+
for content_part in msg['content']:
|
| 163 |
+
if not isinstance(content_part, dict):
|
| 164 |
+
continue
|
| 165 |
+
|
| 166 |
+
content_type = content_part.get('type')
|
| 167 |
+
if content_type in ['image_url', 'image']:
|
| 168 |
+
image_data = content_part.get(content_type)
|
| 169 |
+
assert image_data is not None, f"image data is missing for content part: {content_part}"
|
| 170 |
+
medias.append({
|
| 171 |
+
'type': 'image',
|
| 172 |
+
'image': image_data,
|
| 173 |
+
})
|
| 174 |
+
return medias
|
| 175 |
+
|
| 176 |
+
def apply_chat_template(self, messages, **kwargs):
|
| 177 |
+
return self.tokenizer.apply_chat_template(messages, **kwargs)
|
| 178 |
+
|
| 179 |
+
def batch_decode(self, *args, **kwargs):
|
| 180 |
+
return self.tokenizer.batch_decode(*args, **kwargs)
|
| 181 |
+
|
| 182 |
+
def decode(self, *args, **kwargs):
|
| 183 |
+
return self.tokenizer.decode(*args, **kwargs)
|
| 184 |
+
|
| 185 |
+
@property
|
| 186 |
+
def model_input_names(self):
|
| 187 |
+
return ['input_ids', 'attention_mask', 'pixel_values', 'grid_thws']
|
kimi_k3_vision_processing.py
ADDED
|
@@ -0,0 +1,179 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Image processor class for Kimi-K3.
|
| 2 |
+
"""
|
| 3 |
+
|
| 4 |
+
import json
|
| 5 |
+
from typing import Any, Dict, Optional, Union
|
| 6 |
+
|
| 7 |
+
import numpy as np
|
| 8 |
+
import torch
|
| 9 |
+
from PIL import Image
|
| 10 |
+
from transformers.image_processing_utils import (BaseImageProcessor,
|
| 11 |
+
BatchFeature)
|
| 12 |
+
from transformers.utils import TensorType
|
| 13 |
+
|
| 14 |
+
from .media_utils import (MediaInput, TransparentBgConfig, _to_tensor,
|
| 15 |
+
ensure_media_type, image_to_np, navit_patchify,
|
| 16 |
+
navit_resize_image, normalize)
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
class KimiK3VisionProcessor(BaseImageProcessor):
|
| 20 |
+
model_type = "kimi_k3"
|
| 21 |
+
|
| 22 |
+
def __init__(
|
| 23 |
+
self,
|
| 24 |
+
media_proc_cfg: dict,
|
| 25 |
+
**kwargs,
|
| 26 |
+
):
|
| 27 |
+
super().__init__(**kwargs)
|
| 28 |
+
self.media_proc_cfg = media_proc_cfg
|
| 29 |
+
|
| 30 |
+
@property
|
| 31 |
+
def _transparent_bg_config(self) -> Optional[TransparentBgConfig]:
|
| 32 |
+
cfg = self.media_proc_cfg.get("transparent_bg_config")
|
| 33 |
+
if cfg is None:
|
| 34 |
+
return None
|
| 35 |
+
if isinstance(cfg, TransparentBgConfig):
|
| 36 |
+
return cfg
|
| 37 |
+
return TransparentBgConfig(**cfg)
|
| 38 |
+
|
| 39 |
+
@property
|
| 40 |
+
def _transparent_bg_fill_stage(self) -> str:
|
| 41 |
+
return self.media_proc_cfg.get("transparent_bg_fill_stage",
|
| 42 |
+
"before_resize")
|
| 43 |
+
|
| 44 |
+
def media_tokens_calculator(self, media: MediaInput):
|
| 45 |
+
media = ensure_media_type(
|
| 46 |
+
media,
|
| 47 |
+
transparent_bg_config=self._transparent_bg_config,
|
| 48 |
+
transparent_bg_fill_stage=self._transparent_bg_fill_stage,
|
| 49 |
+
)
|
| 50 |
+
ret = self.get_resize_config(media)
|
| 51 |
+
return ret['num_tokens']
|
| 52 |
+
|
| 53 |
+
@classmethod
|
| 54 |
+
def make_image_prompt(cls, width: int, height: int) -> str:
|
| 55 |
+
"""Build the K3 image placeholder with resolution info."""
|
| 56 |
+
return (f"<|media_begin|>image {width}x{height}"
|
| 57 |
+
f"<|media_content|><|media_pad|><|media_end|>")
|
| 58 |
+
|
| 59 |
+
def get_resize_config(self, media_input: MediaInput) -> dict:
|
| 60 |
+
if media_input['type'] == 'image':
|
| 61 |
+
w, h = media_input['image'].size
|
| 62 |
+
ret = navit_resize_image(
|
| 63 |
+
w, h, self.media_proc_cfg['patch_size'],
|
| 64 |
+
self.media_proc_cfg['merge_kernel_size'],
|
| 65 |
+
self.media_proc_cfg['in_patch_limit'],
|
| 66 |
+
self.media_proc_cfg['patch_limit_on_one_side'],
|
| 67 |
+
self.media_proc_cfg['fixed_output_tokens'])
|
| 68 |
+
return ret
|
| 69 |
+
else:
|
| 70 |
+
raise ValueError("Unsupported type: {}".format(
|
| 71 |
+
media_input['type']))
|
| 72 |
+
|
| 73 |
+
def resize_image(self, image: Image.Image, new_width: int, new_height: int,
|
| 74 |
+
pad_width: int, pad_height: int) -> np.ndarray:
|
| 75 |
+
image_np = image_to_np(
|
| 76 |
+
image,
|
| 77 |
+
(new_width, new_height),
|
| 78 |
+
"resize",
|
| 79 |
+
transparent_bg_config=self._transparent_bg_config,
|
| 80 |
+
transparent_bg_fill_stage=self._transparent_bg_fill_stage,
|
| 81 |
+
)
|
| 82 |
+
image_np = np.pad(
|
| 83 |
+
image_np,
|
| 84 |
+
((0, pad_height), (0, pad_width), (0, 0)),
|
| 85 |
+
mode="constant",
|
| 86 |
+
constant_values=0,
|
| 87 |
+
)
|
| 88 |
+
return image_np
|
| 89 |
+
|
| 90 |
+
def preprocess(
|
| 91 |
+
self,
|
| 92 |
+
medias: list[MediaInput],
|
| 93 |
+
return_tensors: Optional[Union[str, TensorType]] = None,
|
| 94 |
+
) -> BatchFeature:
|
| 95 |
+
"""
|
| 96 |
+
Preprocess a atom vision input (images) into model-ready tensors.
|
| 97 |
+
|
| 98 |
+
Args:
|
| 99 |
+
medias: List of MediaInput.
|
| 100 |
+
return_tensors: Desired output format ('pt', 'np', 'tf', or None).
|
| 101 |
+
|
| 102 |
+
Returns:
|
| 103 |
+
BatchFeature containing 'pixel_values' and 'grid_thws' tensors.
|
| 104 |
+
"""
|
| 105 |
+
if not isinstance(medias, list):
|
| 106 |
+
medias = [medias]
|
| 107 |
+
if medias:
|
| 108 |
+
pixel_values = []
|
| 109 |
+
for item in medias:
|
| 110 |
+
item = ensure_media_type(
|
| 111 |
+
item,
|
| 112 |
+
transparent_bg_config=self._transparent_bg_config,
|
| 113 |
+
transparent_bg_fill_stage=self._transparent_bg_fill_stage,
|
| 114 |
+
)
|
| 115 |
+
resize_config = self.get_resize_config(item)
|
| 116 |
+
new_width, new_height, pad_width, pad_height = resize_config[
|
| 117 |
+
'new_width'], resize_config['new_height'], resize_config[
|
| 118 |
+
'pad_width'], resize_config['pad_height']
|
| 119 |
+
if item['type'] == 'image':
|
| 120 |
+
image = item['image']
|
| 121 |
+
image_np = self.resize_image(image, new_width, new_height,
|
| 122 |
+
pad_width, pad_height)
|
| 123 |
+
pixel_values.append(np.expand_dims(image_np, axis=0))
|
| 124 |
+
else:
|
| 125 |
+
raise ValueError("Unsupported type: {}".format(
|
| 126 |
+
item['type']))
|
| 127 |
+
normalized_pixel_values = []
|
| 128 |
+
image_std_inv = 1.0 / np.array(self.media_proc_cfg['image_std'])
|
| 129 |
+
image_mean = np.array(self.media_proc_cfg['image_mean'])
|
| 130 |
+
for pixels in pixel_values:
|
| 131 |
+
pixels = normalize(pixels, image_mean, image_std_inv)
|
| 132 |
+
pixels_and_thw = navit_patchify(
|
| 133 |
+
pixels,
|
| 134 |
+
self.media_proc_cfg['patch_size'],
|
| 135 |
+
)
|
| 136 |
+
normalized_pixel_values.append(pixels_and_thw)
|
| 137 |
+
|
| 138 |
+
pixel_values = torch.cat([
|
| 139 |
+
_to_tensor(pixel_value['pixel_values'])
|
| 140 |
+
for pixel_value in normalized_pixel_values
|
| 141 |
+
])
|
| 142 |
+
grid_thws = torch.cat([
|
| 143 |
+
_to_tensor(pixel_value['grid_thw'],
|
| 144 |
+
dtype=torch.int64).unsqueeze(0)
|
| 145 |
+
for pixel_value in normalized_pixel_values
|
| 146 |
+
])
|
| 147 |
+
|
| 148 |
+
data = {
|
| 149 |
+
'pixel_values': pixel_values,
|
| 150 |
+
'grid_thws': grid_thws,
|
| 151 |
+
}
|
| 152 |
+
|
| 153 |
+
else:
|
| 154 |
+
data = {}
|
| 155 |
+
|
| 156 |
+
return BatchFeature(data=data, tensor_type=return_tensors)
|
| 157 |
+
|
| 158 |
+
def __repr__(self):
|
| 159 |
+
return f"KimiK3VisionProcessor(media_proc_cfg={self.media_proc_cfg})"
|
| 160 |
+
|
| 161 |
+
def to_dict(self) -> Dict[str, Any]:
|
| 162 |
+
output = super().to_dict()
|
| 163 |
+
output["media_proc_cfg"] = self.media_proc_cfg
|
| 164 |
+
if "media_processor" in output:
|
| 165 |
+
del output["media_processor"]
|
| 166 |
+
return output
|
| 167 |
+
|
| 168 |
+
@classmethod
|
| 169 |
+
def from_dict(cls, config_dict: Dict[str, Any], **kwargs):
|
| 170 |
+
config = config_dict.copy()
|
| 171 |
+
media_proc_cfg = config.pop("media_proc_cfg", {})
|
| 172 |
+
return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs)
|
| 173 |
+
|
| 174 |
+
def to_json_string(self):
|
| 175 |
+
dictionary = self.to_dict()
|
| 176 |
+
for key, value in dictionary.items():
|
| 177 |
+
if hasattr(value, 'tolist'):
|
| 178 |
+
dictionary[key] = value.tolist()
|
| 179 |
+
return json.dumps(dictionary, indent=2, sort_keys=True) + "\n"
|
media_utils.py
ADDED
|
@@ -0,0 +1,376 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import base64
|
| 2 |
+
import functools
|
| 3 |
+
import io
|
| 4 |
+
import math
|
| 5 |
+
from dataclasses import dataclass
|
| 6 |
+
from typing import Literal, TypedDict
|
| 7 |
+
|
| 8 |
+
import numpy as np
|
| 9 |
+
from PIL import Image
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
class ImageInput(TypedDict):
|
| 13 |
+
type: Literal['image']
|
| 14 |
+
image: Image.Image
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
MediaInput = ImageInput
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
@dataclass
|
| 21 |
+
class TransparentBgConfig:
|
| 22 |
+
"""The config of the transparent background."""
|
| 23 |
+
|
| 24 |
+
pattern: Literal["white", "black", "gray", "chessboard"] = "black"
|
| 25 |
+
"""The pattern of the transparent background."""
|
| 26 |
+
|
| 27 |
+
chessboard_square_size: int = 16
|
| 28 |
+
"""The size of the squares in the chessboard background."""
|
| 29 |
+
|
| 30 |
+
chessboard_square_on_top_left: bool = True
|
| 31 |
+
"""Whether to start the chessboard with a white square on the top left."""
|
| 32 |
+
|
| 33 |
+
chessboard_white_value: int = 255
|
| 34 |
+
"""The value of the white pixels in the background."""
|
| 35 |
+
|
| 36 |
+
chessboard_gray_value: int = 200
|
| 37 |
+
"""The value of the gray pixels in the background."""
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
@functools.lru_cache(maxsize=256)
|
| 41 |
+
def _create_chessboard_background(
|
| 42 |
+
height: int,
|
| 43 |
+
width: int,
|
| 44 |
+
square_size: int,
|
| 45 |
+
square_on_top_left: bool,
|
| 46 |
+
white_value: int,
|
| 47 |
+
gray_value: int,
|
| 48 |
+
) -> np.ndarray:
|
| 49 |
+
"""Create a chessboard background."""
|
| 50 |
+
bg = np.ones((height, width, 3), dtype=np.uint8) * white_value
|
| 51 |
+
for y in range(0, height, square_size):
|
| 52 |
+
for x in range(0, width, square_size):
|
| 53 |
+
if (y // square_size + x // square_size) % 2 == (
|
| 54 |
+
1 if square_on_top_left else 0):
|
| 55 |
+
bg[y:y + square_size, x:x + square_size] = gray_value
|
| 56 |
+
return bg
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def fill_transparent_bg_with(
|
| 60 |
+
image: Image.Image,
|
| 61 |
+
transparent_bg_config: TransparentBgConfig | None = None,
|
| 62 |
+
) -> Image.Image:
|
| 63 |
+
"""Composite a (possibly) transparent image onto a configured background.
|
| 64 |
+
|
| 65 |
+
When ``transparent_bg_config`` is ``None``, the image is simply converted
|
| 66 |
+
to RGB (preserving the historical behavior). Otherwise the alpha channel
|
| 67 |
+
is alpha-composited over a background generated according to the config.
|
| 68 |
+
"""
|
| 69 |
+
if transparent_bg_config is None:
|
| 70 |
+
return image.convert("RGB")
|
| 71 |
+
|
| 72 |
+
if image.mode == "RGB":
|
| 73 |
+
return image
|
| 74 |
+
|
| 75 |
+
has_alpha = "A" in image.getbands() or "transparency" in image.info
|
| 76 |
+
if not has_alpha:
|
| 77 |
+
return image.convert("RGB")
|
| 78 |
+
|
| 79 |
+
img = np.array(image.convert("RGBA"))
|
| 80 |
+
height, width = img.shape[:2]
|
| 81 |
+
bg_pattern = transparent_bg_config.pattern
|
| 82 |
+
if bg_pattern == "white":
|
| 83 |
+
bg = np.full((height, width, 3), 255, dtype=np.uint8)
|
| 84 |
+
elif bg_pattern == "black":
|
| 85 |
+
bg = np.zeros((height, width, 3), dtype=np.uint8)
|
| 86 |
+
elif bg_pattern == "gray":
|
| 87 |
+
bg = np.full((height, width, 3), 128, dtype=np.uint8)
|
| 88 |
+
elif bg_pattern == "chessboard":
|
| 89 |
+
bg = _create_chessboard_background(
|
| 90 |
+
height,
|
| 91 |
+
width,
|
| 92 |
+
transparent_bg_config.chessboard_square_size,
|
| 93 |
+
transparent_bg_config.chessboard_square_on_top_left,
|
| 94 |
+
transparent_bg_config.chessboard_white_value,
|
| 95 |
+
transparent_bg_config.chessboard_gray_value,
|
| 96 |
+
)
|
| 97 |
+
else:
|
| 98 |
+
raise ValueError(f"Invalid background pattern: {bg_pattern}")
|
| 99 |
+
|
| 100 |
+
alpha = img[:, :, 3]
|
| 101 |
+
img_rgb = img[:, :, :3]
|
| 102 |
+
alpha_normalized = alpha.astype(np.float32) / 255.0
|
| 103 |
+
alpha_3d = np.stack([alpha_normalized] * 3, axis=2)
|
| 104 |
+
result = alpha_3d * img_rgb + (1 - alpha_3d) * bg
|
| 105 |
+
result = result.astype(np.uint8)
|
| 106 |
+
return Image.fromarray(result)
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
def navit_resize_image(
|
| 110 |
+
width: int,
|
| 111 |
+
height: int,
|
| 112 |
+
patch_size: int,
|
| 113 |
+
merge_kernel_size: int,
|
| 114 |
+
in_patch_limit: int,
|
| 115 |
+
patch_limit_on_one_side: int,
|
| 116 |
+
fixed_output_tokens: int | None,
|
| 117 |
+
):
|
| 118 |
+
# Apply the patch limits.
|
| 119 |
+
s1 = math.sqrt(
|
| 120 |
+
in_patch_limit /
|
| 121 |
+
(max(1.0, width // patch_size) * max(1.0, height // patch_size)))
|
| 122 |
+
s2 = patch_limit_on_one_side * patch_size / width
|
| 123 |
+
s3 = patch_limit_on_one_side * patch_size / height
|
| 124 |
+
scale = min(1.0, s1, s2, s3)
|
| 125 |
+
new_w, new_h = max(1, int(width * scale)), max(1, int(height * scale))
|
| 126 |
+
new_w = min(new_w, patch_limit_on_one_side * patch_size)
|
| 127 |
+
new_h = min(new_h, patch_limit_on_one_side * patch_size)
|
| 128 |
+
|
| 129 |
+
# Calculate the padding to make the height and width divisible by the merge kernel size and patch size.
|
| 130 |
+
factor = merge_kernel_size * patch_size
|
| 131 |
+
|
| 132 |
+
pad_height = (factor - new_h % factor) % factor
|
| 133 |
+
pad_width = (factor - new_w % factor) % factor
|
| 134 |
+
|
| 135 |
+
if fixed_output_tokens is not None:
|
| 136 |
+
num_tokens = fixed_output_tokens
|
| 137 |
+
else:
|
| 138 |
+
# Calculate new dimensions after padding and patching
|
| 139 |
+
token_height = (new_h + pad_height) // factor
|
| 140 |
+
token_width = (new_w + pad_width) // factor
|
| 141 |
+
|
| 142 |
+
assert token_height * merge_kernel_size <= patch_limit_on_one_side, (
|
| 143 |
+
f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}"
|
| 144 |
+
)
|
| 145 |
+
assert token_width * merge_kernel_size <= patch_limit_on_one_side, (
|
| 146 |
+
f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}"
|
| 147 |
+
)
|
| 148 |
+
|
| 149 |
+
num_tokens = token_height * token_width
|
| 150 |
+
return {
|
| 151 |
+
"num_tokens": num_tokens,
|
| 152 |
+
"new_width": new_w,
|
| 153 |
+
"new_height": new_h,
|
| 154 |
+
"pad_width": pad_width,
|
| 155 |
+
"pad_height": pad_height,
|
| 156 |
+
"sampled_nframes": 1,
|
| 157 |
+
}
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
def _to_pil(
|
| 161 |
+
data: str | bytes | Image.Image,
|
| 162 |
+
transparent_bg_config: TransparentBgConfig | None = None,
|
| 163 |
+
to_rgb: bool = True,
|
| 164 |
+
) -> Image.Image:
|
| 165 |
+
"""Load an image and (optionally) composite its transparent background.
|
| 166 |
+
|
| 167 |
+
Args:
|
| 168 |
+
data: A PIL Image, a base64 ``data:`` URL, a file path, or raw bytes.
|
| 169 |
+
transparent_bg_config: The config used to fill the transparent
|
| 170 |
+
background. ``None`` keeps the historical behavior of converting
|
| 171 |
+
to RGB without compositing.
|
| 172 |
+
to_rgb: If ``False`` the image is returned as-is (the
|
| 173 |
+
``transparent_bg_config`` is ignored). The caller is then
|
| 174 |
+
expected to call :func:`fill_transparent_bg_with` later — e.g.
|
| 175 |
+
after a resize.
|
| 176 |
+
"""
|
| 177 |
+
if isinstance(data, Image.Image):
|
| 178 |
+
image = data
|
| 179 |
+
elif isinstance(data, str):
|
| 180 |
+
if data.startswith("data:"):
|
| 181 |
+
raw_base64 = data.split(",")[1]
|
| 182 |
+
image = Image.open(io.BytesIO(base64.b64decode(raw_base64)))
|
| 183 |
+
else:
|
| 184 |
+
image = Image.open(data)
|
| 185 |
+
elif isinstance(data, bytes):
|
| 186 |
+
image = Image.open(io.BytesIO(data))
|
| 187 |
+
else:
|
| 188 |
+
raise ValueError(f"Unsupported data type: {type(data)}")
|
| 189 |
+
|
| 190 |
+
if not to_rgb:
|
| 191 |
+
return image
|
| 192 |
+
|
| 193 |
+
return fill_transparent_bg_with(image, transparent_bg_config)
|
| 194 |
+
|
| 195 |
+
|
| 196 |
+
def ensure_media_type(
|
| 197 |
+
media: MediaInput,
|
| 198 |
+
transparent_bg_config: TransparentBgConfig | None = None,
|
| 199 |
+
transparent_bg_fill_stage: Literal["before_resize",
|
| 200 |
+
"after_resize"] = "before_resize",
|
| 201 |
+
) -> MediaInput:
|
| 202 |
+
if media['type'] == 'image':
|
| 203 |
+
media['image'] = _to_pil(
|
| 204 |
+
media['image'],
|
| 205 |
+
transparent_bg_config=transparent_bg_config,
|
| 206 |
+
to_rgb=transparent_bg_fill_stage == "before_resize",
|
| 207 |
+
)
|
| 208 |
+
return media
|
| 209 |
+
else:
|
| 210 |
+
raise ValueError(f"Unsupported media type: {media['type']}")
|
| 211 |
+
|
| 212 |
+
|
| 213 |
+
def image_to_np(
|
| 214 |
+
image: Image.Image,
|
| 215 |
+
resize_to: tuple[int, int] | None = None,
|
| 216 |
+
mode: str = "resize",
|
| 217 |
+
raise_error_for_ill_resize: bool = True,
|
| 218 |
+
transparent_bg_config: TransparentBgConfig | None = None,
|
| 219 |
+
transparent_bg_fill_stage: Literal["before_resize",
|
| 220 |
+
"after_resize"] = "before_resize",
|
| 221 |
+
) -> np.ndarray:
|
| 222 |
+
"""Convert an image to a numpy array.
|
| 223 |
+
|
| 224 |
+
Args:
|
| 225 |
+
content: The image to convert.
|
| 226 |
+
resize_to: The size to resize the image to.
|
| 227 |
+
mode: The mode to resize the image to.
|
| 228 |
+
raise_error_for_ill_resize: Whether to raise an error for ill-sized resize.
|
| 229 |
+
transparent_bg_config: The config of the transparent background. Only
|
| 230 |
+
used when ``transparent_bg_fill_stage == "after_resize"`` (the
|
| 231 |
+
caller is responsible for filling before resize otherwise).
|
| 232 |
+
transparent_bg_fill_stage: When to composite the transparent
|
| 233 |
+
background — before or after the resize step.
|
| 234 |
+
|
| 235 |
+
Returns:
|
| 236 |
+
A numpy array.
|
| 237 |
+
"""
|
| 238 |
+
assert isinstance(image, Image.Image), "image must be a PIL Image"
|
| 239 |
+
if resize_to is not None:
|
| 240 |
+
if mode == "resize":
|
| 241 |
+
image = image.resize(resize_to, resample=Image.Resampling.BICUBIC)
|
| 242 |
+
if transparent_bg_fill_stage == "after_resize":
|
| 243 |
+
image = fill_transparent_bg_with(image, transparent_bg_config)
|
| 244 |
+
|
| 245 |
+
elif mode == "rescale_and_pad_to_center":
|
| 246 |
+
scale = min(resize_to[0] / image.width,
|
| 247 |
+
resize_to[1] / image.height, 1.0)
|
| 248 |
+
new_width = round(image.width * scale)
|
| 249 |
+
new_height = round(image.height * scale)
|
| 250 |
+
if new_width == 0 or new_height == 0:
|
| 251 |
+
if raise_error_for_ill_resize:
|
| 252 |
+
raise ValueError(
|
| 253 |
+
f"Invalid resize to: {resize_to}, from image size: {image.size}"
|
| 254 |
+
)
|
| 255 |
+
else:
|
| 256 |
+
return np.zeros((resize_to[1], resize_to[0], 3),
|
| 257 |
+
dtype=np.uint8)
|
| 258 |
+
|
| 259 |
+
image = image.resize((new_width, new_height),
|
| 260 |
+
resample=Image.Resampling.BICUBIC)
|
| 261 |
+
if transparent_bg_fill_stage == "after_resize":
|
| 262 |
+
image = fill_transparent_bg_with(image, transparent_bg_config)
|
| 263 |
+
padding_left = (resize_to[0] - new_width) // 2
|
| 264 |
+
padding_right = resize_to[0] - new_width - padding_left
|
| 265 |
+
padding_top = (resize_to[1] - new_height) // 2
|
| 266 |
+
padding_bottom = resize_to[1] - new_height - padding_top
|
| 267 |
+
image = np.asarray(image)
|
| 268 |
+
image = np.pad(
|
| 269 |
+
image,
|
| 270 |
+
((padding_top, padding_bottom), (padding_left, padding_right),
|
| 271 |
+
(0, 0)),
|
| 272 |
+
mode="constant",
|
| 273 |
+
constant_values=0,
|
| 274 |
+
)
|
| 275 |
+
assert image.shape == (resize_to[1], resize_to[0], 3)
|
| 276 |
+
|
| 277 |
+
elif mode == "rescale_and_pad_to_rightbottom":
|
| 278 |
+
scale = min(resize_to[0] / image.width,
|
| 279 |
+
resize_to[1] / image.height, 1.0)
|
| 280 |
+
new_width = round(image.width * scale)
|
| 281 |
+
new_height = round(image.height * scale)
|
| 282 |
+
if new_width == 0 or new_height == 0:
|
| 283 |
+
if raise_error_for_ill_resize:
|
| 284 |
+
raise ValueError(
|
| 285 |
+
f"Invalid resize to: {resize_to}, from image size: {image.size}"
|
| 286 |
+
)
|
| 287 |
+
else:
|
| 288 |
+
return np.zeros((resize_to[1], resize_to[0], 3),
|
| 289 |
+
dtype=np.uint8)
|
| 290 |
+
|
| 291 |
+
image = image.resize((new_width, new_height),
|
| 292 |
+
resample=Image.Resampling.BICUBIC)
|
| 293 |
+
if transparent_bg_fill_stage == "after_resize":
|
| 294 |
+
image = fill_transparent_bg_with(image, transparent_bg_config)
|
| 295 |
+
padding_right = resize_to[0] - new_width
|
| 296 |
+
padding_bottom = resize_to[1] - new_height
|
| 297 |
+
image = np.asarray(image)
|
| 298 |
+
image = np.pad(
|
| 299 |
+
image,
|
| 300 |
+
((0, padding_bottom), (0, padding_right), (0, 0)),
|
| 301 |
+
mode="constant",
|
| 302 |
+
constant_values=0,
|
| 303 |
+
)
|
| 304 |
+
assert image.shape == (resize_to[1], resize_to[0], 3)
|
| 305 |
+
|
| 306 |
+
else:
|
| 307 |
+
raise ValueError(f"Invalid mode: {mode}")
|
| 308 |
+
|
| 309 |
+
if isinstance(image, Image.Image):
|
| 310 |
+
return np.asarray(image)
|
| 311 |
+
else:
|
| 312 |
+
return image
|
| 313 |
+
|
| 314 |
+
|
| 315 |
+
def navit_patchify(pixel_values: np.ndarray,
|
| 316 |
+
patch_size: int) -> dict[str, np.ndarray]:
|
| 317 |
+
"""Reshape the pixel values to a navit shape.
|
| 318 |
+
|
| 319 |
+
Args:
|
| 320 |
+
pixel_values: np.ndarray, shape (t, h, w, c)
|
| 321 |
+
patch_size: int
|
| 322 |
+
|
| 323 |
+
Returns:
|
| 324 |
+
dict[str, np.ndarray]
|
| 325 |
+
- patches: np.ndarray, shape (t * h//patch_size * w//patch_size, c, patch_size, patch_size)
|
| 326 |
+
- grid_thw: np.ndarray, (t, h//patch_size, w//patch_size)
|
| 327 |
+
"""
|
| 328 |
+
T, H, W, C = pixel_values.shape
|
| 329 |
+
assert C == 3, "pixel_values must have 3 channels"
|
| 330 |
+
|
| 331 |
+
patches = pixel_values.reshape(T, H // patch_size, patch_size,
|
| 332 |
+
W // patch_size, patch_size, C)
|
| 333 |
+
# (T, H//patch_size, W//patch_size, C, patch_size, patch_size)
|
| 334 |
+
patches = patches.transpose(0, 1, 3, 5, 2, 4)
|
| 335 |
+
patches = patches.reshape(-1, C, patch_size, patch_size)
|
| 336 |
+
grid_thw = np.array([T, H // patch_size, W // patch_size])
|
| 337 |
+
return {"pixel_values": patches, "grid_thw": grid_thw}
|
| 338 |
+
|
| 339 |
+
|
| 340 |
+
def normalize(x: np.ndarray,
|
| 341 |
+
mean,
|
| 342 |
+
std_inv,
|
| 343 |
+
pixels_dtype: np.dtype = np.float32) -> np.ndarray:
|
| 344 |
+
"""Normalize the image.
|
| 345 |
+
|
| 346 |
+
Args:
|
| 347 |
+
x: The image to normalize. The shape is (..., 3). The dtype is uint8. The range is [0, 255].
|
| 348 |
+
mean: The mean of the image.
|
| 349 |
+
std_inv: The inverse of the std of the image.
|
| 350 |
+
pixels_dtype: The dtype of the image.
|
| 351 |
+
Returns:
|
| 352 |
+
The normalized image. The shape is (..., 3). The dtype is determined by the pixels_dtype.
|
| 353 |
+
"""
|
| 354 |
+
x = (x / 255.0).astype(pixels_dtype)
|
| 355 |
+
x -= mean
|
| 356 |
+
x *= std_inv
|
| 357 |
+
return x
|
| 358 |
+
|
| 359 |
+
|
| 360 |
+
def _to_tensor(data, **kwargs):
|
| 361 |
+
import torch
|
| 362 |
+
|
| 363 |
+
if isinstance(data, np.ndarray):
|
| 364 |
+
return torch.from_numpy(data).to(**kwargs)
|
| 365 |
+
elif isinstance(data, torch.Tensor):
|
| 366 |
+
return data.to(**kwargs)
|
| 367 |
+
elif isinstance(data, list):
|
| 368 |
+
return [_to_tensor(item, **kwargs) for item in data]
|
| 369 |
+
elif isinstance(data, tuple):
|
| 370 |
+
return tuple(_to_tensor(item, **kwargs) for item in data)
|
| 371 |
+
elif isinstance(data, dict):
|
| 372 |
+
return {k: _to_tensor(v, **kwargs) for k, v in data.items()}
|
| 373 |
+
elif data is None:
|
| 374 |
+
return None
|
| 375 |
+
else:
|
| 376 |
+
raise ValueError(f"Unsupported data type: {type(data)}")
|
model-00001-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:ca3e1f926f0da291c3423ed554e2facfaa1d4121c6a343938f4fab9db2f930b3
|
| 3 |
+
size 2341216112
|
model-00002-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:fa4ff4394450e2c9efaea95091b0935ad87b146feb0f7a5aae08c921103390cc
|
| 3 |
+
size 16990911504
|
model-00003-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:9eb36f76e8df07c3c42f2ed7b2fbb82050557efbaf612b84b15104fc5e28ee8e
|
| 3 |
+
size 16990911504
|
model-00004-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:24c89436d95f03baef3dca9f62e8a0fddbb31ae17b260e025cb40048409b008f
|
| 3 |
+
size 16567501776
|
model-00005-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b600311c5d577340df7d35fb2b978735da595d289ae1aaa6c1f91654cf7fa629
|
| 3 |
+
size 16990911504
|
model-00006-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:8db3589e4ea063721ea09e1e13158008f146ad3a61eca208c65cf540cc53ae31
|
| 3 |
+
size 16990911504
|
model-00007-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a23d133889b2485931ea7b64e33c00dc93605f2ddb9593a66af2b86845749d13
|
| 3 |
+
size 16990911504
|
model-00008-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:0ee82aed1bbb5fbacbaccd7c6c20e8e65ea37ca0482d12cca6a73377abf35019
|
| 3 |
+
size 16567501776
|
model-00009-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:cd0d9c51c586a88b953b28299fd5147f5d172da674bf8f34773c2b6f66971301
|
| 3 |
+
size 16990911504
|
model-00010-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:7b8efeb0b4ef2a633f1d81cf5045d789c75a56a051d2d54ecac15c7654b1bbff
|
| 3 |
+
size 16990911504
|
model-00011-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:019d3c8ba61ec44dd4bc51a93c9b2bf6cc46ce8504950dda1d901d75447645c4
|
| 3 |
+
size 16990916912
|
model-00012-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:413ac73803d68d1ed9a65b53352a7f4c48b81a620d4ad827aa986214df09e968
|
| 3 |
+
size 16567507176
|
model-00013-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:25643aac8567945b5ec2a8794edc0c48d3efb3cded56ab69108e88b8575105fa
|
| 3 |
+
size 16990916912
|
model-00014-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:d92a11e29ce9bccec1a60d9ed2dbc38ea0a6f6f171fa5c3ef60c9e406f318edc
|
| 3 |
+
size 16990916912
|
model-00015-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:842cbe51489f73a25396755b6689e368276192e378b4536d6d7c6c83312c0e4a
|
| 3 |
+
size 16990916912
|
model-00016-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:635c8221c13ddaf596dc62054f5b08c379c0b81ae87f228cfce08d3cbbb384e8
|
| 3 |
+
size 16567507176
|
model-00017-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:8fd26354232e24ddf1504b4f51511345073ca71db809ca2923a06bb997c93942
|
| 3 |
+
size 16990916912
|
model-00018-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a6375a0416d791417627abc87b008c4bd835455fc9609dffa9952e27136fd678
|
| 3 |
+
size 16990916912
|
model-00019-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b93ccbac9888d939ca94659215d52ea4da53f636515e1bd17088097d7c01f7d5
|
| 3 |
+
size 16990916912
|
model-00020-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:9c211ec8bc7ce7c7a4148ebb2f074d6ec225e54f5fa611c046119c33a776b71a
|
| 3 |
+
size 16567507176
|
model-00021-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:fea2e7f59d76aaa8365055e7f7022623e680137afdb5a0f46a96a5cb52661a23
|
| 3 |
+
size 16990916912
|
model-00022-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:5b02b63f0783dc51ce33c5f9ed4deb8695206d6b77995e2a8ff985fd90cdd1b3
|
| 3 |
+
size 16990916912
|
model-00023-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b1e1b3eaf6d034c9fd33b8296ba723b876174aef5c8335269d9b96b586eb0621
|
| 3 |
+
size 16990916912
|
model-00024-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:7f23243580226c91605b15fe1fd8d7dd5b24d11942fea36c416969c0b15a362f
|
| 3 |
+
size 16567507176
|
model-00025-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:29be09005f802a7351bdec8ee77b42cecfd492e5ffc25f8c91bda5cedb02688a
|
| 3 |
+
size 16990916912
|
model-00026-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e474b6fc04c8b36cb23cc78f0fd3a00398f26619aa864dbd04e5859f2a39a6bc
|
| 3 |
+
size 16990916912
|
model-00027-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:eec4d4773a5e25c6f6edf6b664b2aaab050ac6d659da5fd46fffe2172c5c34f0
|
| 3 |
+
size 16990916912
|
model-00028-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:4ffe399002a4982c20fe09ce28762b45e3249fd1d76f8ad61a63b31ebbaa7ebb
|
| 3 |
+
size 16567507176
|
model-00029-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e384d165847c9dd92e64442d753a8ca38ecd8efda91862abdac99766c4415a47
|
| 3 |
+
size 16990916912
|
model-00030-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:07b8c5a517cc00b5e59fe1c21888ca86cd8fffb9504007406edb8913379d9a97
|
| 3 |
+
size 16990916912
|
model-00031-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:dcabfea9a0abd0b3e4b2e32807a2ec441d39ce6650988b6c5be657f59db285aa
|
| 3 |
+
size 16990916912
|
model-00032-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:271ba5da377eb639efa324c74070d71c9391c310375990847fd63a652c775096
|
| 3 |
+
size 16567507176
|
model-00033-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:85a0abaa9341b7fee090314bf40b3d6f05cb00c26c8438d762dff5445cfbb3cd
|
| 3 |
+
size 16990916912
|
model-00034-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:1f6a1cd3bc7cf3eefc08cc1be94517fd7c5d59969b5916232b6a065becf3ba00
|
| 3 |
+
size 16990916912
|
model-00035-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:835518f58c230eae24ffaa1d6a54746a845843f78b3963be574890cb3bafc600
|
| 3 |
+
size 16990916912
|
model-00036-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:03e4c43386682c8e07053a96eb841d13282396f0ebc872f80528b22273080825
|
| 3 |
+
size 16567507176
|
model-00037-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:5087848a65cde0681fa4077f5785e933314b37ebf25e31bc0d00860ab1f0ed8e
|
| 3 |
+
size 16990916912
|
model-00038-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:bbf7b2b58b3d703b48615551bc5709bde0b09f7324dd06e014aca397f90963d8
|
| 3 |
+
size 16990916912
|
model-00039-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:5c1ebb9e62a6697f94d931e27cd7a61d5b386300886b408303a5303823e36e15
|
| 3 |
+
size 16990916912
|
model-00040-of-000096.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:27cc21669dc0e3b1f550e2c2c66b1c9601ed79ddeb293e425b46c034b88deae6
|
| 3 |
+
size 16567507176
|