Instructions to use tiny-random/lfm3-omni with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use tiny-random/lfm3-omni with Transformers:
# pip install -U transformers accelerate # Load model directly from transformers import AutoModelForImageTextToText model = AutoModelForImageTextToText.from_pretrained("tiny-random/lfm3-omni", trust_remote_code=True, device_map="auto") - Notebooks
- Google Colab
- Kaggle
Upload folder using huggingface_hub
Browse files- .gitattributes +1 -0
- .meta/create_model.py +229 -0
- .meta/env.json +44 -0
- .meta/model.log +382 -0
- LICENSE +71 -0
- README.md +144 -0
- audio_tokenizer/config.json +36 -0
- audio_tokenizer/configuration_lfm3.py +132 -0
- audio_tokenizer/configuration_liquid_codec.py +152 -0
- audio_tokenizer/lfm3_conv.py +202 -0
- audio_tokenizer/model.safetensors +3 -0
- audio_tokenizer/modeling_lfm3.py +881 -0
- audio_tokenizer/modeling_liquid_codec.py +381 -0
- chat_template.jinja +296 -0
- config.json +207 -0
- configuration_lfm3.py +132 -0
- configuration_lfm3_omni.py +300 -0
- configuration_lfm3_vl.py +211 -0
- configuration_liquid_codec.py +152 -0
- feature_extraction_lfm3_omni.py +202 -0
- generation_config.json +19 -0
- image_processing_lfm3_vl.py +293 -0
- lfm3_conv.py +202 -0
- model.safetensors +3 -0
- modeling_lfm3.py +881 -0
- modeling_lfm3_omni.py +0 -0
- modeling_lfm3_vl.py +965 -0
- modeling_liquid_codec.py +381 -0
- processing_lfm3_omni.py +769 -0
- processing_lfm3_vl.py +248 -0
- processor_config.json +98 -0
- tokenizer.json +3 -0
- tokenizer_config.json +22 -0
- video_preprocessor_config.json +36 -0
- video_processing_lfm3_vl.py +396 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
tokenizer.json filter=lfs diff=lfs merge=lfs -text
|
.meta/create_model.py
ADDED
|
@@ -0,0 +1,229 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
def create_model(source_model_id: str, source_revision: str, save_folder: str) -> str:
|
| 2 |
+
import json
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
|
| 5 |
+
import torch
|
| 6 |
+
from huggingface_hub import hf_hub_download
|
| 7 |
+
from transformers import AutoConfig, AutoModel, AutoModelForImageTextToText, set_seed
|
| 8 |
+
|
| 9 |
+
# The bundled codec hardcodes the full-size backbone dims in lfm_config(); expose them as config fields
|
| 10 |
+
# (source defaults unchanged) and build a small codec through config values only.
|
| 11 |
+
codec_config_patch = """ initializer_range: float = 0.02
|
| 12 |
+
codec_intermediate: int = 2304
|
| 13 |
+
codec_num_layers: int = 8
|
| 14 |
+
codec_num_heads: int = 8
|
| 15 |
+
codec_num_kv_heads: int = 2
|
| 16 |
+
codec_head_dim: int = 64
|
| 17 |
+
"""
|
| 18 |
+
codec_doc_patch = """ initializer_range (`float`, *optional*, defaults to 0.02):
|
| 19 |
+
The standard deviation of the truncated normal initializer of the weights.
|
| 20 |
+
codec_intermediate (`int`, *optional*, defaults to 2304):
|
| 21 |
+
The feed-forward width of each tiny codec LFM3 layer (full-size codec uses 2,304).
|
| 22 |
+
codec_num_layers (`int`, *optional*, defaults to 8):
|
| 23 |
+
The number of LFM3 layers in each tiny codec encoder/decoder stack (full-size codec uses 8).
|
| 24 |
+
codec_num_heads (`int`, *optional*, defaults to 8):
|
| 25 |
+
The number of query attention heads per tiny codec LFM3 layer (full-size codec uses 8).
|
| 26 |
+
codec_num_kv_heads (`int`, *optional*, defaults to 2):
|
| 27 |
+
The number of key/value attention heads per tiny codec LFM3 layer (full-size codec uses 2).
|
| 28 |
+
codec_head_dim (`int`, *optional*, defaults to 64):
|
| 29 |
+
The attention head dimension of each tiny codec LFM3 layer (full-size codec uses 64).
|
| 30 |
+
"""
|
| 31 |
+
codec_lfm_patch = """ return _config_class("lfm3")(
|
| 32 |
+
# the backbones are not sub-configs, so the requested attention implementation is handed over here
|
| 33 |
+
attn_implementation=self._attn_implementation,
|
| 34 |
+
vocab_size=1,
|
| 35 |
+
pad_token_id=None,
|
| 36 |
+
bos_token_id=None,
|
| 37 |
+
eos_token_id=None,
|
| 38 |
+
hidden_size=self.encoder_dim,
|
| 39 |
+
intermediate_size=self.codec_intermediate,
|
| 40 |
+
num_hidden_layers=self.codec_num_layers,
|
| 41 |
+
num_attention_heads=self.codec_num_heads,
|
| 42 |
+
num_key_value_heads=self.codec_num_kv_heads,
|
| 43 |
+
head_dim=self.codec_head_dim,
|
| 44 |
+
max_position_embeddings=131_072,
|
| 45 |
+
norm_eps=1e-5,
|
| 46 |
+
conv_L_cache=7,
|
| 47 |
+
conv_bias=False,
|
| 48 |
+
use_sandwich_norm=True,
|
| 49 |
+
layer_types=[
|
| 50 |
+
attention_type if index in (1, self.codec_num_layers - 1) else "conv"
|
| 51 |
+
for index in range(self.codec_num_layers)
|
| 52 |
+
],
|
| 53 |
+
sliding_window=sliding_window,
|
| 54 |
+
rope_parameters={
|
| 55 |
+
attention_type: {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 1.0}
|
| 56 |
+
},
|
| 57 |
+
use_cache=False,
|
| 58 |
+
tie_word_embeddings=False,
|
| 59 |
+
)
|
| 60 |
+
"""
|
| 61 |
+
|
| 62 |
+
folder = Path(save_folder)
|
| 63 |
+
folder.mkdir(parents=True, exist_ok=True)
|
| 64 |
+
# The Omni implementation is not in a released Transformers; ship the pinned remote code for every component.
|
| 65 |
+
root_assets = (
|
| 66 |
+
"configuration_lfm3.py",
|
| 67 |
+
"configuration_lfm3_vl.py",
|
| 68 |
+
"modeling_lfm3.py",
|
| 69 |
+
"modeling_lfm3_vl.py",
|
| 70 |
+
"lfm3_conv.py",
|
| 71 |
+
"modeling_lfm3_omni.py",
|
| 72 |
+
"configuration_lfm3_omni.py",
|
| 73 |
+
"processing_lfm3_omni.py",
|
| 74 |
+
"processing_lfm3_vl.py",
|
| 75 |
+
"image_processing_lfm3_vl.py",
|
| 76 |
+
"video_processing_lfm3_vl.py",
|
| 77 |
+
"feature_extraction_lfm3_omni.py",
|
| 78 |
+
"modeling_liquid_codec.py",
|
| 79 |
+
"configuration_liquid_codec.py",
|
| 80 |
+
"chat_template.jinja",
|
| 81 |
+
"generation_config.json",
|
| 82 |
+
"processor_config.json",
|
| 83 |
+
"video_preprocessor_config.json",
|
| 84 |
+
"tokenizer.json",
|
| 85 |
+
"tokenizer_config.json",
|
| 86 |
+
"LICENSE",
|
| 87 |
+
)
|
| 88 |
+
downloaded = {
|
| 89 |
+
name: Path(hf_hub_download(source_model_id, name, revision=source_revision)).read_bytes()
|
| 90 |
+
for name in root_assets
|
| 91 |
+
}
|
| 92 |
+
for name, data in downloaded.items():
|
| 93 |
+
(folder / name).write_bytes(data)
|
| 94 |
+
|
| 95 |
+
# The bundled codec hardcodes the full-size backbone dims in lfm_config(); expose them as config fields
|
| 96 |
+
# (source defaults unchanged) and build a four-layer width-64 codec through config values only.
|
| 97 |
+
codec_config_source = downloaded["configuration_liquid_codec.py"].decode()
|
| 98 |
+
marker = " initializer_range: float = 0.02\n"
|
| 99 |
+
assert codec_config_source.count(marker) == 1
|
| 100 |
+
codec_config_source = codec_config_source.replace(marker, codec_config_patch, 1)
|
| 101 |
+
doc_marker = (
|
| 102 |
+
" initializer_range (`float`, *optional*, defaults to 0.02):\n"
|
| 103 |
+
" The standard deviation of the truncated normal initializer of the weights.\n"
|
| 104 |
+
' """\n'
|
| 105 |
+
)
|
| 106 |
+
assert codec_config_source.count(doc_marker) == 1
|
| 107 |
+
codec_config_source = codec_config_source.replace(
|
| 108 |
+
doc_marker,
|
| 109 |
+
codec_doc_patch + ' """\n',
|
| 110 |
+
1,
|
| 111 |
+
)
|
| 112 |
+
start = codec_config_source.index(' return _config_class("lfm3")(')
|
| 113 |
+
end = codec_config_source.index(" )\n", start) + len(" )\n")
|
| 114 |
+
block = codec_config_source[start:end]
|
| 115 |
+
codec_config_source = codec_config_source.replace(block, codec_lfm_patch, 1)
|
| 116 |
+
(folder / "configuration_liquid_codec.py").write_text(codec_config_source)
|
| 117 |
+
(folder / "modeling_liquid_codec.py").write_bytes(downloaded["modeling_liquid_codec.py"])
|
| 118 |
+
|
| 119 |
+
config = json.loads(Path(hf_hub_download(source_model_id, "config.json", revision=source_revision)).read_text())
|
| 120 |
+
# Text backbone: tiny-random/lfm3-vl geometry, source 4:1 Q/KV ratio with four KV heads for TP=4.
|
| 121 |
+
text_config = config["text_config"]
|
| 122 |
+
text_config.update(
|
| 123 |
+
hidden_size=64,
|
| 124 |
+
intermediate_size=128,
|
| 125 |
+
num_hidden_layers=10,
|
| 126 |
+
num_attention_heads=16,
|
| 127 |
+
num_key_value_heads=4,
|
| 128 |
+
head_dim=32,
|
| 129 |
+
max_position_embeddings=128000,
|
| 130 |
+
)
|
| 131 |
+
text_config["layer_types"] = text_config["layer_types"][:10]
|
| 132 |
+
# Vision tower: width 128 keeps the derived head dim 32 with the source 2:1 Q/KV ratio (4 Q, 2 KV heads);
|
| 133 |
+
# six layers = conv/conv prefix + two complete full/conv cycles.
|
| 134 |
+
vision_config = config["vision_config"]
|
| 135 |
+
vision_config.update(
|
| 136 |
+
hidden_size=128, intermediate_size=320, num_hidden_layers=6, num_attention_heads=4, num_key_value_heads=2
|
| 137 |
+
)
|
| 138 |
+
vision_config["layer_types"] = ["conv", "conv", "full_attention", "conv", "full_attention", "conv"]
|
| 139 |
+
config["projector_hidden_size"] = 128
|
| 140 |
+
# 16 kHz causal audio encoder: mel/stride/window contracts fixed, width 64 with two conv/conv/full cycles
|
| 141 |
+
# (two heads over one KV head keeps the derived head dim at 32 and the source 2:1 Q/KV ratio).
|
| 142 |
+
audio_encoder = config["audio_encoder_config"]
|
| 143 |
+
audio_encoder["hidden_size"] = 64
|
| 144 |
+
audio_encoder["sliding_window"] = 32
|
| 145 |
+
audio_encoder["conv_channels"] = 256
|
| 146 |
+
audio_text = audio_encoder["text_config"]
|
| 147 |
+
audio_text.update(
|
| 148 |
+
hidden_size=64,
|
| 149 |
+
intermediate_size=128,
|
| 150 |
+
num_hidden_layers=6,
|
| 151 |
+
num_attention_heads=2,
|
| 152 |
+
num_key_value_heads=1,
|
| 153 |
+
conv_L_cache=3,
|
| 154 |
+
)
|
| 155 |
+
audio_text["layer_types"] = ["conv", "conv", "full_attention", "conv", "conv", "full_attention"]
|
| 156 |
+
# Maskgit depth decoder: three bidirectional passes over the fixed 12 codebook slots; width 32 keeps head dim 32.
|
| 157 |
+
depth = config["depth_config"]
|
| 158 |
+
depth.update(
|
| 159 |
+
hidden_size=32,
|
| 160 |
+
intermediate_size=64,
|
| 161 |
+
num_hidden_layers=3,
|
| 162 |
+
num_attention_heads=2,
|
| 163 |
+
num_key_value_heads=1,
|
| 164 |
+
head_dim=32,
|
| 165 |
+
max_position_embeddings=12,
|
| 166 |
+
)
|
| 167 |
+
(folder / "config.json").write_text(json.dumps(config, indent=2) + "\n")
|
| 168 |
+
|
| 169 |
+
set_seed(42)
|
| 170 |
+
model = AutoModelForImageTextToText.from_config(
|
| 171 |
+
AutoConfig.from_pretrained(folder, trust_remote_code=True), trust_remote_code=True, dtype=torch.bfloat16
|
| 172 |
+
)
|
| 173 |
+
model.save_pretrained(folder)
|
| 174 |
+
representation = str(model)
|
| 175 |
+
|
| 176 |
+
# Bundled 24 kHz LiquidCodec in its audio_tokenizer/ component directory.
|
| 177 |
+
codec_dir = folder / "audio_tokenizer"
|
| 178 |
+
codec_dir.mkdir(parents=True, exist_ok=True)
|
| 179 |
+
codec_config = json.loads(
|
| 180 |
+
Path(hf_hub_download(source_model_id, "audio_tokenizer/config.json", revision=source_revision)).read_text()
|
| 181 |
+
)
|
| 182 |
+
codec_config.update(
|
| 183 |
+
encoder_dim=64,
|
| 184 |
+
encoder_conv_dim=32,
|
| 185 |
+
quantizer_dim=48,
|
| 186 |
+
codec_intermediate=128,
|
| 187 |
+
codec_num_layers=4,
|
| 188 |
+
codec_num_heads=2,
|
| 189 |
+
codec_num_kv_heads=1,
|
| 190 |
+
codec_head_dim=32,
|
| 191 |
+
)
|
| 192 |
+
# Remote code lives next to the component checkpoint so AutoConfig/AutoModel resolve from this subfolder.
|
| 193 |
+
codec_config["auto_map"] = {
|
| 194 |
+
"AutoConfig": "configuration_liquid_codec.LiquidCodecConfig",
|
| 195 |
+
"AutoModel": "modeling_liquid_codec.LiquidCodecModel",
|
| 196 |
+
}
|
| 197 |
+
(codec_dir / "config.json").write_text(json.dumps(codec_config, indent=2) + "\n")
|
| 198 |
+
for name in (
|
| 199 |
+
"modeling_liquid_codec.py",
|
| 200 |
+
"configuration_liquid_codec.py",
|
| 201 |
+
"modeling_lfm3.py",
|
| 202 |
+
"configuration_lfm3.py",
|
| 203 |
+
"lfm3_conv.py",
|
| 204 |
+
):
|
| 205 |
+
(codec_dir / name).write_bytes((folder / name).read_bytes())
|
| 206 |
+
set_seed(42)
|
| 207 |
+
codec = AutoModel.from_config(
|
| 208 |
+
AutoConfig.from_pretrained(codec_dir, trust_remote_code=True), trust_remote_code=True, dtype=torch.bfloat16
|
| 209 |
+
)
|
| 210 |
+
codec.save_pretrained(codec_dir)
|
| 211 |
+
|
| 212 |
+
# save_pretrained rewrites some files with this environment's metadata; restore pinned source bytes and the
|
| 213 |
+
# deliberate config edits so only intended fields differ from the source.
|
| 214 |
+
(folder / "config.json").write_text(json.dumps(config, indent=2) + "\n")
|
| 215 |
+
(folder / "generation_config.json").write_bytes(downloaded["generation_config.json"])
|
| 216 |
+
(codec_dir / "config.json").write_text(json.dumps(codec_config, indent=2) + "\n")
|
| 217 |
+
|
| 218 |
+
from safetensors import safe_open
|
| 219 |
+
|
| 220 |
+
with safe_open(codec_dir / "model.safetensors", framework="pt") as tensors:
|
| 221 |
+
assert tensors.get_tensor("decoder.istft.window").dtype == torch.float32
|
| 222 |
+
return representation
|
| 223 |
+
|
| 224 |
+
|
| 225 |
+
create_model(
|
| 226 |
+
source_model_id="LiquidAI/LFM3-Omni-2B",
|
| 227 |
+
source_revision="21d75123f86c1f7ffb4bef1ea202710c4e7ebe67",
|
| 228 |
+
save_folder="/tmp/lfm3_omni",
|
| 229 |
+
)
|
.meta/env.json
ADDED
|
@@ -0,0 +1,44 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"source_model": {
|
| 3 |
+
"name": "LiquidAI/LFM3-Omni-2B",
|
| 4 |
+
"commit_id": "21d75123f86c1f7ffb4bef1ea202710c4e7ebe67"
|
| 5 |
+
},
|
| 6 |
+
"environments": {
|
| 7 |
+
"creation": {
|
| 8 |
+
"python": "3.12.13",
|
| 9 |
+
"packages": {
|
| 10 |
+
"huggingface_hub": "2.3.0",
|
| 11 |
+
"numpy": "2.5.2",
|
| 12 |
+
"pillow": "12.3.0",
|
| 13 |
+
"safetensors": "0.8.0",
|
| 14 |
+
"tokenizers": "0.23.3",
|
| 15 |
+
"torch": "2.13.0+cu126",
|
| 16 |
+
"torchvision": "0.28.0+cu126",
|
| 17 |
+
"transformers": "5.22.0.dev0"
|
| 18 |
+
},
|
| 19 |
+
"cuda": "12.6",
|
| 20 |
+
"hardware": {
|
| 21 |
+
"gpu_model": "NVIDIA H20"
|
| 22 |
+
}
|
| 23 |
+
},
|
| 24 |
+
"transformers": {
|
| 25 |
+
"python": "3.12.13",
|
| 26 |
+
"packages": {
|
| 27 |
+
"huggingface_hub": "2.3.0",
|
| 28 |
+
"numpy": "2.5.2",
|
| 29 |
+
"pillow": "12.3.0",
|
| 30 |
+
"safetensors": "0.8.0",
|
| 31 |
+
"soundfile": "0.13.1",
|
| 32 |
+
"tokenizers": "0.23.3",
|
| 33 |
+
"torch": "2.13.0+cu126",
|
| 34 |
+
"torchvision": "0.28.0+cu126",
|
| 35 |
+
"transformers": "5.22.0.dev0",
|
| 36 |
+
"triton": "3.7.1"
|
| 37 |
+
},
|
| 38 |
+
"cuda": "12.6",
|
| 39 |
+
"hardware": {
|
| 40 |
+
"gpu_model": "NVIDIA H20"
|
| 41 |
+
}
|
| 42 |
+
}
|
| 43 |
+
}
|
| 44 |
+
}
|
.meta/model.log
ADDED
|
@@ -0,0 +1,382 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Lfm3OmniForConditionalGeneration(
|
| 2 |
+
(model): Lfm3OmniModel(
|
| 3 |
+
(language_model): Lfm3Model(
|
| 4 |
+
(embed_tokens): Embedding(128000, 64, padding_idx=124893)
|
| 5 |
+
(layers): ModuleList(
|
| 6 |
+
(0-1): 2 x Lfm3DecoderLayer(
|
| 7 |
+
(conv): Lfm3ShortConv(
|
| 8 |
+
(conv): Conv1d(64, 64, kernel_size=(7,), stride=(1,), padding=(6,), groups=64, bias=False)
|
| 9 |
+
(in_proj): Linear(in_features=64, out_features=192, bias=False)
|
| 10 |
+
(out_proj): Linear(in_features=64, out_features=64, bias=False)
|
| 11 |
+
)
|
| 12 |
+
(feed_forward): Lfm3MLP(
|
| 13 |
+
(w1): Linear(in_features=64, out_features=128, bias=False)
|
| 14 |
+
(w3): Linear(in_features=64, out_features=128, bias=False)
|
| 15 |
+
(w2): Linear(in_features=128, out_features=64, bias=False)
|
| 16 |
+
)
|
| 17 |
+
(operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 18 |
+
(ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 19 |
+
(post_operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 20 |
+
(post_ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 21 |
+
)
|
| 22 |
+
(2): Lfm3DecoderLayer(
|
| 23 |
+
(self_attn): Lfm3Attention(
|
| 24 |
+
(q_proj): Linear(in_features=64, out_features=512, bias=False)
|
| 25 |
+
(k_proj): Linear(in_features=64, out_features=128, bias=False)
|
| 26 |
+
(v_proj): Linear(in_features=64, out_features=128, bias=False)
|
| 27 |
+
(o_proj): Linear(in_features=512, out_features=64, bias=False)
|
| 28 |
+
(q_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 29 |
+
(k_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 30 |
+
(gate_proj): Linear(in_features=64, out_features=16, bias=False)
|
| 31 |
+
)
|
| 32 |
+
(feed_forward): Lfm3MLP(
|
| 33 |
+
(w1): Linear(in_features=64, out_features=128, bias=False)
|
| 34 |
+
(w3): Linear(in_features=64, out_features=128, bias=False)
|
| 35 |
+
(w2): Linear(in_features=128, out_features=64, bias=False)
|
| 36 |
+
)
|
| 37 |
+
(operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 38 |
+
(ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 39 |
+
(post_operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 40 |
+
(post_ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 41 |
+
)
|
| 42 |
+
(3): Lfm3DecoderLayer(
|
| 43 |
+
(conv): Lfm3ShortConv(
|
| 44 |
+
(conv): Conv1d(64, 64, kernel_size=(7,), stride=(1,), padding=(6,), groups=64, bias=False)
|
| 45 |
+
(in_proj): Linear(in_features=64, out_features=192, bias=False)
|
| 46 |
+
(out_proj): Linear(in_features=64, out_features=64, bias=False)
|
| 47 |
+
)
|
| 48 |
+
(feed_forward): Lfm3MLP(
|
| 49 |
+
(w1): Linear(in_features=64, out_features=128, bias=False)
|
| 50 |
+
(w3): Linear(in_features=64, out_features=128, bias=False)
|
| 51 |
+
(w2): Linear(in_features=128, out_features=64, bias=False)
|
| 52 |
+
)
|
| 53 |
+
(operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 54 |
+
(ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 55 |
+
(post_operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 56 |
+
(post_ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 57 |
+
)
|
| 58 |
+
(4): Lfm3DecoderLayer(
|
| 59 |
+
(self_attn): Lfm3Attention(
|
| 60 |
+
(q_proj): Linear(in_features=64, out_features=512, bias=False)
|
| 61 |
+
(k_proj): Linear(in_features=64, out_features=128, bias=False)
|
| 62 |
+
(v_proj): Linear(in_features=64, out_features=128, bias=False)
|
| 63 |
+
(o_proj): Linear(in_features=512, out_features=64, bias=False)
|
| 64 |
+
(q_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 65 |
+
(k_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 66 |
+
(gate_proj): Linear(in_features=64, out_features=16, bias=False)
|
| 67 |
+
)
|
| 68 |
+
(feed_forward): Lfm3MLP(
|
| 69 |
+
(w1): Linear(in_features=64, out_features=128, bias=False)
|
| 70 |
+
(w3): Linear(in_features=64, out_features=128, bias=False)
|
| 71 |
+
(w2): Linear(in_features=128, out_features=64, bias=False)
|
| 72 |
+
)
|
| 73 |
+
(operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 74 |
+
(ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 75 |
+
(post_operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 76 |
+
(post_ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 77 |
+
)
|
| 78 |
+
(5): Lfm3DecoderLayer(
|
| 79 |
+
(conv): Lfm3ShortConv(
|
| 80 |
+
(conv): Conv1d(64, 64, kernel_size=(7,), stride=(1,), padding=(6,), groups=64, bias=False)
|
| 81 |
+
(in_proj): Linear(in_features=64, out_features=192, bias=False)
|
| 82 |
+
(out_proj): Linear(in_features=64, out_features=64, bias=False)
|
| 83 |
+
)
|
| 84 |
+
(feed_forward): Lfm3MLP(
|
| 85 |
+
(w1): Linear(in_features=64, out_features=128, bias=False)
|
| 86 |
+
(w3): Linear(in_features=64, out_features=128, bias=False)
|
| 87 |
+
(w2): Linear(in_features=128, out_features=64, bias=False)
|
| 88 |
+
)
|
| 89 |
+
(operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 90 |
+
(ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 91 |
+
(post_operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 92 |
+
(post_ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 93 |
+
)
|
| 94 |
+
(6): Lfm3DecoderLayer(
|
| 95 |
+
(self_attn): Lfm3Attention(
|
| 96 |
+
(q_proj): Linear(in_features=64, out_features=512, bias=False)
|
| 97 |
+
(k_proj): Linear(in_features=64, out_features=128, bias=False)
|
| 98 |
+
(v_proj): Linear(in_features=64, out_features=128, bias=False)
|
| 99 |
+
(o_proj): Linear(in_features=512, out_features=64, bias=False)
|
| 100 |
+
(q_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 101 |
+
(k_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 102 |
+
(gate_proj): Linear(in_features=64, out_features=16, bias=False)
|
| 103 |
+
)
|
| 104 |
+
(feed_forward): Lfm3MLP(
|
| 105 |
+
(w1): Linear(in_features=64, out_features=128, bias=False)
|
| 106 |
+
(w3): Linear(in_features=64, out_features=128, bias=False)
|
| 107 |
+
(w2): Linear(in_features=128, out_features=64, bias=False)
|
| 108 |
+
)
|
| 109 |
+
(operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 110 |
+
(ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 111 |
+
(post_operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 112 |
+
(post_ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 113 |
+
)
|
| 114 |
+
(7): Lfm3DecoderLayer(
|
| 115 |
+
(conv): Lfm3ShortConv(
|
| 116 |
+
(conv): Conv1d(64, 64, kernel_size=(7,), stride=(1,), padding=(6,), groups=64, bias=False)
|
| 117 |
+
(in_proj): Linear(in_features=64, out_features=192, bias=False)
|
| 118 |
+
(out_proj): Linear(in_features=64, out_features=64, bias=False)
|
| 119 |
+
)
|
| 120 |
+
(feed_forward): Lfm3MLP(
|
| 121 |
+
(w1): Linear(in_features=64, out_features=128, bias=False)
|
| 122 |
+
(w3): Linear(in_features=64, out_features=128, bias=False)
|
| 123 |
+
(w2): Linear(in_features=128, out_features=64, bias=False)
|
| 124 |
+
)
|
| 125 |
+
(operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 126 |
+
(ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 127 |
+
(post_operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 128 |
+
(post_ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 129 |
+
)
|
| 130 |
+
(8): Lfm3DecoderLayer(
|
| 131 |
+
(self_attn): Lfm3Attention(
|
| 132 |
+
(q_proj): Linear(in_features=64, out_features=512, bias=False)
|
| 133 |
+
(k_proj): Linear(in_features=64, out_features=128, bias=False)
|
| 134 |
+
(v_proj): Linear(in_features=64, out_features=128, bias=False)
|
| 135 |
+
(o_proj): Linear(in_features=512, out_features=64, bias=False)
|
| 136 |
+
(q_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 137 |
+
(k_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 138 |
+
(gate_proj): Linear(in_features=64, out_features=16, bias=False)
|
| 139 |
+
)
|
| 140 |
+
(feed_forward): Lfm3MLP(
|
| 141 |
+
(w1): Linear(in_features=64, out_features=128, bias=False)
|
| 142 |
+
(w3): Linear(in_features=64, out_features=128, bias=False)
|
| 143 |
+
(w2): Linear(in_features=128, out_features=64, bias=False)
|
| 144 |
+
)
|
| 145 |
+
(operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 146 |
+
(ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 147 |
+
(post_operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 148 |
+
(post_ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 149 |
+
)
|
| 150 |
+
(9): Lfm3DecoderLayer(
|
| 151 |
+
(conv): Lfm3ShortConv(
|
| 152 |
+
(conv): Conv1d(64, 64, kernel_size=(7,), stride=(1,), padding=(6,), groups=64, bias=False)
|
| 153 |
+
(in_proj): Linear(in_features=64, out_features=192, bias=False)
|
| 154 |
+
(out_proj): Linear(in_features=64, out_features=64, bias=False)
|
| 155 |
+
)
|
| 156 |
+
(feed_forward): Lfm3MLP(
|
| 157 |
+
(w1): Linear(in_features=64, out_features=128, bias=False)
|
| 158 |
+
(w3): Linear(in_features=64, out_features=128, bias=False)
|
| 159 |
+
(w2): Linear(in_features=128, out_features=64, bias=False)
|
| 160 |
+
)
|
| 161 |
+
(operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 162 |
+
(ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 163 |
+
(post_operator_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 164 |
+
(post_ffn_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 165 |
+
)
|
| 166 |
+
)
|
| 167 |
+
(rotary_emb): Lfm3RotaryEmbedding()
|
| 168 |
+
(embedding_norm): Lfm3RMSNorm((64,), eps=1e-05)
|
| 169 |
+
)
|
| 170 |
+
(audio_encoder): Lfm3OmniEncoder(
|
| 171 |
+
(downsample): Lfm3OmniDownsampler(
|
| 172 |
+
(conv_in): Conv2d(1, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
|
| 173 |
+
(conv_mid): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=256)
|
| 174 |
+
(lin_mid): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))
|
| 175 |
+
(conv_end): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=256)
|
| 176 |
+
(lin_end): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))
|
| 177 |
+
(lin_out): Conv1d(4096, 64, kernel_size=(1,), stride=(1,))
|
| 178 |
+
(act): SiLU()
|
| 179 |
+
)
|
| 180 |
+
(lfm): Lfm2Model(
|
| 181 |
+
(embed_tokens): None
|
| 182 |
+
(layers): ModuleList(
|
| 183 |
+
(0-1): 2 x Lfm2DecoderLayer(
|
| 184 |
+
(conv): Lfm2ShortConv(
|
| 185 |
+
(conv): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(2,), groups=64, bias=False)
|
| 186 |
+
(in_proj): Linear(in_features=64, out_features=192, bias=False)
|
| 187 |
+
(out_proj): Linear(in_features=64, out_features=64, bias=False)
|
| 188 |
+
)
|
| 189 |
+
(feed_forward): Lfm2MLP(
|
| 190 |
+
(w1): Linear(in_features=64, out_features=256, bias=False)
|
| 191 |
+
(w3): Linear(in_features=64, out_features=256, bias=False)
|
| 192 |
+
(w2): Linear(in_features=256, out_features=64, bias=False)
|
| 193 |
+
)
|
| 194 |
+
(operator_norm): Lfm2RMSNorm((64,), eps=1e-05)
|
| 195 |
+
(ffn_norm): Lfm2RMSNorm((64,), eps=1e-05)
|
| 196 |
+
)
|
| 197 |
+
(2): Lfm2DecoderLayer(
|
| 198 |
+
(self_attn): Lfm2Attention(
|
| 199 |
+
(q_proj): Linear(in_features=64, out_features=64, bias=False)
|
| 200 |
+
(k_proj): Linear(in_features=64, out_features=32, bias=False)
|
| 201 |
+
(v_proj): Linear(in_features=64, out_features=32, bias=False)
|
| 202 |
+
(out_proj): Linear(in_features=64, out_features=64, bias=False)
|
| 203 |
+
(q_layernorm): Lfm2RMSNorm((32,), eps=1e-05)
|
| 204 |
+
(k_layernorm): Lfm2RMSNorm((32,), eps=1e-05)
|
| 205 |
+
)
|
| 206 |
+
(feed_forward): Lfm2MLP(
|
| 207 |
+
(w1): Linear(in_features=64, out_features=256, bias=False)
|
| 208 |
+
(w3): Linear(in_features=64, out_features=256, bias=False)
|
| 209 |
+
(w2): Linear(in_features=256, out_features=64, bias=False)
|
| 210 |
+
)
|
| 211 |
+
(operator_norm): Lfm2RMSNorm((64,), eps=1e-05)
|
| 212 |
+
(ffn_norm): Lfm2RMSNorm((64,), eps=1e-05)
|
| 213 |
+
)
|
| 214 |
+
(3-4): 2 x Lfm2DecoderLayer(
|
| 215 |
+
(conv): Lfm2ShortConv(
|
| 216 |
+
(conv): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(2,), groups=64, bias=False)
|
| 217 |
+
(in_proj): Linear(in_features=64, out_features=192, bias=False)
|
| 218 |
+
(out_proj): Linear(in_features=64, out_features=64, bias=False)
|
| 219 |
+
)
|
| 220 |
+
(feed_forward): Lfm2MLP(
|
| 221 |
+
(w1): Linear(in_features=64, out_features=256, bias=False)
|
| 222 |
+
(w3): Linear(in_features=64, out_features=256, bias=False)
|
| 223 |
+
(w2): Linear(in_features=256, out_features=64, bias=False)
|
| 224 |
+
)
|
| 225 |
+
(operator_norm): Lfm2RMSNorm((64,), eps=1e-05)
|
| 226 |
+
(ffn_norm): Lfm2RMSNorm((64,), eps=1e-05)
|
| 227 |
+
)
|
| 228 |
+
(5): Lfm2DecoderLayer(
|
| 229 |
+
(self_attn): Lfm2Attention(
|
| 230 |
+
(q_proj): Linear(in_features=64, out_features=64, bias=False)
|
| 231 |
+
(k_proj): Linear(in_features=64, out_features=32, bias=False)
|
| 232 |
+
(v_proj): Linear(in_features=64, out_features=32, bias=False)
|
| 233 |
+
(out_proj): Linear(in_features=64, out_features=64, bias=False)
|
| 234 |
+
(q_layernorm): Lfm2RMSNorm((32,), eps=1e-05)
|
| 235 |
+
(k_layernorm): Lfm2RMSNorm((32,), eps=1e-05)
|
| 236 |
+
)
|
| 237 |
+
(feed_forward): Lfm2MLP(
|
| 238 |
+
(w1): Linear(in_features=64, out_features=256, bias=False)
|
| 239 |
+
(w3): Linear(in_features=64, out_features=256, bias=False)
|
| 240 |
+
(w2): Linear(in_features=256, out_features=64, bias=False)
|
| 241 |
+
)
|
| 242 |
+
(operator_norm): Lfm2RMSNorm((64,), eps=1e-05)
|
| 243 |
+
(ffn_norm): Lfm2RMSNorm((64,), eps=1e-05)
|
| 244 |
+
)
|
| 245 |
+
)
|
| 246 |
+
(rotary_emb): Lfm2RotaryEmbedding()
|
| 247 |
+
(embedding_norm): Lfm2RMSNorm((64,), eps=1e-05)
|
| 248 |
+
)
|
| 249 |
+
)
|
| 250 |
+
(audio_projector): Lfm3OmniProjector(
|
| 251 |
+
(layer_norm): LayerNorm((64,), eps=1e-05, elementwise_affine=True, bias=True)
|
| 252 |
+
(linear_1): Linear(in_features=64, out_features=64, bias=True)
|
| 253 |
+
(act): GELU(approximate='none')
|
| 254 |
+
(linear_2): Linear(in_features=64, out_features=64, bias=True)
|
| 255 |
+
)
|
| 256 |
+
(codebook_embedding): Embedding(24204, 64)
|
| 257 |
+
(depth_decoder): Lfm3OmniDepthDecoder(
|
| 258 |
+
(depth_linear): Linear(in_features=64, out_features=384, bias=True)
|
| 259 |
+
(model): Lfm3Model(
|
| 260 |
+
(embed_tokens): None
|
| 261 |
+
(layers): ModuleList(
|
| 262 |
+
(0-2): 3 x Lfm3DecoderLayer(
|
| 263 |
+
(self_attn): Lfm3Attention(
|
| 264 |
+
(q_proj): Linear(in_features=32, out_features=64, bias=False)
|
| 265 |
+
(k_proj): Linear(in_features=32, out_features=32, bias=False)
|
| 266 |
+
(v_proj): Linear(in_features=32, out_features=32, bias=False)
|
| 267 |
+
(o_proj): Linear(in_features=64, out_features=32, bias=False)
|
| 268 |
+
(q_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 269 |
+
(k_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 270 |
+
(gate_proj): Linear(in_features=32, out_features=2, bias=False)
|
| 271 |
+
)
|
| 272 |
+
(feed_forward): Lfm3MLP(
|
| 273 |
+
(w1): Linear(in_features=32, out_features=64, bias=False)
|
| 274 |
+
(w3): Linear(in_features=32, out_features=64, bias=False)
|
| 275 |
+
(w2): Linear(in_features=64, out_features=32, bias=False)
|
| 276 |
+
)
|
| 277 |
+
(operator_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 278 |
+
(ffn_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 279 |
+
)
|
| 280 |
+
)
|
| 281 |
+
(rotary_emb): Lfm3RotaryEmbedding()
|
| 282 |
+
(embedding_norm): Lfm3RMSNorm((32,), eps=1e-05)
|
| 283 |
+
)
|
| 284 |
+
(depth_embeddings): ModuleList(
|
| 285 |
+
(0-11): 12 x Lfm3OmniCodebookEmbedding(
|
| 286 |
+
(embedding): Embedding(2017, 32)
|
| 287 |
+
(embedding_norm): Lfm3OmniRMSNorm()
|
| 288 |
+
)
|
| 289 |
+
)
|
| 290 |
+
)
|
| 291 |
+
(vision_tower): Lfm3VlVisionModel(
|
| 292 |
+
(patchifier): Linear(in_features=1536, out_features=128, bias=False)
|
| 293 |
+
(embedding_norm): Lfm3VlRMSNorm((128,), eps=1e-05)
|
| 294 |
+
(rotary_pos_emb): Lfm3VlVisionRotaryEmbedding()
|
| 295 |
+
(layers): ModuleList(
|
| 296 |
+
(0-1): 2 x Lfm3VlVisionLayer(
|
| 297 |
+
(conv): Lfm3VlVisionConv2d(
|
| 298 |
+
(in_proj): Linear(in_features=128, out_features=384, bias=False)
|
| 299 |
+
(spatial_conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=128, bias=False)
|
| 300 |
+
(out_proj): Linear(in_features=128, out_features=128, bias=False)
|
| 301 |
+
)
|
| 302 |
+
(feed_forward): Lfm3VlMLP(
|
| 303 |
+
(w1): Linear(in_features=128, out_features=320, bias=False)
|
| 304 |
+
(w3): Linear(in_features=128, out_features=320, bias=False)
|
| 305 |
+
(w2): Linear(in_features=320, out_features=128, bias=False)
|
| 306 |
+
)
|
| 307 |
+
(operator_norm): Lfm3VlRMSNorm((128,), eps=1e-05)
|
| 308 |
+
(ffn_norm): Lfm3VlRMSNorm((128,), eps=1e-05)
|
| 309 |
+
)
|
| 310 |
+
(2): Lfm3VlVisionLayer(
|
| 311 |
+
(self_attn): Lfm3VlVisionAttention(
|
| 312 |
+
(q_proj): Linear(in_features=128, out_features=128, bias=False)
|
| 313 |
+
(k_proj): Linear(in_features=128, out_features=64, bias=False)
|
| 314 |
+
(v_proj): Linear(in_features=128, out_features=64, bias=False)
|
| 315 |
+
(out_proj): Linear(in_features=128, out_features=128, bias=False)
|
| 316 |
+
(q_layernorm): Lfm3VlRMSNorm((32,), eps=1e-05)
|
| 317 |
+
(k_layernorm): Lfm3VlRMSNorm((32,), eps=1e-05)
|
| 318 |
+
)
|
| 319 |
+
(feed_forward): Lfm3VlMLP(
|
| 320 |
+
(w1): Linear(in_features=128, out_features=320, bias=False)
|
| 321 |
+
(w3): Linear(in_features=128, out_features=320, bias=False)
|
| 322 |
+
(w2): Linear(in_features=320, out_features=128, bias=False)
|
| 323 |
+
)
|
| 324 |
+
(operator_norm): Lfm3VlRMSNorm((128,), eps=1e-05)
|
| 325 |
+
(ffn_norm): Lfm3VlRMSNorm((128,), eps=1e-05)
|
| 326 |
+
)
|
| 327 |
+
(3): Lfm3VlVisionLayer(
|
| 328 |
+
(conv): Lfm3VlVisionConv2d(
|
| 329 |
+
(in_proj): Linear(in_features=128, out_features=384, bias=False)
|
| 330 |
+
(spatial_conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=128, bias=False)
|
| 331 |
+
(out_proj): Linear(in_features=128, out_features=128, bias=False)
|
| 332 |
+
)
|
| 333 |
+
(feed_forward): Lfm3VlMLP(
|
| 334 |
+
(w1): Linear(in_features=128, out_features=320, bias=False)
|
| 335 |
+
(w3): Linear(in_features=128, out_features=320, bias=False)
|
| 336 |
+
(w2): Linear(in_features=320, out_features=128, bias=False)
|
| 337 |
+
)
|
| 338 |
+
(operator_norm): Lfm3VlRMSNorm((128,), eps=1e-05)
|
| 339 |
+
(ffn_norm): Lfm3VlRMSNorm((128,), eps=1e-05)
|
| 340 |
+
)
|
| 341 |
+
(4): Lfm3VlVisionLayer(
|
| 342 |
+
(self_attn): Lfm3VlVisionAttention(
|
| 343 |
+
(q_proj): Linear(in_features=128, out_features=128, bias=False)
|
| 344 |
+
(k_proj): Linear(in_features=128, out_features=64, bias=False)
|
| 345 |
+
(v_proj): Linear(in_features=128, out_features=64, bias=False)
|
| 346 |
+
(out_proj): Linear(in_features=128, out_features=128, bias=False)
|
| 347 |
+
(q_layernorm): Lfm3VlRMSNorm((32,), eps=1e-05)
|
| 348 |
+
(k_layernorm): Lfm3VlRMSNorm((32,), eps=1e-05)
|
| 349 |
+
)
|
| 350 |
+
(feed_forward): Lfm3VlMLP(
|
| 351 |
+
(w1): Linear(in_features=128, out_features=320, bias=False)
|
| 352 |
+
(w3): Linear(in_features=128, out_features=320, bias=False)
|
| 353 |
+
(w2): Linear(in_features=320, out_features=128, bias=False)
|
| 354 |
+
)
|
| 355 |
+
(operator_norm): Lfm3VlRMSNorm((128,), eps=1e-05)
|
| 356 |
+
(ffn_norm): Lfm3VlRMSNorm((128,), eps=1e-05)
|
| 357 |
+
)
|
| 358 |
+
(5): Lfm3VlVisionLayer(
|
| 359 |
+
(conv): Lfm3VlVisionConv2d(
|
| 360 |
+
(in_proj): Linear(in_features=128, out_features=384, bias=False)
|
| 361 |
+
(spatial_conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=128, bias=False)
|
| 362 |
+
(out_proj): Linear(in_features=128, out_features=128, bias=False)
|
| 363 |
+
)
|
| 364 |
+
(feed_forward): Lfm3VlMLP(
|
| 365 |
+
(w1): Linear(in_features=128, out_features=320, bias=False)
|
| 366 |
+
(w3): Linear(in_features=128, out_features=320, bias=False)
|
| 367 |
+
(w2): Linear(in_features=320, out_features=128, bias=False)
|
| 368 |
+
)
|
| 369 |
+
(operator_norm): Lfm3VlRMSNorm((128,), eps=1e-05)
|
| 370 |
+
(ffn_norm): Lfm3VlRMSNorm((128,), eps=1e-05)
|
| 371 |
+
)
|
| 372 |
+
)
|
| 373 |
+
)
|
| 374 |
+
(multi_modal_projector): Lfm3OmniMultiModalProjector(
|
| 375 |
+
(layer_norm): LayerNorm((128,), eps=1e-06, elementwise_affine=True, bias=True)
|
| 376 |
+
(linear_1): Linear(in_features=512, out_features=128, bias=True)
|
| 377 |
+
(act): GELUActivation()
|
| 378 |
+
(linear_2): Linear(in_features=128, out_features=64, bias=True)
|
| 379 |
+
)
|
| 380 |
+
)
|
| 381 |
+
(lm_head): Linear(in_features=64, out_features=128000, bias=False)
|
| 382 |
+
)
|
LICENSE
ADDED
|
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
LFM Open License v1.0
|
| 2 |
+
|
| 3 |
+
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
| 4 |
+
|
| 5 |
+
1. Definitions.
|
| 6 |
+
|
| 7 |
+
"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by this document.
|
| 8 |
+
|
| 9 |
+
"Licensor" shall mean Liquid AI, Inc.
|
| 10 |
+
|
| 11 |
+
"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
|
| 12 |
+
|
| 13 |
+
"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
|
| 14 |
+
|
| 15 |
+
"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
|
| 16 |
+
|
| 17 |
+
"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
|
| 18 |
+
|
| 19 |
+
"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work.
|
| 20 |
+
|
| 21 |
+
"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
|
| 22 |
+
|
| 23 |
+
"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
|
| 24 |
+
|
| 25 |
+
"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
|
| 26 |
+
|
| 27 |
+
"Commercial Use" shall mean any use of the Work for direct or indirect commercial advantage or monetary compensation.
|
| 28 |
+
|
| 29 |
+
"Qualified Non-Profit Organization" shall mean a Legal Entity that is organized and operated exclusively for religious, charitable, scientific, testing for public safety, literary, or educational purposes, and which is exempt from federal income tax under Section 501(c)(3) of the United States Internal Revenue Code of 1986, as amended, or any equivalent non-profit or charitable organization in a foreign jurisdiction.
|
| 30 |
+
|
| 31 |
+
"Non-Commercial or Research Purposes" shall mean purposes that do not involve any use of the Work or a Derivative Work for Commercial Use.
|
| 32 |
+
|
| 33 |
+
"Threshold" shall mean annual revenue of 10 million United States dollars ($10,000,000) or more.
|
| 34 |
+
|
| 35 |
+
2. Grant of Copyright License. Subject to the terms and conditions of this License, including the Commercial Use limitation set forth in Section 5, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
|
| 36 |
+
|
| 37 |
+
3. Grant of Patent License. Subject to the terms and conditions of this License, including the Commercial Use limitation set forth in Section 5, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
|
| 38 |
+
|
| 39 |
+
4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
|
| 40 |
+
|
| 41 |
+
(a) You must give any other recipients of the Work or Derivative Works a copy of this License; and
|
| 42 |
+
|
| 43 |
+
(b) You must cause any modified files to carry prominent notices stating that You changed the files; and
|
| 44 |
+
|
| 45 |
+
(c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and
|
| 46 |
+
|
| 47 |
+
(d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License.
|
| 48 |
+
|
| 49 |
+
You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License.
|
| 50 |
+
|
| 51 |
+
5. Commercial Use Limitation.
|
| 52 |
+
|
| 53 |
+
(a) The rights granted under this License for Commercial Use are conditioned upon You or Your Legal Entity not exceeding the Threshold.
|
| 54 |
+
|
| 55 |
+
(b) Any Commercial Use of the Work or a Derivative Work by a Legal Entity that exceeds the Threshold is not licensed under this Agreement.
|
| 56 |
+
|
| 57 |
+
(c) The Threshold shall not apply to a Qualified Non-Profit Organization's use of the Work or a Derivative Work for Non-Commercial or Research Purposes.
|
| 58 |
+
|
| 59 |
+
6. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
|
| 60 |
+
|
| 61 |
+
7. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except for the reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
|
| 62 |
+
|
| 63 |
+
8. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
|
| 64 |
+
|
| 65 |
+
9. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
|
| 66 |
+
|
| 67 |
+
10. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
|
| 68 |
+
|
| 69 |
+
11. Termination. This License will terminate automatically and immediately if You fail to comply with any of its terms and conditions. Upon termination, You must cease all use of the Work and any Derivative Works and delete all copies in Your possession.
|
| 70 |
+
|
| 71 |
+
END OF TERMS AND CONDITIONS
|
README.md
ADDED
|
@@ -0,0 +1,144 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
library_name: transformers
|
| 3 |
+
pipeline_tag: any-to-any
|
| 4 |
+
base_model:
|
| 5 |
+
- LiquidAI/LFM3-Omni-2B
|
| 6 |
+
---
|
| 7 |
+
|
| 8 |
+
This is a randomly initialized tiny model based on [LiquidAI/LFM3-Omni-2B](https://huggingface.co/LiquidAI/LFM3-Omni-2B) for testing and debugging.
|
| 9 |
+
|
| 10 |
+
| Feature | Source | This repo |
|
| 11 |
+
|---|---|---|
|
| 12 |
+
| Checkpoint tensor elements | 2.24B BF16, 1.28K F32 | 13.91M BF16, 1.28K F32 |
|
| 13 |
+
| Safetensors weight files | 4.48 GB (2 files) | 27.88 MB (2 files) |
|
| 14 |
+
| Text layers / cycle | 24 / two-conv prefix then full/conv/sliding/conv | 10 / same prefix and two complete cycles |
|
| 15 |
+
| Text hidden / MLP width | 2,048 / 7,680 | 64 / 128 |
|
| 16 |
+
| Text Q / KV heads / head dim | 16 / 4 / 128 | 16 / 4 / 32 |
|
| 17 |
+
| Audio encoder layers / width | 17 / 768 | 6 / 64 (two conv/conv/full cycles) |
|
| 18 |
+
| Audio mel bins / stride / window | 128 / 8 / 32 frames | Same |
|
| 19 |
+
| Depth decoder layers / width | 3 / 1,024 | 3 / 32 |
|
| 20 |
+
| Codebooks / codebook size / maskgit | 12 / 2,017 / [4, 4, 4] | Same |
|
| 21 |
+
| Vision layers / width / Q-KV heads | 14 / 1,024 / 16-8 | 6 / 128 / 4-2 |
|
| 22 |
+
| Patch / temporal patch / merge | 16 / 2 / 2 | Same |
|
| 23 |
+
| Projector MLP hidden / downsample | 4,096 / 2 | 128 / 2 |
|
| 24 |
+
| Codec sample rate / frame samples | 24,000 / 1,920 | Same |
|
| 25 |
+
| Codec STFT n_fft / hop / frame stack | 1,280 / 320 / 6 | Same |
|
| 26 |
+
| Codec FSQ levels / codebooks | [8, 7, 6, 6] / 12 | Same |
|
| 27 |
+
| Codec LFM3 layers / width | 8 + 8 / 512 | 4 + 4 / 64 |
|
| 28 |
+
| Vocabulary rows / tying | 128,000 / tied | Same |
|
| 29 |
+
| Stored precision | BF16; codec ISTFT window FP32 | Same |
|
| 30 |
+
|
| 31 |
+
Text and vision geometry follows [tiny-random/lfm3-vl](https://huggingface.co/tiny-random/lfm3-vl). The pinned upstream implementation is included for Transformers loading, including the separate `audio_tokenizer/` LiquidCodec component; the only runtime change lets the codec LFM3 backbone dimensions come from its config (source defaults preserved). The 16 kHz feature extractor and 24 kHz codec STFT/FSQ contracts are unchanged. This repo provides a Transformers example with image, speech in and speech out; no serving-backend validation is claimed.
|
| 32 |
+
|
| 33 |
+
| File path | Size |
|
| 34 |
+
|------|------|
|
| 35 |
+
| audio_tokenizer/model.safetensors | 1.0MB |
|
| 36 |
+
| model.safetensors | 26.9MB |
|
| 37 |
+
|
| 38 |
+
### Example usage
|
| 39 |
+
|
| 40 |
+
#### Transformers (✅ verified)
|
| 41 |
+
|
| 42 |
+
Load the included pinned upstream implementation with `trust_remote_code=True`. The example sends a 16 kHz spoken question together with an image and four predecoded video frames, answers as text, then synthesizes 24 kHz speech through the bundled LiquidCodec. The Triton short-conv path requires CUDA.
|
| 43 |
+
|
| 44 |
+
```python
|
| 45 |
+
def run_transformers(model_id: str):
|
| 46 |
+
import numpy as np
|
| 47 |
+
import torch
|
| 48 |
+
from PIL import Image
|
| 49 |
+
from transformers import AutoModelForImageTextToText, AutoProcessor
|
| 50 |
+
|
| 51 |
+
model_id = str(model_id)
|
| 52 |
+
processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
|
| 53 |
+
model = (
|
| 54 |
+
AutoModelForImageTextToText.from_pretrained(
|
| 55 |
+
model_id, trust_remote_code=True, dtype=torch.bfloat16, attn_implementation="sdpa"
|
| 56 |
+
)
|
| 57 |
+
.eval()
|
| 58 |
+
.cuda()
|
| 59 |
+
)
|
| 60 |
+
codec = processor.audio_tokenizer.to(device=model.device).eval()
|
| 61 |
+
|
| 62 |
+
image = Image.fromarray(np.random.randint(0, 255, (256, 256, 3), dtype=np.uint8), "RGB")
|
| 63 |
+
frames = np.stack(
|
| 64 |
+
[
|
| 65 |
+
np.asarray(Image.new("RGB", (256, 256), color))
|
| 66 |
+
for color in ((220, 30, 30), (30, 220, 30), (30, 30, 220), (220, 220, 30))
|
| 67 |
+
]
|
| 68 |
+
)
|
| 69 |
+
sample_rate = 16000
|
| 70 |
+
speech = (np.sin(2 * np.pi * 440 * np.arange(sample_rate) / sample_rate) * 0.1).astype(np.float32)
|
| 71 |
+
conversation = [
|
| 72 |
+
{"role": "system", "content": "Respond in text."},
|
| 73 |
+
{
|
| 74 |
+
"role": "user",
|
| 75 |
+
"content": [
|
| 76 |
+
{"type": "image", "image": image},
|
| 77 |
+
{"type": "audio", "audio": speech},
|
| 78 |
+
{"type": "video", "video": frames},
|
| 79 |
+
{"type": "text", "text": "Transcribe the question, describe the image and the video."},
|
| 80 |
+
],
|
| 81 |
+
},
|
| 82 |
+
]
|
| 83 |
+
inputs = processor.apply_chat_template(
|
| 84 |
+
conversation,
|
| 85 |
+
add_generation_prompt=True,
|
| 86 |
+
tokenize=True,
|
| 87 |
+
return_dict=True,
|
| 88 |
+
return_tensors="pt",
|
| 89 |
+
sampling_rate=sample_rate,
|
| 90 |
+
processor_kwargs={"video_metadata": [{"total_num_frames": 4, "fps": 2, "frames_indices": [0, 1, 2, 3]}]},
|
| 91 |
+
).to(model.device)
|
| 92 |
+
assert (
|
| 93 |
+
inputs["input_features"].shape[-1] > 0
|
| 94 |
+
and inputs["pixel_values"].shape[0] > 0
|
| 95 |
+
and inputs["pixel_values_videos"].shape[0] > 0
|
| 96 |
+
)
|
| 97 |
+
with torch.no_grad():
|
| 98 |
+
answered = model.generate(**inputs, max_new_tokens=32, generation_mode="text")
|
| 99 |
+
text = processor.decode(
|
| 100 |
+
answered.sequences[0, inputs["input_ids"].shape[1] : answered.sequence_lengths[0]],
|
| 101 |
+
skip_special_tokens=True,
|
| 102 |
+
)
|
| 103 |
+
print(text)
|
| 104 |
+
|
| 105 |
+
tts_conversation = [
|
| 106 |
+
{"role": "system", "content": "Perform TTS in American English using a female voice."},
|
| 107 |
+
{"role": "user", "content": "The quick brown fox jumps over the lazy dog."},
|
| 108 |
+
]
|
| 109 |
+
tts_inputs = processor.apply_chat_template(
|
| 110 |
+
tts_conversation,
|
| 111 |
+
add_generation_prompt=True,
|
| 112 |
+
tokenize=True,
|
| 113 |
+
return_dict=True,
|
| 114 |
+
return_tensors="pt",
|
| 115 |
+
).to(model.device)
|
| 116 |
+
with torch.no_grad():
|
| 117 |
+
spoken = model.generate(
|
| 118 |
+
**tts_inputs,
|
| 119 |
+
max_new_tokens=64,
|
| 120 |
+
generation_mode="audio",
|
| 121 |
+
audio_temperature=0.4,
|
| 122 |
+
audio_repetition_penalty=1.2,
|
| 123 |
+
)
|
| 124 |
+
assert spoken.audio_codes.shape[-1] > 0
|
| 125 |
+
waveform = codec.decode(spoken.audio_codes[:1, :, : spoken.audio_lengths[0]]).audio_values[0]
|
| 126 |
+
assert waveform.numel() > 0 and torch.isfinite(waveform).all()
|
| 127 |
+
print(waveform.shape)
|
| 128 |
+
|
| 129 |
+
|
| 130 |
+
if __name__ == "__main__":
|
| 131 |
+
run_transformers(model_id="scratch/lfm3_omni")
|
| 132 |
+
```
|
| 133 |
+
|
| 134 |
+
### Creating this repo
|
| 135 |
+
|
| 136 |
+
- Code: [.meta/create_model.py](.meta/create_model.py)
|
| 137 |
+
- Model structure: [.meta/model.log](.meta/model.log)
|
| 138 |
+
- Environment: [.meta/env.json](.meta/env.json)
|
| 139 |
+
|
| 140 |
+
### Change log
|
| 141 |
+
|
| 142 |
+
- 2026-11-03: Initial version. Text backbone follows tiny-random/lfm3-vl (residual 64 / MLP 128, ten layers with the two-conv prefix and two full/conv/sliding/conv cycles, 16 Q / 4 KV heads at head dim 32); vision keeps the source 2:1 Q/KV ratio at width 128 (4 Q / 2 KV heads, head dim 32) with six conv/full-cadence layers. The 16 kHz causal audio encoder keeps 128 mel bins, stride-8 convolutions, the 32-frame window and a two-group conv/conv/full schedule at width 64; the 12-codebook maskgit depth decoder keeps the three full-attention passes and 2,017-entry codebooks at width 32. The bundled 24 kHz LiquidCodec keeps its STFT contract (1,280 FFT, hop 320, frame stack 6), 12 FSQ codebooks [8, 7, 6, 6] and FP32 synthesis window, with a four-conv/attention-layer LFM3 stack at width 64 via a pinned, localized runtime config patch.
|
| 143 |
+
|
| 144 |
+
*This repo is generated by GPT-6.1-Sol.*
|
audio_tokenizer/config.json
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"LiquidCodecModel"
|
| 4 |
+
],
|
| 5 |
+
"decoder_window": 30,
|
| 6 |
+
"dtype": "bfloat16",
|
| 7 |
+
"encoder_conv_dim": 32,
|
| 8 |
+
"encoder_dim": 64,
|
| 9 |
+
"frame_samples": 1920,
|
| 10 |
+
"fsq_levels": [
|
| 11 |
+
8,
|
| 12 |
+
7,
|
| 13 |
+
6,
|
| 14 |
+
6
|
| 15 |
+
],
|
| 16 |
+
"initializer_range": 0.02,
|
| 17 |
+
"model_type": "liquid_codec",
|
| 18 |
+
"n_q": 12,
|
| 19 |
+
"quantizer_dim": 48,
|
| 20 |
+
"sample_rate": 24000,
|
| 21 |
+
"sampling_rate": 24000,
|
| 22 |
+
"stft_frame_stack": 6,
|
| 23 |
+
"stft_hop_length": 320,
|
| 24 |
+
"stft_n_fft": 1280,
|
| 25 |
+
"stft_win_length": 1280,
|
| 26 |
+
"transformers_version": "5.19.0.dev0",
|
| 27 |
+
"codec_intermediate": 128,
|
| 28 |
+
"codec_num_layers": 4,
|
| 29 |
+
"codec_num_heads": 2,
|
| 30 |
+
"codec_num_kv_heads": 1,
|
| 31 |
+
"codec_head_dim": 32,
|
| 32 |
+
"auto_map": {
|
| 33 |
+
"AutoConfig": "configuration_liquid_codec.LiquidCodecConfig",
|
| 34 |
+
"AutoModel": "modeling_liquid_codec.LiquidCodecModel"
|
| 35 |
+
}
|
| 36 |
+
}
|
audio_tokenizer/configuration_lfm3.py
ADDED
|
@@ -0,0 +1,132 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
from huggingface_hub.dataclasses import strict
|
| 15 |
+
|
| 16 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 17 |
+
from transformers.utils import auto_docstring
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
@auto_docstring(checkpoint="LiquidAI/LFM3-1.7B-Preview")
|
| 21 |
+
@strict
|
| 22 |
+
class Lfm3Config(PreTrainedConfig):
|
| 23 |
+
r"""
|
| 24 |
+
rope_parameters (`dict[str, RopeParameters]`, *optional*):
|
| 25 |
+
RoPE parameters keyed by attention layer type, for example `{"full_attention": {"rope_type": "default",
|
| 26 |
+
"rope_theta": 1000000.0, "partial_rotary_factor": 0.5}, "sliding_attention": {"rope_type": "default",
|
| 27 |
+
"rope_theta": 10000.0}}`. `conv` layers do not use RoPE and have no entry. Full-attention layers rotate the
|
| 28 |
+
first `partial_rotary_factor * head_dim` dimensions and leave the rest unrotated.
|
| 29 |
+
conv_bias (`bool`, *optional*, defaults to `False`):
|
| 30 |
+
Whether to use bias in the conv layers.
|
| 31 |
+
conv_L_cache (`int`, *optional*, defaults to 7):
|
| 32 |
+
L_cache dim in the conv layers.
|
| 33 |
+
layer_types (`list[str]`, *optional*):
|
| 34 |
+
Per-layer operator type, each one of `"full_attention"`, `"sliding_attention"` or `"conv"`. Defaults to the
|
| 35 |
+
LFM3 pattern of a full-attention layer every fourth layer starting from the third, a sliding-attention layer
|
| 36 |
+
in between, and gated short convolutions elsewhere.
|
| 37 |
+
sliding_window (`int`, *optional*, defaults to 256):
|
| 38 |
+
Sliding-window size of the `sliding_attention` layers.
|
| 39 |
+
use_sandwich_norm (`bool`, *optional*, defaults to `True`):
|
| 40 |
+
Whether to apply an RMSNorm to the operator and feed-forward outputs before they enter the residual stream.
|
| 41 |
+
|
| 42 |
+
```python
|
| 43 |
+
>>> from transformers import Lfm3Model, Lfm3Config
|
| 44 |
+
|
| 45 |
+
>>> # Initializing a LFM3 model
|
| 46 |
+
>>> configuration = Lfm3Config()
|
| 47 |
+
|
| 48 |
+
>>> # Initializing a model from the LFM3-1.7B style configuration
|
| 49 |
+
>>> model = Lfm3Model(configuration)
|
| 50 |
+
|
| 51 |
+
>>> # Accessing the model configuration
|
| 52 |
+
>>> configuration = model.config
|
| 53 |
+
```
|
| 54 |
+
"""
|
| 55 |
+
|
| 56 |
+
model_type = "lfm3"
|
| 57 |
+
keys_to_ignore_at_inference = ["past_key_values"]
|
| 58 |
+
default_theta = 1000000.0
|
| 59 |
+
|
| 60 |
+
vocab_size: int = 128_000
|
| 61 |
+
hidden_size: int = 2048
|
| 62 |
+
intermediate_size: int = 7680
|
| 63 |
+
num_hidden_layers: int = 24
|
| 64 |
+
num_attention_heads: int = 16
|
| 65 |
+
num_key_value_heads: int = 4
|
| 66 |
+
max_position_embeddings: int = 128_000
|
| 67 |
+
initializer_range: float = 0.02
|
| 68 |
+
norm_eps: float = 0.00001
|
| 69 |
+
use_cache: bool = True
|
| 70 |
+
pad_token_id: int | None = 124_893
|
| 71 |
+
bos_token_id: int | None = 124_894
|
| 72 |
+
eos_token_id: int | list[int] | None = 124_900
|
| 73 |
+
tie_word_embeddings: bool = True
|
| 74 |
+
rope_parameters: dict | None = None
|
| 75 |
+
conv_bias: bool = False
|
| 76 |
+
conv_L_cache: int = 7
|
| 77 |
+
layer_types: list[str] | None = None
|
| 78 |
+
_attention_layer_types = ("full_attention", "sliding_attention")
|
| 79 |
+
head_dim: int | None = None
|
| 80 |
+
sliding_window: int | None = 256
|
| 81 |
+
use_sandwich_norm: bool = True
|
| 82 |
+
|
| 83 |
+
def __post_init__(self, **kwargs):
|
| 84 |
+
if self.head_dim is None:
|
| 85 |
+
self.head_dim = self.hidden_size // self.num_attention_heads
|
| 86 |
+
if self.rope_parameters is None:
|
| 87 |
+
self.rope_parameters = {
|
| 88 |
+
"full_attention": {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 0.5},
|
| 89 |
+
"sliding_attention": {"rope_type": "default", "rope_theta": 10_000.0},
|
| 90 |
+
}
|
| 91 |
+
if self.layer_types is None:
|
| 92 |
+
self.layer_types = [
|
| 93 |
+
"full_attention" if i % 4 == 2 else "sliding_attention" if i % 4 == 0 and i > 0 else "conv"
|
| 94 |
+
for i in range(self.num_hidden_layers)
|
| 95 |
+
]
|
| 96 |
+
|
| 97 |
+
self.tie_word_embeddings = kwargs.pop("tie_embedding", self.tie_word_embeddings)
|
| 98 |
+
super().__post_init__(**kwargs)
|
| 99 |
+
|
| 100 |
+
def validate_architecture(self):
|
| 101 |
+
"""Part of ``@strict``-powered validation."""
|
| 102 |
+
if len(self.layer_types) != self.num_hidden_layers:
|
| 103 |
+
raise ValueError(
|
| 104 |
+
f"`layer_types` has {len(self.layer_types)} entries but `num_hidden_layers` is {self.num_hidden_layers}."
|
| 105 |
+
)
|
| 106 |
+
allowed_layer_types = (*self._attention_layer_types, "conv")
|
| 107 |
+
invalid_layer_types = [layer_type for layer_type in self.layer_types if layer_type not in allowed_layer_types]
|
| 108 |
+
if invalid_layer_types:
|
| 109 |
+
raise ValueError(f"`layer_types` entries must be one of {allowed_layer_types}; got {invalid_layer_types}.")
|
| 110 |
+
if "sliding_attention" in self.layer_types and self.sliding_window is None:
|
| 111 |
+
raise ValueError("`layer_types` has `sliding_attention` layers but `sliding_window` is None.")
|
| 112 |
+
missing_rope = [
|
| 113 |
+
layer_type
|
| 114 |
+
for layer_type in self._attention_layer_types
|
| 115 |
+
if layer_type in self.layer_types and layer_type not in self.rope_parameters
|
| 116 |
+
]
|
| 117 |
+
if missing_rope:
|
| 118 |
+
raise ValueError(f"`rope_parameters` has no entry for the {missing_rope} layers in `layer_types`.")
|
| 119 |
+
for layer_type, rope_params in self.rope_parameters.items():
|
| 120 |
+
if not isinstance(rope_params, dict):
|
| 121 |
+
continue
|
| 122 |
+
if (
|
| 123 |
+
rope_params.get("rope_type", "default") in ("dynamic", "longrope")
|
| 124 |
+
and rope_params.get("partial_rotary_factor", 1.0) < 1
|
| 125 |
+
):
|
| 126 |
+
raise ValueError(
|
| 127 |
+
f"`{rope_params['rope_type']}` RoPE recomputes its frequencies during the forward pass over the "
|
| 128 |
+
f"rotated dimension only and cannot keep LFM3's partial rotation on `{layer_type}` layers."
|
| 129 |
+
)
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
__all__ = ["Lfm3Config"]
|
audio_tokenizer/configuration_liquid_codec.py
ADDED
|
@@ -0,0 +1,152 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
import math
|
| 16 |
+
|
| 17 |
+
from huggingface_hub.dataclasses import strict
|
| 18 |
+
|
| 19 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 20 |
+
from transformers.models.auto import CONFIG_MAPPING
|
| 21 |
+
from transformers.utils import auto_docstring
|
| 22 |
+
|
| 23 |
+
from .configuration_lfm3 import Lfm3Config
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def _config_class(model_type):
|
| 27 |
+
# the LFM3 configs ship with the checkpoint instead of the transformers registry
|
| 28 |
+
remote = {"lfm3": Lfm3Config}
|
| 29 |
+
return remote.get(model_type) or CONFIG_MAPPING[model_type]
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
@auto_docstring(checkpoint="LiquidAI/LFM3-Audio-2B-Preview")
|
| 33 |
+
@strict
|
| 34 |
+
class LiquidCodecConfig(PreTrainedConfig):
|
| 35 |
+
r"""
|
| 36 |
+
Configuration of the [`LiquidCodecModel`] bundled with LFM3-Audio-2B and LFM3-Omni-2B in their
|
| 37 |
+
`audio_tokenizer/` directory.
|
| 38 |
+
|
| 39 |
+
encoder_dim (`int`, *optional*, defaults to 512):
|
| 40 |
+
The width of the encoder and decoder LFM3 models.
|
| 41 |
+
encoder_conv_dim (`int`, *optional*, defaults to 256):
|
| 42 |
+
The number of channels of the STFT front end.
|
| 43 |
+
stft_n_fft (`int`, *optional*, defaults to 1280):
|
| 44 |
+
The FFT size of the analysis and synthesis transforms.
|
| 45 |
+
stft_hop_length (`int`, *optional*, defaults to 320):
|
| 46 |
+
The hop of the analysis and synthesis transforms.
|
| 47 |
+
stft_win_length (`int`, *optional*, defaults to 1280):
|
| 48 |
+
The window length of the analysis and synthesis transforms.
|
| 49 |
+
stft_frame_stack (`int`, *optional*, defaults to 6):
|
| 50 |
+
The number of STFT frames per codec frame.
|
| 51 |
+
quantizer_dim (`int`, *optional*, defaults to 256):
|
| 52 |
+
The width of the latent handed to the quantizer.
|
| 53 |
+
fsq_levels (`list[int]`, *optional*, defaults to `[8, 7, 6, 6]`):
|
| 54 |
+
The number of levels of each finite-scalar-quantization dimension; their product is the codebook size.
|
| 55 |
+
n_q (`int`, *optional*, defaults to 12):
|
| 56 |
+
The number of codebooks.
|
| 57 |
+
decoder_window (`int`, *optional*, defaults to 30):
|
| 58 |
+
The sliding attention window of the decoder LFM3 model, in 75 Hz frames.
|
| 59 |
+
sampling_rate (`int`, *optional*, defaults to 24000):
|
| 60 |
+
The sampling rate of the audio.
|
| 61 |
+
initializer_range (`float`, *optional*, defaults to 0.02):
|
| 62 |
+
The standard deviation of the truncated normal initializer of the weights.
|
| 63 |
+
codec_intermediate (`int`, *optional*, defaults to 2304):
|
| 64 |
+
The feed-forward width of each tiny codec LFM3 layer (full-size codec uses 2,304).
|
| 65 |
+
codec_num_layers (`int`, *optional*, defaults to 8):
|
| 66 |
+
The number of LFM3 layers in each tiny codec encoder/decoder stack (full-size codec uses 8).
|
| 67 |
+
codec_num_heads (`int`, *optional*, defaults to 8):
|
| 68 |
+
The number of query attention heads per tiny codec LFM3 layer (full-size codec uses 8).
|
| 69 |
+
codec_num_kv_heads (`int`, *optional*, defaults to 2):
|
| 70 |
+
The number of key/value attention heads per tiny codec LFM3 layer (full-size codec uses 2).
|
| 71 |
+
codec_head_dim (`int`, *optional*, defaults to 64):
|
| 72 |
+
The attention head dimension of each tiny codec LFM3 layer (full-size codec uses 64).
|
| 73 |
+
"""
|
| 74 |
+
|
| 75 |
+
model_type = "liquid_codec"
|
| 76 |
+
|
| 77 |
+
encoder_dim: int = 512
|
| 78 |
+
encoder_conv_dim: int = 256
|
| 79 |
+
stft_n_fft: int = 1280
|
| 80 |
+
stft_hop_length: int = 320
|
| 81 |
+
stft_win_length: int = 1280
|
| 82 |
+
stft_frame_stack: int = 6
|
| 83 |
+
quantizer_dim: int = 256
|
| 84 |
+
fsq_levels: list[int] | None = None
|
| 85 |
+
n_q: int = 12
|
| 86 |
+
decoder_window: int = 30
|
| 87 |
+
sampling_rate: int = 24_000
|
| 88 |
+
initializer_range: float = 0.02
|
| 89 |
+
codec_intermediate: int = 2304
|
| 90 |
+
codec_num_layers: int = 8
|
| 91 |
+
codec_num_heads: int = 8
|
| 92 |
+
codec_num_kv_heads: int = 2
|
| 93 |
+
codec_head_dim: int = 64
|
| 94 |
+
|
| 95 |
+
def __post_init__(self, **kwargs):
|
| 96 |
+
# the released configs nest the fields under `model` and dump the two LFM3 sub-configs
|
| 97 |
+
dropped = ("encoder_lfm_config", "decoder_lfm_config", "n_output_samples")
|
| 98 |
+
for key, value in kwargs.pop("model", {}).items():
|
| 99 |
+
if key in dropped:
|
| 100 |
+
continue
|
| 101 |
+
if key not in self.__dataclass_fields__:
|
| 102 |
+
raise ValueError(f"Unknown Liquid codec setting `{key}`")
|
| 103 |
+
setattr(self, key, value)
|
| 104 |
+
for key in dropped:
|
| 105 |
+
kwargs.pop(key, None)
|
| 106 |
+
if self.fsq_levels is None:
|
| 107 |
+
self.fsq_levels = [8, 7, 6, 6]
|
| 108 |
+
self.fsq_levels = list(self.fsq_levels)
|
| 109 |
+
super().__post_init__(**kwargs)
|
| 110 |
+
|
| 111 |
+
@property
|
| 112 |
+
def codebook_size(self) -> int:
|
| 113 |
+
return math.prod(self.fsq_levels)
|
| 114 |
+
|
| 115 |
+
@property
|
| 116 |
+
def frame_size(self) -> int:
|
| 117 |
+
return self.stft_hop_length * self.stft_frame_stack
|
| 118 |
+
|
| 119 |
+
def lfm_config(self, sliding_window: int | None = None) -> PreTrainedConfig:
|
| 120 |
+
attention_type = "full_attention" if sliding_window is None else "sliding_attention"
|
| 121 |
+
return _config_class("lfm3")(
|
| 122 |
+
# the backbones are not sub-configs, so the requested attention implementation is handed over here
|
| 123 |
+
attn_implementation=self._attn_implementation,
|
| 124 |
+
vocab_size=1,
|
| 125 |
+
pad_token_id=None,
|
| 126 |
+
bos_token_id=None,
|
| 127 |
+
eos_token_id=None,
|
| 128 |
+
hidden_size=self.encoder_dim,
|
| 129 |
+
intermediate_size=self.codec_intermediate,
|
| 130 |
+
num_hidden_layers=self.codec_num_layers,
|
| 131 |
+
num_attention_heads=self.codec_num_heads,
|
| 132 |
+
num_key_value_heads=self.codec_num_kv_heads,
|
| 133 |
+
head_dim=self.codec_head_dim,
|
| 134 |
+
max_position_embeddings=131_072,
|
| 135 |
+
norm_eps=1e-5,
|
| 136 |
+
conv_L_cache=7,
|
| 137 |
+
conv_bias=False,
|
| 138 |
+
use_sandwich_norm=True,
|
| 139 |
+
layer_types=[
|
| 140 |
+
attention_type if index in (1, self.codec_num_layers - 1) else "conv"
|
| 141 |
+
for index in range(self.codec_num_layers)
|
| 142 |
+
],
|
| 143 |
+
sliding_window=sliding_window,
|
| 144 |
+
rope_parameters={
|
| 145 |
+
attention_type: {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 1.0}
|
| 146 |
+
},
|
| 147 |
+
use_cache=False,
|
| 148 |
+
tie_word_embeddings=False,
|
| 149 |
+
)
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
__all__ = ["LiquidCodecConfig"]
|
audio_tokenizer/lfm3_conv.py
ADDED
|
@@ -0,0 +1,202 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
"""Inference kernels for LFM3's width-7 depthwise convolution."""
|
| 15 |
+
|
| 16 |
+
import torch
|
| 17 |
+
|
| 18 |
+
from transformers.utils import is_triton_available
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
if is_triton_available():
|
| 22 |
+
import triton
|
| 23 |
+
import triton.language as tl
|
| 24 |
+
|
| 25 |
+
# the length and strides are runtime arguments, or every prompt length would compile its own kernel
|
| 26 |
+
@triton.jit(do_not_specialize=["T", "X0", "X1", "X2", "S0", "S1", "ST"])
|
| 27 |
+
def _prefill(
|
| 28 |
+
X,
|
| 29 |
+
W,
|
| 30 |
+
BIAS,
|
| 31 |
+
SEQ,
|
| 32 |
+
Y,
|
| 33 |
+
T,
|
| 34 |
+
D: tl.constexpr,
|
| 35 |
+
X0,
|
| 36 |
+
X1,
|
| 37 |
+
X2,
|
| 38 |
+
S0,
|
| 39 |
+
S1,
|
| 40 |
+
ST,
|
| 41 |
+
HAS_BIAS: tl.constexpr,
|
| 42 |
+
PACKED: tl.constexpr,
|
| 43 |
+
BLOCK: tl.constexpr,
|
| 44 |
+
):
|
| 45 |
+
# int64 offsets: batch * channels * length can exceed 2**31 elements
|
| 46 |
+
b, d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
|
| 47 |
+
t = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
|
| 48 |
+
acc = tl.full((BLOCK,), 0, tl.float32)
|
| 49 |
+
if PACKED:
|
| 50 |
+
doc = tl.load(SEQ + b * S0 + (t - T + ST) * S1, (t < T) & (t >= T - ST), other=-1)
|
| 51 |
+
for k in tl.static_range(7):
|
| 52 |
+
pos = t + k - 6
|
| 53 |
+
valid = (t < T) & (pos >= 0)
|
| 54 |
+
if PACKED:
|
| 55 |
+
prev = tl.load(SEQ + b * S0 + (pos - T + ST) * S1, valid & (pos >= T - ST), other=-1)
|
| 56 |
+
valid = valid & (doc == prev)
|
| 57 |
+
x = tl.load(X + b * X0 + d * X1 + pos * X2, valid, other=0).to(tl.float32)
|
| 58 |
+
w = tl.load(W + d * 7 + k).to(tl.float32)
|
| 59 |
+
acc += x * w
|
| 60 |
+
if HAS_BIAS:
|
| 61 |
+
acc += tl.load(BIAS + d).to(tl.float32)
|
| 62 |
+
tl.store(Y + (b * D + d) * T + t, acc, t < T)
|
| 63 |
+
|
| 64 |
+
@triton.jit
|
| 65 |
+
def _decode(
|
| 66 |
+
X,
|
| 67 |
+
W,
|
| 68 |
+
BIAS,
|
| 69 |
+
STATE,
|
| 70 |
+
Y,
|
| 71 |
+
D: tl.constexpr,
|
| 72 |
+
K: tl.constexpr,
|
| 73 |
+
X0: tl.constexpr,
|
| 74 |
+
X1: tl.constexpr,
|
| 75 |
+
S0: tl.constexpr,
|
| 76 |
+
S1: tl.constexpr,
|
| 77 |
+
S2: tl.constexpr,
|
| 78 |
+
HAS_BIAS: tl.constexpr,
|
| 79 |
+
BLOCK: tl.constexpr,
|
| 80 |
+
):
|
| 81 |
+
b = tl.program_id(0).to(tl.int64)
|
| 82 |
+
d = tl.program_id(1).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
|
| 83 |
+
k = tl.arange(0, 8)
|
| 84 |
+
# Load the complete state before writing the shifted window.
|
| 85 |
+
old = tl.load(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, (d[:, None] < D) & (k[None, :] < K), other=0)
|
| 86 |
+
x = tl.load(X + b * X0 + d * X1, d < D, other=0)
|
| 87 |
+
shifted = tl.gather(old, tl.broadcast_to(tl.minimum(k + 1, 7)[None, :], (BLOCK, 8)), 1)
|
| 88 |
+
shifted = tl.where(k[None, :] == K - 1, x[:, None], shifted)
|
| 89 |
+
# A cache may store either K=6 previous inputs or K=7 including the oldest discarded input.
|
| 90 |
+
positions = k + K - 6
|
| 91 |
+
values = tl.gather(old, tl.broadcast_to(tl.minimum(positions, 7)[None, :], (BLOCK, 8)), 1)
|
| 92 |
+
values = tl.where(k[None, :] == 6, x[:, None], values)
|
| 93 |
+
weights = tl.load(W + d[:, None] * 7 + k[None, :], (d[:, None] < D) & (k[None, :] < 7), other=0)
|
| 94 |
+
acc = tl.sum(values.to(tl.float32) * weights.to(tl.float32), 1)
|
| 95 |
+
if HAS_BIAS:
|
| 96 |
+
acc += tl.load(BIAS + d, d < D, other=0).to(tl.float32)
|
| 97 |
+
tl.store(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, shifted, (d[:, None] < D) & (k[None, :] < K))
|
| 98 |
+
tl.store(Y + b * D + d, acc, d < D)
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def can_use_lfm3_conv(hidden_states, weight):
|
| 102 |
+
return (
|
| 103 |
+
is_triton_available()
|
| 104 |
+
and hasattr(tl, "gather") # the decode kernel needs Triton >= 3.3
|
| 105 |
+
and hidden_states.is_cuda
|
| 106 |
+
and not torch.is_grad_enabled()
|
| 107 |
+
and hidden_states.dtype == weight.dtype
|
| 108 |
+
and hidden_states.dtype in (torch.float16, torch.bfloat16, torch.float32)
|
| 109 |
+
and weight.shape[-1] == 7
|
| 110 |
+
)
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
def lfm3_causal_conv1d(hidden_states, weight, bias=None, seq_idx=None):
|
| 114 |
+
batch, channels, length = hidden_states.shape
|
| 115 |
+
# the kernels index every tensor by the batch and channel of `hidden_states` and cannot bound the others
|
| 116 |
+
if weight.shape != (channels, 7) or (bias is not None and bias.shape != (channels,)):
|
| 117 |
+
raise ValueError(
|
| 118 |
+
f"`weight` must have shape ({channels}, 7) and `bias` ({channels},), got {tuple(weight.shape)}"
|
| 119 |
+
)
|
| 120 |
+
if seq_idx is not None and (
|
| 121 |
+
seq_idx.ndim != 2 or seq_idx.shape[0] not in (1, batch) or seq_idx.dtype not in (torch.int32, torch.int64)
|
| 122 |
+
):
|
| 123 |
+
raise ValueError(
|
| 124 |
+
f"`seq_idx` must be an int32 or int64 tensor of shape (1, length) or ({batch}, length), got "
|
| 125 |
+
f"{seq_idx.dtype} {tuple(seq_idx.shape)}"
|
| 126 |
+
)
|
| 127 |
+
if any(t is not None and t.device != hidden_states.device for t in (weight, bias, seq_idx)):
|
| 128 |
+
raise ValueError(
|
| 129 |
+
f"`weight`, `bias` and `seq_idx` must be on the device of `hidden_states` ({hidden_states.device})"
|
| 130 |
+
)
|
| 131 |
+
output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
|
| 132 |
+
# Triton launches on the current device; the model may live on another one (`device_map`, `.to("cuda:1")`).
|
| 133 |
+
with torch.cuda.device(hidden_states.device):
|
| 134 |
+
_prefill[(batch, channels, triton.cdiv(length, 128))](
|
| 135 |
+
hidden_states,
|
| 136 |
+
# the kernels index `bias` and `weight` densely
|
| 137 |
+
weight.contiguous(),
|
| 138 |
+
bias.contiguous() if bias is not None else None,
|
| 139 |
+
seq_idx,
|
| 140 |
+
output,
|
| 141 |
+
length,
|
| 142 |
+
channels,
|
| 143 |
+
*hidden_states.stride(),
|
| 144 |
+
*(
|
| 145 |
+
(0 if seq_idx.shape[0] == 1 else seq_idx.stride(0), seq_idx.stride(1))
|
| 146 |
+
if seq_idx is not None
|
| 147 |
+
else (0, 0)
|
| 148 |
+
),
|
| 149 |
+
seq_idx.shape[-1] if seq_idx is not None else 0,
|
| 150 |
+
bias is not None,
|
| 151 |
+
seq_idx is not None,
|
| 152 |
+
128,
|
| 153 |
+
enable_fp_fusion=False,
|
| 154 |
+
)
|
| 155 |
+
return output
|
| 156 |
+
|
| 157 |
+
|
| 158 |
+
def lfm3_causal_conv1d_update(hidden_states, conv_state, weight, bias=None):
|
| 159 |
+
batch, channels, length = hidden_states.shape
|
| 160 |
+
# the kernel addresses one position per row and bounds nothing by the other tensors' shapes
|
| 161 |
+
if length != 1:
|
| 162 |
+
raise ValueError(f"`hidden_states` must hold one position per row, got {tuple(hidden_states.shape)}")
|
| 163 |
+
# the 8-wide tile holds a 6- or 7-wide state (see `positions` in `_decode`)
|
| 164 |
+
if (
|
| 165 |
+
conv_state.ndim != 3
|
| 166 |
+
or conv_state.shape[:2] != (batch, channels)
|
| 167 |
+
or conv_state.shape[-1] not in (6, 7)
|
| 168 |
+
or weight.shape != (channels, 7)
|
| 169 |
+
):
|
| 170 |
+
raise ValueError(
|
| 171 |
+
f"`conv_state` must have shape ({batch}, {channels}, 6 or 7) and `weight` ({channels}, 7), got "
|
| 172 |
+
f"{tuple(conv_state.shape)} and {tuple(weight.shape)}"
|
| 173 |
+
)
|
| 174 |
+
if bias is not None and bias.shape != (channels,):
|
| 175 |
+
raise ValueError(f"`bias` must have shape ({channels},), got {tuple(bias.shape)}")
|
| 176 |
+
# the kernel writes the state in place: rows or channels sharing memory (an `expand`ed cache) would race
|
| 177 |
+
extent = 0
|
| 178 |
+
for stride, size in sorted(zip(conv_state.stride(), conv_state.shape)):
|
| 179 |
+
if size > 1 and stride <= extent:
|
| 180 |
+
raise ValueError(f"`conv_state` must not overlap itself, got strides {conv_state.stride()}")
|
| 181 |
+
extent += stride * (size - 1)
|
| 182 |
+
if any(t is not None and t.device != hidden_states.device for t in (conv_state, weight, bias)):
|
| 183 |
+
raise ValueError(
|
| 184 |
+
f"`conv_state`, `weight` and `bias` must be on the device of `hidden_states` ({hidden_states.device})"
|
| 185 |
+
)
|
| 186 |
+
output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
|
| 187 |
+
with torch.cuda.device(hidden_states.device):
|
| 188 |
+
_decode[(batch, triton.cdiv(channels, 32))](
|
| 189 |
+
hidden_states,
|
| 190 |
+
weight.contiguous(),
|
| 191 |
+
bias.contiguous() if bias is not None else None,
|
| 192 |
+
conv_state,
|
| 193 |
+
output,
|
| 194 |
+
channels,
|
| 195 |
+
conv_state.shape[-1],
|
| 196 |
+
*hidden_states.stride()[:2],
|
| 197 |
+
*conv_state.stride(),
|
| 198 |
+
bias is not None,
|
| 199 |
+
32,
|
| 200 |
+
enable_fp_fusion=False,
|
| 201 |
+
)
|
| 202 |
+
return output
|
audio_tokenizer/model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a50b0605c16f09ee9479da791a249d135b8713e975b66d680d6e328ed8717887
|
| 3 |
+
size 982140
|
audio_tokenizer/modeling_lfm3.py
ADDED
|
@@ -0,0 +1,881 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
import copy
|
| 15 |
+
from collections.abc import Callable
|
| 16 |
+
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn.functional as F
|
| 19 |
+
from torch import nn
|
| 20 |
+
|
| 21 |
+
from transformers import initialization as init
|
| 22 |
+
from transformers.activations import ACT2FN
|
| 23 |
+
from transformers.cache_utils import Cache, DynamicCache
|
| 24 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 25 |
+
from transformers.generation import GenerationMixin
|
| 26 |
+
from transformers.integrations import (
|
| 27 |
+
use_kernel_forward_from_hub,
|
| 28 |
+
use_kernel_func_from_hub_with_fallback,
|
| 29 |
+
use_kernelized_func,
|
| 30 |
+
)
|
| 31 |
+
from transformers.integrations.accelerate import force_accelerate_hooks
|
| 32 |
+
from transformers.masking_utils import (
|
| 33 |
+
create_causal_mask,
|
| 34 |
+
create_recurrent_attention_mask,
|
| 35 |
+
create_sliding_window_causal_mask,
|
| 36 |
+
)
|
| 37 |
+
from transformers.modeling_layers import GradientCheckpointingLayer
|
| 38 |
+
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
|
| 39 |
+
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
|
| 40 |
+
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
|
| 41 |
+
from transformers.processing_utils import Unpack
|
| 42 |
+
from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
|
| 43 |
+
from transformers.utils.generic import merge_with_config_defaults
|
| 44 |
+
from transformers.utils.import_utils import is_tracing
|
| 45 |
+
from transformers.utils.output_capturing import capture_outputs
|
| 46 |
+
|
| 47 |
+
from .configuration_lfm3 import Lfm3Config
|
| 48 |
+
from .lfm3_conv import can_use_lfm3_conv, lfm3_causal_conv1d, lfm3_causal_conv1d_update
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
@use_kernel_forward_from_hub("RMSNorm")
|
| 52 |
+
class Lfm3RMSNorm(nn.Module):
|
| 53 |
+
def __init__(self, hidden_size, eps: float = 1e-6) -> None:
|
| 54 |
+
"""
|
| 55 |
+
Lfm3RMSNorm is equivalent to T5LayerNorm
|
| 56 |
+
"""
|
| 57 |
+
super().__init__()
|
| 58 |
+
self.weight = nn.Parameter(torch.ones(hidden_size))
|
| 59 |
+
self.variance_epsilon = eps
|
| 60 |
+
|
| 61 |
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 62 |
+
input_dtype = hidden_states.dtype
|
| 63 |
+
hidden_states = hidden_states.to(torch.float32)
|
| 64 |
+
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
| 65 |
+
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
| 66 |
+
return self.weight * hidden_states.to(input_dtype)
|
| 67 |
+
|
| 68 |
+
def extra_repr(self):
|
| 69 |
+
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
class Lfm3RotaryEmbedding(nn.Module):
|
| 73 |
+
"""
|
| 74 |
+
Rotary embedding per attention layer type, like Gemma 3; `conv` layers register no buffers. Partial RoPE keeps the
|
| 75 |
+
frequency progression of the full head dimension and zeroes the tail instead of shrinking the rotated dimension.
|
| 76 |
+
"""
|
| 77 |
+
|
| 78 |
+
def __init__(self, config: Lfm3Config, device=None):
|
| 79 |
+
super().__init__()
|
| 80 |
+
self.max_seq_len_cached = config.max_position_embeddings
|
| 81 |
+
self.original_max_seq_len = config.max_position_embeddings
|
| 82 |
+
self.config = config
|
| 83 |
+
self.layer_types = sorted(set(config.layer_types) & set(config.rope_parameters))
|
| 84 |
+
self.rope_type = {}
|
| 85 |
+
for layer_type in self.layer_types:
|
| 86 |
+
rope_params = self.config.rope_parameters[layer_type]
|
| 87 |
+
self.rope_type[layer_type] = rope_params["rope_type"]
|
| 88 |
+
curr_inv_freq, curr_attention_scaling = self.compute_rope_parameters(config, device, layer_type=layer_type)
|
| 89 |
+
setattr(self, f"{layer_type}_inv_freq", nn.Buffer(curr_inv_freq, persistent=False))
|
| 90 |
+
setattr(self, f"{layer_type}_original_inv_freq", nn.Buffer(curr_inv_freq.clone(), persistent=False))
|
| 91 |
+
setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling)
|
| 92 |
+
|
| 93 |
+
@staticmethod
|
| 94 |
+
def compute_default_rope_parameters(
|
| 95 |
+
config: Lfm3Config, device=None, layer_type: str | None = None, **kwargs
|
| 96 |
+
) -> tuple[torch.Tensor, float]:
|
| 97 |
+
"""
|
| 98 |
+
Inverse frequencies of the default RoPE with LFM3's zero-tail partial rotation: the first
|
| 99 |
+
`partial_rotary_factor * head_dim // 2` frequencies follow the progression of the full head dimension and the
|
| 100 |
+
rest are zero, leaving those dimensions unrotated.
|
| 101 |
+
|
| 102 |
+
Returns:
|
| 103 |
+
Tuple of (`torch.Tensor`, `float`): the inverse frequencies and the attention scaling factor (always 1).
|
| 104 |
+
"""
|
| 105 |
+
base = config.rope_parameters[layer_type]["rope_theta"]
|
| 106 |
+
partial_rotary_factor = config.rope_parameters[layer_type].get("partial_rotary_factor", 1.0)
|
| 107 |
+
dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
|
| 108 |
+
rope_angles = int(partial_rotary_factor * dim) // 2
|
| 109 |
+
inv_freq = 1.0 / (base ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.float) / dim))
|
| 110 |
+
inv_freq = torch.cat((inv_freq, torch.zeros(dim // 2 - rope_angles, dtype=torch.float)))
|
| 111 |
+
return inv_freq.to(device), 1.0
|
| 112 |
+
|
| 113 |
+
@torch.no_grad()
|
| 114 |
+
@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
|
| 115 |
+
def forward(self, x, position_ids, layer_type):
|
| 116 |
+
inv_freq = getattr(self, f"{layer_type}_inv_freq")
|
| 117 |
+
attention_scaling = getattr(self, f"{layer_type}_attention_scaling")
|
| 118 |
+
|
| 119 |
+
freqs = position_ids[..., None].float() * inv_freq.to(device=x.device, dtype=torch.float)
|
| 120 |
+
emb = torch.cat((freqs, freqs), dim=-1)
|
| 121 |
+
cos = emb.cos() * attention_scaling
|
| 122 |
+
sin = emb.sin() * attention_scaling
|
| 123 |
+
|
| 124 |
+
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
| 125 |
+
|
| 126 |
+
@classmethod
|
| 127 |
+
def compute_rope_parameters(
|
| 128 |
+
cls, config: Lfm3Config, device=None, layer_type: str | None = None
|
| 129 |
+
) -> tuple[torch.Tensor, float]:
|
| 130 |
+
"""Inverse frequencies of any RoPE type in the zero-tail layout: scaled types are evaluated over the full head
|
| 131 |
+
dimension and truncated to the rotated prefix."""
|
| 132 |
+
rope_params = config.rope_parameters[layer_type]
|
| 133 |
+
if rope_params["rope_type"] == "default":
|
| 134 |
+
return cls.compute_default_rope_parameters(config, device, layer_type=layer_type)
|
| 135 |
+
full_config = copy.copy(config)
|
| 136 |
+
full_config.rope_parameters = {
|
| 137 |
+
**config.rope_parameters,
|
| 138 |
+
layer_type: {**rope_params, "partial_rotary_factor": 1.0},
|
| 139 |
+
}
|
| 140 |
+
inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_params["rope_type"]](
|
| 141 |
+
full_config, device, layer_type=layer_type
|
| 142 |
+
)
|
| 143 |
+
dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
|
| 144 |
+
rope_angles = int(rope_params.get("partial_rotary_factor", 1.0) * dim) // 2
|
| 145 |
+
inv_freq = torch.cat((inv_freq[:rope_angles], inv_freq.new_zeros(dim // 2 - rope_angles)))
|
| 146 |
+
return inv_freq, attention_factor
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
class Lfm3MLP(nn.Module):
|
| 150 |
+
def __init__(self, config: Lfm3Config):
|
| 151 |
+
super().__init__()
|
| 152 |
+
self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 153 |
+
self.w3 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 154 |
+
self.w2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
| 155 |
+
|
| 156 |
+
def forward(self, x):
|
| 157 |
+
return self.w2(F.silu(self.w1(x)) * self.w3(x))
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
def rotate_half(x):
|
| 161 |
+
"""Rotates half the hidden dims of the input."""
|
| 162 |
+
x1 = x[..., : x.shape[-1] // 2]
|
| 163 |
+
x2 = x[..., x.shape[-1] // 2 :]
|
| 164 |
+
return torch.cat((-x2, x1), dim=-1)
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
| 168 |
+
"""
|
| 169 |
+
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
| 170 |
+
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
| 171 |
+
"""
|
| 172 |
+
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
| 173 |
+
if n_rep == 1:
|
| 174 |
+
return hidden_states
|
| 175 |
+
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
| 176 |
+
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
def eager_attention_forward(
|
| 180 |
+
module: nn.Module,
|
| 181 |
+
query: torch.Tensor,
|
| 182 |
+
key: torch.Tensor,
|
| 183 |
+
value: torch.Tensor,
|
| 184 |
+
attention_mask: torch.Tensor | None,
|
| 185 |
+
scaling: float,
|
| 186 |
+
dropout: float = 0.0,
|
| 187 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 188 |
+
):
|
| 189 |
+
key_states = repeat_kv(key, module.num_key_value_groups)
|
| 190 |
+
value_states = repeat_kv(value, module.num_key_value_groups)
|
| 191 |
+
|
| 192 |
+
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
|
| 193 |
+
if attention_mask is not None:
|
| 194 |
+
attn_weights = attn_weights + attention_mask
|
| 195 |
+
|
| 196 |
+
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
|
| 197 |
+
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
|
| 198 |
+
attn_output = torch.matmul(attn_weights, value_states)
|
| 199 |
+
attn_output = attn_output.transpose(1, 2).contiguous()
|
| 200 |
+
|
| 201 |
+
return attn_output, attn_weights
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
|
| 205 |
+
"""Applies RoPE to the query and key tensors in float32, like the reference implementation."""
|
| 206 |
+
cos = cos.unsqueeze(unsqueeze_dim).float()
|
| 207 |
+
sin = sin.unsqueeze(unsqueeze_dim).float()
|
| 208 |
+
q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
|
| 209 |
+
k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
|
| 210 |
+
return q_embed.to(q.dtype), k_embed.to(k.dtype)
|
| 211 |
+
|
| 212 |
+
|
| 213 |
+
class Lfm3Attention(nn.Module):
|
| 214 |
+
"""Llama attention with QK RMSNorm, a sigmoid output gate per head, a sliding window on `sliding_attention`
|
| 215 |
+
layers and RoPE in float32."""
|
| 216 |
+
|
| 217 |
+
def __init__(self, config: Lfm3Config, layer_idx: int):
|
| 218 |
+
super().__init__()
|
| 219 |
+
self.config = config
|
| 220 |
+
self.layer_idx = layer_idx
|
| 221 |
+
self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
| 222 |
+
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
|
| 223 |
+
self.scaling = self.head_dim**-0.5
|
| 224 |
+
self.is_causal = True
|
| 225 |
+
self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
|
| 226 |
+
self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
|
| 227 |
+
self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
|
| 228 |
+
self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
|
| 229 |
+
self.q_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
|
| 230 |
+
self.k_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
|
| 231 |
+
self.gate_proj = nn.Linear(config.hidden_size, config.num_attention_heads, bias=False)
|
| 232 |
+
self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
|
| 233 |
+
|
| 234 |
+
def forward(
|
| 235 |
+
self,
|
| 236 |
+
hidden_states: torch.Tensor,
|
| 237 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
| 238 |
+
attention_mask: torch.Tensor | None,
|
| 239 |
+
past_key_values: Cache | None = None,
|
| 240 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 241 |
+
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
| 242 |
+
input_shape = hidden_states.shape[:-1]
|
| 243 |
+
hidden_shape = (*input_shape, -1, self.head_dim)
|
| 244 |
+
|
| 245 |
+
query_states = self.q_norm(self.q_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
|
| 246 |
+
key_states = self.k_norm(self.k_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
|
| 247 |
+
value_states = self.v_proj(hidden_states).view(*hidden_shape).transpose(1, 2)
|
| 248 |
+
|
| 249 |
+
cos, sin = position_embeddings
|
| 250 |
+
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
| 251 |
+
|
| 252 |
+
if past_key_values is not None:
|
| 253 |
+
key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
|
| 254 |
+
|
| 255 |
+
attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
|
| 256 |
+
self.config._attn_implementation, eager_attention_forward
|
| 257 |
+
)
|
| 258 |
+
|
| 259 |
+
attn_output, attn_weights = attention_interface(
|
| 260 |
+
self,
|
| 261 |
+
query_states,
|
| 262 |
+
key_states,
|
| 263 |
+
value_states,
|
| 264 |
+
attention_mask,
|
| 265 |
+
dropout=0.0,
|
| 266 |
+
scaling=self.scaling,
|
| 267 |
+
sliding_window=self.sliding_window,
|
| 268 |
+
**kwargs,
|
| 269 |
+
)
|
| 270 |
+
gate = torch.sigmoid(self.gate_proj(hidden_states)).unsqueeze(-1)
|
| 271 |
+
attn_output = (attn_output * gate).reshape(*input_shape, -1).contiguous()
|
| 272 |
+
attn_output = self.o_proj(attn_output)
|
| 273 |
+
return attn_output, attn_weights
|
| 274 |
+
|
| 275 |
+
|
| 276 |
+
def apply_mask_to_padding_states(hidden_states, attention_mask):
|
| 277 |
+
"""
|
| 278 |
+
Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
|
| 279 |
+
"""
|
| 280 |
+
# NOTE: attention mask is a 2D boolean tensor
|
| 281 |
+
if attention_mask is not None:
|
| 282 |
+
dtype = hidden_states.dtype
|
| 283 |
+
hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
|
| 284 |
+
|
| 285 |
+
return hidden_states
|
| 286 |
+
|
| 287 |
+
|
| 288 |
+
@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
|
| 289 |
+
def causal_conv1d_update(
|
| 290 |
+
hidden_states: torch.Tensor,
|
| 291 |
+
conv_state: torch.Tensor,
|
| 292 |
+
weight: nn.Parameter,
|
| 293 |
+
bias: nn.Parameter | None = None,
|
| 294 |
+
activation: str | None = None,
|
| 295 |
+
):
|
| 296 |
+
_, hidden_size, seq_len = hidden_states.shape
|
| 297 |
+
state_len = conv_state.shape[-1]
|
| 298 |
+
|
| 299 |
+
hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
|
| 300 |
+
conv_state.copy_(hidden_states_new[:, :, -state_len:])
|
| 301 |
+
out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
|
| 302 |
+
out = out[:, :, -seq_len:]
|
| 303 |
+
if activation is not None:
|
| 304 |
+
out = ACT2FN[activation](out)
|
| 305 |
+
return out.to(hidden_states.dtype)
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
|
| 309 |
+
def causal_conv1d_fn(
|
| 310 |
+
hidden_states: torch.Tensor,
|
| 311 |
+
weight: nn.Parameter,
|
| 312 |
+
bias: nn.Parameter | None = None,
|
| 313 |
+
activation: str | None = None,
|
| 314 |
+
**kwargs,
|
| 315 |
+
):
|
| 316 |
+
_, hidden_size, seq_len = hidden_states.shape
|
| 317 |
+
padding = weight.shape[-1] - 1
|
| 318 |
+
|
| 319 |
+
out = F.conv1d(
|
| 320 |
+
hidden_states.to(weight.dtype),
|
| 321 |
+
weight=weight.unsqueeze(1),
|
| 322 |
+
bias=bias,
|
| 323 |
+
padding=padding,
|
| 324 |
+
groups=hidden_size,
|
| 325 |
+
)[:, :, :seq_len]
|
| 326 |
+
if activation is not None:
|
| 327 |
+
out = ACT2FN[activation](out)
|
| 328 |
+
return out.to(hidden_states.dtype)
|
| 329 |
+
|
| 330 |
+
|
| 331 |
+
@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
|
| 332 |
+
class Lfm3ShortConv(nn.Module):
|
| 333 |
+
"""LFM2 gated short convolution with packed-document boundaries (`seq_idx`), the padding mask re-applied after a
|
| 334 |
+
biased `in_proj`, and width-7 Triton inference kernels with a PyTorch fallback."""
|
| 335 |
+
|
| 336 |
+
def __init__(
|
| 337 |
+
self,
|
| 338 |
+
config: Lfm3Config,
|
| 339 |
+
layer_idx: int,
|
| 340 |
+
):
|
| 341 |
+
super().__init__()
|
| 342 |
+
self.config = config
|
| 343 |
+
self.layer_idx = layer_idx
|
| 344 |
+
self.conv_kernel_size = config.conv_L_cache
|
| 345 |
+
self.bias = config.conv_bias
|
| 346 |
+
|
| 347 |
+
self.conv = nn.Conv1d(
|
| 348 |
+
in_channels=config.hidden_size,
|
| 349 |
+
out_channels=config.hidden_size,
|
| 350 |
+
kernel_size=self.conv_kernel_size,
|
| 351 |
+
groups=config.hidden_size,
|
| 352 |
+
bias=self.bias,
|
| 353 |
+
padding=self.conv_kernel_size - 1,
|
| 354 |
+
)
|
| 355 |
+
self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=self.bias)
|
| 356 |
+
self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=self.bias)
|
| 357 |
+
|
| 358 |
+
self.layer_type = config.layer_types[layer_idx]
|
| 359 |
+
|
| 360 |
+
@force_accelerate_hooks("conv")
|
| 361 |
+
def forward(
|
| 362 |
+
self,
|
| 363 |
+
hidden_states: torch.Tensor,
|
| 364 |
+
past_key_values: Cache | None = None,
|
| 365 |
+
attention_mask: torch.Tensor | None = None,
|
| 366 |
+
seq_idx: torch.IntTensor | None = None,
|
| 367 |
+
):
|
| 368 |
+
seq_len = hidden_states.shape[1]
|
| 369 |
+
|
| 370 |
+
hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
|
| 371 |
+
BCx = self.in_proj(hidden_states).transpose(-1, -2)
|
| 372 |
+
B, C, x = BCx.chunk(3, dim=-2)
|
| 373 |
+
hidden_states = B * x
|
| 374 |
+
if attention_mask is not None:
|
| 375 |
+
# with `conv_bias` the projection makes the padded positions non-zero again
|
| 376 |
+
hidden_states = hidden_states * attention_mask[:, None, :].to(hidden_states.dtype)
|
| 377 |
+
|
| 378 |
+
use_precomputed_states = past_key_values is not None and past_key_values.has_previous_state(self.layer_idx)
|
| 379 |
+
use_kernels = 2 <= self.conv_kernel_size <= 4
|
| 380 |
+
use_wide_kernel = can_use_lfm3_conv(hidden_states, self.conv.weight)
|
| 381 |
+
|
| 382 |
+
if use_precomputed_states and seq_len == 1 and not past_key_values.layers[self.layer_idx].record_past:
|
| 383 |
+
conv_state = past_key_values.layers[self.layer_idx].conv_states[0]
|
| 384 |
+
if use_wide_kernel and conv_state.shape[-1] in (6, 7):
|
| 385 |
+
hidden_states = lfm3_causal_conv1d_update(
|
| 386 |
+
hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
|
| 387 |
+
)
|
| 388 |
+
elif use_kernels:
|
| 389 |
+
hidden_states = causal_conv1d_update(
|
| 390 |
+
hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
|
| 391 |
+
)
|
| 392 |
+
else:
|
| 393 |
+
hidden_states = torch.cat([conv_state, hidden_states], dim=-1)
|
| 394 |
+
conv_state.copy_(hidden_states[:, :, -conv_state.shape[-1] :])
|
| 395 |
+
hidden_states = self._conv(hidden_states)[:, :, -seq_len:]
|
| 396 |
+
else:
|
| 397 |
+
if past_key_values is not None and seq_idx is not None:
|
| 398 |
+
# `seq_idx` is only set on an empty cache, and the decode steps read the cached window without it:
|
| 399 |
+
# only the last document (after any positional left padding) may enter the window
|
| 400 |
+
last_document = (seq_idx == seq_idx[:, -1:])[:, None].to(hidden_states.dtype)
|
| 401 |
+
past_key_values.update_conv_state(
|
| 402 |
+
hidden_states * last_document, self.layer_idx, conv_kernel_size=self.conv_kernel_size
|
| 403 |
+
)
|
| 404 |
+
elif past_key_values is not None:
|
| 405 |
+
hidden_states = past_key_values.update_conv_state(
|
| 406 |
+
hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
|
| 407 |
+
)
|
| 408 |
+
|
| 409 |
+
if use_wide_kernel:
|
| 410 |
+
hidden_states = lfm3_causal_conv1d(hidden_states, self.conv.weight.squeeze(1), self.conv.bias, seq_idx)
|
| 411 |
+
elif seq_idx is not None:
|
| 412 |
+
hidden_states = self._segmented_conv(hidden_states, seq_idx)
|
| 413 |
+
elif use_kernels:
|
| 414 |
+
hidden_states = causal_conv1d_fn(hidden_states, self.conv.weight.squeeze(1), self.conv.bias)
|
| 415 |
+
else:
|
| 416 |
+
hidden_states = self._conv(hidden_states)
|
| 417 |
+
|
| 418 |
+
# Drop the additional previous states
|
| 419 |
+
if past_key_values is not None:
|
| 420 |
+
hidden_states = hidden_states[:, :, -seq_len:]
|
| 421 |
+
|
| 422 |
+
y = C * hidden_states
|
| 423 |
+
y = y.transpose(-1, -2).contiguous()
|
| 424 |
+
y = self.out_proj(y)
|
| 425 |
+
return y
|
| 426 |
+
|
| 427 |
+
def _conv(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 428 |
+
"""Causal depthwise conv in torch: the `causal_conv1d` kernels only accept widths 2 to 4."""
|
| 429 |
+
conv_out = self.conv(hidden_states.to(self.conv.weight.dtype))[..., : hidden_states.shape[-1]]
|
| 430 |
+
return conv_out.to(hidden_states.dtype)
|
| 431 |
+
|
| 432 |
+
def _segmented_conv(self, hidden_states: torch.Tensor, seq_idx: torch.Tensor) -> torch.Tensor:
|
| 433 |
+
"""Causal depthwise conv that only reads positions sharing the current token's document index."""
|
| 434 |
+
kernel_size = self.conv_kernel_size
|
| 435 |
+
seq_idx = F.pad(seq_idx, (hidden_states.shape[-1] - seq_idx.shape[-1] + kernel_size - 1, 0), value=-1)
|
| 436 |
+
windows = F.pad(hidden_states, (kernel_size - 1, 0)).unfold(-1, kernel_size, 1)
|
| 437 |
+
seq_idx_windows = seq_idx.unfold(-1, kernel_size, 1)
|
| 438 |
+
same_document = (seq_idx_windows == seq_idx_windows[..., -1:]).unsqueeze(1)
|
| 439 |
+
conv_out = torch.einsum("bdtk,dk->bdt", windows * same_document.to(windows.dtype), self.conv.weight.squeeze(1))
|
| 440 |
+
if self.bias:
|
| 441 |
+
conv_out = conv_out + self.conv.bias[None, :, None]
|
| 442 |
+
return conv_out.to(hidden_states.dtype)
|
| 443 |
+
|
| 444 |
+
|
| 445 |
+
class Lfm3DecoderLayer(GradientCheckpointingLayer):
|
| 446 |
+
def __init__(self, config: Lfm3Config, layer_idx: int):
|
| 447 |
+
super().__init__()
|
| 448 |
+
self.is_attention_layer = config.layer_types[layer_idx] in config._attention_layer_types
|
| 449 |
+
self.use_sandwich_norm = config.use_sandwich_norm
|
| 450 |
+
|
| 451 |
+
if self.is_attention_layer:
|
| 452 |
+
self.self_attn = Lfm3Attention(config, layer_idx)
|
| 453 |
+
else:
|
| 454 |
+
self.conv = Lfm3ShortConv(config, layer_idx)
|
| 455 |
+
self.feed_forward = Lfm3MLP(config)
|
| 456 |
+
self.operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 457 |
+
self.ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 458 |
+
if self.use_sandwich_norm:
|
| 459 |
+
self.post_operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 460 |
+
self.post_ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 461 |
+
|
| 462 |
+
def forward(
|
| 463 |
+
self,
|
| 464 |
+
hidden_states: torch.Tensor,
|
| 465 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
| 466 |
+
attention_mask: torch.Tensor | None = None,
|
| 467 |
+
position_ids: torch.LongTensor | None = None,
|
| 468 |
+
past_key_values: Cache | None = None,
|
| 469 |
+
seq_idx: torch.IntTensor | None = None,
|
| 470 |
+
**kwargs,
|
| 471 |
+
) -> torch.Tensor:
|
| 472 |
+
residual = hidden_states
|
| 473 |
+
if self.is_attention_layer:
|
| 474 |
+
hidden_states, _ = self.self_attn(
|
| 475 |
+
hidden_states=self.operator_norm(hidden_states),
|
| 476 |
+
position_embeddings=position_embeddings,
|
| 477 |
+
attention_mask=attention_mask,
|
| 478 |
+
position_ids=position_ids,
|
| 479 |
+
past_key_values=past_key_values,
|
| 480 |
+
**kwargs,
|
| 481 |
+
)
|
| 482 |
+
else:
|
| 483 |
+
hidden_states = self.conv(
|
| 484 |
+
hidden_states=self.operator_norm(hidden_states),
|
| 485 |
+
past_key_values=past_key_values,
|
| 486 |
+
attention_mask=attention_mask,
|
| 487 |
+
seq_idx=seq_idx,
|
| 488 |
+
)
|
| 489 |
+
if self.use_sandwich_norm:
|
| 490 |
+
hidden_states = self.post_operator_norm(hidden_states)
|
| 491 |
+
hidden_states = hidden_states + residual
|
| 492 |
+
|
| 493 |
+
ffn_out = self.feed_forward(self.ffn_norm(hidden_states))
|
| 494 |
+
if self.use_sandwich_norm:
|
| 495 |
+
ffn_out = self.post_ffn_norm(ffn_out)
|
| 496 |
+
hidden_states = hidden_states + ffn_out
|
| 497 |
+
|
| 498 |
+
return hidden_states
|
| 499 |
+
|
| 500 |
+
|
| 501 |
+
def create_masks_by_layer_type(
|
| 502 |
+
config: PreTrainedConfig,
|
| 503 |
+
inputs_embeds: torch.Tensor,
|
| 504 |
+
attention_mask: torch.Tensor | None,
|
| 505 |
+
past_key_values: Cache | None,
|
| 506 |
+
position_ids: torch.Tensor | None = None,
|
| 507 |
+
) -> dict:
|
| 508 |
+
"""The cache sizes each mask from the layer it is asked about, and the generic mask utilities ask about the first
|
| 509 |
+
non-sliding layer, which for a leading `conv` layer redirects to the first attention layer of any type: name the
|
| 510 |
+
first layer of each mask's own type instead."""
|
| 511 |
+
config = config.get_text_config()
|
| 512 |
+
mask_kwargs = {
|
| 513 |
+
"config": config,
|
| 514 |
+
"inputs_embeds": inputs_embeds,
|
| 515 |
+
"attention_mask": attention_mask,
|
| 516 |
+
"past_key_values": past_key_values,
|
| 517 |
+
"position_ids": position_ids,
|
| 518 |
+
}
|
| 519 |
+
layer_types, causal_mask_mapping = config.layer_types, {}
|
| 520 |
+
if "full_attention" in layer_types:
|
| 521 |
+
layer_idx = layer_types.index("full_attention")
|
| 522 |
+
causal_mask_mapping["full_attention"] = create_causal_mask(**mask_kwargs, layer_idx=layer_idx)
|
| 523 |
+
if "sliding_attention" in layer_types:
|
| 524 |
+
layer_idx = layer_types.index("sliding_attention")
|
| 525 |
+
causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(
|
| 526 |
+
**mask_kwargs, layer_idx=layer_idx
|
| 527 |
+
)
|
| 528 |
+
if "conv" in layer_types:
|
| 529 |
+
causal_mask_mapping["conv"] = create_recurrent_attention_mask(**mask_kwargs)
|
| 530 |
+
return causal_mask_mapping
|
| 531 |
+
|
| 532 |
+
|
| 533 |
+
@auto_docstring
|
| 534 |
+
class Lfm3PreTrainedModel(PreTrainedModel):
|
| 535 |
+
config: Lfm3Config
|
| 536 |
+
base_model_prefix = "model"
|
| 537 |
+
supports_gradient_checkpointing = True
|
| 538 |
+
_no_split_modules = ["Lfm3DecoderLayer"]
|
| 539 |
+
_skip_keys_device_placement = ["past_key_values"]
|
| 540 |
+
_supports_flash_attn = True
|
| 541 |
+
_supports_sdpa = True
|
| 542 |
+
_supports_flex_attn = True
|
| 543 |
+
|
| 544 |
+
_can_compile_fullgraph = True
|
| 545 |
+
_supports_attention_backend = True
|
| 546 |
+
_can_record_outputs = {
|
| 547 |
+
"hidden_states": Lfm3DecoderLayer,
|
| 548 |
+
"attentions": Lfm3Attention,
|
| 549 |
+
}
|
| 550 |
+
|
| 551 |
+
@torch.no_grad()
|
| 552 |
+
def _init_weights(self, module):
|
| 553 |
+
super()._init_weights(module)
|
| 554 |
+
if isinstance(module, Lfm3RotaryEmbedding):
|
| 555 |
+
for layer_type in module.layer_types:
|
| 556 |
+
curr_inv_freq, _ = module.compute_rope_parameters(module.config, layer_type=layer_type)
|
| 557 |
+
init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
|
| 558 |
+
init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
|
| 559 |
+
|
| 560 |
+
@staticmethod
|
| 561 |
+
def create_masks_for_generate(
|
| 562 |
+
config: PreTrainedConfig,
|
| 563 |
+
inputs_embeds: torch.Tensor,
|
| 564 |
+
attention_mask: torch.Tensor | None,
|
| 565 |
+
past_key_values: Cache | None,
|
| 566 |
+
position_ids: torch.Tensor | None = None,
|
| 567 |
+
**kwargs,
|
| 568 |
+
) -> dict:
|
| 569 |
+
return create_masks_by_layer_type(config, inputs_embeds, attention_mask, past_key_values, position_ids)
|
| 570 |
+
|
| 571 |
+
|
| 572 |
+
def packed_under_padding(position_ids: torch.Tensor, padding_mask: torch.Tensor, filled) -> torch.Tensor:
|
| 573 |
+
"""Padding repeats position 0, so packing is a restart between two kept tokens, or a kept position 0 on a filled
|
| 574 |
+
cache whose row holds a kept token; attention under a padding mask cannot isolate it."""
|
| 575 |
+
restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
|
| 576 |
+
valid = padding_mask[:, -restarts.shape[-1] :].bool()
|
| 577 |
+
packed = restarts[:, 1:] & valid[:, 1:] & valid[:, :-1]
|
| 578 |
+
if filled:
|
| 579 |
+
# a row whose cached tokens are all padding (chunked prefill of a left-padded batch) has no document to pack
|
| 580 |
+
# against; a mask covering only the chunk keeps the strict rule
|
| 581 |
+
cached = padding_mask[:, : padding_mask.shape[-1] - restarts.shape[-1]]
|
| 582 |
+
cached_valid = cached.bool().any(-1, keepdim=True) if cached.shape[-1] else torch.ones_like(valid[:, :1])
|
| 583 |
+
packed = packed | (valid[:, :1] & (position_ids[:, :1] == 0) & cached_valid)
|
| 584 |
+
return packed
|
| 585 |
+
|
| 586 |
+
|
| 587 |
+
@auto_docstring
|
| 588 |
+
class Lfm3Model(Lfm3PreTrainedModel):
|
| 589 |
+
"""LFM2 model with sliding-attention layers as a third layer type: the mask and the rotary embedding are built per
|
| 590 |
+
layer type, and packed documents (`position_ids` restarting) reset the short-conv window through `seq_idx`."""
|
| 591 |
+
|
| 592 |
+
def __init__(self, config: Lfm3Config):
|
| 593 |
+
super().__init__(config)
|
| 594 |
+
self.padding_idx = config.pad_token_id
|
| 595 |
+
self.vocab_size = config.vocab_size
|
| 596 |
+
|
| 597 |
+
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
| 598 |
+
self.layers = nn.ModuleList(
|
| 599 |
+
[Lfm3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
| 600 |
+
)
|
| 601 |
+
self.rotary_emb = Lfm3RotaryEmbedding(config=config)
|
| 602 |
+
self.gradient_checkpointing = False
|
| 603 |
+
self.embedding_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 604 |
+
|
| 605 |
+
# Initialize weights and apply final processing
|
| 606 |
+
self.post_init()
|
| 607 |
+
|
| 608 |
+
@merge_with_config_defaults
|
| 609 |
+
@capture_outputs
|
| 610 |
+
@auto_docstring
|
| 611 |
+
def forward(
|
| 612 |
+
self,
|
| 613 |
+
input_ids: torch.LongTensor | None = None,
|
| 614 |
+
attention_mask: torch.Tensor | None = None,
|
| 615 |
+
position_ids: torch.LongTensor | None = None,
|
| 616 |
+
past_key_values: Cache | None = None,
|
| 617 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 618 |
+
use_cache: bool | None = None,
|
| 619 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 620 |
+
) -> BaseModelOutputWithPast:
|
| 621 |
+
r"""
|
| 622 |
+
attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)` or `dict`, *optional*):
|
| 623 |
+
Mask to avoid performing attention on padding token indices (1 keeps a token, 0 masks it), or the mask
|
| 624 |
+
dict of `create_masks_for_generate`. A mask dict is used as given: `generate` refuses packed
|
| 625 |
+
`position_ids` before its prefill.
|
| 626 |
+
"""
|
| 627 |
+
if (input_ids is None) ^ (inputs_embeds is not None):
|
| 628 |
+
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
| 629 |
+
# `seq_idx` is a shared kwarg (`DataCollatorWithFlattening(return_seq_idx=True)`, read by Lfm2); LFM3 derives
|
| 630 |
+
# the documents from the restarting `position_ids` that come with it
|
| 631 |
+
if kwargs.pop("seq_idx", None) is not None and position_ids is None:
|
| 632 |
+
raise ValueError("Pass packed documents as `position_ids` restarting at 0; LFM3 does not read `seq_idx`")
|
| 633 |
+
|
| 634 |
+
if inputs_embeds is None:
|
| 635 |
+
inputs_embeds = self.embed_tokens(input_ids)
|
| 636 |
+
|
| 637 |
+
if use_cache and past_key_values is None:
|
| 638 |
+
past_key_values = DynamicCache(config=self.config)
|
| 639 |
+
filled = past_key_values is not None and (
|
| 640 |
+
past_key_values.get_seq_length() > 0
|
| 641 |
+
if any(t != "conv" for t in self.config.layer_types)
|
| 642 |
+
else past_key_values.has_previous_state()
|
| 643 |
+
)
|
| 644 |
+
packed_cache = getattr(past_key_values, "lfm3_packed", None)
|
| 645 |
+
if isinstance(packed_cache, torch.Tensor) or packed_cache:
|
| 646 |
+
# Reset layers have no history to mix. Ordinary caches carry a Python `False` and skip this, so the decode
|
| 647 |
+
# steps do no tensor work here; only a compiled packed prefill leaves a tensor flag, and `filled` stays
|
| 648 |
+
# tensor-valued for static caches.
|
| 649 |
+
packed_message = (
|
| 650 |
+
"A cache filled by a packed-document forward cannot be continued: its attention spans the documents"
|
| 651 |
+
)
|
| 652 |
+
packed_cache = packed_cache & filled if isinstance(packed_cache, torch.Tensor) else filled
|
| 653 |
+
if isinstance(packed_cache, torch.Tensor):
|
| 654 |
+
torch_compilable_check(~packed_cache, packed_message)
|
| 655 |
+
if not is_tracing(packed_cache):
|
| 656 |
+
past_key_values.lfm3_packed = bool(packed_cache)
|
| 657 |
+
elif packed_cache:
|
| 658 |
+
raise ValueError(packed_message)
|
| 659 |
+
else:
|
| 660 |
+
past_key_values.lfm3_packed = False
|
| 661 |
+
|
| 662 |
+
# The mask utilities detect packing only with `attention_mask=None`, so a multi-token forward with explicit
|
| 663 |
+
# positions drops the all-ones mask tokenizers emit (a data-dependent check, hence not when traced).
|
| 664 |
+
explicit_positions = position_ids is not None and inputs_embeds.shape[1] > 1
|
| 665 |
+
if (
|
| 666 |
+
explicit_positions
|
| 667 |
+
and isinstance(attention_mask, torch.Tensor)
|
| 668 |
+
and attention_mask.ndim == 2
|
| 669 |
+
and not is_tracing(position_ids)
|
| 670 |
+
and bool(attention_mask.all())
|
| 671 |
+
):
|
| 672 |
+
attention_mask = None
|
| 673 |
+
seq_idx = None
|
| 674 |
+
if position_ids is None:
|
| 675 |
+
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
| 676 |
+
position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
|
| 677 |
+
position_ids = position_ids.unsqueeze(0)
|
| 678 |
+
elif explicit_positions and not isinstance(attention_mask, dict):
|
| 679 |
+
# a mask dict comes from `generate`, which refused packing before its prefill
|
| 680 |
+
if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 2:
|
| 681 |
+
# data-dependent: a traced forward skips it (`generate` refuses packing eagerly before its prefill)
|
| 682 |
+
if not is_tracing(position_ids):
|
| 683 |
+
torch_compilable_check(
|
| 684 |
+
~packed_under_padding(position_ids, attention_mask, filled),
|
| 685 |
+
"Packed documents (`position_ids` restarting) cannot be combined with an attention mask: "
|
| 686 |
+
"attention would not isolate the documents. Pass `attention_mask=None` without a static cache.",
|
| 687 |
+
)
|
| 688 |
+
else:
|
| 689 |
+
restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
|
| 690 |
+
document_ids = restarts.cumsum(-1)
|
| 691 |
+
if is_tracing(position_ids) or bool((document_ids[..., -1] > 0).any()):
|
| 692 |
+
seq_idx = document_ids.to(torch.int32)
|
| 693 |
+
# Packed documents are prefilled in one forward on an empty cache (cached tokens would join the first
|
| 694 |
+
# document), and their cache is never continued
|
| 695 |
+
if filled:
|
| 696 |
+
torch_compilable_check(
|
| 697 |
+
(~restarts[:, 1:]).all() & (position_ids[:, 0] != 0).all(),
|
| 698 |
+
"Packed documents must be prefilled in one forward; the cache already holds tokens",
|
| 699 |
+
)
|
| 700 |
+
# Ordinary continuation must retain the convolution window from the cached document.
|
| 701 |
+
seq_idx = None
|
| 702 |
+
if attention_mask is None:
|
| 703 |
+
# without a mask only the cacheless packed-sequence mask isolates repeated zero positions (left
|
| 704 |
+
# padding): a cache-sized mask and every later decode step would attend them, so they count too
|
| 705 |
+
packed = (document_ids[..., -1] > 0).any()
|
| 706 |
+
else:
|
| 707 |
+
# a prepared 4D mask is the caller's and carries the padding (repeated zero positions), but may
|
| 708 |
+
# or may not isolate documents, so the cache records the packing and is never continued
|
| 709 |
+
steps = torch.diff(position_ids, dim=-1)
|
| 710 |
+
packed = ((steps != 1) & ~((steps == 0) & (position_ids[:, :-1] == 0))).any()
|
| 711 |
+
if past_key_values is not None:
|
| 712 |
+
# resolved here, where the prefill already syncs, so the decode steps carry a Python flag: no
|
| 713 |
+
# per-step host sync, data-dependent op or CUDA-graph output on the cache
|
| 714 |
+
past_key_values.lfm3_packed = packed if is_tracing(position_ids) else bool(packed)
|
| 715 |
+
|
| 716 |
+
if not isinstance(causal_mask_mapping := attention_mask, dict):
|
| 717 |
+
mask_kwargs = {
|
| 718 |
+
"config": self.config,
|
| 719 |
+
"inputs_embeds": inputs_embeds,
|
| 720 |
+
"attention_mask": attention_mask,
|
| 721 |
+
"past_key_values": past_key_values,
|
| 722 |
+
"position_ids": position_ids,
|
| 723 |
+
}
|
| 724 |
+
# The mask utilities only isolate packed documents without a cache, and an empty cache adds nothing.
|
| 725 |
+
layer_types = self.config.layer_types
|
| 726 |
+
if seq_idx is not None and past_key_values is not None:
|
| 727 |
+
attention_layer = next((i for i, t in enumerate(layer_types) if t != "conv"), None)
|
| 728 |
+
if attention_layer is not None and past_key_values.get_seq_length(attention_layer) == 0:
|
| 729 |
+
if past_key_values.is_compileable:
|
| 730 |
+
# static layers always return their full length, so the masks are sized by the cache, which
|
| 731 |
+
# isolates neither documents nor positional-only left padding
|
| 732 |
+
torch_compilable_check(
|
| 733 |
+
~packed,
|
| 734 |
+
"Packed documents, or left padding without an attention mask, cannot be prefilled into a "
|
| 735 |
+
"static cache: its masks are sized by the cache. Pass the padding as `attention_mask`, or "
|
| 736 |
+
"use a `DynamicCache` for packed documents",
|
| 737 |
+
)
|
| 738 |
+
else:
|
| 739 |
+
mask_kwargs["past_key_values"] = None
|
| 740 |
+
causal_mask_mapping = create_masks_by_layer_type(**mask_kwargs)
|
| 741 |
+
|
| 742 |
+
hidden_states = inputs_embeds
|
| 743 |
+
position_embeddings = {
|
| 744 |
+
layer_type: self.rotary_emb(hidden_states, position_ids=position_ids, layer_type=layer_type)
|
| 745 |
+
for layer_type in self.rotary_emb.layer_types
|
| 746 |
+
}
|
| 747 |
+
|
| 748 |
+
for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
|
| 749 |
+
layer_type = self.config.layer_types[i]
|
| 750 |
+
hidden_states = decoder_layer(
|
| 751 |
+
hidden_states,
|
| 752 |
+
attention_mask=causal_mask_mapping[layer_type],
|
| 753 |
+
position_embeddings=position_embeddings.get(layer_type),
|
| 754 |
+
position_ids=position_ids,
|
| 755 |
+
past_key_values=past_key_values,
|
| 756 |
+
seq_idx=seq_idx,
|
| 757 |
+
**kwargs,
|
| 758 |
+
)
|
| 759 |
+
|
| 760 |
+
hidden_states = self.embedding_norm(hidden_states)
|
| 761 |
+
|
| 762 |
+
return BaseModelOutputWithPast(
|
| 763 |
+
last_hidden_state=hidden_states,
|
| 764 |
+
past_key_values=past_key_values,
|
| 765 |
+
)
|
| 766 |
+
|
| 767 |
+
|
| 768 |
+
def reject_packed_generation(config: PreTrainedConfig, input_ids, generation_config, model_kwargs, is_first_iteration):
|
| 769 |
+
"""`generate` decodes from the cache its prefill fills, and no mask it builds isolates packed documents
|
| 770 |
+
(`position_ids` restarting): the prompt is refused before the prefill, one forward or every chunk, writes to it."""
|
| 771 |
+
position_ids, attention_mask = model_kwargs.get("position_ids"), model_kwargs.get("attention_mask")
|
| 772 |
+
cache = model_kwargs.get("past_key_values")
|
| 773 |
+
if cache is None or position_ids is None or position_ids.ndim != 2:
|
| 774 |
+
return
|
| 775 |
+
if attention_mask is not None and attention_mask.ndim != 2:
|
| 776 |
+
return
|
| 777 |
+
has_attention = any(t != "conv" for t in config.get_text_config().layer_types)
|
| 778 |
+
filled = bool(cache.get_seq_length() > 0 if has_attention else cache.has_previous_state())
|
| 779 |
+
inputs_embeds = model_kwargs.get("inputs_embeds") if is_first_iteration else None
|
| 780 |
+
length = input_ids.shape[1] if inputs_embeds is None else inputs_embeds.shape[1]
|
| 781 |
+
if generation_config.prefill_chunk_size is not None:
|
| 782 |
+
# the chunks run from the start of the prompt
|
| 783 |
+
position_ids = position_ids[:, :length]
|
| 784 |
+
else:
|
| 785 |
+
if filled and has_attention:
|
| 786 |
+
# the slice `_prefill` feeds: a prompt that spans the cache is cut to its uncached part
|
| 787 |
+
mask_length = getattr(generation_config, "_mask_length", -1)
|
| 788 |
+
if attention_mask is not None:
|
| 789 |
+
mask_length = attention_mask.shape[1]
|
| 790 |
+
if inputs_embeds is not None or mask_length == length:
|
| 791 |
+
length -= int(cache.get_seq_length())
|
| 792 |
+
position_ids = position_ids[:, position_ids.shape[1] - length :]
|
| 793 |
+
if attention_mask is None:
|
| 794 |
+
attention_mask = torch.ones_like(position_ids)
|
| 795 |
+
if bool(packed_under_padding(position_ids, attention_mask, filled).any()):
|
| 796 |
+
raise ValueError(
|
| 797 |
+
"`generate` cannot decode packed documents (`position_ids` restarting, or starting at 0 on a cache that "
|
| 798 |
+
"holds tokens): attention would span them. Generate each document in its own row of a left-padded batch."
|
| 799 |
+
)
|
| 800 |
+
|
| 801 |
+
|
| 802 |
+
@auto_docstring
|
| 803 |
+
class Lfm3ForCausalLM(Lfm3PreTrainedModel, GenerationMixin):
|
| 804 |
+
_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
|
| 805 |
+
_tp_plan = {"lm_head": "colwise_gather_output"}
|
| 806 |
+
_pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
|
| 807 |
+
_fsdp_plan = {"lm_head": "keep_full_weight"}
|
| 808 |
+
|
| 809 |
+
def __init__(self, config):
|
| 810 |
+
super().__init__(config)
|
| 811 |
+
self.model = Lfm3Model(config)
|
| 812 |
+
self.vocab_size = config.vocab_size
|
| 813 |
+
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
| 814 |
+
|
| 815 |
+
# Initialize weights and apply final processing
|
| 816 |
+
self.post_init()
|
| 817 |
+
|
| 818 |
+
@can_return_tuple
|
| 819 |
+
@auto_docstring
|
| 820 |
+
def forward(
|
| 821 |
+
self,
|
| 822 |
+
input_ids: torch.LongTensor | None = None,
|
| 823 |
+
attention_mask: torch.Tensor | None = None,
|
| 824 |
+
position_ids: torch.LongTensor | None = None,
|
| 825 |
+
past_key_values: Cache | None = None,
|
| 826 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 827 |
+
labels: torch.LongTensor | None = None,
|
| 828 |
+
use_cache: bool | None = None,
|
| 829 |
+
logits_to_keep: int | torch.Tensor = 0,
|
| 830 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 831 |
+
) -> CausalLMOutputWithPast:
|
| 832 |
+
r"""
|
| 833 |
+
Example:
|
| 834 |
+
|
| 835 |
+
```python
|
| 836 |
+
>>> from transformers import AutoTokenizer, Lfm3ForCausalLM
|
| 837 |
+
|
| 838 |
+
>>> model = Lfm3ForCausalLM.from_pretrained("LiquidAI/LFM3-1.7B-Preview")
|
| 839 |
+
>>> tokenizer = AutoTokenizer.from_pretrained("LiquidAI/LFM3-1.7B-Preview")
|
| 840 |
+
|
| 841 |
+
>>> inputs = tokenizer.apply_chat_template(
|
| 842 |
+
... [{"role": "user", "content": "What is C. elegans?"}], add_generation_prompt=True, return_tensors="pt"
|
| 843 |
+
... )
|
| 844 |
+
>>> generate_ids = model.generate(**inputs, max_new_tokens=64)
|
| 845 |
+
>>> tokenizer.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True)
|
| 846 |
+
```"""
|
| 847 |
+
outputs: BaseModelOutputWithPast = self.model(
|
| 848 |
+
input_ids=input_ids,
|
| 849 |
+
attention_mask=attention_mask,
|
| 850 |
+
position_ids=position_ids,
|
| 851 |
+
past_key_values=past_key_values,
|
| 852 |
+
inputs_embeds=inputs_embeds,
|
| 853 |
+
use_cache=use_cache,
|
| 854 |
+
**kwargs,
|
| 855 |
+
)
|
| 856 |
+
|
| 857 |
+
hidden_states = outputs.last_hidden_state
|
| 858 |
+
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
|
| 859 |
+
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
| 860 |
+
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
| 861 |
+
|
| 862 |
+
loss = None
|
| 863 |
+
if labels is not None:
|
| 864 |
+
loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
|
| 865 |
+
|
| 866 |
+
return CausalLMOutputWithPast(
|
| 867 |
+
loss=loss,
|
| 868 |
+
logits=logits,
|
| 869 |
+
past_key_values=outputs.past_key_values,
|
| 870 |
+
hidden_states=outputs.hidden_states,
|
| 871 |
+
attentions=outputs.attentions,
|
| 872 |
+
)
|
| 873 |
+
|
| 874 |
+
def _prefill(self, input_ids, generation_config, model_kwargs, is_first_iteration=True):
|
| 875 |
+
# every prefill passes here once before its first forward: a check per chunk would leave the earlier chunks
|
| 876 |
+
# in the cache, and a dynamic cache would only refuse packing at the first decode step
|
| 877 |
+
reject_packed_generation(self.config, input_ids, generation_config, model_kwargs, is_first_iteration)
|
| 878 |
+
return super()._prefill(input_ids, generation_config, model_kwargs, is_first_iteration=is_first_iteration)
|
| 879 |
+
|
| 880 |
+
|
| 881 |
+
__all__ = ["Lfm3ForCausalLM", "Lfm3Model", "Lfm3PreTrainedModel"]
|
audio_tokenizer/modeling_liquid_codec.py
ADDED
|
@@ -0,0 +1,381 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
import math
|
| 16 |
+
from dataclasses import dataclass
|
| 17 |
+
|
| 18 |
+
import torch
|
| 19 |
+
import torch.nn.functional as F
|
| 20 |
+
from torch import nn
|
| 21 |
+
|
| 22 |
+
from transformers import initialization as init
|
| 23 |
+
from transformers.cache_utils import DynamicCache
|
| 24 |
+
from transformers.conversion_mapping import register_checkpoint_conversion_mapping
|
| 25 |
+
from transformers.core_model_loading import WeightRenaming
|
| 26 |
+
from transformers.modeling_utils import PreTrainedAudioTokenizerBase
|
| 27 |
+
from transformers.utils import ModelOutput, auto_docstring, can_return_tuple, torch_compilable_check
|
| 28 |
+
from transformers.utils.import_utils import is_tracing
|
| 29 |
+
|
| 30 |
+
from .configuration_liquid_codec import LiquidCodecConfig
|
| 31 |
+
from .modeling_lfm3 import Lfm3Model
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
@auto_docstring
|
| 35 |
+
@dataclass
|
| 36 |
+
class LiquidCodecEncoderOutput(ModelOutput):
|
| 37 |
+
r"""
|
| 38 |
+
audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, num_frames)`):
|
| 39 |
+
The code of every codebook of every frame.
|
| 40 |
+
"""
|
| 41 |
+
|
| 42 |
+
audio_codes: torch.LongTensor = None
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
@auto_docstring
|
| 46 |
+
@dataclass
|
| 47 |
+
class LiquidCodecDecoderOutput(ModelOutput):
|
| 48 |
+
r"""
|
| 49 |
+
audio_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`):
|
| 50 |
+
The decoded 24 kHz audio in float32.
|
| 51 |
+
"""
|
| 52 |
+
|
| 53 |
+
audio_values: torch.FloatTensor = None
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
@auto_docstring
|
| 57 |
+
@dataclass
|
| 58 |
+
class LiquidCodecOutput(ModelOutput):
|
| 59 |
+
r"""
|
| 60 |
+
audio_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`):
|
| 61 |
+
The reconstructed audio, cut to the input length.
|
| 62 |
+
audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, num_frames)`):
|
| 63 |
+
The code of every codebook of every frame.
|
| 64 |
+
"""
|
| 65 |
+
|
| 66 |
+
audio_values: torch.FloatTensor = None
|
| 67 |
+
audio_codes: torch.LongTensor = None
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
class LiquidCodecStreamingCache:
|
| 71 |
+
"""Carries the decoder state between chunks: the LFM3 cache and the unfinished overlap-add tail."""
|
| 72 |
+
|
| 73 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 74 |
+
self.past_key_values = DynamicCache(config=config.lfm_config(config.decoder_window))
|
| 75 |
+
self.audio: torch.Tensor | None = None
|
| 76 |
+
self.envelope: torch.Tensor | None = None
|
| 77 |
+
self.started = False
|
| 78 |
+
self.finished = False
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
class LiquidCodecSTFT(nn.Module):
|
| 82 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 83 |
+
super().__init__()
|
| 84 |
+
self.n_fft = config.stft_n_fft
|
| 85 |
+
self.hop_length = config.stft_hop_length
|
| 86 |
+
self.win_length = config.stft_win_length
|
| 87 |
+
self.window = nn.Buffer(torch.hann_window(self.win_length, dtype=torch.float32), persistent=False)
|
| 88 |
+
self.num_features = 3 * (self.n_fft // 2 + 1)
|
| 89 |
+
|
| 90 |
+
def forward(self, audio: torch.Tensor) -> torch.Tensor:
|
| 91 |
+
total_pad = self.win_length - self.hop_length
|
| 92 |
+
audio = F.pad(audio.float(), (total_pad // 2, total_pad - total_pad // 2), mode="reflect")
|
| 93 |
+
spec = torch.stft(
|
| 94 |
+
audio,
|
| 95 |
+
n_fft=self.n_fft,
|
| 96 |
+
hop_length=self.hop_length,
|
| 97 |
+
win_length=self.win_length,
|
| 98 |
+
window=self.window.float(),
|
| 99 |
+
center=False,
|
| 100 |
+
return_complex=True,
|
| 101 |
+
)
|
| 102 |
+
magnitude = spec.abs()
|
| 103 |
+
phase_denominator = magnitude.clamp_min(1e-5)
|
| 104 |
+
return torch.cat((torch.log1p(magnitude), spec.real / phase_denominator, spec.imag / phase_denominator), dim=1)
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
class LiquidCodecFrontEnd(nn.Module):
|
| 108 |
+
def __init__(self, config: LiquidCodecConfig, num_features: int):
|
| 109 |
+
super().__init__()
|
| 110 |
+
self.frame_stack = config.stft_frame_stack
|
| 111 |
+
self.in_proj = nn.Conv1d(num_features, config.encoder_conv_dim, kernel_size=1)
|
| 112 |
+
self.conv = nn.Conv1d(config.encoder_conv_dim, config.encoder_conv_dim, kernel_size=3, padding=2)
|
| 113 |
+
self.out = nn.Conv1d(
|
| 114 |
+
config.encoder_conv_dim, config.encoder_dim, kernel_size=self.frame_stack, stride=self.frame_stack
|
| 115 |
+
)
|
| 116 |
+
self.act = nn.SiLU()
|
| 117 |
+
|
| 118 |
+
def forward(self, features: torch.Tensor) -> torch.Tensor:
|
| 119 |
+
hidden_states = self.act(self.in_proj(features.to(self.in_proj.weight.dtype)))
|
| 120 |
+
hidden_states = self.act(self.conv(hidden_states)[:, :, : hidden_states.shape[-1]])
|
| 121 |
+
return self.out(hidden_states)
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
class LiquidCodecLfm(nn.Module):
|
| 125 |
+
def __init__(self, config: LiquidCodecConfig, sliding_window: int | None = None):
|
| 126 |
+
super().__init__()
|
| 127 |
+
self.lfm = Lfm3Model(config.lfm_config(sliding_window))
|
| 128 |
+
self.lfm.register_module("embed_tokens", None)
|
| 129 |
+
# the projections run in the model dtype but the rotary tables stay in float32; a named hook keeps the
|
| 130 |
+
# module picklable
|
| 131 |
+
self.lfm.rotary_emb.register_forward_pre_hook(self._rotary_in_float32)
|
| 132 |
+
|
| 133 |
+
@staticmethod
|
| 134 |
+
def _rotary_in_float32(module, args):
|
| 135 |
+
return (args[0].float(), *args[1:])
|
| 136 |
+
|
| 137 |
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 138 |
+
return self.lfm(inputs_embeds=hidden_states, use_cache=False).last_hidden_state
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
class LiquidCodecQuantizer(nn.Module):
|
| 142 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 143 |
+
super().__init__()
|
| 144 |
+
levels = list(config.fsq_levels)
|
| 145 |
+
self.num_codebooks = config.n_q
|
| 146 |
+
self.codebook_dim = len(levels)
|
| 147 |
+
self.levels = nn.Buffer(torch.tensor(levels, dtype=torch.int32), persistent=False)
|
| 148 |
+
self.basis = nn.Buffer(
|
| 149 |
+
torch.tensor([math.prod(levels[:index]) for index in range(self.codebook_dim)], dtype=torch.int32),
|
| 150 |
+
persistent=False,
|
| 151 |
+
)
|
| 152 |
+
effective_dim = self.codebook_dim * self.num_codebooks
|
| 153 |
+
if config.quantizer_dim == effective_dim:
|
| 154 |
+
self.project_in = nn.Identity()
|
| 155 |
+
self.project_out = nn.Identity()
|
| 156 |
+
else:
|
| 157 |
+
self.project_in = nn.Linear(config.quantizer_dim, effective_dim)
|
| 158 |
+
self.project_out = nn.Linear(effective_dim, config.quantizer_dim)
|
| 159 |
+
|
| 160 |
+
def _bound(self, latents: torch.Tensor, eps: float = 1e-3) -> torch.Tensor:
|
| 161 |
+
half_levels = (self.levels - 1) * (1 + eps) / 2
|
| 162 |
+
offset = torch.where(self.levels % 2 == 0, 0.5, 0.0)
|
| 163 |
+
shift = torch.atanh(offset / half_levels)
|
| 164 |
+
return torch.tanh(latents + shift) * half_levels - offset
|
| 165 |
+
|
| 166 |
+
def encode(self, latents: torch.Tensor) -> torch.Tensor:
|
| 167 |
+
# the quantizer runs outside `forward`, so no `device_map` hook moves its inputs (here and in `decode`) to the
|
| 168 |
+
# device of its buffers
|
| 169 |
+
latents = self.project_in(latents.movedim(1, -1).contiguous()).float().to(self.levels.device)
|
| 170 |
+
if not is_tracing(latents):
|
| 171 |
+
# finite audio far above full scale overflows the float32 STFT, and NaN latents round to INT64_MIN codes
|
| 172 |
+
torch_compilable_check(
|
| 173 |
+
torch.isfinite(latents),
|
| 174 |
+
"The audio overflows the codec (non-finite latents); expected samples in [-1, 1]",
|
| 175 |
+
)
|
| 176 |
+
latents = latents.reshape(*latents.shape[:-1], self.num_codebooks, self.codebook_dim)
|
| 177 |
+
half_width = self.levels // 2
|
| 178 |
+
codes = self._bound(latents).round() / half_width
|
| 179 |
+
indices = ((codes * half_width + half_width) * self.basis).sum(-1).round().long()
|
| 180 |
+
return indices.transpose(1, 2)
|
| 181 |
+
|
| 182 |
+
def decode(self, codes: torch.Tensor) -> torch.Tensor:
|
| 183 |
+
indices = codes.to(self.basis.device).transpose(1, 2).unsqueeze(-1)
|
| 184 |
+
half_width = self.levels // 2
|
| 185 |
+
latents = ((indices // self.basis) % self.levels - half_width) / half_width
|
| 186 |
+
latents = latents.reshape(*latents.shape[:-2], -1)
|
| 187 |
+
if isinstance(self.project_out, nn.Linear):
|
| 188 |
+
latents = latents.to(self.project_out.weight.dtype)
|
| 189 |
+
return self.project_out(latents).movedim(-1, 1)
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
class LiquidCodecISTFT(nn.Module):
|
| 193 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 194 |
+
super().__init__()
|
| 195 |
+
self.n_fft = config.stft_n_fft
|
| 196 |
+
self.hop_length = config.stft_hop_length
|
| 197 |
+
self.win_length = config.stft_win_length
|
| 198 |
+
self.window = nn.Buffer(torch.hann_window(self.win_length, dtype=torch.float32))
|
| 199 |
+
|
| 200 |
+
def forward(self, spec: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
| 201 |
+
num_frames = spec.shape[-1]
|
| 202 |
+
frames = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") * self.window[None, :, None]
|
| 203 |
+
output_size = (num_frames - 1) * self.hop_length + self.win_length
|
| 204 |
+
fold = {"output_size": (1, output_size), "kernel_size": (1, self.win_length), "stride": (1, self.hop_length)}
|
| 205 |
+
audio = F.fold(frames, **fold)[:, 0, 0]
|
| 206 |
+
envelope = F.fold(self.window.square().expand(1, num_frames, -1).transpose(1, 2), **fold)[0, 0, 0]
|
| 207 |
+
return audio, envelope
|
| 208 |
+
|
| 209 |
+
|
| 210 |
+
class LiquidCodecDecoder(nn.Module):
|
| 211 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 212 |
+
super().__init__()
|
| 213 |
+
self.frame_stack = config.stft_frame_stack
|
| 214 |
+
self.lfm = LiquidCodecLfm(config, sliding_window=config.decoder_window).lfm
|
| 215 |
+
self.lin = nn.Linear(config.encoder_dim, 2 * (config.stft_n_fft // 2 + 1))
|
| 216 |
+
self.istft = LiquidCodecISTFT(config)
|
| 217 |
+
|
| 218 |
+
def forward(
|
| 219 |
+
self, embeddings: torch.Tensor, past_key_values: DynamicCache | None = None
|
| 220 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 221 |
+
hidden_states = self.lfm(
|
| 222 |
+
inputs_embeds=F.interpolate(
|
| 223 |
+
embeddings.mT, embeddings.shape[1] * self.frame_stack, mode="nearest-exact"
|
| 224 |
+
).mT,
|
| 225 |
+
past_key_values=past_key_values,
|
| 226 |
+
use_cache=past_key_values is not None,
|
| 227 |
+
).last_hidden_state
|
| 228 |
+
log_magnitude, angle = torch.chunk(self.lin(hidden_states).mT.contiguous(), 2, 1)
|
| 229 |
+
return self.istft(torch.polar(log_magnitude.float().exp(), angle.float()))
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
@auto_docstring(
|
| 233 |
+
custom_intro="The Liquid Codec neural audio codec: 24 kHz audio to 12 codebooks at 12.5 frames per second and back."
|
| 234 |
+
)
|
| 235 |
+
class LiquidCodecModel(PreTrainedAudioTokenizerBase):
|
| 236 |
+
config: LiquidCodecConfig
|
| 237 |
+
base_model_prefix = "codec"
|
| 238 |
+
main_input_name = "input_values"
|
| 239 |
+
# the analysis window is a non-persistent buffer cast in the forward; only the synthesis window is a weight
|
| 240 |
+
_keep_in_fp32_modules_strict = ["istft"]
|
| 241 |
+
_supports_flash_attn = False
|
| 242 |
+
_supports_sdpa = True
|
| 243 |
+
|
| 244 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 245 |
+
super().__init__(config)
|
| 246 |
+
self.stft = LiquidCodecSTFT(config)
|
| 247 |
+
self.frontend = LiquidCodecFrontEnd(config, self.stft.num_features)
|
| 248 |
+
self.encoder = LiquidCodecLfm(config)
|
| 249 |
+
self.to_quantizer = nn.Linear(config.encoder_dim, config.quantizer_dim)
|
| 250 |
+
self.quantizer = LiquidCodecQuantizer(config)
|
| 251 |
+
self.decoder = LiquidCodecDecoder(config)
|
| 252 |
+
self.from_quantizer = nn.Linear(config.quantizer_dim, config.encoder_dim)
|
| 253 |
+
self.post_init()
|
| 254 |
+
|
| 255 |
+
@torch.no_grad()
|
| 256 |
+
def _init_weights(self, module):
|
| 257 |
+
super()._init_weights(module)
|
| 258 |
+
if isinstance(module, (LiquidCodecSTFT, LiquidCodecISTFT)):
|
| 259 |
+
init.copy_(module.window, torch.hann_window(module.win_length, dtype=torch.float32))
|
| 260 |
+
elif isinstance(module, LiquidCodecQuantizer):
|
| 261 |
+
levels = list(self.config.fsq_levels)
|
| 262 |
+
init.copy_(module.levels, torch.tensor(levels, dtype=torch.int32))
|
| 263 |
+
init.copy_(
|
| 264 |
+
module.basis,
|
| 265 |
+
torch.tensor([math.prod(levels[:index]) for index in range(len(levels))], dtype=torch.int32),
|
| 266 |
+
)
|
| 267 |
+
|
| 268 |
+
@can_return_tuple
|
| 269 |
+
def encode(self, input_values: torch.FloatTensor, *, return_dict: bool | None = None) -> LiquidCodecEncoderOutput:
|
| 270 |
+
r"""
|
| 271 |
+
Non-finite audio, or audio that overflows the codec, is refused eagerly; a compiled `encode` skips the check.
|
| 272 |
+
|
| 273 |
+
input_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`):
|
| 274 |
+
24 kHz audio. The tail is zero-padded to a whole codec frame of 1920 samples.
|
| 275 |
+
return_dict (`bool`, *optional*):
|
| 276 |
+
Whether to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
| 277 |
+
"""
|
| 278 |
+
if input_values.ndim != 2 or input_values.shape[-1] == 0:
|
| 279 |
+
raise ValueError(
|
| 280 |
+
f"Expected non-empty audio of shape (batch_size, num_samples), got {tuple(input_values.shape)}"
|
| 281 |
+
)
|
| 282 |
+
if not is_tracing(input_values):
|
| 283 |
+
# a non-finite sample would round to an INT64_MIN code that only fails far downstream; the tensor-valued
|
| 284 |
+
# check would break a `fullgraph` compile or a CUDA-graph capture
|
| 285 |
+
torch_compilable_check(torch.isfinite(input_values), "`input_values` must be finite")
|
| 286 |
+
frame_size = self.config.frame_size
|
| 287 |
+
num_samples = -(input_values.shape[-1] // -frame_size) * frame_size
|
| 288 |
+
input_values = F.pad(input_values, (0, num_samples - input_values.shape[-1]))
|
| 289 |
+
hidden_states = self.frontend(self.stft(input_values)).mT
|
| 290 |
+
latents = self.to_quantizer(self.encoder(hidden_states)).mT
|
| 291 |
+
return LiquidCodecEncoderOutput(audio_codes=self.quantizer.encode(latents))
|
| 292 |
+
|
| 293 |
+
@can_return_tuple
|
| 294 |
+
def decode(
|
| 295 |
+
self,
|
| 296 |
+
audio_codes: torch.LongTensor,
|
| 297 |
+
streaming_cache: LiquidCodecStreamingCache | None = None,
|
| 298 |
+
finalize: bool = True,
|
| 299 |
+
*,
|
| 300 |
+
return_dict: bool | None = None,
|
| 301 |
+
) -> LiquidCodecDecoderOutput:
|
| 302 |
+
r"""
|
| 303 |
+
audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, num_frames)`):
|
| 304 |
+
The codes to decode; every frame becomes 1920 samples.
|
| 305 |
+
streaming_cache ([`LiquidCodecStreamingCache`], *optional*):
|
| 306 |
+
The state carried between consecutive chunks of one stream. The chunks decode to the same audio as the
|
| 307 |
+
concatenated codes.
|
| 308 |
+
finalize (`bool`, *optional*, defaults to `True`):
|
| 309 |
+
Whether this is the last chunk of the stream. Until then the trailing samples that later frames still
|
| 310 |
+
contribute to are held back in the cache.
|
| 311 |
+
return_dict (`bool`, *optional*):
|
| 312 |
+
Whether to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
| 313 |
+
"""
|
| 314 |
+
if audio_codes.ndim != 3 or audio_codes.shape[1] != self.config.n_q:
|
| 315 |
+
raise ValueError(
|
| 316 |
+
f"Expected codes of shape (batch_size, {self.config.n_q}, num_frames), got {tuple(audio_codes.shape)}"
|
| 317 |
+
)
|
| 318 |
+
torch_compilable_check(
|
| 319 |
+
(audio_codes >= 0) & (audio_codes < self.config.codebook_size),
|
| 320 |
+
f"Codes must be in [0, {self.config.codebook_size})",
|
| 321 |
+
)
|
| 322 |
+
if streaming_cache is not None and streaming_cache.finished:
|
| 323 |
+
raise ValueError("This stream is finalized; start with a new cache")
|
| 324 |
+
trim = (self.config.stft_win_length - self.config.stft_hop_length) // 2
|
| 325 |
+
tail = 2 * trim
|
| 326 |
+
batch_size = audio_codes.shape[0]
|
| 327 |
+
if not audio_codes.shape[-1] and not (finalize and streaming_cache is not None and streaming_cache.started):
|
| 328 |
+
# no frames and no held-back tail to flush: the overlap-add state must stay untouched
|
| 329 |
+
if streaming_cache is not None and finalize:
|
| 330 |
+
streaming_cache.finished = True
|
| 331 |
+
return LiquidCodecDecoderOutput(audio_values=audio_codes.new_zeros((batch_size, 0), dtype=torch.float32))
|
| 332 |
+
if audio_codes.shape[-1]:
|
| 333 |
+
# `quantizer.decode` returns float32 latents when its projections are `nn.Identity`
|
| 334 |
+
latents = self.quantizer.decode(audio_codes).mT.to(self.from_quantizer.weight.dtype)
|
| 335 |
+
embeddings = self.from_quantizer(latents)
|
| 336 |
+
past_key_values = None if streaming_cache is None else streaming_cache.past_key_values
|
| 337 |
+
audio, envelope = self.decoder(embeddings, past_key_values)
|
| 338 |
+
else:
|
| 339 |
+
audio = self.decoder.istft.window.new_zeros(batch_size, tail)
|
| 340 |
+
envelope = self.decoder.istft.window.new_zeros(tail)
|
| 341 |
+
if streaming_cache is None:
|
| 342 |
+
audio, envelope = audio[:, trim:-trim], envelope[trim:-trim]
|
| 343 |
+
else:
|
| 344 |
+
if streaming_cache.audio is not None:
|
| 345 |
+
audio = torch.cat([audio[:, :tail] + streaming_cache.audio, audio[:, tail:]], dim=1)
|
| 346 |
+
envelope = torch.cat([envelope[:tail] + streaming_cache.envelope, envelope[tail:]])
|
| 347 |
+
if not streaming_cache.started:
|
| 348 |
+
audio, envelope = audio[:, trim:], envelope[trim:]
|
| 349 |
+
streaming_cache.started = True
|
| 350 |
+
if finalize:
|
| 351 |
+
audio, envelope = audio[:, :-trim], envelope[:-trim]
|
| 352 |
+
streaming_cache.audio, streaming_cache.envelope = None, None
|
| 353 |
+
streaming_cache.finished = True
|
| 354 |
+
else:
|
| 355 |
+
streaming_cache.audio, streaming_cache.envelope = audio[:, -tail:], envelope[-tail:]
|
| 356 |
+
audio, envelope = audio[:, :-tail], envelope[:-tail]
|
| 357 |
+
return LiquidCodecDecoderOutput(audio_values=(audio / envelope).float())
|
| 358 |
+
|
| 359 |
+
@can_return_tuple
|
| 360 |
+
@auto_docstring
|
| 361 |
+
def forward(self, input_values: torch.FloatTensor) -> LiquidCodecOutput:
|
| 362 |
+
r"""
|
| 363 |
+
input_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`):
|
| 364 |
+
24 kHz audio.
|
| 365 |
+
"""
|
| 366 |
+
audio_codes = self.encode(input_values, return_dict=True).audio_codes
|
| 367 |
+
audio = self.decode(audio_codes, return_dict=True).audio_values
|
| 368 |
+
return LiquidCodecOutput(audio_values=audio[..., : input_values.shape[-1]], audio_codes=audio_codes)
|
| 369 |
+
|
| 370 |
+
|
| 371 |
+
__all__ = ["LiquidCodecModel", "LiquidCodecStreamingCache"]
|
| 372 |
+
|
| 373 |
+
|
| 374 |
+
register_checkpoint_conversion_mapping(
|
| 375 |
+
"liquid_codec",
|
| 376 |
+
[
|
| 377 |
+
WeightRenaming(source_patterns=r"^_frontend\.", target_patterns="frontend."),
|
| 378 |
+
WeightRenaming(source_patterns=r"^quantizer\.fsq\.", target_patterns="quantizer."),
|
| 379 |
+
],
|
| 380 |
+
overwrite=True,
|
| 381 |
+
)
|
chat_template.jinja
ADDED
|
@@ -0,0 +1,296 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{{- bos_token -}}
|
| 2 |
+
{%- set preserve_thinking = preserve_thinking | default(false) -%}
|
| 3 |
+
{%- set enable_thinking = enable_thinking | default(false) -%}
|
| 4 |
+
|
| 5 |
+
{%- macro render_value(v) -%}
|
| 6 |
+
{%- if v is string -%}
|
| 7 |
+
{{- v -}}
|
| 8 |
+
{%- elif v is boolean -%}
|
| 9 |
+
{{- "true" if v else "false" -}}
|
| 10 |
+
{%- elif v is none -%}
|
| 11 |
+
{{- "null" -}}
|
| 12 |
+
{%- elif v is mapping or (v is sequence and v is not string) -%}
|
| 13 |
+
{{- v | tojson -}}
|
| 14 |
+
{%- else -%}
|
| 15 |
+
{{- v -}}
|
| 16 |
+
{%- endif -%}
|
| 17 |
+
{%- endmacro -%}
|
| 18 |
+
|
| 19 |
+
{%- macro parse_content(content) -%}
|
| 20 |
+
{%- if content is string -%}
|
| 21 |
+
{{- content -}}
|
| 22 |
+
{%- elif content is mapping and content.get("type") == "image" -%}
|
| 23 |
+
{{- "<|vision_start|><|image|><|vision_end|>" -}}
|
| 24 |
+
{%- elif content is mapping and content.get("type") == "video" -%}
|
| 25 |
+
{{- "<|vision_start|><|video|><|vision_end|>" -}}
|
| 26 |
+
{%- elif content is mapping and content.get("type") in ["audio", "audio_url", "input_audio"] -%}
|
| 27 |
+
{{- "<|vision_start|><|audio|><|vision_end|>" -}}
|
| 28 |
+
{%- elif content is mapping and content.get("type") == "audio_out" -%}
|
| 29 |
+
{{- "<|audio_out|>" -}}
|
| 30 |
+
{%- elif content is mapping and content.get("type") == "text" -%}
|
| 31 |
+
{{- (content.get("text") or "") | string -}}
|
| 32 |
+
{%- elif content is mapping -%}
|
| 33 |
+
{{- content | tojson -}}
|
| 34 |
+
{%- elif content is iterable -%}
|
| 35 |
+
{%- set _ns = namespace(result="") -%}
|
| 36 |
+
{%- for item in content -%}
|
| 37 |
+
{%- if item is string -%}
|
| 38 |
+
{%- set _ns.result = _ns.result + item -%}
|
| 39 |
+
{%- elif item is mapping and item.get("type") == "image" -%}
|
| 40 |
+
{%- set _ns.result = _ns.result + "<|vision_start|><|image|><|vision_end|>" -%}
|
| 41 |
+
{%- elif item is mapping and item.get("type") == "video" -%}
|
| 42 |
+
{%- set _ns.result = _ns.result + "<|vision_start|><|video|><|vision_end|>" -%}
|
| 43 |
+
{%- elif item is mapping and item.get("type") in ["audio", "audio_url", "input_audio"] -%}
|
| 44 |
+
{%- set _ns.result = _ns.result + "<|vision_start|><|audio|><|vision_end|>" -%}
|
| 45 |
+
{%- elif item is mapping and item.get("type") == "audio_out" -%}
|
| 46 |
+
{%- set _ns.result = _ns.result + "<|audio_out|>" -%}
|
| 47 |
+
{%- elif item is mapping and item.get("type") == "text" -%}
|
| 48 |
+
{%- set _ns.result = _ns.result + ((item.get("text") or "") | string) -%}
|
| 49 |
+
{%- else -%}
|
| 50 |
+
{%- set _ns.result = _ns.result + (item | tojson) -%}
|
| 51 |
+
{%- endif -%}
|
| 52 |
+
{%- endfor -%}
|
| 53 |
+
{{- _ns.result -}}
|
| 54 |
+
{%- endif -%}
|
| 55 |
+
{%- endmacro -%}
|
| 56 |
+
|
| 57 |
+
{%- macro render_tool_calls(tool_calls) -%}
|
| 58 |
+
{%- for tool_call in tool_calls -%}
|
| 59 |
+
{%- if not loop.first -%}{{- "\n\n" -}}{%- endif -%}
|
| 60 |
+
{%- set func = tool_call["function"] if "function" in tool_call else tool_call -%}
|
| 61 |
+
{%- set func_args = func.get("arguments") -%}
|
| 62 |
+
{%- if func_args is string and (func_args | trim) not in ["", "{}", "null"] -%}
|
| 63 |
+
{{- raise_exception("Tool call arguments must be a mapping, got a JSON-encoded string: parse arguments with json.loads() before applying the chat template") -}}
|
| 64 |
+
{%- endif -%}
|
| 65 |
+
{{- "<tool_call>\n<function=" + func["name"] + ">\n" -}}
|
| 66 |
+
{%- if func_args is mapping -%}
|
| 67 |
+
{%- for arg_name, arg_value in func_args.items() -%}
|
| 68 |
+
{{- "<parameter=" + arg_name + ">\n" + render_value(arg_value) + "\n</parameter>\n" -}}
|
| 69 |
+
{%- endfor -%}
|
| 70 |
+
{%- endif -%}
|
| 71 |
+
{{- "</function>\n</tool_call>" -}}
|
| 72 |
+
{%- endfor -%}
|
| 73 |
+
{%- endmacro -%}
|
| 74 |
+
|
| 75 |
+
{%- macro tool_to_json(tool) -%}
|
| 76 |
+
{%- set fn = tool.function if tool.function is defined else tool -%}
|
| 77 |
+
{%- set ns_t = namespace(first=true) -%}
|
| 78 |
+
{{- "{" -}}
|
| 79 |
+
{%- for k, v in fn.items() -%}
|
| 80 |
+
{%- if k != "defer_loading" -%}
|
| 81 |
+
{%- if not ns_t.first -%}{{- ", " -}}{%- endif -%}
|
| 82 |
+
{%- set ns_t.first = false -%}
|
| 83 |
+
{{- k | tojson -}}{{- ": " -}}{{- v | tojson -}}
|
| 84 |
+
{%- endif -%}
|
| 85 |
+
{%- endfor -%}
|
| 86 |
+
{{- "}" -}}
|
| 87 |
+
{%- endmacro -%}
|
| 88 |
+
|
| 89 |
+
{#- a tool_reference content part (vLLM, SGLang) names a tool from `tools`,
|
| 90 |
+
usually a deferred one; it renders as that tool's definition -#}
|
| 91 |
+
{%- macro is_tool_reference(item) -%}
|
| 92 |
+
{{- "1" if item is mapping and item.get("type") == "tool_reference" else "" -}}
|
| 93 |
+
{%- endmacro -%}
|
| 94 |
+
|
| 95 |
+
{%- macro referenced_tool_to_json(name) -%}
|
| 96 |
+
{%- set _ns = namespace(found=none) -%}
|
| 97 |
+
{%- for tool in tools or [] -%}
|
| 98 |
+
{%- set fn = tool.function if tool.function is defined else tool -%}
|
| 99 |
+
{%- if fn.name == name -%}
|
| 100 |
+
{%- set _ns.found = fn -%}
|
| 101 |
+
{%- endif -%}
|
| 102 |
+
{%- endfor -%}
|
| 103 |
+
{%- if _ns.found is none -%}
|
| 104 |
+
{{- raise_exception("tool_reference names a tool that is not in tools: " + name | string) -}}
|
| 105 |
+
{%- endif -%}
|
| 106 |
+
{{- tool_to_json(_ns.found) -}}
|
| 107 |
+
{%- endmacro -%}
|
| 108 |
+
|
| 109 |
+
{%- macro render_tool_response(content) -%}
|
| 110 |
+
{#- content that already carries the tag (pre-formatted by the caller) is
|
| 111 |
+
passed through as-is instead of being wrapped a second time -#}
|
| 112 |
+
{%- if content is string and "<tool_response>" in content -%}
|
| 113 |
+
{{- content -}}
|
| 114 |
+
{%- else -%}
|
| 115 |
+
{{- "<tool_response>\n" -}}
|
| 116 |
+
{%- if content is string -%}
|
| 117 |
+
{{- content -}}
|
| 118 |
+
{%- elif content is mapping -%}
|
| 119 |
+
{{- parse_content(content) -}}
|
| 120 |
+
{%- elif content is iterable -%}
|
| 121 |
+
{%- for item in content -%}
|
| 122 |
+
{%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
|
| 123 |
+
{%- if is_tool_reference(item) -%}
|
| 124 |
+
{#- consecutive references share one <additional_tools> block -#}
|
| 125 |
+
{%- if loop.first or not is_tool_reference(loop.previtem) -%}
|
| 126 |
+
{{- "<additional_tools>\n" -}}
|
| 127 |
+
{%- endif -%}
|
| 128 |
+
{{- referenced_tool_to_json(item.name) -}}
|
| 129 |
+
{%- if loop.last or not is_tool_reference(loop.nextitem) -%}
|
| 130 |
+
{{- "\n</additional_tools>" -}}
|
| 131 |
+
{%- endif -%}
|
| 132 |
+
{%- elif item is mapping and item.result is defined and item.result is mapping and item.result.tools is defined -%}
|
| 133 |
+
{{- "<additional_tools>\n" -}}
|
| 134 |
+
{%- set _tl = item.result.tools -%}
|
| 135 |
+
{%- if _tl is string -%}
|
| 136 |
+
{{- _tl -}}
|
| 137 |
+
{%- elif _tl is mapping -%}
|
| 138 |
+
{{- tool_to_json(_tl) -}}
|
| 139 |
+
{%- else -%}
|
| 140 |
+
{%- for t in _tl -%}
|
| 141 |
+
{%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
|
| 142 |
+
{{- tool_to_json(t) -}}
|
| 143 |
+
{%- endfor -%}
|
| 144 |
+
{%- endif -%}
|
| 145 |
+
{{- "\n</additional_tools>" -}}
|
| 146 |
+
{%- else -%}
|
| 147 |
+
{{- parse_content(item) -}}
|
| 148 |
+
{%- endif -%}
|
| 149 |
+
{%- endfor -%}
|
| 150 |
+
{%- endif -%}
|
| 151 |
+
{{- "\n</tool_response>" -}}
|
| 152 |
+
{%- endif -%}
|
| 153 |
+
{%- endmacro -%}
|
| 154 |
+
|
| 155 |
+
{%- set ns = namespace(system_prompt="", last_user_index=-1) -%}
|
| 156 |
+
{%- if messages and messages[0]["role"] == "system" -%}
|
| 157 |
+
{%- if messages[0].get("content") -%}
|
| 158 |
+
{%- set ns.system_prompt = parse_content(messages[0]["content"]) -%}
|
| 159 |
+
{%- endif -%}
|
| 160 |
+
{%- set messages = messages[1:] -%}
|
| 161 |
+
{%- endif -%}
|
| 162 |
+
{%- if tools -%}
|
| 163 |
+
{%- set tools_ns = namespace(lines=[]) -%}
|
| 164 |
+
{%- for tool in tools -%}
|
| 165 |
+
{%- set fn = tool.function if tool.function is defined else tool -%}
|
| 166 |
+
{%- set _dyn = (fn.defer_loading is defined and fn.defer_loading) or (tool.defer_loading is defined and tool.defer_loading) -%}
|
| 167 |
+
{%- if not _dyn -%}
|
| 168 |
+
{%- set tools_ns.lines = tools_ns.lines + [tool_to_json(fn)] -%}
|
| 169 |
+
{%- endif -%}
|
| 170 |
+
{%- endfor -%}
|
| 171 |
+
{%- if tools_ns.lines -%}
|
| 172 |
+
{%- set ns.system_prompt = ns.system_prompt + ("\n\n" if ns.system_prompt else "") + "<tools>\n" + (tools_ns.lines | join("\n")) + "\n</tools>" -%}
|
| 173 |
+
{%- endif -%}
|
| 174 |
+
{%- endif -%}
|
| 175 |
+
{%- if ns.system_prompt -%}
|
| 176 |
+
{{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}}
|
| 177 |
+
{%- endif -%}
|
| 178 |
+
{%- for message in messages -%}
|
| 179 |
+
{%- if message["role"] == "user" -%}
|
| 180 |
+
{%- set ns.last_user_index = loop.index0 -%}
|
| 181 |
+
{%- endif -%}
|
| 182 |
+
{%- endfor -%}
|
| 183 |
+
{#- vLLM's and SGLang's Anthropic /v1/messages adapters split one tool_result that carries
|
| 184 |
+
tool_reference blocks into consecutive tool messages with the same tool_call_id (text first,
|
| 185 |
+
references second). Merge such runs into one tool turn so they render exactly like a single
|
| 186 |
+
OpenAI-format tool message whose content holds the same items. -#}
|
| 187 |
+
{%- set carry = namespace(items=[]) -%}
|
| 188 |
+
{%- for message in messages -%}
|
| 189 |
+
{%- set _next = messages[loop.index0 + 1] if not loop.last else none -%}
|
| 190 |
+
{%- set _merge = message.role == "tool" and not message.get("additional_tools")
|
| 191 |
+
and _next is not none and _next.role == "tool" and not _next.get("additional_tools")
|
| 192 |
+
and message.get("tool_call_id") and _next.get("tool_call_id") == message.get("tool_call_id") -%}
|
| 193 |
+
{%- set _items = [] -%}
|
| 194 |
+
{%- if message.role == "tool" and not message.get("additional_tools") -%}
|
| 195 |
+
{%- set _c = message.get("content") -%}
|
| 196 |
+
{%- if _c is string -%}
|
| 197 |
+
{%- set _items = [{"type": "text", "text": _c}] if _c else [] -%}
|
| 198 |
+
{%- elif _c is mapping -%}
|
| 199 |
+
{%- set _items = [_c] -%}
|
| 200 |
+
{%- elif _c -%}
|
| 201 |
+
{%- set _items = _c | list -%}
|
| 202 |
+
{%- endif -%}
|
| 203 |
+
{%- endif -%}
|
| 204 |
+
{%- if _merge -%}
|
| 205 |
+
{%- set carry.items = carry.items + _items -%}
|
| 206 |
+
{%- else -%}
|
| 207 |
+
{{- "<|im_start|>" + message.role + "\n" -}}
|
| 208 |
+
{%- if message.role == "assistant" -%}
|
| 209 |
+
{%- generation -%}
|
| 210 |
+
{%- set thinking = message.thinking if message.thinking is defined else message.reasoning if message.reasoning is defined else message.reasoning_content -%}
|
| 211 |
+
{%- set keep_thinking = preserve_thinking or loop.index0 > ns.last_user_index -%}
|
| 212 |
+
{%- if thinking is defined and keep_thinking -%}
|
| 213 |
+
{{- "<think>" + (thinking or "") + "</think>" -}}
|
| 214 |
+
{%- elif keep_thinking and "</think>" not in (message.get("content") if message.get("content") is string else "") -%}
|
| 215 |
+
{{- "<think></think>" -}}
|
| 216 |
+
{%- endif -%}
|
| 217 |
+
{%- set _cfm_tag = "CONTINUE_FINAL_MESSAGE_TAG " -%}
|
| 218 |
+
{%- set _has_cfm = false -%}
|
| 219 |
+
{%- set content = "" -%}
|
| 220 |
+
{%- if message.get("content") -%}
|
| 221 |
+
{%- set content = parse_content(message.content) -%}
|
| 222 |
+
{%- endif -%}
|
| 223 |
+
{%- if not keep_thinking -%}
|
| 224 |
+
{%- if "</think>" in content -%}
|
| 225 |
+
{%- set content = content.split("</think>")[-1] | trim -%}
|
| 226 |
+
{%- endif -%}
|
| 227 |
+
{%- set content = "<think></think>" + content -%}
|
| 228 |
+
{%- endif -%}
|
| 229 |
+
{%- if content.endswith(_cfm_tag) -%}
|
| 230 |
+
{%- set _has_cfm = true -%}
|
| 231 |
+
{%- set _trunc_len = (content | length) - (_cfm_tag | length) -%}
|
| 232 |
+
{%- set content = content[:_trunc_len] -%}
|
| 233 |
+
{%- endif -%}
|
| 234 |
+
{{- content -}}
|
| 235 |
+
{%- if message.tool_calls -%}
|
| 236 |
+
{{- render_tool_calls(message.tool_calls) -}}
|
| 237 |
+
{%- endif -%}
|
| 238 |
+
{%- if _has_cfm -%}
|
| 239 |
+
{{- _cfm_tag -}}
|
| 240 |
+
{%- endif -%}
|
| 241 |
+
{{- "<|im_end|>\n" -}}
|
| 242 |
+
{%- endgeneration -%}
|
| 243 |
+
{%- elif message.role == "tool" -%}
|
| 244 |
+
{%- if message.get("additional_tools") -%}
|
| 245 |
+
{%- set _at = message["additional_tools"] -%}
|
| 246 |
+
{#- definitions may live in additional_tools itself, or in content
|
| 247 |
+
with additional_tools used as a plain truthy flag -#}
|
| 248 |
+
{%- if _at is mapping or (_at is sequence and _at is not string) -%}
|
| 249 |
+
{%- set _src = _at -%}
|
| 250 |
+
{%- else -%}
|
| 251 |
+
{%- set _src = message.get("content") -%}
|
| 252 |
+
{%- endif -%}
|
| 253 |
+
{#- nothing to reveal: skip the wrapper entirely rather than
|
| 254 |
+
emitting an empty <additional_tools> block -#}
|
| 255 |
+
{%- if _src is string and "<tool_response>" in _src -%}
|
| 256 |
+
{{- _src -}}
|
| 257 |
+
{%- elif _src is string and "<additional_tools>" in _src -%}
|
| 258 |
+
{{- "<tool_response>\n" + _src + "\n</tool_response>" -}}
|
| 259 |
+
{%- elif _src -%}
|
| 260 |
+
{{- "<tool_response>\n<additional_tools>\n" -}}
|
| 261 |
+
{%- if _src is string -%}
|
| 262 |
+
{{- _src -}}
|
| 263 |
+
{%- elif _src is mapping -%}
|
| 264 |
+
{{- tool_to_json(_src) -}}
|
| 265 |
+
{%- else -%}
|
| 266 |
+
{%- for t in _src -%}
|
| 267 |
+
{%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
|
| 268 |
+
{{- tool_to_json(t) -}}
|
| 269 |
+
{%- endfor -%}
|
| 270 |
+
{%- endif -%}
|
| 271 |
+
{{- "\n</additional_tools>\n</tool_response>" -}}
|
| 272 |
+
{%- endif -%}
|
| 273 |
+
{%- elif carry.items -%}
|
| 274 |
+
{{- render_tool_response(carry.items + _items) -}}
|
| 275 |
+
{%- set carry.items = [] -%}
|
| 276 |
+
{%- else -%}
|
| 277 |
+
{#- an empty or missing result still gets the wrapper, so every tool turn has the same shape -#}
|
| 278 |
+
{{- render_tool_response(message.get("content") or "") -}}
|
| 279 |
+
{%- endif -%}
|
| 280 |
+
{{- "<|im_end|>\n" -}}
|
| 281 |
+
{%- else %}
|
| 282 |
+
{%- if message.get("content") -%}
|
| 283 |
+
{{- parse_content(message["content"]) -}}
|
| 284 |
+
{%- endif -%}
|
| 285 |
+
{{- "<|im_end|>\n" -}}
|
| 286 |
+
{%- endif %}
|
| 287 |
+
{%- endif -%}
|
| 288 |
+
{%- endfor -%}
|
| 289 |
+
{%- if add_generation_prompt -%}
|
| 290 |
+
{{- "<|im_start|>assistant\n" -}}
|
| 291 |
+
{%- if enable_thinking -%}
|
| 292 |
+
{{- "<think>" -}}
|
| 293 |
+
{%- else -%}
|
| 294 |
+
{{- "<think></think>" -}}
|
| 295 |
+
{%- endif -%}
|
| 296 |
+
{%- endif -%}
|
config.json
ADDED
|
@@ -0,0 +1,207 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"Lfm3OmniForConditionalGeneration"
|
| 4 |
+
],
|
| 5 |
+
"audio_encoder_config": {
|
| 6 |
+
"conv_channels": 256,
|
| 7 |
+
"hidden_size": 64,
|
| 8 |
+
"initializer_range": 0.02,
|
| 9 |
+
"model_type": "lfm3_omni_encoder",
|
| 10 |
+
"num_mel_bins": 128,
|
| 11 |
+
"sliding_window": 32,
|
| 12 |
+
"text_config": {
|
| 13 |
+
"_name_or_path": "",
|
| 14 |
+
"architectures": null,
|
| 15 |
+
"block_auto_adjust_ff_dim": true,
|
| 16 |
+
"block_ffn_dim_multiplier": 1.0,
|
| 17 |
+
"block_multiple_of": 256,
|
| 18 |
+
"bos_token_id": 1,
|
| 19 |
+
"chunk_size_feed_forward": 0,
|
| 20 |
+
"conv_L_cache": 3,
|
| 21 |
+
"conv_bias": false,
|
| 22 |
+
"dtype": null,
|
| 23 |
+
"eos_token_id": null,
|
| 24 |
+
"full_attn_idxs": null,
|
| 25 |
+
"hidden_size": 64,
|
| 26 |
+
"id2label": {
|
| 27 |
+
"0": "LABEL_0",
|
| 28 |
+
"1": "LABEL_1"
|
| 29 |
+
},
|
| 30 |
+
"initializer_range": 0.02,
|
| 31 |
+
"intermediate_size": 128,
|
| 32 |
+
"is_encoder_decoder": false,
|
| 33 |
+
"label2id": {
|
| 34 |
+
"LABEL_0": 0,
|
| 35 |
+
"LABEL_1": 1
|
| 36 |
+
},
|
| 37 |
+
"layer_types": [
|
| 38 |
+
"conv",
|
| 39 |
+
"conv",
|
| 40 |
+
"full_attention",
|
| 41 |
+
"conv",
|
| 42 |
+
"conv",
|
| 43 |
+
"full_attention"
|
| 44 |
+
],
|
| 45 |
+
"max_position_embeddings": 128000,
|
| 46 |
+
"model_type": "lfm2",
|
| 47 |
+
"norm_eps": 1e-05,
|
| 48 |
+
"num_attention_heads": 2,
|
| 49 |
+
"num_hidden_layers": 6,
|
| 50 |
+
"num_key_value_heads": 1,
|
| 51 |
+
"output_attentions": false,
|
| 52 |
+
"output_hidden_states": false,
|
| 53 |
+
"pad_token_id": 0,
|
| 54 |
+
"problem_type": null,
|
| 55 |
+
"return_dict": true,
|
| 56 |
+
"rope_parameters": {
|
| 57 |
+
"rope_theta": 1000000.0,
|
| 58 |
+
"rope_type": "default"
|
| 59 |
+
},
|
| 60 |
+
"tie_word_embeddings": true,
|
| 61 |
+
"use_cache": true,
|
| 62 |
+
"vocab_size": 2
|
| 63 |
+
}
|
| 64 |
+
},
|
| 65 |
+
"audio_eos_code": 2016,
|
| 66 |
+
"audio_loss_weight": 1.0,
|
| 67 |
+
"audio_out_token_id": 125032,
|
| 68 |
+
"audio_start_token_id": 125012,
|
| 69 |
+
"audio_token_id": 125031,
|
| 70 |
+
"codebook_size": 2017,
|
| 71 |
+
"depth_config": {
|
| 72 |
+
"bos_token_id": null,
|
| 73 |
+
"conv_L_cache": 7,
|
| 74 |
+
"conv_bias": false,
|
| 75 |
+
"eos_token_id": null,
|
| 76 |
+
"head_dim": 32,
|
| 77 |
+
"hidden_size": 32,
|
| 78 |
+
"initializer_range": 0.02,
|
| 79 |
+
"intermediate_size": 64,
|
| 80 |
+
"layer_types": [
|
| 81 |
+
"full_attention",
|
| 82 |
+
"full_attention",
|
| 83 |
+
"full_attention"
|
| 84 |
+
],
|
| 85 |
+
"max_position_embeddings": 12,
|
| 86 |
+
"model_type": "lfm3",
|
| 87 |
+
"norm_eps": 1e-05,
|
| 88 |
+
"num_attention_heads": 2,
|
| 89 |
+
"num_hidden_layers": 3,
|
| 90 |
+
"num_key_value_heads": 1,
|
| 91 |
+
"pad_token_id": null,
|
| 92 |
+
"rope_parameters": {
|
| 93 |
+
"full_attention": {
|
| 94 |
+
"partial_rotary_factor": 1.0,
|
| 95 |
+
"rope_theta": 1000000.0,
|
| 96 |
+
"rope_type": "default"
|
| 97 |
+
}
|
| 98 |
+
},
|
| 99 |
+
"sliding_window": 256,
|
| 100 |
+
"tie_word_embeddings": true,
|
| 101 |
+
"use_cache": false,
|
| 102 |
+
"use_sandwich_norm": false,
|
| 103 |
+
"vocab_size": 1
|
| 104 |
+
},
|
| 105 |
+
"downsample_factor": 2,
|
| 106 |
+
"image_token_id": 125030,
|
| 107 |
+
"initializer_range": 0.02,
|
| 108 |
+
"interleaved_n_audio": 12,
|
| 109 |
+
"interleaved_n_text": 6,
|
| 110 |
+
"maskgit_schedule": [
|
| 111 |
+
4,
|
| 112 |
+
4,
|
| 113 |
+
4
|
| 114 |
+
],
|
| 115 |
+
"mixed_start_token_id": 125015,
|
| 116 |
+
"model_type": "lfm3_omni",
|
| 117 |
+
"num_codebooks": 12,
|
| 118 |
+
"projector_bias": true,
|
| 119 |
+
"projector_hidden_act": "gelu",
|
| 120 |
+
"projector_hidden_size": 128,
|
| 121 |
+
"projector_layer_norm_eps": 1e-06,
|
| 122 |
+
"text_config": {
|
| 123 |
+
"bos_token_id": 124894,
|
| 124 |
+
"conv_L_cache": 7,
|
| 125 |
+
"conv_bias": false,
|
| 126 |
+
"eos_token_id": 124900,
|
| 127 |
+
"head_dim": 32,
|
| 128 |
+
"hidden_size": 64,
|
| 129 |
+
"initializer_range": 0.02,
|
| 130 |
+
"intermediate_size": 128,
|
| 131 |
+
"layer_types": [
|
| 132 |
+
"conv",
|
| 133 |
+
"conv",
|
| 134 |
+
"full_attention",
|
| 135 |
+
"conv",
|
| 136 |
+
"sliding_attention",
|
| 137 |
+
"conv",
|
| 138 |
+
"full_attention",
|
| 139 |
+
"conv",
|
| 140 |
+
"sliding_attention",
|
| 141 |
+
"conv"
|
| 142 |
+
],
|
| 143 |
+
"max_position_embeddings": 128000,
|
| 144 |
+
"model_type": "lfm3",
|
| 145 |
+
"norm_eps": 1e-05,
|
| 146 |
+
"num_attention_heads": 16,
|
| 147 |
+
"num_hidden_layers": 10,
|
| 148 |
+
"num_key_value_heads": 4,
|
| 149 |
+
"pad_token_id": 124893,
|
| 150 |
+
"rope_parameters": {
|
| 151 |
+
"full_attention": {
|
| 152 |
+
"partial_rotary_factor": 0.5,
|
| 153 |
+
"rope_theta": 1000000.0,
|
| 154 |
+
"rope_type": "default"
|
| 155 |
+
},
|
| 156 |
+
"sliding_attention": {
|
| 157 |
+
"rope_theta": 10000,
|
| 158 |
+
"rope_type": "default"
|
| 159 |
+
}
|
| 160 |
+
},
|
| 161 |
+
"sliding_window": 256,
|
| 162 |
+
"tie_word_embeddings": true,
|
| 163 |
+
"use_cache": true,
|
| 164 |
+
"use_sandwich_norm": true,
|
| 165 |
+
"vocab_size": 128000
|
| 166 |
+
},
|
| 167 |
+
"text_end_token_id": 125014,
|
| 168 |
+
"text_loss_weight": 1.0,
|
| 169 |
+
"tie_word_embeddings": true,
|
| 170 |
+
"transformers_version": "5.19.0.dev0",
|
| 171 |
+
"video_token_id": 125029,
|
| 172 |
+
"vision_config": {
|
| 173 |
+
"conv_kernel_size": 3,
|
| 174 |
+
"dtype": "bfloat16",
|
| 175 |
+
"hidden_size": 128,
|
| 176 |
+
"in_channels": 3,
|
| 177 |
+
"initializer_range": 0.02,
|
| 178 |
+
"intermediate_size": 320,
|
| 179 |
+
"layer_types": [
|
| 180 |
+
"conv",
|
| 181 |
+
"conv",
|
| 182 |
+
"full_attention",
|
| 183 |
+
"conv",
|
| 184 |
+
"full_attention",
|
| 185 |
+
"conv"
|
| 186 |
+
],
|
| 187 |
+
"model_type": "lfm3_vl_vision",
|
| 188 |
+
"norm_eps": 1e-05,
|
| 189 |
+
"num_attention_heads": 4,
|
| 190 |
+
"num_hidden_layers": 6,
|
| 191 |
+
"num_key_value_heads": 2,
|
| 192 |
+
"patch_size": 16,
|
| 193 |
+
"rope_parameters": {
|
| 194 |
+
"rope_theta": 1000000.0,
|
| 195 |
+
"rope_type": "axial"
|
| 196 |
+
},
|
| 197 |
+
"spatial_merge_size": 2,
|
| 198 |
+
"temporal_patch_size": 2
|
| 199 |
+
},
|
| 200 |
+
"auto_map": {
|
| 201 |
+
"AutoConfig": "configuration_lfm3_omni.Lfm3OmniConfig",
|
| 202 |
+
"AutoModel": "modeling_lfm3_omni.Lfm3OmniModel",
|
| 203 |
+
"AutoModelForImageTextToText": "modeling_lfm3_omni.Lfm3OmniForConditionalGeneration",
|
| 204 |
+
"AutoModelForMultimodalLM": "modeling_lfm3_omni.Lfm3OmniForConditionalGeneration",
|
| 205 |
+
"AutoModelForSeq2SeqLM": "modeling_lfm3_omni.Lfm3OmniForConditionalGeneration"
|
| 206 |
+
}
|
| 207 |
+
}
|
configuration_lfm3.py
ADDED
|
@@ -0,0 +1,132 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
from huggingface_hub.dataclasses import strict
|
| 15 |
+
|
| 16 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 17 |
+
from transformers.utils import auto_docstring
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
@auto_docstring(checkpoint="LiquidAI/LFM3-1.7B-Preview")
|
| 21 |
+
@strict
|
| 22 |
+
class Lfm3Config(PreTrainedConfig):
|
| 23 |
+
r"""
|
| 24 |
+
rope_parameters (`dict[str, RopeParameters]`, *optional*):
|
| 25 |
+
RoPE parameters keyed by attention layer type, for example `{"full_attention": {"rope_type": "default",
|
| 26 |
+
"rope_theta": 1000000.0, "partial_rotary_factor": 0.5}, "sliding_attention": {"rope_type": "default",
|
| 27 |
+
"rope_theta": 10000.0}}`. `conv` layers do not use RoPE and have no entry. Full-attention layers rotate the
|
| 28 |
+
first `partial_rotary_factor * head_dim` dimensions and leave the rest unrotated.
|
| 29 |
+
conv_bias (`bool`, *optional*, defaults to `False`):
|
| 30 |
+
Whether to use bias in the conv layers.
|
| 31 |
+
conv_L_cache (`int`, *optional*, defaults to 7):
|
| 32 |
+
L_cache dim in the conv layers.
|
| 33 |
+
layer_types (`list[str]`, *optional*):
|
| 34 |
+
Per-layer operator type, each one of `"full_attention"`, `"sliding_attention"` or `"conv"`. Defaults to the
|
| 35 |
+
LFM3 pattern of a full-attention layer every fourth layer starting from the third, a sliding-attention layer
|
| 36 |
+
in between, and gated short convolutions elsewhere.
|
| 37 |
+
sliding_window (`int`, *optional*, defaults to 256):
|
| 38 |
+
Sliding-window size of the `sliding_attention` layers.
|
| 39 |
+
use_sandwich_norm (`bool`, *optional*, defaults to `True`):
|
| 40 |
+
Whether to apply an RMSNorm to the operator and feed-forward outputs before they enter the residual stream.
|
| 41 |
+
|
| 42 |
+
```python
|
| 43 |
+
>>> from transformers import Lfm3Model, Lfm3Config
|
| 44 |
+
|
| 45 |
+
>>> # Initializing a LFM3 model
|
| 46 |
+
>>> configuration = Lfm3Config()
|
| 47 |
+
|
| 48 |
+
>>> # Initializing a model from the LFM3-1.7B style configuration
|
| 49 |
+
>>> model = Lfm3Model(configuration)
|
| 50 |
+
|
| 51 |
+
>>> # Accessing the model configuration
|
| 52 |
+
>>> configuration = model.config
|
| 53 |
+
```
|
| 54 |
+
"""
|
| 55 |
+
|
| 56 |
+
model_type = "lfm3"
|
| 57 |
+
keys_to_ignore_at_inference = ["past_key_values"]
|
| 58 |
+
default_theta = 1000000.0
|
| 59 |
+
|
| 60 |
+
vocab_size: int = 128_000
|
| 61 |
+
hidden_size: int = 2048
|
| 62 |
+
intermediate_size: int = 7680
|
| 63 |
+
num_hidden_layers: int = 24
|
| 64 |
+
num_attention_heads: int = 16
|
| 65 |
+
num_key_value_heads: int = 4
|
| 66 |
+
max_position_embeddings: int = 128_000
|
| 67 |
+
initializer_range: float = 0.02
|
| 68 |
+
norm_eps: float = 0.00001
|
| 69 |
+
use_cache: bool = True
|
| 70 |
+
pad_token_id: int | None = 124_893
|
| 71 |
+
bos_token_id: int | None = 124_894
|
| 72 |
+
eos_token_id: int | list[int] | None = 124_900
|
| 73 |
+
tie_word_embeddings: bool = True
|
| 74 |
+
rope_parameters: dict | None = None
|
| 75 |
+
conv_bias: bool = False
|
| 76 |
+
conv_L_cache: int = 7
|
| 77 |
+
layer_types: list[str] | None = None
|
| 78 |
+
_attention_layer_types = ("full_attention", "sliding_attention")
|
| 79 |
+
head_dim: int | None = None
|
| 80 |
+
sliding_window: int | None = 256
|
| 81 |
+
use_sandwich_norm: bool = True
|
| 82 |
+
|
| 83 |
+
def __post_init__(self, **kwargs):
|
| 84 |
+
if self.head_dim is None:
|
| 85 |
+
self.head_dim = self.hidden_size // self.num_attention_heads
|
| 86 |
+
if self.rope_parameters is None:
|
| 87 |
+
self.rope_parameters = {
|
| 88 |
+
"full_attention": {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 0.5},
|
| 89 |
+
"sliding_attention": {"rope_type": "default", "rope_theta": 10_000.0},
|
| 90 |
+
}
|
| 91 |
+
if self.layer_types is None:
|
| 92 |
+
self.layer_types = [
|
| 93 |
+
"full_attention" if i % 4 == 2 else "sliding_attention" if i % 4 == 0 and i > 0 else "conv"
|
| 94 |
+
for i in range(self.num_hidden_layers)
|
| 95 |
+
]
|
| 96 |
+
|
| 97 |
+
self.tie_word_embeddings = kwargs.pop("tie_embedding", self.tie_word_embeddings)
|
| 98 |
+
super().__post_init__(**kwargs)
|
| 99 |
+
|
| 100 |
+
def validate_architecture(self):
|
| 101 |
+
"""Part of ``@strict``-powered validation."""
|
| 102 |
+
if len(self.layer_types) != self.num_hidden_layers:
|
| 103 |
+
raise ValueError(
|
| 104 |
+
f"`layer_types` has {len(self.layer_types)} entries but `num_hidden_layers` is {self.num_hidden_layers}."
|
| 105 |
+
)
|
| 106 |
+
allowed_layer_types = (*self._attention_layer_types, "conv")
|
| 107 |
+
invalid_layer_types = [layer_type for layer_type in self.layer_types if layer_type not in allowed_layer_types]
|
| 108 |
+
if invalid_layer_types:
|
| 109 |
+
raise ValueError(f"`layer_types` entries must be one of {allowed_layer_types}; got {invalid_layer_types}.")
|
| 110 |
+
if "sliding_attention" in self.layer_types and self.sliding_window is None:
|
| 111 |
+
raise ValueError("`layer_types` has `sliding_attention` layers but `sliding_window` is None.")
|
| 112 |
+
missing_rope = [
|
| 113 |
+
layer_type
|
| 114 |
+
for layer_type in self._attention_layer_types
|
| 115 |
+
if layer_type in self.layer_types and layer_type not in self.rope_parameters
|
| 116 |
+
]
|
| 117 |
+
if missing_rope:
|
| 118 |
+
raise ValueError(f"`rope_parameters` has no entry for the {missing_rope} layers in `layer_types`.")
|
| 119 |
+
for layer_type, rope_params in self.rope_parameters.items():
|
| 120 |
+
if not isinstance(rope_params, dict):
|
| 121 |
+
continue
|
| 122 |
+
if (
|
| 123 |
+
rope_params.get("rope_type", "default") in ("dynamic", "longrope")
|
| 124 |
+
and rope_params.get("partial_rotary_factor", 1.0) < 1
|
| 125 |
+
):
|
| 126 |
+
raise ValueError(
|
| 127 |
+
f"`{rope_params['rope_type']}` RoPE recomputes its frequencies during the forward pass over the "
|
| 128 |
+
f"rotated dimension only and cannot keep LFM3's partial rotation on `{layer_type}` layers."
|
| 129 |
+
)
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
__all__ = ["Lfm3Config"]
|
configuration_lfm3_omni.py
ADDED
|
@@ -0,0 +1,300 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
from huggingface_hub.dataclasses import strict
|
| 15 |
+
|
| 16 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 17 |
+
from transformers.models.auto import CONFIG_MAPPING, AutoConfig
|
| 18 |
+
from transformers.utils import auto_docstring
|
| 19 |
+
|
| 20 |
+
from .configuration_lfm3 import Lfm3Config
|
| 21 |
+
from .configuration_lfm3_vl import Lfm3VlVisionConfig
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def _config_class(model_type):
|
| 25 |
+
# the LFM3 configs ship with the checkpoint instead of the transformers registry
|
| 26 |
+
remote = {"lfm3": Lfm3Config, "lfm3_vl_vision": Lfm3VlVisionConfig, "lfm3_omni_encoder": Lfm3OmniEncoderConfig}
|
| 27 |
+
return remote.get(model_type) or CONFIG_MAPPING[model_type]
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
# `dtype` is dropped so that the top-level dtype governs every component of a legacy export
|
| 31 |
+
_LEGACY_SUBCONFIG_DROPPED_KEYS = (
|
| 32 |
+
"architectures",
|
| 33 |
+
"transformers_version",
|
| 34 |
+
"is_causal",
|
| 35 |
+
"auto_map",
|
| 36 |
+
"_name_or_path",
|
| 37 |
+
"dtype",
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def _subconfig_dict(values: dict, default_model_type: str) -> dict:
|
| 42 |
+
values = {key: value for key, value in values.items() if key not in _LEGACY_SUBCONFIG_DROPPED_KEYS}
|
| 43 |
+
values.setdefault("model_type", default_model_type)
|
| 44 |
+
return values
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
@auto_docstring(checkpoint="LiquidAI/LFM3-Omni-2B-Preview")
|
| 48 |
+
@strict
|
| 49 |
+
class Lfm3OmniEncoderConfig(PreTrainedConfig):
|
| 50 |
+
r"""
|
| 51 |
+
text_config (`Lfm2Config` or `dict`, *optional*):
|
| 52 |
+
The configuration of the LFM2 model run over the downsampled spectrogram. Defaults to the 17-layer,
|
| 53 |
+
768-wide model of LFM3-Audio.
|
| 54 |
+
num_mel_bins (`int`, *optional*, defaults to 128):
|
| 55 |
+
The number of mel bins of the input features.
|
| 56 |
+
conv_channels (`int`, *optional*, defaults to 256):
|
| 57 |
+
The number of channels of the downsampling convolutions.
|
| 58 |
+
hidden_size (`int`, *optional*, defaults to 768):
|
| 59 |
+
The width of the encoder output, equal to `text_config.hidden_size`.
|
| 60 |
+
sliding_window (`int`, *optional*, defaults to 32):
|
| 61 |
+
The number of past frames each frame attends to, itself included.
|
| 62 |
+
initializer_range (`float`, *optional*, defaults to 0.02):
|
| 63 |
+
The standard deviation of the truncated normal initializer of the convolutions.
|
| 64 |
+
"""
|
| 65 |
+
|
| 66 |
+
model_type = "lfm3_omni_encoder"
|
| 67 |
+
sub_configs = {"text_config": AutoConfig}
|
| 68 |
+
|
| 69 |
+
text_config: dict | PreTrainedConfig | None = None
|
| 70 |
+
num_mel_bins: int = 128
|
| 71 |
+
conv_channels: int = 256
|
| 72 |
+
hidden_size: int = 768
|
| 73 |
+
sliding_window: int = 32
|
| 74 |
+
initializer_range: float = 0.02
|
| 75 |
+
|
| 76 |
+
def __post_init__(self, **kwargs):
|
| 77 |
+
if isinstance(self.text_config, dict):
|
| 78 |
+
# the encoder consumes embeddings only: no end-of-sequence id in its two-token vocabulary
|
| 79 |
+
self.text_config = {"eos_token_id": None, **_subconfig_dict(self.text_config, "lfm2")}
|
| 80 |
+
self.text_config = _config_class(self.text_config["model_type"])(**self.text_config)
|
| 81 |
+
elif self.text_config is None:
|
| 82 |
+
self.text_config = _config_class("lfm2")(
|
| 83 |
+
vocab_size=2,
|
| 84 |
+
hidden_size=self.hidden_size,
|
| 85 |
+
intermediate_size=2560,
|
| 86 |
+
num_hidden_layers=17,
|
| 87 |
+
num_attention_heads=8,
|
| 88 |
+
num_key_value_heads=4,
|
| 89 |
+
layer_types=(["conv", "conv", "full_attention"] * 5) + ["full_attention", "full_attention"],
|
| 90 |
+
conv_L_cache=3,
|
| 91 |
+
rope_parameters={"rope_type": "default", "rope_theta": 1_000_000.0},
|
| 92 |
+
pad_token_id=0,
|
| 93 |
+
bos_token_id=1,
|
| 94 |
+
eos_token_id=None,
|
| 95 |
+
)
|
| 96 |
+
super().__post_init__(**kwargs)
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
@auto_docstring(checkpoint="LiquidAI/LFM3-Omni-2B-Preview")
|
| 100 |
+
@strict
|
| 101 |
+
class Lfm3OmniConfig(PreTrainedConfig):
|
| 102 |
+
r"""
|
| 103 |
+
This is the configuration class to store the configuration of a [`Lfm3OmniForConditionalGeneration`].
|
| 104 |
+
|
| 105 |
+
text_config (`Lfm3Config` or `dict`, *optional*):
|
| 106 |
+
The configuration of the LFM3 language model.
|
| 107 |
+
audio_encoder_config (`Lfm3OmniEncoderConfig` or `dict`, *optional*):
|
| 108 |
+
The configuration of the speech encoder.
|
| 109 |
+
depth_config (`Lfm3Config` or `dict`, *optional*):
|
| 110 |
+
The configuration of the depth decoder, a bidirectional LFM3 model over the codebook axis of one speech
|
| 111 |
+
frame. Defaults to the 3-layer, 1024-wide model of LFM3-Audio.
|
| 112 |
+
num_codebooks (`int`, *optional*, defaults to 12):
|
| 113 |
+
The number of codebooks per speech frame.
|
| 114 |
+
codebook_size (`int`, *optional*, defaults to 2017):
|
| 115 |
+
The size of each codebook, the codec's codes plus the end-of-audio code.
|
| 116 |
+
audio_eos_code (`int`, *optional*, defaults to 2016):
|
| 117 |
+
The end-of-audio code. A frame holding it in every codebook ends a speech segment.
|
| 118 |
+
maskgit_schedule (`list[int]`, *optional*, defaults to `[4, 4, 4]`):
|
| 119 |
+
The number of codebooks revealed by each pass of the depth decoder when sampling a frame.
|
| 120 |
+
interleaved_n_text (`int`, *optional*, defaults to 6):
|
| 121 |
+
The number of text tokens per block of the interleaved generation mode.
|
| 122 |
+
interleaved_n_audio (`int`, *optional*, defaults to 12):
|
| 123 |
+
The number of speech frames per block of the interleaved generation mode.
|
| 124 |
+
audio_token_id (`int`, *optional*, defaults to 125031):
|
| 125 |
+
The id of the placeholder token standing for one encoded input speech frame.
|
| 126 |
+
audio_out_token_id (`int`, *optional*, defaults to 125032):
|
| 127 |
+
The id of the placeholder token standing for one output speech frame.
|
| 128 |
+
audio_start_token_id (`int`, *optional*, defaults to 125012):
|
| 129 |
+
The id of `<|audio_start|>`, which switches generation to speech until an end-of-audio frame.
|
| 130 |
+
text_end_token_id (`int`, *optional*, defaults to 125014):
|
| 131 |
+
The id of `<|text_end|>`, which ends the text of an interleaved answer.
|
| 132 |
+
mixed_start_token_id (`int`, *optional*, defaults to 125015):
|
| 133 |
+
The id of `<|mixed_start|>`, which switches generation to interleaved text and speech.
|
| 134 |
+
tie_word_embeddings (`bool`, *optional*, defaults to `True`):
|
| 135 |
+
Whether to tie the language model head to the input embeddings.
|
| 136 |
+
text_loss_weight (`float`, *optional*, defaults to 1.0):
|
| 137 |
+
The weight of the text loss in the training loss.
|
| 138 |
+
audio_loss_weight (`float`, *optional*, defaults to 1.0):
|
| 139 |
+
The weight of the speech loss in the training loss.
|
| 140 |
+
initializer_range (`float`, *optional*, defaults to 0.02):
|
| 141 |
+
The standard deviation of the truncated normal initializer of the projector and depth decoder.
|
| 142 |
+
|
| 143 |
+
vision_config (`Lfm3VlVisionConfig` or `dict`, *optional*):
|
| 144 |
+
The configuration of the image encoder.
|
| 145 |
+
video_token_id (`int`, *optional*, defaults to 125029):
|
| 146 |
+
The id of the placeholder token standing for one encoded video patch group.
|
| 147 |
+
image_token_id (`int`, *optional*, defaults to 125030):
|
| 148 |
+
The placeholder for one projected image token.
|
| 149 |
+
projector_hidden_size (`int`, *optional*, defaults to 4096):
|
| 150 |
+
The intermediate dimension of the image projector.
|
| 151 |
+
projector_hidden_act (`str`, *optional*, defaults to `"gelu"`):
|
| 152 |
+
The image projector activation.
|
| 153 |
+
projector_bias (`bool`, *optional*, defaults to `True`):
|
| 154 |
+
Whether the image projector linear layers have biases.
|
| 155 |
+
projector_layer_norm_eps (`float`, *optional*, defaults to 1e-06):
|
| 156 |
+
The epsilon of the image projector layer normalization.
|
| 157 |
+
downsample_factor (`int`, *optional*, defaults to 2):
|
| 158 |
+
The side length of each image patch block merged into one language token.
|
| 159 |
+
"""
|
| 160 |
+
|
| 161 |
+
model_type = "lfm3_omni"
|
| 162 |
+
sub_configs = {
|
| 163 |
+
"text_config": AutoConfig,
|
| 164 |
+
"audio_encoder_config": Lfm3OmniEncoderConfig,
|
| 165 |
+
"depth_config": AutoConfig,
|
| 166 |
+
"vision_config": AutoConfig,
|
| 167 |
+
}
|
| 168 |
+
|
| 169 |
+
text_config: dict | PreTrainedConfig | None = None
|
| 170 |
+
audio_encoder_config: dict | PreTrainedConfig | None = None
|
| 171 |
+
depth_config: dict | PreTrainedConfig | None = None
|
| 172 |
+
num_codebooks: int = 12
|
| 173 |
+
codebook_size: int = 2017
|
| 174 |
+
audio_eos_code: int = 2016
|
| 175 |
+
maskgit_schedule: list[int] | None = None
|
| 176 |
+
interleaved_n_text: int = 6
|
| 177 |
+
interleaved_n_audio: int = 12
|
| 178 |
+
audio_token_id: int = 125_031
|
| 179 |
+
audio_out_token_id: int = 125_032
|
| 180 |
+
audio_start_token_id: int = 125_012
|
| 181 |
+
text_end_token_id: int = 125_014
|
| 182 |
+
mixed_start_token_id: int = 125_015
|
| 183 |
+
tie_word_embeddings: bool = True
|
| 184 |
+
text_loss_weight: float = 1.0
|
| 185 |
+
audio_loss_weight: float = 1.0
|
| 186 |
+
initializer_range: float = 0.02
|
| 187 |
+
|
| 188 |
+
# the ids `generate` and `transcribe_live` index the text logits with
|
| 189 |
+
_token_id_fields = (
|
| 190 |
+
"audio_token_id",
|
| 191 |
+
"audio_out_token_id",
|
| 192 |
+
"audio_start_token_id",
|
| 193 |
+
"text_end_token_id",
|
| 194 |
+
"mixed_start_token_id",
|
| 195 |
+
)
|
| 196 |
+
vision_config: dict | PreTrainedConfig | None = None
|
| 197 |
+
image_token_id: int = 125_030
|
| 198 |
+
video_token_id: int = 125_029
|
| 199 |
+
projector_hidden_size: int = 4096
|
| 200 |
+
projector_hidden_act: str = "gelu"
|
| 201 |
+
projector_bias: bool = True
|
| 202 |
+
projector_layer_norm_eps: float = 1e-6
|
| 203 |
+
downsample_factor: int = 2
|
| 204 |
+
|
| 205 |
+
def __post_init__(self, **kwargs):
|
| 206 |
+
if isinstance(self.vision_config, dict):
|
| 207 |
+
self.vision_config = _config_class("lfm3_vl_vision")(**self.vision_config)
|
| 208 |
+
elif self.vision_config is None:
|
| 209 |
+
self.vision_config = _config_class("lfm3_vl_vision")()
|
| 210 |
+
if "lfm" in kwargs or "format_version" in kwargs:
|
| 211 |
+
raise ValueError(
|
| 212 |
+
"LFM3-Omni reads the canonical checkpoint layout only; training exports (`lfm` / `format_version`) "
|
| 213 |
+
"are not supported, load a canonical Hugging Face checkpoint."
|
| 214 |
+
)
|
| 215 |
+
if isinstance(self.text_config, dict):
|
| 216 |
+
self.text_config = _subconfig_dict(self.text_config, "lfm3")
|
| 217 |
+
self.text_config = _config_class(self.text_config["model_type"])(**self.text_config)
|
| 218 |
+
elif self.text_config is None:
|
| 219 |
+
self.text_config = _config_class("lfm3")()
|
| 220 |
+
self.validate_vision_fields()
|
| 221 |
+
if isinstance(self.audio_encoder_config, dict):
|
| 222 |
+
self.audio_encoder_config = self.sub_configs["audio_encoder_config"](**self.audio_encoder_config)
|
| 223 |
+
elif self.audio_encoder_config is None:
|
| 224 |
+
self.audio_encoder_config = self.sub_configs["audio_encoder_config"]()
|
| 225 |
+
if isinstance(self.depth_config, dict):
|
| 226 |
+
# the depth decoder consumes embeddings only; its token ids are placeholders
|
| 227 |
+
self.depth_config = {
|
| 228 |
+
"pad_token_id": None,
|
| 229 |
+
"bos_token_id": None,
|
| 230 |
+
"eos_token_id": None,
|
| 231 |
+
**_subconfig_dict(self.depth_config, "lfm3"),
|
| 232 |
+
}
|
| 233 |
+
self.depth_config = _config_class(self.depth_config["model_type"])(**self.depth_config)
|
| 234 |
+
elif self.depth_config is None:
|
| 235 |
+
self.depth_config = _config_class("lfm3")(
|
| 236 |
+
vocab_size=1,
|
| 237 |
+
pad_token_id=None,
|
| 238 |
+
hidden_size=1024,
|
| 239 |
+
intermediate_size=2816,
|
| 240 |
+
num_hidden_layers=3,
|
| 241 |
+
num_attention_heads=32,
|
| 242 |
+
num_key_value_heads=8,
|
| 243 |
+
head_dim=32,
|
| 244 |
+
max_position_embeddings=self.num_codebooks,
|
| 245 |
+
layer_types=["full_attention"] * 3,
|
| 246 |
+
rope_parameters={
|
| 247 |
+
"full_attention": {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 1.0}
|
| 248 |
+
},
|
| 249 |
+
sliding_window=None,
|
| 250 |
+
use_sandwich_norm=False,
|
| 251 |
+
use_cache=False,
|
| 252 |
+
bos_token_id=None,
|
| 253 |
+
eos_token_id=None,
|
| 254 |
+
)
|
| 255 |
+
steps = kwargs.pop("maskgit_steps", 3)
|
| 256 |
+
if self.maskgit_schedule is None:
|
| 257 |
+
if not 1 <= steps <= self.num_codebooks:
|
| 258 |
+
raise ValueError(
|
| 259 |
+
f"maskgit_steps must be between 1 and num_codebooks={self.num_codebooks}, got {steps}"
|
| 260 |
+
)
|
| 261 |
+
self.maskgit_schedule = [self.num_codebooks // steps] * steps
|
| 262 |
+
self.maskgit_schedule[-1] += self.num_codebooks - sum(self.maskgit_schedule)
|
| 263 |
+
self.maskgit_schedule = list(self.maskgit_schedule)
|
| 264 |
+
self._validate_speech_fields()
|
| 265 |
+
super().__post_init__(**kwargs)
|
| 266 |
+
|
| 267 |
+
def validate_speech_fields(self):
|
| 268 |
+
"""Part of `@strict`-powered validation, so `validate()` and `save_pretrained` refuse later edits too."""
|
| 269 |
+
self._validate_speech_fields()
|
| 270 |
+
|
| 271 |
+
def _validate_speech_fields(self):
|
| 272 |
+
if min(self.maskgit_schedule, default=0) < 1 or sum(self.maskgit_schedule) != self.num_codebooks:
|
| 273 |
+
raise ValueError(
|
| 274 |
+
f"maskgit_schedule {self.maskgit_schedule} must be positive counts summing to "
|
| 275 |
+
f"num_codebooks={self.num_codebooks}"
|
| 276 |
+
)
|
| 277 |
+
# the end-of-audio code indexes every codebook; a code outside them could never be sampled or embedded
|
| 278 |
+
if not 0 <= self.audio_eos_code < self.codebook_size:
|
| 279 |
+
raise ValueError(
|
| 280 |
+
f"audio_eos_code={self.audio_eos_code} must lie in [0, codebook_size={self.codebook_size})"
|
| 281 |
+
)
|
| 282 |
+
if self.interleaved_n_text < 1 or self.interleaved_n_audio < 1:
|
| 283 |
+
raise ValueError("interleaved_n_text and interleaved_n_audio must be positive")
|
| 284 |
+
vocab_size = self.text_config.vocab_size
|
| 285 |
+
for name in self._token_id_fields:
|
| 286 |
+
if not 0 <= getattr(self, name) < vocab_size:
|
| 287 |
+
raise ValueError(f"{name}={getattr(self, name)} must lie in [0, text_config.vocab_size={vocab_size})")
|
| 288 |
+
|
| 289 |
+
def validate_vision_fields(self):
|
| 290 |
+
"""Part of `@strict`-powered validation, so `validate()` and `save_pretrained` refuse later edits too."""
|
| 291 |
+
if self.downsample_factor != self.vision_config.spatial_merge_size:
|
| 292 |
+
raise ValueError("downsample_factor must match vision_config.spatial_merge_size")
|
| 293 |
+
# an `inputs_embeds` forward embeds both placeholder ids to find them
|
| 294 |
+
vocab_size = self.text_config.vocab_size
|
| 295 |
+
for name in ("image_token_id", "video_token_id"):
|
| 296 |
+
if not 0 <= getattr(self, name) < vocab_size:
|
| 297 |
+
raise ValueError(f"{name}={getattr(self, name)} must lie in [0, text_config.vocab_size={vocab_size})")
|
| 298 |
+
|
| 299 |
+
|
| 300 |
+
__all__ = ["Lfm3OmniConfig", "Lfm3OmniEncoderConfig"]
|
configuration_lfm3_vl.py
ADDED
|
@@ -0,0 +1,211 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
from huggingface_hub.dataclasses import strict
|
| 15 |
+
|
| 16 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 17 |
+
from transformers.models.auto import CONFIG_MAPPING, AutoConfig
|
| 18 |
+
from transformers.utils import auto_docstring
|
| 19 |
+
|
| 20 |
+
from .configuration_lfm3 import Lfm3Config
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def _config_class(model_type):
|
| 24 |
+
# the LFM3 configs ship with the checkpoint instead of the transformers registry
|
| 25 |
+
remote = {"lfm3": Lfm3Config, "lfm3_vl_vision": Lfm3VlVisionConfig}
|
| 26 |
+
return remote.get(model_type) or CONFIG_MAPPING[model_type]
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
@auto_docstring(checkpoint="LiquidAI/LFM3-VL-1.9B-Preview")
|
| 30 |
+
@strict
|
| 31 |
+
class Lfm3VlVisionConfig(PreTrainedConfig):
|
| 32 |
+
r"""
|
| 33 |
+
layer_types (`list[str]`, *optional*):
|
| 34 |
+
Per-layer operator type, each one of `"full_attention"` or `"conv"`. Defaults to two gated convolution
|
| 35 |
+
layers followed by alternating attention and convolution layers.
|
| 36 |
+
patch_size (`int`, *optional*, defaults to 16):
|
| 37 |
+
The spatial size of each image patch.
|
| 38 |
+
temporal_patch_size (`int`, *optional*, defaults to 2):
|
| 39 |
+
The number of video frames packed into one patch.
|
| 40 |
+
spatial_merge_size (`int`, *optional*, defaults to 2):
|
| 41 |
+
The side of the square block of patches merged into one language token.
|
| 42 |
+
conv_kernel_size (`int`, *optional*, defaults to 3):
|
| 43 |
+
The kernel size of the depthwise 2D convolution layers.
|
| 44 |
+
norm_eps (`float`, *optional*, defaults to 1e-05):
|
| 45 |
+
The epsilon used by the RMS normalization layers.
|
| 46 |
+
"""
|
| 47 |
+
|
| 48 |
+
model_type = "lfm3_vl_vision"
|
| 49 |
+
base_config_key = "vision_config"
|
| 50 |
+
default_theta = 1_000_000.0
|
| 51 |
+
default_rope_type = "axial"
|
| 52 |
+
_attention_layer_types = ("full_attention",)
|
| 53 |
+
|
| 54 |
+
hidden_size: int = 1024
|
| 55 |
+
intermediate_size: int = 2560
|
| 56 |
+
num_hidden_layers: int = 14
|
| 57 |
+
num_attention_heads: int = 16
|
| 58 |
+
num_key_value_heads: int = 8
|
| 59 |
+
layer_types: list[str] | None = None
|
| 60 |
+
patch_size: int = 16
|
| 61 |
+
temporal_patch_size: int = 2
|
| 62 |
+
in_channels: int = 3
|
| 63 |
+
spatial_merge_size: int = 2
|
| 64 |
+
conv_kernel_size: int = 3
|
| 65 |
+
norm_eps: float = 1e-5
|
| 66 |
+
initializer_range: float = 0.02
|
| 67 |
+
rope_parameters: dict | None = None
|
| 68 |
+
|
| 69 |
+
def __post_init__(self, **kwargs):
|
| 70 |
+
if self.layer_types is None:
|
| 71 |
+
self.layer_types = ["conv" if i < 2 or i % 2 else "full_attention" for i in range(self.num_hidden_layers)]
|
| 72 |
+
super().__post_init__(**kwargs)
|
| 73 |
+
|
| 74 |
+
def validate_architecture(self):
|
| 75 |
+
"""Part of ``@strict``-powered validation."""
|
| 76 |
+
if len(self.layer_types) != self.num_hidden_layers:
|
| 77 |
+
raise ValueError(
|
| 78 |
+
f"`layer_types` has {len(self.layer_types)} entries but `num_hidden_layers` is {self.num_hidden_layers}."
|
| 79 |
+
)
|
| 80 |
+
allowed_layer_types = (*self._attention_layer_types, "conv")
|
| 81 |
+
invalid_layer_types = [layer_type for layer_type in self.layer_types if layer_type not in allowed_layer_types]
|
| 82 |
+
if invalid_layer_types:
|
| 83 |
+
raise ValueError(f"`layer_types` entries must be one of {allowed_layer_types}; got {invalid_layer_types}.")
|
| 84 |
+
if self.hidden_size % self.num_attention_heads:
|
| 85 |
+
raise ValueError("`hidden_size` must be a multiple of `num_attention_heads`.")
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
# The first checkpoints spell every text and `vision_*` field at the top level of the config
|
| 89 |
+
_LEGACY_TEXT_KEYS = (
|
| 90 |
+
"vocab_size",
|
| 91 |
+
"hidden_size",
|
| 92 |
+
"intermediate_size",
|
| 93 |
+
"num_hidden_layers",
|
| 94 |
+
"num_attention_heads",
|
| 95 |
+
"num_key_value_heads",
|
| 96 |
+
"head_dim",
|
| 97 |
+
"layer_types",
|
| 98 |
+
"conv_bias",
|
| 99 |
+
"conv_L_cache",
|
| 100 |
+
"sliding_window",
|
| 101 |
+
"max_position_embeddings",
|
| 102 |
+
"use_sandwich_norm",
|
| 103 |
+
"use_cache",
|
| 104 |
+
)
|
| 105 |
+
_LEGACY_VISION_KEYS = {
|
| 106 |
+
"vision_hidden_size": "hidden_size",
|
| 107 |
+
"vision_intermediate_size": "intermediate_size",
|
| 108 |
+
"vision_num_hidden_layers": "num_hidden_layers",
|
| 109 |
+
"vision_num_attention_heads": "num_attention_heads",
|
| 110 |
+
"vision_num_key_value_heads": "num_key_value_heads",
|
| 111 |
+
"vision_layer_types": "layer_types",
|
| 112 |
+
"vision_patch_size": "patch_size",
|
| 113 |
+
"vision_temporal_patch_size": "temporal_patch_size",
|
| 114 |
+
"vision_num_channels": "in_channels",
|
| 115 |
+
"vision_spatial_merge_size": "spatial_merge_size",
|
| 116 |
+
}
|
| 117 |
+
_LEGACY_DROPPED_KEYS = ("vision_patch_dim", "vision_max_position_embeddings")
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
@auto_docstring(checkpoint="LiquidAI/LFM3-VL-1.9B-Preview")
|
| 121 |
+
@strict
|
| 122 |
+
class Lfm3VlConfig(PreTrainedConfig):
|
| 123 |
+
r"""
|
| 124 |
+
video_token_id (`int`, *optional*, defaults to 125029):
|
| 125 |
+
The id of the placeholder token for video patches.
|
| 126 |
+
vision_start_token_id (`int`, *optional*, defaults to 125009):
|
| 127 |
+
The id of the token opening an image or a video frame.
|
| 128 |
+
vision_end_token_id (`int`, *optional*, defaults to 125010):
|
| 129 |
+
The id of the token closing an image or a video frame.
|
| 130 |
+
projector_layer_norm_eps (`float`, *optional*, defaults to 1e-06):
|
| 131 |
+
The epsilon of the layer normalization applied to the vision features before projection.
|
| 132 |
+
downsample_factor (`int`, *optional*, defaults to 2):
|
| 133 |
+
The side of the square block of vision tokens merged by the projector, equal to the vision encoder's
|
| 134 |
+
`spatial_merge_size`.
|
| 135 |
+
"""
|
| 136 |
+
|
| 137 |
+
model_type = "lfm3_vl"
|
| 138 |
+
sub_configs = {"text_config": AutoConfig, "vision_config": Lfm3VlVisionConfig}
|
| 139 |
+
|
| 140 |
+
text_config: dict | PreTrainedConfig | None = None
|
| 141 |
+
vision_config: dict | PreTrainedConfig | None = None
|
| 142 |
+
image_token_id: int = 125_030
|
| 143 |
+
video_token_id: int = 125_029
|
| 144 |
+
vision_start_token_id: int = 125_009
|
| 145 |
+
vision_end_token_id: int = 125_010
|
| 146 |
+
projector_hidden_act: str = "gelu"
|
| 147 |
+
projector_hidden_size: int = 4096
|
| 148 |
+
projector_bias: bool = True
|
| 149 |
+
projector_layer_norm_eps: float = 1e-6
|
| 150 |
+
downsample_factor: int = 2
|
| 151 |
+
tie_word_embeddings: bool = True
|
| 152 |
+
|
| 153 |
+
def __post_init__(self, **kwargs):
|
| 154 |
+
if "vision_hidden_size" in kwargs:
|
| 155 |
+
kwargs = self._convert_legacy_flat_config(kwargs)
|
| 156 |
+
|
| 157 |
+
if isinstance(self.vision_config, dict):
|
| 158 |
+
self.vision_config = self.sub_configs["vision_config"](**self.vision_config)
|
| 159 |
+
elif self.vision_config is None:
|
| 160 |
+
self.vision_config = self.sub_configs["vision_config"]()
|
| 161 |
+
|
| 162 |
+
if isinstance(self.text_config, dict):
|
| 163 |
+
self.text_config = {"model_type": "lfm3", **self.text_config}
|
| 164 |
+
self.text_config = _config_class(self.text_config["model_type"])(**self.text_config)
|
| 165 |
+
elif self.text_config is None:
|
| 166 |
+
self.text_config = _config_class("lfm3")()
|
| 167 |
+
self.validate_vision_fields()
|
| 168 |
+
|
| 169 |
+
super().__post_init__(**kwargs)
|
| 170 |
+
|
| 171 |
+
def validate_vision_fields(self):
|
| 172 |
+
"""Part of `@strict`-powered validation, so `validate()` and `save_pretrained` refuse later edits too."""
|
| 173 |
+
if self.downsample_factor != self.vision_config.spatial_merge_size:
|
| 174 |
+
raise ValueError("downsample_factor must match vision_config.spatial_merge_size")
|
| 175 |
+
# an `inputs_embeds` forward embeds both placeholder ids to find them
|
| 176 |
+
vocab_size = self.text_config.vocab_size
|
| 177 |
+
for name in ("image_token_id", "video_token_id"):
|
| 178 |
+
if not 0 <= getattr(self, name) < vocab_size:
|
| 179 |
+
raise ValueError(f"{name}={getattr(self, name)} must lie in [0, text_config.vocab_size={vocab_size})")
|
| 180 |
+
|
| 181 |
+
def _convert_legacy_flat_config(self, kwargs: dict) -> dict:
|
| 182 |
+
text_config = {key: kwargs.pop(key) for key in _LEGACY_TEXT_KEYS if key in kwargs}
|
| 183 |
+
text_config["norm_eps"] = kwargs.pop("rms_norm_eps", 1e-5)
|
| 184 |
+
text_config["rope_parameters"] = {
|
| 185 |
+
"full_attention": {
|
| 186 |
+
"rope_type": "default",
|
| 187 |
+
"rope_theta": kwargs.pop("full_attention_rope_theta", 1_000_000.0),
|
| 188 |
+
"partial_rotary_factor": kwargs.pop("partial_rotary_factor", 0.5),
|
| 189 |
+
},
|
| 190 |
+
"sliding_attention": {
|
| 191 |
+
"rope_type": "default",
|
| 192 |
+
"rope_theta": kwargs.pop("sliding_attention_rope_theta", 10_000.0),
|
| 193 |
+
},
|
| 194 |
+
}
|
| 195 |
+
for key in ("pad_token_id", "bos_token_id", "eos_token_id"):
|
| 196 |
+
if kwargs.get(key) is not None:
|
| 197 |
+
text_config[key] = kwargs[key]
|
| 198 |
+
text_config["tie_word_embeddings"] = self.tie_word_embeddings
|
| 199 |
+
|
| 200 |
+
vision_config = {new: kwargs.pop(old) for old, new in _LEGACY_VISION_KEYS.items() if old in kwargs}
|
| 201 |
+
vision_config["norm_eps"] = text_config["norm_eps"]
|
| 202 |
+
vision_config["rope_parameters"] = {"rope_type": "axial", "rope_theta": kwargs.pop("vision_rope_theta", 1e6)}
|
| 203 |
+
for key in _LEGACY_DROPPED_KEYS:
|
| 204 |
+
kwargs.pop(key, None)
|
| 205 |
+
|
| 206 |
+
self.text_config = text_config
|
| 207 |
+
self.vision_config = vision_config
|
| 208 |
+
return kwargs
|
| 209 |
+
|
| 210 |
+
|
| 211 |
+
__all__ = ["Lfm3VlConfig", "Lfm3VlVisionConfig"]
|
configuration_liquid_codec.py
ADDED
|
@@ -0,0 +1,152 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
import math
|
| 16 |
+
|
| 17 |
+
from huggingface_hub.dataclasses import strict
|
| 18 |
+
|
| 19 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 20 |
+
from transformers.models.auto import CONFIG_MAPPING
|
| 21 |
+
from transformers.utils import auto_docstring
|
| 22 |
+
|
| 23 |
+
from .configuration_lfm3 import Lfm3Config
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def _config_class(model_type):
|
| 27 |
+
# the LFM3 configs ship with the checkpoint instead of the transformers registry
|
| 28 |
+
remote = {"lfm3": Lfm3Config}
|
| 29 |
+
return remote.get(model_type) or CONFIG_MAPPING[model_type]
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
@auto_docstring(checkpoint="LiquidAI/LFM3-Audio-2B-Preview")
|
| 33 |
+
@strict
|
| 34 |
+
class LiquidCodecConfig(PreTrainedConfig):
|
| 35 |
+
r"""
|
| 36 |
+
Configuration of the [`LiquidCodecModel`] bundled with LFM3-Audio-2B and LFM3-Omni-2B in their
|
| 37 |
+
`audio_tokenizer/` directory.
|
| 38 |
+
|
| 39 |
+
encoder_dim (`int`, *optional*, defaults to 512):
|
| 40 |
+
The width of the encoder and decoder LFM3 models.
|
| 41 |
+
encoder_conv_dim (`int`, *optional*, defaults to 256):
|
| 42 |
+
The number of channels of the STFT front end.
|
| 43 |
+
stft_n_fft (`int`, *optional*, defaults to 1280):
|
| 44 |
+
The FFT size of the analysis and synthesis transforms.
|
| 45 |
+
stft_hop_length (`int`, *optional*, defaults to 320):
|
| 46 |
+
The hop of the analysis and synthesis transforms.
|
| 47 |
+
stft_win_length (`int`, *optional*, defaults to 1280):
|
| 48 |
+
The window length of the analysis and synthesis transforms.
|
| 49 |
+
stft_frame_stack (`int`, *optional*, defaults to 6):
|
| 50 |
+
The number of STFT frames per codec frame.
|
| 51 |
+
quantizer_dim (`int`, *optional*, defaults to 256):
|
| 52 |
+
The width of the latent handed to the quantizer.
|
| 53 |
+
fsq_levels (`list[int]`, *optional*, defaults to `[8, 7, 6, 6]`):
|
| 54 |
+
The number of levels of each finite-scalar-quantization dimension; their product is the codebook size.
|
| 55 |
+
n_q (`int`, *optional*, defaults to 12):
|
| 56 |
+
The number of codebooks.
|
| 57 |
+
decoder_window (`int`, *optional*, defaults to 30):
|
| 58 |
+
The sliding attention window of the decoder LFM3 model, in 75 Hz frames.
|
| 59 |
+
sampling_rate (`int`, *optional*, defaults to 24000):
|
| 60 |
+
The sampling rate of the audio.
|
| 61 |
+
initializer_range (`float`, *optional*, defaults to 0.02):
|
| 62 |
+
The standard deviation of the truncated normal initializer of the weights.
|
| 63 |
+
codec_intermediate (`int`, *optional*, defaults to 2304):
|
| 64 |
+
The feed-forward width of each tiny codec LFM3 layer (full-size codec uses 2,304).
|
| 65 |
+
codec_num_layers (`int`, *optional*, defaults to 8):
|
| 66 |
+
The number of LFM3 layers in each tiny codec encoder/decoder stack (full-size codec uses 8).
|
| 67 |
+
codec_num_heads (`int`, *optional*, defaults to 8):
|
| 68 |
+
The number of query attention heads per tiny codec LFM3 layer (full-size codec uses 8).
|
| 69 |
+
codec_num_kv_heads (`int`, *optional*, defaults to 2):
|
| 70 |
+
The number of key/value attention heads per tiny codec LFM3 layer (full-size codec uses 2).
|
| 71 |
+
codec_head_dim (`int`, *optional*, defaults to 64):
|
| 72 |
+
The attention head dimension of each tiny codec LFM3 layer (full-size codec uses 64).
|
| 73 |
+
"""
|
| 74 |
+
|
| 75 |
+
model_type = "liquid_codec"
|
| 76 |
+
|
| 77 |
+
encoder_dim: int = 512
|
| 78 |
+
encoder_conv_dim: int = 256
|
| 79 |
+
stft_n_fft: int = 1280
|
| 80 |
+
stft_hop_length: int = 320
|
| 81 |
+
stft_win_length: int = 1280
|
| 82 |
+
stft_frame_stack: int = 6
|
| 83 |
+
quantizer_dim: int = 256
|
| 84 |
+
fsq_levels: list[int] | None = None
|
| 85 |
+
n_q: int = 12
|
| 86 |
+
decoder_window: int = 30
|
| 87 |
+
sampling_rate: int = 24_000
|
| 88 |
+
initializer_range: float = 0.02
|
| 89 |
+
codec_intermediate: int = 2304
|
| 90 |
+
codec_num_layers: int = 8
|
| 91 |
+
codec_num_heads: int = 8
|
| 92 |
+
codec_num_kv_heads: int = 2
|
| 93 |
+
codec_head_dim: int = 64
|
| 94 |
+
|
| 95 |
+
def __post_init__(self, **kwargs):
|
| 96 |
+
# the released configs nest the fields under `model` and dump the two LFM3 sub-configs
|
| 97 |
+
dropped = ("encoder_lfm_config", "decoder_lfm_config", "n_output_samples")
|
| 98 |
+
for key, value in kwargs.pop("model", {}).items():
|
| 99 |
+
if key in dropped:
|
| 100 |
+
continue
|
| 101 |
+
if key not in self.__dataclass_fields__:
|
| 102 |
+
raise ValueError(f"Unknown Liquid codec setting `{key}`")
|
| 103 |
+
setattr(self, key, value)
|
| 104 |
+
for key in dropped:
|
| 105 |
+
kwargs.pop(key, None)
|
| 106 |
+
if self.fsq_levels is None:
|
| 107 |
+
self.fsq_levels = [8, 7, 6, 6]
|
| 108 |
+
self.fsq_levels = list(self.fsq_levels)
|
| 109 |
+
super().__post_init__(**kwargs)
|
| 110 |
+
|
| 111 |
+
@property
|
| 112 |
+
def codebook_size(self) -> int:
|
| 113 |
+
return math.prod(self.fsq_levels)
|
| 114 |
+
|
| 115 |
+
@property
|
| 116 |
+
def frame_size(self) -> int:
|
| 117 |
+
return self.stft_hop_length * self.stft_frame_stack
|
| 118 |
+
|
| 119 |
+
def lfm_config(self, sliding_window: int | None = None) -> PreTrainedConfig:
|
| 120 |
+
attention_type = "full_attention" if sliding_window is None else "sliding_attention"
|
| 121 |
+
return _config_class("lfm3")(
|
| 122 |
+
# the backbones are not sub-configs, so the requested attention implementation is handed over here
|
| 123 |
+
attn_implementation=self._attn_implementation,
|
| 124 |
+
vocab_size=1,
|
| 125 |
+
pad_token_id=None,
|
| 126 |
+
bos_token_id=None,
|
| 127 |
+
eos_token_id=None,
|
| 128 |
+
hidden_size=self.encoder_dim,
|
| 129 |
+
intermediate_size=self.codec_intermediate,
|
| 130 |
+
num_hidden_layers=self.codec_num_layers,
|
| 131 |
+
num_attention_heads=self.codec_num_heads,
|
| 132 |
+
num_key_value_heads=self.codec_num_kv_heads,
|
| 133 |
+
head_dim=self.codec_head_dim,
|
| 134 |
+
max_position_embeddings=131_072,
|
| 135 |
+
norm_eps=1e-5,
|
| 136 |
+
conv_L_cache=7,
|
| 137 |
+
conv_bias=False,
|
| 138 |
+
use_sandwich_norm=True,
|
| 139 |
+
layer_types=[
|
| 140 |
+
attention_type if index in (1, self.codec_num_layers - 1) else "conv"
|
| 141 |
+
for index in range(self.codec_num_layers)
|
| 142 |
+
],
|
| 143 |
+
sliding_window=sliding_window,
|
| 144 |
+
rope_parameters={
|
| 145 |
+
attention_type: {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 1.0}
|
| 146 |
+
},
|
| 147 |
+
use_cache=False,
|
| 148 |
+
tie_word_embeddings=False,
|
| 149 |
+
)
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
__all__ = ["LiquidCodecConfig"]
|
feature_extraction_lfm3_omni.py
ADDED
|
@@ -0,0 +1,202 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
import numpy as np
|
| 16 |
+
|
| 17 |
+
from transformers.audio_utils import mel_filter_bank
|
| 18 |
+
from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
|
| 19 |
+
from transformers.feature_extraction_utils import BatchFeature
|
| 20 |
+
from transformers.utils import TensorType, is_torch_available, logging
|
| 21 |
+
|
| 22 |
+
from .modeling_lfm3_omni import Lfm3OmniFeatureStreamingCache
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
if is_torch_available():
|
| 26 |
+
import torch
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
logger = logging.get_logger(__name__)
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
class Lfm3OmniFeatureExtractor(SequenceFeatureExtractor):
|
| 33 |
+
"""The LFM3-Audio log-mel front end; canonical checkpoints bundle its configuration files."""
|
| 34 |
+
|
| 35 |
+
model_input_names = ["input_features", "feature_attention_mask"]
|
| 36 |
+
|
| 37 |
+
def __init__(
|
| 38 |
+
self,
|
| 39 |
+
feature_size: int = 128,
|
| 40 |
+
sampling_rate: int = 16000,
|
| 41 |
+
n_fft: int = 400,
|
| 42 |
+
hop_length: int = 160,
|
| 43 |
+
win_length: int = 400,
|
| 44 |
+
log_floor: float = -12.0,
|
| 45 |
+
normalization_offset: float = 4.0,
|
| 46 |
+
normalization_divisor: float = 4.0,
|
| 47 |
+
encoder_stride: int = 8,
|
| 48 |
+
padding_value: float = 0.0,
|
| 49 |
+
**kwargs,
|
| 50 |
+
):
|
| 51 |
+
super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs)
|
| 52 |
+
self.n_fft = n_fft
|
| 53 |
+
self.hop_length = hop_length
|
| 54 |
+
self.win_length = win_length
|
| 55 |
+
self.log_floor = log_floor
|
| 56 |
+
self.normalization_offset = normalization_offset
|
| 57 |
+
self.normalization_divisor = normalization_divisor
|
| 58 |
+
self.encoder_stride = encoder_stride
|
| 59 |
+
self.mel_filters = mel_filter_bank(
|
| 60 |
+
num_frequency_bins=n_fft // 2 + 1,
|
| 61 |
+
num_mel_filters=feature_size,
|
| 62 |
+
min_frequency=0.0,
|
| 63 |
+
max_frequency=sampling_rate / 2,
|
| 64 |
+
sampling_rate=sampling_rate,
|
| 65 |
+
norm="slaney",
|
| 66 |
+
mel_scale="slaney",
|
| 67 |
+
)
|
| 68 |
+
|
| 69 |
+
def num_frames(self, num_samples: int) -> int:
|
| 70 |
+
padded = self.pad_length(num_samples)
|
| 71 |
+
return 1 + (padded - self.n_fft) // self.hop_length
|
| 72 |
+
|
| 73 |
+
def pad_length(self, num_samples: int) -> int:
|
| 74 |
+
if num_samples < self.n_fft:
|
| 75 |
+
return self.n_fft
|
| 76 |
+
return num_samples + (-(num_samples - self.n_fft) % self.hop_length)
|
| 77 |
+
|
| 78 |
+
def num_encoder_tokens(self, num_frames: int) -> int:
|
| 79 |
+
return -(num_frames // -self.encoder_stride)
|
| 80 |
+
|
| 81 |
+
def _log_mel(self, waveform: torch.Tensor) -> torch.Tensor:
|
| 82 |
+
window = torch.hann_window(self.win_length, device=waveform.device)
|
| 83 |
+
spectrogram = torch.stft(
|
| 84 |
+
waveform,
|
| 85 |
+
n_fft=self.n_fft,
|
| 86 |
+
hop_length=self.hop_length,
|
| 87 |
+
win_length=self.win_length,
|
| 88 |
+
window=window,
|
| 89 |
+
center=False,
|
| 90 |
+
return_complex=True,
|
| 91 |
+
)
|
| 92 |
+
power = spectrogram.abs().square()
|
| 93 |
+
mel_filters = torch.from_numpy(self.mel_filters).to(device=waveform.device, dtype=power.dtype)
|
| 94 |
+
log_mel = (mel_filters.T @ power).log10().clamp(min=self.log_floor)
|
| 95 |
+
return (log_mel + self.normalization_offset) / self.normalization_divisor
|
| 96 |
+
|
| 97 |
+
@torch.no_grad()
|
| 98 |
+
def stream(
|
| 99 |
+
self,
|
| 100 |
+
raw_speech: np.ndarray | torch.Tensor | list[float],
|
| 101 |
+
cache: Lfm3OmniFeatureStreamingCache | None = None,
|
| 102 |
+
is_final: bool = False,
|
| 103 |
+
sampling_rate: int | None = None,
|
| 104 |
+
) -> tuple[BatchFeature, Lfm3OmniFeatureStreamingCache]:
|
| 105 |
+
"""Extract completed windows from a mono chunk, retaining overlap. Finalization pads the last window."""
|
| 106 |
+
if sampling_rate is not None and sampling_rate != self.sampling_rate:
|
| 107 |
+
raise ValueError(f"Expected {self.sampling_rate} Hz, got {sampling_rate} Hz")
|
| 108 |
+
if cache is None:
|
| 109 |
+
cache = Lfm3OmniFeatureStreamingCache()
|
| 110 |
+
if cache.finished:
|
| 111 |
+
raise ValueError("This stream is finalized; start with a new cache")
|
| 112 |
+
chunk = torch.as_tensor(raw_speech, dtype=torch.float32)
|
| 113 |
+
if chunk.ndim != 1:
|
| 114 |
+
raise ValueError("stream expects one mono waveform")
|
| 115 |
+
samples = chunk if cache.samples is None else torch.cat((cache.samples, chunk))
|
| 116 |
+
num_samples = cache.num_samples + chunk.numel()
|
| 117 |
+
if is_final:
|
| 118 |
+
count = self.num_frames(num_samples) - cache.num_frames
|
| 119 |
+
if count:
|
| 120 |
+
size = (count - 1) * self.hop_length + self.n_fft
|
| 121 |
+
samples = torch.nn.functional.pad(samples, (0, size - samples.numel()))
|
| 122 |
+
else:
|
| 123 |
+
count = max(0, 1 + (samples.numel() - self.n_fft) // self.hop_length)
|
| 124 |
+
features = self._log_mel(samples) if count else samples.new_empty((self.feature_size, 0))
|
| 125 |
+
# the cache changes only once the chunk's features exist, so a failed call can be retried
|
| 126 |
+
cache.samples = samples[count * self.hop_length :].clone() if not is_final else None
|
| 127 |
+
cache.num_samples = num_samples
|
| 128 |
+
cache.num_frames += count
|
| 129 |
+
cache.finished = is_final
|
| 130 |
+
return BatchFeature(
|
| 131 |
+
{
|
| 132 |
+
"input_features": features[None],
|
| 133 |
+
"feature_attention_mask": torch.ones((1, count), device=chunk.device, dtype=torch.long),
|
| 134 |
+
}
|
| 135 |
+
), cache
|
| 136 |
+
|
| 137 |
+
def __call__(
|
| 138 |
+
self,
|
| 139 |
+
raw_speech: np.ndarray | torch.Tensor | list[np.ndarray] | list[torch.Tensor] | list[float],
|
| 140 |
+
sampling_rate: int | None = None,
|
| 141 |
+
return_tensors: str | TensorType | None = None,
|
| 142 |
+
**kwargs,
|
| 143 |
+
) -> BatchFeature:
|
| 144 |
+
r"""
|
| 145 |
+
raw_speech (`np.ndarray`, `torch.Tensor`, `list[np.ndarray]`, `list[torch.Tensor]`):
|
| 146 |
+
One mono clip or a list of mono clips, as float samples at `sampling_rate`.
|
| 147 |
+
sampling_rate (`int`, *optional*):
|
| 148 |
+
The sampling rate of `raw_speech`, checked against the extractor's.
|
| 149 |
+
return_tensors (`str` or [`~utils.TensorType`], *optional*):
|
| 150 |
+
If set, will return tensors instead of numpy arrays.
|
| 151 |
+
"""
|
| 152 |
+
if kwargs:
|
| 153 |
+
raise TypeError(
|
| 154 |
+
f"{self.__class__.__name__} does not implement {sorted(kwargs)}: clips are always right-padded to "
|
| 155 |
+
"whole frames and `feature_attention_mask` is always returned"
|
| 156 |
+
)
|
| 157 |
+
if sampling_rate is not None and sampling_rate != self.sampling_rate:
|
| 158 |
+
raise ValueError(
|
| 159 |
+
f"{self.__class__.__name__} expects audio sampled at {self.sampling_rate} Hz, got {sampling_rate} Hz"
|
| 160 |
+
)
|
| 161 |
+
if sampling_rate is None:
|
| 162 |
+
logger.warning_once(
|
| 163 |
+
f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()` "
|
| 164 |
+
"to avoid silent errors."
|
| 165 |
+
)
|
| 166 |
+
if isinstance(raw_speech, (list, tuple)) and len(raw_speech) == 0:
|
| 167 |
+
raise ValueError(f"{self.__class__.__name__} needs at least one clip")
|
| 168 |
+
is_batched = isinstance(raw_speech, (list, tuple)) and hasattr(raw_speech[0], "__len__")
|
| 169 |
+
is_batched = is_batched or (isinstance(raw_speech, (np.ndarray, torch.Tensor)) and raw_speech.ndim == 2)
|
| 170 |
+
clips = list(raw_speech) if is_batched else [raw_speech]
|
| 171 |
+
clips = [
|
| 172 |
+
clip.detach().cpu().float()
|
| 173 |
+
if isinstance(clip, torch.Tensor)
|
| 174 |
+
else torch.as_tensor(np.asarray(clip, dtype=np.float32))
|
| 175 |
+
for clip in clips
|
| 176 |
+
]
|
| 177 |
+
for clip in clips:
|
| 178 |
+
if clip.ndim != 1:
|
| 179 |
+
raise ValueError(
|
| 180 |
+
f"{self.__class__.__name__} expects mono clips, got a clip of shape {tuple(clip.shape)}"
|
| 181 |
+
)
|
| 182 |
+
|
| 183 |
+
features = []
|
| 184 |
+
for clip in clips:
|
| 185 |
+
clip = torch.nn.functional.pad(clip, (0, self.pad_length(clip.shape[0]) - clip.shape[0]))
|
| 186 |
+
features.append(self._log_mel(clip))
|
| 187 |
+
lengths = [feature.shape[-1] for feature in features]
|
| 188 |
+
max_length = max(lengths)
|
| 189 |
+
input_features = torch.full((len(features), self.feature_size, max_length), float(self.padding_value))
|
| 190 |
+
feature_attention_mask = torch.zeros((len(features), max_length), dtype=torch.long)
|
| 191 |
+
for index, feature in enumerate(features):
|
| 192 |
+
input_features[index, :, : feature.shape[-1]] = feature
|
| 193 |
+
feature_attention_mask[index, : feature.shape[-1]] = 1
|
| 194 |
+
|
| 195 |
+
data = {
|
| 196 |
+
"input_features": input_features.numpy(),
|
| 197 |
+
"feature_attention_mask": feature_attention_mask.numpy(),
|
| 198 |
+
}
|
| 199 |
+
return BatchFeature(data=data, tensor_type=return_tensors)
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
__all__ = ["Lfm3OmniFeatureExtractor"]
|
generation_config.json
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token_id": 124894,
|
| 3 |
+
"eos_token_id": [
|
| 4 |
+
124900,
|
| 5 |
+
125014
|
| 6 |
+
],
|
| 7 |
+
"pad_token_id": 124893,
|
| 8 |
+
"use_cache": true,
|
| 9 |
+
"output_attentions": false,
|
| 10 |
+
"output_hidden_states": false,
|
| 11 |
+
"do_sample": false,
|
| 12 |
+
"temperature": 0,
|
| 13 |
+
"top_k": 1,
|
| 14 |
+
"repetition_penalty": 1.0,
|
| 15 |
+
"audio_temperature": 0.4,
|
| 16 |
+
"audio_top_k": null,
|
| 17 |
+
"audio_repetition_penalty": 1.2,
|
| 18 |
+
"audio_repetition_window": 64
|
| 19 |
+
}
|
image_processing_lfm3_vl.py
ADDED
|
@@ -0,0 +1,293 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
import math
|
| 15 |
+
import warnings
|
| 16 |
+
from collections.abc import Iterable
|
| 17 |
+
|
| 18 |
+
import torch
|
| 19 |
+
from torchvision.transforms.v2 import functional as tvF
|
| 20 |
+
|
| 21 |
+
from transformers.image_processing_backends import TorchvisionBackend
|
| 22 |
+
from transformers.image_processing_utils import BatchFeature
|
| 23 |
+
from transformers.image_transforms import group_images_by_shape, reorder_images
|
| 24 |
+
from transformers.image_utils import ImageInput, PILImageResampling, SizeDict
|
| 25 |
+
from transformers.processing_utils import ImagesKwargs, Unpack
|
| 26 |
+
from transformers.utils import TensorType, auto_docstring
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
class Lfm3VlImageProcessorKwargs(ImagesKwargs, total=False):
|
| 30 |
+
r"""
|
| 31 |
+
min_pixels (`int`, *optional*, defaults to `56 * 56`):
|
| 32 |
+
The min pixels of the image to resize the image.
|
| 33 |
+
max_pixels (`int`, *optional*, defaults to `28 * 28 * 1280`):
|
| 34 |
+
The max pixels of the image to resize the image.
|
| 35 |
+
patch_size (`int`, *optional*, defaults to 14):
|
| 36 |
+
The spatial patch size of the vision encoder.
|
| 37 |
+
temporal_patch_size (`int`, *optional*, defaults to 2):
|
| 38 |
+
The temporal patch size of the vision encoder.
|
| 39 |
+
merge_size (`int`, *optional*, defaults to 2):
|
| 40 |
+
The merge size of the vision encoder to llm encoder.
|
| 41 |
+
"""
|
| 42 |
+
|
| 43 |
+
min_pixels: int
|
| 44 |
+
max_pixels: int
|
| 45 |
+
patch_size: int
|
| 46 |
+
temporal_patch_size: int
|
| 47 |
+
merge_size: int
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def smart_resize(
|
| 51 |
+
height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280
|
| 52 |
+
):
|
| 53 |
+
"""Rescales the image so that the following conditions are met:
|
| 54 |
+
|
| 55 |
+
1. Both dimensions (height and width) are divisible by 'factor'.
|
| 56 |
+
|
| 57 |
+
2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
|
| 58 |
+
|
| 59 |
+
3. The aspect ratio of the image is maintained as closely as possible.
|
| 60 |
+
|
| 61 |
+
"""
|
| 62 |
+
if max(height, width) / min(height, width) > 200:
|
| 63 |
+
raise ValueError(
|
| 64 |
+
f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
|
| 65 |
+
)
|
| 66 |
+
h_bar = round(height / factor) * factor
|
| 67 |
+
w_bar = round(width / factor) * factor
|
| 68 |
+
if h_bar * w_bar > max_pixels:
|
| 69 |
+
beta = math.sqrt((height * width) / max_pixels)
|
| 70 |
+
h_bar = max(factor, math.floor(height / beta / factor) * factor)
|
| 71 |
+
w_bar = max(factor, math.floor(width / beta / factor) * factor)
|
| 72 |
+
elif h_bar * w_bar < min_pixels:
|
| 73 |
+
beta = math.sqrt(min_pixels / (height * width))
|
| 74 |
+
h_bar = math.ceil(height * beta / factor) * factor
|
| 75 |
+
w_bar = math.ceil(width * beta / factor) * factor
|
| 76 |
+
return h_bar, w_bar
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
@auto_docstring
|
| 80 |
+
class Lfm3VlImageProcessor(TorchvisionBackend):
|
| 81 |
+
do_resize = True
|
| 82 |
+
resample = PILImageResampling.BICUBIC
|
| 83 |
+
size = {"shortest_edge": 65_536, "longest_edge": 16_777_216}
|
| 84 |
+
default_to_square = False
|
| 85 |
+
do_rescale = True
|
| 86 |
+
do_normalize = True
|
| 87 |
+
image_mean = [0.5, 0.5, 0.5]
|
| 88 |
+
image_std = [0.5, 0.5, 0.5]
|
| 89 |
+
do_convert_rgb = True
|
| 90 |
+
patch_size = 16
|
| 91 |
+
temporal_patch_size = 2
|
| 92 |
+
merge_size = 2
|
| 93 |
+
valid_kwargs = Lfm3VlImageProcessorKwargs
|
| 94 |
+
model_input_names = ["pixel_values", "image_grid_thw"]
|
| 95 |
+
|
| 96 |
+
def __init__(self, **kwargs: Unpack[Lfm3VlImageProcessorKwargs]):
|
| 97 |
+
# backward compatibility: override size with min_pixels and max_pixels if they are provided
|
| 98 |
+
size = kwargs.pop("size", None)
|
| 99 |
+
size = dict(self.size) if size is None else size
|
| 100 |
+
if (min_pixels := kwargs.pop("min_pixels", None)) is not None:
|
| 101 |
+
size["shortest_edge"] = min_pixels
|
| 102 |
+
size.pop("min_pixels", None)
|
| 103 |
+
if (max_pixels := kwargs.pop("max_pixels", None)) is not None:
|
| 104 |
+
size["longest_edge"] = max_pixels
|
| 105 |
+
size.pop("max_pixels", None)
|
| 106 |
+
super().__init__(size=size, **kwargs)
|
| 107 |
+
|
| 108 |
+
def _standardize_kwargs(
|
| 109 |
+
self,
|
| 110 |
+
size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
|
| 111 |
+
min_pixels: int | None = None,
|
| 112 |
+
max_pixels: int | None = None,
|
| 113 |
+
**kwargs,
|
| 114 |
+
) -> dict:
|
| 115 |
+
if min_pixels is not None or max_pixels is not None:
|
| 116 |
+
warnings.warn(
|
| 117 |
+
"Passing `min_pixels` and `max_pixels` to a processor call is deprecated and will be removed in v5.23. "
|
| 118 |
+
"Pass in `size={'longest_edge': xxx, 'shortest_edge': xxx} to override the target size.`",
|
| 119 |
+
FutureWarning,
|
| 120 |
+
)
|
| 121 |
+
|
| 122 |
+
size_dict = dict(size) if isinstance(size, (dict, SizeDict)) else {}
|
| 123 |
+
if min_pixels is not None:
|
| 124 |
+
size_dict["shortest_edge"] = min_pixels
|
| 125 |
+
if max_pixels is not None:
|
| 126 |
+
size_dict["longest_edge"] = max_pixels
|
| 127 |
+
size = SizeDict(**size_dict)
|
| 128 |
+
return super()._standardize_kwargs(size=size, **kwargs)
|
| 129 |
+
|
| 130 |
+
@auto_docstring
|
| 131 |
+
def preprocess(
|
| 132 |
+
self,
|
| 133 |
+
images: ImageInput,
|
| 134 |
+
**kwargs: Unpack[Lfm3VlImageProcessorKwargs],
|
| 135 |
+
) -> BatchFeature:
|
| 136 |
+
return super().preprocess(images, **kwargs)
|
| 137 |
+
|
| 138 |
+
def resize(
|
| 139 |
+
self,
|
| 140 |
+
images: "torch.Tensor",
|
| 141 |
+
size: SizeDict,
|
| 142 |
+
resample: "PILImageResampling | tvF.InterpolationMode | int | None",
|
| 143 |
+
factor: int,
|
| 144 |
+
**kwargs,
|
| 145 |
+
) -> "torch.Tensor":
|
| 146 |
+
"""Resize dynamically based on input image aspect ratio."""
|
| 147 |
+
if not size.shortest_edge or not size.longest_edge:
|
| 148 |
+
raise ValueError(f"`size` dict must contain 'shortest_edge' and 'longest_edge' keys but got {size}.")
|
| 149 |
+
|
| 150 |
+
height, width = images.shape[-2:]
|
| 151 |
+
resized_height, resized_width = smart_resize(
|
| 152 |
+
height,
|
| 153 |
+
width,
|
| 154 |
+
factor=factor,
|
| 155 |
+
min_pixels=size.shortest_edge,
|
| 156 |
+
max_pixels=size.longest_edge,
|
| 157 |
+
)
|
| 158 |
+
return super().resize(
|
| 159 |
+
image=images,
|
| 160 |
+
size=SizeDict(height=resized_height, width=resized_width),
|
| 161 |
+
resample=resample,
|
| 162 |
+
)
|
| 163 |
+
|
| 164 |
+
def patchify(
|
| 165 |
+
self,
|
| 166 |
+
images: "torch.Tensor",
|
| 167 |
+
patch_size: int,
|
| 168 |
+
merge_size: int,
|
| 169 |
+
temporal_patch_size: int,
|
| 170 |
+
) -> tuple["torch.Tensor", int, int]:
|
| 171 |
+
"""Patchifies each image into flat layout of shape (`seq_len`, `patch_dim`) so we can concat dynamically shaped pixels."""
|
| 172 |
+
batch_size, channel, resized_height, resized_width = images.shape
|
| 173 |
+
grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
|
| 174 |
+
patches = images.reshape(
|
| 175 |
+
batch_size,
|
| 176 |
+
channel,
|
| 177 |
+
grid_h // merge_size,
|
| 178 |
+
merge_size,
|
| 179 |
+
patch_size,
|
| 180 |
+
grid_w // merge_size,
|
| 181 |
+
merge_size,
|
| 182 |
+
patch_size,
|
| 183 |
+
)
|
| 184 |
+
# Reorder dimensions to group grid and patch information for subsequent flattening.
|
| 185 |
+
# [batch, grid_h/merge, grid_w/merge, merge, merge, channel, patch, patch]
|
| 186 |
+
patches = patches.permute(0, 2, 5, 3, 6, 1, 4, 7)
|
| 187 |
+
|
| 188 |
+
flatten_patches = (
|
| 189 |
+
patches.unsqueeze(6)
|
| 190 |
+
.expand(-1, -1, -1, -1, -1, -1, temporal_patch_size, -1, -1)
|
| 191 |
+
.reshape(
|
| 192 |
+
batch_size,
|
| 193 |
+
grid_h * grid_w,
|
| 194 |
+
channel * temporal_patch_size * patch_size * patch_size,
|
| 195 |
+
)
|
| 196 |
+
)
|
| 197 |
+
return flatten_patches, grid_h, grid_w
|
| 198 |
+
|
| 199 |
+
def _preprocess(
|
| 200 |
+
self,
|
| 201 |
+
images: list["torch.Tensor"],
|
| 202 |
+
do_resize: bool,
|
| 203 |
+
size: SizeDict,
|
| 204 |
+
resample: "PILImageResampling | tvF.InterpolationMode | int | None",
|
| 205 |
+
do_rescale: bool,
|
| 206 |
+
rescale_factor: float,
|
| 207 |
+
do_normalize: bool,
|
| 208 |
+
image_mean: float | list[float] | None,
|
| 209 |
+
image_std: float | list[float] | None,
|
| 210 |
+
patch_size: int,
|
| 211 |
+
temporal_patch_size: int,
|
| 212 |
+
merge_size: int,
|
| 213 |
+
disable_grouping: bool | None,
|
| 214 |
+
return_tensors: str | TensorType | None,
|
| 215 |
+
**kwargs,
|
| 216 |
+
) -> BatchFeature:
|
| 217 |
+
grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
|
| 218 |
+
resized_images_grouped = {}
|
| 219 |
+
for shape, stacked_images in grouped_images.items():
|
| 220 |
+
if do_resize:
|
| 221 |
+
stacked_images = self.resize(
|
| 222 |
+
images=stacked_images,
|
| 223 |
+
size=size,
|
| 224 |
+
resample=resample,
|
| 225 |
+
factor=patch_size * merge_size,
|
| 226 |
+
)
|
| 227 |
+
resized_images_grouped[shape] = stacked_images
|
| 228 |
+
resized_images = reorder_images(resized_images_grouped, grouped_images_index)
|
| 229 |
+
|
| 230 |
+
grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
|
| 231 |
+
processed_images_grouped = {}
|
| 232 |
+
processed_grids = {}
|
| 233 |
+
for shape, stacked_images in grouped_images.items():
|
| 234 |
+
stacked_images = self.rescale_and_normalize(
|
| 235 |
+
stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
|
| 236 |
+
)
|
| 237 |
+
patches, grid_h, grid_w = self.patchify(
|
| 238 |
+
stacked_images,
|
| 239 |
+
patch_size=patch_size,
|
| 240 |
+
merge_size=merge_size,
|
| 241 |
+
temporal_patch_size=temporal_patch_size,
|
| 242 |
+
)
|
| 243 |
+
|
| 244 |
+
processed_images_grouped[shape] = patches
|
| 245 |
+
processed_grids[shape] = [[1, grid_h, grid_w]] * len(stacked_images)
|
| 246 |
+
|
| 247 |
+
processed_images = reorder_images(processed_images_grouped, grouped_images_index)
|
| 248 |
+
processed_grids_ordered = reorder_images(processed_grids, grouped_images_index)
|
| 249 |
+
pixel_values = processed_images[0] if len(processed_images) == 1 else torch.cat(processed_images, dim=0)
|
| 250 |
+
image_grid_thw = torch.tensor(processed_grids_ordered, dtype=torch.long)
|
| 251 |
+
|
| 252 |
+
return BatchFeature(
|
| 253 |
+
data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, tensor_type=return_tensors
|
| 254 |
+
)
|
| 255 |
+
|
| 256 |
+
def get_number_of_image_patches(self, height: int, width: int, images_kwargs: dict | None = None) -> int:
|
| 257 |
+
"""
|
| 258 |
+
A utility that returns number of image patches for a given image size.
|
| 259 |
+
|
| 260 |
+
Note: Do not remove this method! It is used by vLLM to infer the number of patches and placeholders
|
| 261 |
+
without an image input.
|
| 262 |
+
|
| 263 |
+
Args:
|
| 264 |
+
height (`int`):
|
| 265 |
+
Height of the input image.
|
| 266 |
+
width (`int`):
|
| 267 |
+
Width of the input image.
|
| 268 |
+
images_kwargs (`dict`, *optional*)
|
| 269 |
+
Any kwargs to override defaults of the image processor.
|
| 270 |
+
Returns:
|
| 271 |
+
`int`: Number of image patches per image.
|
| 272 |
+
"""
|
| 273 |
+
# `size` is the supported override (`min_pixels` / `max_pixels` are deprecated): count with the budget the
|
| 274 |
+
# preprocessing applies
|
| 275 |
+
images_kwargs = dict(images_kwargs or {})
|
| 276 |
+
size = images_kwargs.pop("size", None) or self.size
|
| 277 |
+
images_kwargs.setdefault("min_pixels", size["shortest_edge"])
|
| 278 |
+
images_kwargs.setdefault("max_pixels", size["longest_edge"])
|
| 279 |
+
images_kwargs = images_kwargs or {}
|
| 280 |
+
min_pixels = images_kwargs["min_pixels"] if "min_pixels" in images_kwargs else self.size["shortest_edge"]
|
| 281 |
+
max_pixels = images_kwargs["max_pixels"] if "max_pixels" in images_kwargs else self.size["longest_edge"]
|
| 282 |
+
patch_size = images_kwargs.get("patch_size", self.patch_size)
|
| 283 |
+
merge_size = images_kwargs.get("merge_size", self.merge_size)
|
| 284 |
+
|
| 285 |
+
factor = patch_size * merge_size
|
| 286 |
+
resized_height, resized_width = smart_resize(
|
| 287 |
+
height, width, factor, min_pixels=min_pixels, max_pixels=max_pixels
|
| 288 |
+
)
|
| 289 |
+
grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
|
| 290 |
+
return grid_h * grid_w
|
| 291 |
+
|
| 292 |
+
|
| 293 |
+
__all__ = ["Lfm3VlImageProcessor"]
|
lfm3_conv.py
ADDED
|
@@ -0,0 +1,202 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
"""Inference kernels for LFM3's width-7 depthwise convolution."""
|
| 15 |
+
|
| 16 |
+
import torch
|
| 17 |
+
|
| 18 |
+
from transformers.utils import is_triton_available
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
if is_triton_available():
|
| 22 |
+
import triton
|
| 23 |
+
import triton.language as tl
|
| 24 |
+
|
| 25 |
+
# the length and strides are runtime arguments, or every prompt length would compile its own kernel
|
| 26 |
+
@triton.jit(do_not_specialize=["T", "X0", "X1", "X2", "S0", "S1", "ST"])
|
| 27 |
+
def _prefill(
|
| 28 |
+
X,
|
| 29 |
+
W,
|
| 30 |
+
BIAS,
|
| 31 |
+
SEQ,
|
| 32 |
+
Y,
|
| 33 |
+
T,
|
| 34 |
+
D: tl.constexpr,
|
| 35 |
+
X0,
|
| 36 |
+
X1,
|
| 37 |
+
X2,
|
| 38 |
+
S0,
|
| 39 |
+
S1,
|
| 40 |
+
ST,
|
| 41 |
+
HAS_BIAS: tl.constexpr,
|
| 42 |
+
PACKED: tl.constexpr,
|
| 43 |
+
BLOCK: tl.constexpr,
|
| 44 |
+
):
|
| 45 |
+
# int64 offsets: batch * channels * length can exceed 2**31 elements
|
| 46 |
+
b, d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
|
| 47 |
+
t = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
|
| 48 |
+
acc = tl.full((BLOCK,), 0, tl.float32)
|
| 49 |
+
if PACKED:
|
| 50 |
+
doc = tl.load(SEQ + b * S0 + (t - T + ST) * S1, (t < T) & (t >= T - ST), other=-1)
|
| 51 |
+
for k in tl.static_range(7):
|
| 52 |
+
pos = t + k - 6
|
| 53 |
+
valid = (t < T) & (pos >= 0)
|
| 54 |
+
if PACKED:
|
| 55 |
+
prev = tl.load(SEQ + b * S0 + (pos - T + ST) * S1, valid & (pos >= T - ST), other=-1)
|
| 56 |
+
valid = valid & (doc == prev)
|
| 57 |
+
x = tl.load(X + b * X0 + d * X1 + pos * X2, valid, other=0).to(tl.float32)
|
| 58 |
+
w = tl.load(W + d * 7 + k).to(tl.float32)
|
| 59 |
+
acc += x * w
|
| 60 |
+
if HAS_BIAS:
|
| 61 |
+
acc += tl.load(BIAS + d).to(tl.float32)
|
| 62 |
+
tl.store(Y + (b * D + d) * T + t, acc, t < T)
|
| 63 |
+
|
| 64 |
+
@triton.jit
|
| 65 |
+
def _decode(
|
| 66 |
+
X,
|
| 67 |
+
W,
|
| 68 |
+
BIAS,
|
| 69 |
+
STATE,
|
| 70 |
+
Y,
|
| 71 |
+
D: tl.constexpr,
|
| 72 |
+
K: tl.constexpr,
|
| 73 |
+
X0: tl.constexpr,
|
| 74 |
+
X1: tl.constexpr,
|
| 75 |
+
S0: tl.constexpr,
|
| 76 |
+
S1: tl.constexpr,
|
| 77 |
+
S2: tl.constexpr,
|
| 78 |
+
HAS_BIAS: tl.constexpr,
|
| 79 |
+
BLOCK: tl.constexpr,
|
| 80 |
+
):
|
| 81 |
+
b = tl.program_id(0).to(tl.int64)
|
| 82 |
+
d = tl.program_id(1).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
|
| 83 |
+
k = tl.arange(0, 8)
|
| 84 |
+
# Load the complete state before writing the shifted window.
|
| 85 |
+
old = tl.load(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, (d[:, None] < D) & (k[None, :] < K), other=0)
|
| 86 |
+
x = tl.load(X + b * X0 + d * X1, d < D, other=0)
|
| 87 |
+
shifted = tl.gather(old, tl.broadcast_to(tl.minimum(k + 1, 7)[None, :], (BLOCK, 8)), 1)
|
| 88 |
+
shifted = tl.where(k[None, :] == K - 1, x[:, None], shifted)
|
| 89 |
+
# A cache may store either K=6 previous inputs or K=7 including the oldest discarded input.
|
| 90 |
+
positions = k + K - 6
|
| 91 |
+
values = tl.gather(old, tl.broadcast_to(tl.minimum(positions, 7)[None, :], (BLOCK, 8)), 1)
|
| 92 |
+
values = tl.where(k[None, :] == 6, x[:, None], values)
|
| 93 |
+
weights = tl.load(W + d[:, None] * 7 + k[None, :], (d[:, None] < D) & (k[None, :] < 7), other=0)
|
| 94 |
+
acc = tl.sum(values.to(tl.float32) * weights.to(tl.float32), 1)
|
| 95 |
+
if HAS_BIAS:
|
| 96 |
+
acc += tl.load(BIAS + d, d < D, other=0).to(tl.float32)
|
| 97 |
+
tl.store(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, shifted, (d[:, None] < D) & (k[None, :] < K))
|
| 98 |
+
tl.store(Y + b * D + d, acc, d < D)
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def can_use_lfm3_conv(hidden_states, weight):
|
| 102 |
+
return (
|
| 103 |
+
is_triton_available()
|
| 104 |
+
and hasattr(tl, "gather") # the decode kernel needs Triton >= 3.3
|
| 105 |
+
and hidden_states.is_cuda
|
| 106 |
+
and not torch.is_grad_enabled()
|
| 107 |
+
and hidden_states.dtype == weight.dtype
|
| 108 |
+
and hidden_states.dtype in (torch.float16, torch.bfloat16, torch.float32)
|
| 109 |
+
and weight.shape[-1] == 7
|
| 110 |
+
)
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
def lfm3_causal_conv1d(hidden_states, weight, bias=None, seq_idx=None):
|
| 114 |
+
batch, channels, length = hidden_states.shape
|
| 115 |
+
# the kernels index every tensor by the batch and channel of `hidden_states` and cannot bound the others
|
| 116 |
+
if weight.shape != (channels, 7) or (bias is not None and bias.shape != (channels,)):
|
| 117 |
+
raise ValueError(
|
| 118 |
+
f"`weight` must have shape ({channels}, 7) and `bias` ({channels},), got {tuple(weight.shape)}"
|
| 119 |
+
)
|
| 120 |
+
if seq_idx is not None and (
|
| 121 |
+
seq_idx.ndim != 2 or seq_idx.shape[0] not in (1, batch) or seq_idx.dtype not in (torch.int32, torch.int64)
|
| 122 |
+
):
|
| 123 |
+
raise ValueError(
|
| 124 |
+
f"`seq_idx` must be an int32 or int64 tensor of shape (1, length) or ({batch}, length), got "
|
| 125 |
+
f"{seq_idx.dtype} {tuple(seq_idx.shape)}"
|
| 126 |
+
)
|
| 127 |
+
if any(t is not None and t.device != hidden_states.device for t in (weight, bias, seq_idx)):
|
| 128 |
+
raise ValueError(
|
| 129 |
+
f"`weight`, `bias` and `seq_idx` must be on the device of `hidden_states` ({hidden_states.device})"
|
| 130 |
+
)
|
| 131 |
+
output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
|
| 132 |
+
# Triton launches on the current device; the model may live on another one (`device_map`, `.to("cuda:1")`).
|
| 133 |
+
with torch.cuda.device(hidden_states.device):
|
| 134 |
+
_prefill[(batch, channels, triton.cdiv(length, 128))](
|
| 135 |
+
hidden_states,
|
| 136 |
+
# the kernels index `bias` and `weight` densely
|
| 137 |
+
weight.contiguous(),
|
| 138 |
+
bias.contiguous() if bias is not None else None,
|
| 139 |
+
seq_idx,
|
| 140 |
+
output,
|
| 141 |
+
length,
|
| 142 |
+
channels,
|
| 143 |
+
*hidden_states.stride(),
|
| 144 |
+
*(
|
| 145 |
+
(0 if seq_idx.shape[0] == 1 else seq_idx.stride(0), seq_idx.stride(1))
|
| 146 |
+
if seq_idx is not None
|
| 147 |
+
else (0, 0)
|
| 148 |
+
),
|
| 149 |
+
seq_idx.shape[-1] if seq_idx is not None else 0,
|
| 150 |
+
bias is not None,
|
| 151 |
+
seq_idx is not None,
|
| 152 |
+
128,
|
| 153 |
+
enable_fp_fusion=False,
|
| 154 |
+
)
|
| 155 |
+
return output
|
| 156 |
+
|
| 157 |
+
|
| 158 |
+
def lfm3_causal_conv1d_update(hidden_states, conv_state, weight, bias=None):
|
| 159 |
+
batch, channels, length = hidden_states.shape
|
| 160 |
+
# the kernel addresses one position per row and bounds nothing by the other tensors' shapes
|
| 161 |
+
if length != 1:
|
| 162 |
+
raise ValueError(f"`hidden_states` must hold one position per row, got {tuple(hidden_states.shape)}")
|
| 163 |
+
# the 8-wide tile holds a 6- or 7-wide state (see `positions` in `_decode`)
|
| 164 |
+
if (
|
| 165 |
+
conv_state.ndim != 3
|
| 166 |
+
or conv_state.shape[:2] != (batch, channels)
|
| 167 |
+
or conv_state.shape[-1] not in (6, 7)
|
| 168 |
+
or weight.shape != (channels, 7)
|
| 169 |
+
):
|
| 170 |
+
raise ValueError(
|
| 171 |
+
f"`conv_state` must have shape ({batch}, {channels}, 6 or 7) and `weight` ({channels}, 7), got "
|
| 172 |
+
f"{tuple(conv_state.shape)} and {tuple(weight.shape)}"
|
| 173 |
+
)
|
| 174 |
+
if bias is not None and bias.shape != (channels,):
|
| 175 |
+
raise ValueError(f"`bias` must have shape ({channels},), got {tuple(bias.shape)}")
|
| 176 |
+
# the kernel writes the state in place: rows or channels sharing memory (an `expand`ed cache) would race
|
| 177 |
+
extent = 0
|
| 178 |
+
for stride, size in sorted(zip(conv_state.stride(), conv_state.shape)):
|
| 179 |
+
if size > 1 and stride <= extent:
|
| 180 |
+
raise ValueError(f"`conv_state` must not overlap itself, got strides {conv_state.stride()}")
|
| 181 |
+
extent += stride * (size - 1)
|
| 182 |
+
if any(t is not None and t.device != hidden_states.device for t in (conv_state, weight, bias)):
|
| 183 |
+
raise ValueError(
|
| 184 |
+
f"`conv_state`, `weight` and `bias` must be on the device of `hidden_states` ({hidden_states.device})"
|
| 185 |
+
)
|
| 186 |
+
output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
|
| 187 |
+
with torch.cuda.device(hidden_states.device):
|
| 188 |
+
_decode[(batch, triton.cdiv(channels, 32))](
|
| 189 |
+
hidden_states,
|
| 190 |
+
weight.contiguous(),
|
| 191 |
+
bias.contiguous() if bias is not None else None,
|
| 192 |
+
conv_state,
|
| 193 |
+
output,
|
| 194 |
+
channels,
|
| 195 |
+
conv_state.shape[-1],
|
| 196 |
+
*hidden_states.stride()[:2],
|
| 197 |
+
*conv_state.stride(),
|
| 198 |
+
bias is not None,
|
| 199 |
+
32,
|
| 200 |
+
enable_fp_fusion=False,
|
| 201 |
+
)
|
| 202 |
+
return output
|
model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:09f7375869ba0697a0ad7e6af917441b1f4556f83ad6647cf3bc6f671c589732
|
| 3 |
+
size 26901336
|
modeling_lfm3.py
ADDED
|
@@ -0,0 +1,881 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
import copy
|
| 15 |
+
from collections.abc import Callable
|
| 16 |
+
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn.functional as F
|
| 19 |
+
from torch import nn
|
| 20 |
+
|
| 21 |
+
from transformers import initialization as init
|
| 22 |
+
from transformers.activations import ACT2FN
|
| 23 |
+
from transformers.cache_utils import Cache, DynamicCache
|
| 24 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 25 |
+
from transformers.generation import GenerationMixin
|
| 26 |
+
from transformers.integrations import (
|
| 27 |
+
use_kernel_forward_from_hub,
|
| 28 |
+
use_kernel_func_from_hub_with_fallback,
|
| 29 |
+
use_kernelized_func,
|
| 30 |
+
)
|
| 31 |
+
from transformers.integrations.accelerate import force_accelerate_hooks
|
| 32 |
+
from transformers.masking_utils import (
|
| 33 |
+
create_causal_mask,
|
| 34 |
+
create_recurrent_attention_mask,
|
| 35 |
+
create_sliding_window_causal_mask,
|
| 36 |
+
)
|
| 37 |
+
from transformers.modeling_layers import GradientCheckpointingLayer
|
| 38 |
+
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
|
| 39 |
+
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
|
| 40 |
+
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
|
| 41 |
+
from transformers.processing_utils import Unpack
|
| 42 |
+
from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
|
| 43 |
+
from transformers.utils.generic import merge_with_config_defaults
|
| 44 |
+
from transformers.utils.import_utils import is_tracing
|
| 45 |
+
from transformers.utils.output_capturing import capture_outputs
|
| 46 |
+
|
| 47 |
+
from .configuration_lfm3 import Lfm3Config
|
| 48 |
+
from .lfm3_conv import can_use_lfm3_conv, lfm3_causal_conv1d, lfm3_causal_conv1d_update
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
@use_kernel_forward_from_hub("RMSNorm")
|
| 52 |
+
class Lfm3RMSNorm(nn.Module):
|
| 53 |
+
def __init__(self, hidden_size, eps: float = 1e-6) -> None:
|
| 54 |
+
"""
|
| 55 |
+
Lfm3RMSNorm is equivalent to T5LayerNorm
|
| 56 |
+
"""
|
| 57 |
+
super().__init__()
|
| 58 |
+
self.weight = nn.Parameter(torch.ones(hidden_size))
|
| 59 |
+
self.variance_epsilon = eps
|
| 60 |
+
|
| 61 |
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 62 |
+
input_dtype = hidden_states.dtype
|
| 63 |
+
hidden_states = hidden_states.to(torch.float32)
|
| 64 |
+
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
| 65 |
+
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
| 66 |
+
return self.weight * hidden_states.to(input_dtype)
|
| 67 |
+
|
| 68 |
+
def extra_repr(self):
|
| 69 |
+
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
class Lfm3RotaryEmbedding(nn.Module):
|
| 73 |
+
"""
|
| 74 |
+
Rotary embedding per attention layer type, like Gemma 3; `conv` layers register no buffers. Partial RoPE keeps the
|
| 75 |
+
frequency progression of the full head dimension and zeroes the tail instead of shrinking the rotated dimension.
|
| 76 |
+
"""
|
| 77 |
+
|
| 78 |
+
def __init__(self, config: Lfm3Config, device=None):
|
| 79 |
+
super().__init__()
|
| 80 |
+
self.max_seq_len_cached = config.max_position_embeddings
|
| 81 |
+
self.original_max_seq_len = config.max_position_embeddings
|
| 82 |
+
self.config = config
|
| 83 |
+
self.layer_types = sorted(set(config.layer_types) & set(config.rope_parameters))
|
| 84 |
+
self.rope_type = {}
|
| 85 |
+
for layer_type in self.layer_types:
|
| 86 |
+
rope_params = self.config.rope_parameters[layer_type]
|
| 87 |
+
self.rope_type[layer_type] = rope_params["rope_type"]
|
| 88 |
+
curr_inv_freq, curr_attention_scaling = self.compute_rope_parameters(config, device, layer_type=layer_type)
|
| 89 |
+
setattr(self, f"{layer_type}_inv_freq", nn.Buffer(curr_inv_freq, persistent=False))
|
| 90 |
+
setattr(self, f"{layer_type}_original_inv_freq", nn.Buffer(curr_inv_freq.clone(), persistent=False))
|
| 91 |
+
setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling)
|
| 92 |
+
|
| 93 |
+
@staticmethod
|
| 94 |
+
def compute_default_rope_parameters(
|
| 95 |
+
config: Lfm3Config, device=None, layer_type: str | None = None, **kwargs
|
| 96 |
+
) -> tuple[torch.Tensor, float]:
|
| 97 |
+
"""
|
| 98 |
+
Inverse frequencies of the default RoPE with LFM3's zero-tail partial rotation: the first
|
| 99 |
+
`partial_rotary_factor * head_dim // 2` frequencies follow the progression of the full head dimension and the
|
| 100 |
+
rest are zero, leaving those dimensions unrotated.
|
| 101 |
+
|
| 102 |
+
Returns:
|
| 103 |
+
Tuple of (`torch.Tensor`, `float`): the inverse frequencies and the attention scaling factor (always 1).
|
| 104 |
+
"""
|
| 105 |
+
base = config.rope_parameters[layer_type]["rope_theta"]
|
| 106 |
+
partial_rotary_factor = config.rope_parameters[layer_type].get("partial_rotary_factor", 1.0)
|
| 107 |
+
dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
|
| 108 |
+
rope_angles = int(partial_rotary_factor * dim) // 2
|
| 109 |
+
inv_freq = 1.0 / (base ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.float) / dim))
|
| 110 |
+
inv_freq = torch.cat((inv_freq, torch.zeros(dim // 2 - rope_angles, dtype=torch.float)))
|
| 111 |
+
return inv_freq.to(device), 1.0
|
| 112 |
+
|
| 113 |
+
@torch.no_grad()
|
| 114 |
+
@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
|
| 115 |
+
def forward(self, x, position_ids, layer_type):
|
| 116 |
+
inv_freq = getattr(self, f"{layer_type}_inv_freq")
|
| 117 |
+
attention_scaling = getattr(self, f"{layer_type}_attention_scaling")
|
| 118 |
+
|
| 119 |
+
freqs = position_ids[..., None].float() * inv_freq.to(device=x.device, dtype=torch.float)
|
| 120 |
+
emb = torch.cat((freqs, freqs), dim=-1)
|
| 121 |
+
cos = emb.cos() * attention_scaling
|
| 122 |
+
sin = emb.sin() * attention_scaling
|
| 123 |
+
|
| 124 |
+
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
| 125 |
+
|
| 126 |
+
@classmethod
|
| 127 |
+
def compute_rope_parameters(
|
| 128 |
+
cls, config: Lfm3Config, device=None, layer_type: str | None = None
|
| 129 |
+
) -> tuple[torch.Tensor, float]:
|
| 130 |
+
"""Inverse frequencies of any RoPE type in the zero-tail layout: scaled types are evaluated over the full head
|
| 131 |
+
dimension and truncated to the rotated prefix."""
|
| 132 |
+
rope_params = config.rope_parameters[layer_type]
|
| 133 |
+
if rope_params["rope_type"] == "default":
|
| 134 |
+
return cls.compute_default_rope_parameters(config, device, layer_type=layer_type)
|
| 135 |
+
full_config = copy.copy(config)
|
| 136 |
+
full_config.rope_parameters = {
|
| 137 |
+
**config.rope_parameters,
|
| 138 |
+
layer_type: {**rope_params, "partial_rotary_factor": 1.0},
|
| 139 |
+
}
|
| 140 |
+
inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_params["rope_type"]](
|
| 141 |
+
full_config, device, layer_type=layer_type
|
| 142 |
+
)
|
| 143 |
+
dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
|
| 144 |
+
rope_angles = int(rope_params.get("partial_rotary_factor", 1.0) * dim) // 2
|
| 145 |
+
inv_freq = torch.cat((inv_freq[:rope_angles], inv_freq.new_zeros(dim // 2 - rope_angles)))
|
| 146 |
+
return inv_freq, attention_factor
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
class Lfm3MLP(nn.Module):
|
| 150 |
+
def __init__(self, config: Lfm3Config):
|
| 151 |
+
super().__init__()
|
| 152 |
+
self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 153 |
+
self.w3 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 154 |
+
self.w2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
| 155 |
+
|
| 156 |
+
def forward(self, x):
|
| 157 |
+
return self.w2(F.silu(self.w1(x)) * self.w3(x))
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
def rotate_half(x):
|
| 161 |
+
"""Rotates half the hidden dims of the input."""
|
| 162 |
+
x1 = x[..., : x.shape[-1] // 2]
|
| 163 |
+
x2 = x[..., x.shape[-1] // 2 :]
|
| 164 |
+
return torch.cat((-x2, x1), dim=-1)
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
| 168 |
+
"""
|
| 169 |
+
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
| 170 |
+
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
| 171 |
+
"""
|
| 172 |
+
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
| 173 |
+
if n_rep == 1:
|
| 174 |
+
return hidden_states
|
| 175 |
+
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
| 176 |
+
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
def eager_attention_forward(
|
| 180 |
+
module: nn.Module,
|
| 181 |
+
query: torch.Tensor,
|
| 182 |
+
key: torch.Tensor,
|
| 183 |
+
value: torch.Tensor,
|
| 184 |
+
attention_mask: torch.Tensor | None,
|
| 185 |
+
scaling: float,
|
| 186 |
+
dropout: float = 0.0,
|
| 187 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 188 |
+
):
|
| 189 |
+
key_states = repeat_kv(key, module.num_key_value_groups)
|
| 190 |
+
value_states = repeat_kv(value, module.num_key_value_groups)
|
| 191 |
+
|
| 192 |
+
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
|
| 193 |
+
if attention_mask is not None:
|
| 194 |
+
attn_weights = attn_weights + attention_mask
|
| 195 |
+
|
| 196 |
+
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
|
| 197 |
+
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
|
| 198 |
+
attn_output = torch.matmul(attn_weights, value_states)
|
| 199 |
+
attn_output = attn_output.transpose(1, 2).contiguous()
|
| 200 |
+
|
| 201 |
+
return attn_output, attn_weights
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
|
| 205 |
+
"""Applies RoPE to the query and key tensors in float32, like the reference implementation."""
|
| 206 |
+
cos = cos.unsqueeze(unsqueeze_dim).float()
|
| 207 |
+
sin = sin.unsqueeze(unsqueeze_dim).float()
|
| 208 |
+
q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
|
| 209 |
+
k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
|
| 210 |
+
return q_embed.to(q.dtype), k_embed.to(k.dtype)
|
| 211 |
+
|
| 212 |
+
|
| 213 |
+
class Lfm3Attention(nn.Module):
|
| 214 |
+
"""Llama attention with QK RMSNorm, a sigmoid output gate per head, a sliding window on `sliding_attention`
|
| 215 |
+
layers and RoPE in float32."""
|
| 216 |
+
|
| 217 |
+
def __init__(self, config: Lfm3Config, layer_idx: int):
|
| 218 |
+
super().__init__()
|
| 219 |
+
self.config = config
|
| 220 |
+
self.layer_idx = layer_idx
|
| 221 |
+
self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
| 222 |
+
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
|
| 223 |
+
self.scaling = self.head_dim**-0.5
|
| 224 |
+
self.is_causal = True
|
| 225 |
+
self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
|
| 226 |
+
self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
|
| 227 |
+
self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
|
| 228 |
+
self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
|
| 229 |
+
self.q_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
|
| 230 |
+
self.k_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
|
| 231 |
+
self.gate_proj = nn.Linear(config.hidden_size, config.num_attention_heads, bias=False)
|
| 232 |
+
self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
|
| 233 |
+
|
| 234 |
+
def forward(
|
| 235 |
+
self,
|
| 236 |
+
hidden_states: torch.Tensor,
|
| 237 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
| 238 |
+
attention_mask: torch.Tensor | None,
|
| 239 |
+
past_key_values: Cache | None = None,
|
| 240 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 241 |
+
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
| 242 |
+
input_shape = hidden_states.shape[:-1]
|
| 243 |
+
hidden_shape = (*input_shape, -1, self.head_dim)
|
| 244 |
+
|
| 245 |
+
query_states = self.q_norm(self.q_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
|
| 246 |
+
key_states = self.k_norm(self.k_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
|
| 247 |
+
value_states = self.v_proj(hidden_states).view(*hidden_shape).transpose(1, 2)
|
| 248 |
+
|
| 249 |
+
cos, sin = position_embeddings
|
| 250 |
+
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
| 251 |
+
|
| 252 |
+
if past_key_values is not None:
|
| 253 |
+
key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
|
| 254 |
+
|
| 255 |
+
attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
|
| 256 |
+
self.config._attn_implementation, eager_attention_forward
|
| 257 |
+
)
|
| 258 |
+
|
| 259 |
+
attn_output, attn_weights = attention_interface(
|
| 260 |
+
self,
|
| 261 |
+
query_states,
|
| 262 |
+
key_states,
|
| 263 |
+
value_states,
|
| 264 |
+
attention_mask,
|
| 265 |
+
dropout=0.0,
|
| 266 |
+
scaling=self.scaling,
|
| 267 |
+
sliding_window=self.sliding_window,
|
| 268 |
+
**kwargs,
|
| 269 |
+
)
|
| 270 |
+
gate = torch.sigmoid(self.gate_proj(hidden_states)).unsqueeze(-1)
|
| 271 |
+
attn_output = (attn_output * gate).reshape(*input_shape, -1).contiguous()
|
| 272 |
+
attn_output = self.o_proj(attn_output)
|
| 273 |
+
return attn_output, attn_weights
|
| 274 |
+
|
| 275 |
+
|
| 276 |
+
def apply_mask_to_padding_states(hidden_states, attention_mask):
|
| 277 |
+
"""
|
| 278 |
+
Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
|
| 279 |
+
"""
|
| 280 |
+
# NOTE: attention mask is a 2D boolean tensor
|
| 281 |
+
if attention_mask is not None:
|
| 282 |
+
dtype = hidden_states.dtype
|
| 283 |
+
hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
|
| 284 |
+
|
| 285 |
+
return hidden_states
|
| 286 |
+
|
| 287 |
+
|
| 288 |
+
@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
|
| 289 |
+
def causal_conv1d_update(
|
| 290 |
+
hidden_states: torch.Tensor,
|
| 291 |
+
conv_state: torch.Tensor,
|
| 292 |
+
weight: nn.Parameter,
|
| 293 |
+
bias: nn.Parameter | None = None,
|
| 294 |
+
activation: str | None = None,
|
| 295 |
+
):
|
| 296 |
+
_, hidden_size, seq_len = hidden_states.shape
|
| 297 |
+
state_len = conv_state.shape[-1]
|
| 298 |
+
|
| 299 |
+
hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
|
| 300 |
+
conv_state.copy_(hidden_states_new[:, :, -state_len:])
|
| 301 |
+
out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
|
| 302 |
+
out = out[:, :, -seq_len:]
|
| 303 |
+
if activation is not None:
|
| 304 |
+
out = ACT2FN[activation](out)
|
| 305 |
+
return out.to(hidden_states.dtype)
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
|
| 309 |
+
def causal_conv1d_fn(
|
| 310 |
+
hidden_states: torch.Tensor,
|
| 311 |
+
weight: nn.Parameter,
|
| 312 |
+
bias: nn.Parameter | None = None,
|
| 313 |
+
activation: str | None = None,
|
| 314 |
+
**kwargs,
|
| 315 |
+
):
|
| 316 |
+
_, hidden_size, seq_len = hidden_states.shape
|
| 317 |
+
padding = weight.shape[-1] - 1
|
| 318 |
+
|
| 319 |
+
out = F.conv1d(
|
| 320 |
+
hidden_states.to(weight.dtype),
|
| 321 |
+
weight=weight.unsqueeze(1),
|
| 322 |
+
bias=bias,
|
| 323 |
+
padding=padding,
|
| 324 |
+
groups=hidden_size,
|
| 325 |
+
)[:, :, :seq_len]
|
| 326 |
+
if activation is not None:
|
| 327 |
+
out = ACT2FN[activation](out)
|
| 328 |
+
return out.to(hidden_states.dtype)
|
| 329 |
+
|
| 330 |
+
|
| 331 |
+
@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
|
| 332 |
+
class Lfm3ShortConv(nn.Module):
|
| 333 |
+
"""LFM2 gated short convolution with packed-document boundaries (`seq_idx`), the padding mask re-applied after a
|
| 334 |
+
biased `in_proj`, and width-7 Triton inference kernels with a PyTorch fallback."""
|
| 335 |
+
|
| 336 |
+
def __init__(
|
| 337 |
+
self,
|
| 338 |
+
config: Lfm3Config,
|
| 339 |
+
layer_idx: int,
|
| 340 |
+
):
|
| 341 |
+
super().__init__()
|
| 342 |
+
self.config = config
|
| 343 |
+
self.layer_idx = layer_idx
|
| 344 |
+
self.conv_kernel_size = config.conv_L_cache
|
| 345 |
+
self.bias = config.conv_bias
|
| 346 |
+
|
| 347 |
+
self.conv = nn.Conv1d(
|
| 348 |
+
in_channels=config.hidden_size,
|
| 349 |
+
out_channels=config.hidden_size,
|
| 350 |
+
kernel_size=self.conv_kernel_size,
|
| 351 |
+
groups=config.hidden_size,
|
| 352 |
+
bias=self.bias,
|
| 353 |
+
padding=self.conv_kernel_size - 1,
|
| 354 |
+
)
|
| 355 |
+
self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=self.bias)
|
| 356 |
+
self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=self.bias)
|
| 357 |
+
|
| 358 |
+
self.layer_type = config.layer_types[layer_idx]
|
| 359 |
+
|
| 360 |
+
@force_accelerate_hooks("conv")
|
| 361 |
+
def forward(
|
| 362 |
+
self,
|
| 363 |
+
hidden_states: torch.Tensor,
|
| 364 |
+
past_key_values: Cache | None = None,
|
| 365 |
+
attention_mask: torch.Tensor | None = None,
|
| 366 |
+
seq_idx: torch.IntTensor | None = None,
|
| 367 |
+
):
|
| 368 |
+
seq_len = hidden_states.shape[1]
|
| 369 |
+
|
| 370 |
+
hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
|
| 371 |
+
BCx = self.in_proj(hidden_states).transpose(-1, -2)
|
| 372 |
+
B, C, x = BCx.chunk(3, dim=-2)
|
| 373 |
+
hidden_states = B * x
|
| 374 |
+
if attention_mask is not None:
|
| 375 |
+
# with `conv_bias` the projection makes the padded positions non-zero again
|
| 376 |
+
hidden_states = hidden_states * attention_mask[:, None, :].to(hidden_states.dtype)
|
| 377 |
+
|
| 378 |
+
use_precomputed_states = past_key_values is not None and past_key_values.has_previous_state(self.layer_idx)
|
| 379 |
+
use_kernels = 2 <= self.conv_kernel_size <= 4
|
| 380 |
+
use_wide_kernel = can_use_lfm3_conv(hidden_states, self.conv.weight)
|
| 381 |
+
|
| 382 |
+
if use_precomputed_states and seq_len == 1 and not past_key_values.layers[self.layer_idx].record_past:
|
| 383 |
+
conv_state = past_key_values.layers[self.layer_idx].conv_states[0]
|
| 384 |
+
if use_wide_kernel and conv_state.shape[-1] in (6, 7):
|
| 385 |
+
hidden_states = lfm3_causal_conv1d_update(
|
| 386 |
+
hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
|
| 387 |
+
)
|
| 388 |
+
elif use_kernels:
|
| 389 |
+
hidden_states = causal_conv1d_update(
|
| 390 |
+
hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
|
| 391 |
+
)
|
| 392 |
+
else:
|
| 393 |
+
hidden_states = torch.cat([conv_state, hidden_states], dim=-1)
|
| 394 |
+
conv_state.copy_(hidden_states[:, :, -conv_state.shape[-1] :])
|
| 395 |
+
hidden_states = self._conv(hidden_states)[:, :, -seq_len:]
|
| 396 |
+
else:
|
| 397 |
+
if past_key_values is not None and seq_idx is not None:
|
| 398 |
+
# `seq_idx` is only set on an empty cache, and the decode steps read the cached window without it:
|
| 399 |
+
# only the last document (after any positional left padding) may enter the window
|
| 400 |
+
last_document = (seq_idx == seq_idx[:, -1:])[:, None].to(hidden_states.dtype)
|
| 401 |
+
past_key_values.update_conv_state(
|
| 402 |
+
hidden_states * last_document, self.layer_idx, conv_kernel_size=self.conv_kernel_size
|
| 403 |
+
)
|
| 404 |
+
elif past_key_values is not None:
|
| 405 |
+
hidden_states = past_key_values.update_conv_state(
|
| 406 |
+
hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
|
| 407 |
+
)
|
| 408 |
+
|
| 409 |
+
if use_wide_kernel:
|
| 410 |
+
hidden_states = lfm3_causal_conv1d(hidden_states, self.conv.weight.squeeze(1), self.conv.bias, seq_idx)
|
| 411 |
+
elif seq_idx is not None:
|
| 412 |
+
hidden_states = self._segmented_conv(hidden_states, seq_idx)
|
| 413 |
+
elif use_kernels:
|
| 414 |
+
hidden_states = causal_conv1d_fn(hidden_states, self.conv.weight.squeeze(1), self.conv.bias)
|
| 415 |
+
else:
|
| 416 |
+
hidden_states = self._conv(hidden_states)
|
| 417 |
+
|
| 418 |
+
# Drop the additional previous states
|
| 419 |
+
if past_key_values is not None:
|
| 420 |
+
hidden_states = hidden_states[:, :, -seq_len:]
|
| 421 |
+
|
| 422 |
+
y = C * hidden_states
|
| 423 |
+
y = y.transpose(-1, -2).contiguous()
|
| 424 |
+
y = self.out_proj(y)
|
| 425 |
+
return y
|
| 426 |
+
|
| 427 |
+
def _conv(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 428 |
+
"""Causal depthwise conv in torch: the `causal_conv1d` kernels only accept widths 2 to 4."""
|
| 429 |
+
conv_out = self.conv(hidden_states.to(self.conv.weight.dtype))[..., : hidden_states.shape[-1]]
|
| 430 |
+
return conv_out.to(hidden_states.dtype)
|
| 431 |
+
|
| 432 |
+
def _segmented_conv(self, hidden_states: torch.Tensor, seq_idx: torch.Tensor) -> torch.Tensor:
|
| 433 |
+
"""Causal depthwise conv that only reads positions sharing the current token's document index."""
|
| 434 |
+
kernel_size = self.conv_kernel_size
|
| 435 |
+
seq_idx = F.pad(seq_idx, (hidden_states.shape[-1] - seq_idx.shape[-1] + kernel_size - 1, 0), value=-1)
|
| 436 |
+
windows = F.pad(hidden_states, (kernel_size - 1, 0)).unfold(-1, kernel_size, 1)
|
| 437 |
+
seq_idx_windows = seq_idx.unfold(-1, kernel_size, 1)
|
| 438 |
+
same_document = (seq_idx_windows == seq_idx_windows[..., -1:]).unsqueeze(1)
|
| 439 |
+
conv_out = torch.einsum("bdtk,dk->bdt", windows * same_document.to(windows.dtype), self.conv.weight.squeeze(1))
|
| 440 |
+
if self.bias:
|
| 441 |
+
conv_out = conv_out + self.conv.bias[None, :, None]
|
| 442 |
+
return conv_out.to(hidden_states.dtype)
|
| 443 |
+
|
| 444 |
+
|
| 445 |
+
class Lfm3DecoderLayer(GradientCheckpointingLayer):
|
| 446 |
+
def __init__(self, config: Lfm3Config, layer_idx: int):
|
| 447 |
+
super().__init__()
|
| 448 |
+
self.is_attention_layer = config.layer_types[layer_idx] in config._attention_layer_types
|
| 449 |
+
self.use_sandwich_norm = config.use_sandwich_norm
|
| 450 |
+
|
| 451 |
+
if self.is_attention_layer:
|
| 452 |
+
self.self_attn = Lfm3Attention(config, layer_idx)
|
| 453 |
+
else:
|
| 454 |
+
self.conv = Lfm3ShortConv(config, layer_idx)
|
| 455 |
+
self.feed_forward = Lfm3MLP(config)
|
| 456 |
+
self.operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 457 |
+
self.ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 458 |
+
if self.use_sandwich_norm:
|
| 459 |
+
self.post_operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 460 |
+
self.post_ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 461 |
+
|
| 462 |
+
def forward(
|
| 463 |
+
self,
|
| 464 |
+
hidden_states: torch.Tensor,
|
| 465 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
| 466 |
+
attention_mask: torch.Tensor | None = None,
|
| 467 |
+
position_ids: torch.LongTensor | None = None,
|
| 468 |
+
past_key_values: Cache | None = None,
|
| 469 |
+
seq_idx: torch.IntTensor | None = None,
|
| 470 |
+
**kwargs,
|
| 471 |
+
) -> torch.Tensor:
|
| 472 |
+
residual = hidden_states
|
| 473 |
+
if self.is_attention_layer:
|
| 474 |
+
hidden_states, _ = self.self_attn(
|
| 475 |
+
hidden_states=self.operator_norm(hidden_states),
|
| 476 |
+
position_embeddings=position_embeddings,
|
| 477 |
+
attention_mask=attention_mask,
|
| 478 |
+
position_ids=position_ids,
|
| 479 |
+
past_key_values=past_key_values,
|
| 480 |
+
**kwargs,
|
| 481 |
+
)
|
| 482 |
+
else:
|
| 483 |
+
hidden_states = self.conv(
|
| 484 |
+
hidden_states=self.operator_norm(hidden_states),
|
| 485 |
+
past_key_values=past_key_values,
|
| 486 |
+
attention_mask=attention_mask,
|
| 487 |
+
seq_idx=seq_idx,
|
| 488 |
+
)
|
| 489 |
+
if self.use_sandwich_norm:
|
| 490 |
+
hidden_states = self.post_operator_norm(hidden_states)
|
| 491 |
+
hidden_states = hidden_states + residual
|
| 492 |
+
|
| 493 |
+
ffn_out = self.feed_forward(self.ffn_norm(hidden_states))
|
| 494 |
+
if self.use_sandwich_norm:
|
| 495 |
+
ffn_out = self.post_ffn_norm(ffn_out)
|
| 496 |
+
hidden_states = hidden_states + ffn_out
|
| 497 |
+
|
| 498 |
+
return hidden_states
|
| 499 |
+
|
| 500 |
+
|
| 501 |
+
def create_masks_by_layer_type(
|
| 502 |
+
config: PreTrainedConfig,
|
| 503 |
+
inputs_embeds: torch.Tensor,
|
| 504 |
+
attention_mask: torch.Tensor | None,
|
| 505 |
+
past_key_values: Cache | None,
|
| 506 |
+
position_ids: torch.Tensor | None = None,
|
| 507 |
+
) -> dict:
|
| 508 |
+
"""The cache sizes each mask from the layer it is asked about, and the generic mask utilities ask about the first
|
| 509 |
+
non-sliding layer, which for a leading `conv` layer redirects to the first attention layer of any type: name the
|
| 510 |
+
first layer of each mask's own type instead."""
|
| 511 |
+
config = config.get_text_config()
|
| 512 |
+
mask_kwargs = {
|
| 513 |
+
"config": config,
|
| 514 |
+
"inputs_embeds": inputs_embeds,
|
| 515 |
+
"attention_mask": attention_mask,
|
| 516 |
+
"past_key_values": past_key_values,
|
| 517 |
+
"position_ids": position_ids,
|
| 518 |
+
}
|
| 519 |
+
layer_types, causal_mask_mapping = config.layer_types, {}
|
| 520 |
+
if "full_attention" in layer_types:
|
| 521 |
+
layer_idx = layer_types.index("full_attention")
|
| 522 |
+
causal_mask_mapping["full_attention"] = create_causal_mask(**mask_kwargs, layer_idx=layer_idx)
|
| 523 |
+
if "sliding_attention" in layer_types:
|
| 524 |
+
layer_idx = layer_types.index("sliding_attention")
|
| 525 |
+
causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(
|
| 526 |
+
**mask_kwargs, layer_idx=layer_idx
|
| 527 |
+
)
|
| 528 |
+
if "conv" in layer_types:
|
| 529 |
+
causal_mask_mapping["conv"] = create_recurrent_attention_mask(**mask_kwargs)
|
| 530 |
+
return causal_mask_mapping
|
| 531 |
+
|
| 532 |
+
|
| 533 |
+
@auto_docstring
|
| 534 |
+
class Lfm3PreTrainedModel(PreTrainedModel):
|
| 535 |
+
config: Lfm3Config
|
| 536 |
+
base_model_prefix = "model"
|
| 537 |
+
supports_gradient_checkpointing = True
|
| 538 |
+
_no_split_modules = ["Lfm3DecoderLayer"]
|
| 539 |
+
_skip_keys_device_placement = ["past_key_values"]
|
| 540 |
+
_supports_flash_attn = True
|
| 541 |
+
_supports_sdpa = True
|
| 542 |
+
_supports_flex_attn = True
|
| 543 |
+
|
| 544 |
+
_can_compile_fullgraph = True
|
| 545 |
+
_supports_attention_backend = True
|
| 546 |
+
_can_record_outputs = {
|
| 547 |
+
"hidden_states": Lfm3DecoderLayer,
|
| 548 |
+
"attentions": Lfm3Attention,
|
| 549 |
+
}
|
| 550 |
+
|
| 551 |
+
@torch.no_grad()
|
| 552 |
+
def _init_weights(self, module):
|
| 553 |
+
super()._init_weights(module)
|
| 554 |
+
if isinstance(module, Lfm3RotaryEmbedding):
|
| 555 |
+
for layer_type in module.layer_types:
|
| 556 |
+
curr_inv_freq, _ = module.compute_rope_parameters(module.config, layer_type=layer_type)
|
| 557 |
+
init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
|
| 558 |
+
init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
|
| 559 |
+
|
| 560 |
+
@staticmethod
|
| 561 |
+
def create_masks_for_generate(
|
| 562 |
+
config: PreTrainedConfig,
|
| 563 |
+
inputs_embeds: torch.Tensor,
|
| 564 |
+
attention_mask: torch.Tensor | None,
|
| 565 |
+
past_key_values: Cache | None,
|
| 566 |
+
position_ids: torch.Tensor | None = None,
|
| 567 |
+
**kwargs,
|
| 568 |
+
) -> dict:
|
| 569 |
+
return create_masks_by_layer_type(config, inputs_embeds, attention_mask, past_key_values, position_ids)
|
| 570 |
+
|
| 571 |
+
|
| 572 |
+
def packed_under_padding(position_ids: torch.Tensor, padding_mask: torch.Tensor, filled) -> torch.Tensor:
|
| 573 |
+
"""Padding repeats position 0, so packing is a restart between two kept tokens, or a kept position 0 on a filled
|
| 574 |
+
cache whose row holds a kept token; attention under a padding mask cannot isolate it."""
|
| 575 |
+
restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
|
| 576 |
+
valid = padding_mask[:, -restarts.shape[-1] :].bool()
|
| 577 |
+
packed = restarts[:, 1:] & valid[:, 1:] & valid[:, :-1]
|
| 578 |
+
if filled:
|
| 579 |
+
# a row whose cached tokens are all padding (chunked prefill of a left-padded batch) has no document to pack
|
| 580 |
+
# against; a mask covering only the chunk keeps the strict rule
|
| 581 |
+
cached = padding_mask[:, : padding_mask.shape[-1] - restarts.shape[-1]]
|
| 582 |
+
cached_valid = cached.bool().any(-1, keepdim=True) if cached.shape[-1] else torch.ones_like(valid[:, :1])
|
| 583 |
+
packed = packed | (valid[:, :1] & (position_ids[:, :1] == 0) & cached_valid)
|
| 584 |
+
return packed
|
| 585 |
+
|
| 586 |
+
|
| 587 |
+
@auto_docstring
|
| 588 |
+
class Lfm3Model(Lfm3PreTrainedModel):
|
| 589 |
+
"""LFM2 model with sliding-attention layers as a third layer type: the mask and the rotary embedding are built per
|
| 590 |
+
layer type, and packed documents (`position_ids` restarting) reset the short-conv window through `seq_idx`."""
|
| 591 |
+
|
| 592 |
+
def __init__(self, config: Lfm3Config):
|
| 593 |
+
super().__init__(config)
|
| 594 |
+
self.padding_idx = config.pad_token_id
|
| 595 |
+
self.vocab_size = config.vocab_size
|
| 596 |
+
|
| 597 |
+
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
| 598 |
+
self.layers = nn.ModuleList(
|
| 599 |
+
[Lfm3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
| 600 |
+
)
|
| 601 |
+
self.rotary_emb = Lfm3RotaryEmbedding(config=config)
|
| 602 |
+
self.gradient_checkpointing = False
|
| 603 |
+
self.embedding_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 604 |
+
|
| 605 |
+
# Initialize weights and apply final processing
|
| 606 |
+
self.post_init()
|
| 607 |
+
|
| 608 |
+
@merge_with_config_defaults
|
| 609 |
+
@capture_outputs
|
| 610 |
+
@auto_docstring
|
| 611 |
+
def forward(
|
| 612 |
+
self,
|
| 613 |
+
input_ids: torch.LongTensor | None = None,
|
| 614 |
+
attention_mask: torch.Tensor | None = None,
|
| 615 |
+
position_ids: torch.LongTensor | None = None,
|
| 616 |
+
past_key_values: Cache | None = None,
|
| 617 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 618 |
+
use_cache: bool | None = None,
|
| 619 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 620 |
+
) -> BaseModelOutputWithPast:
|
| 621 |
+
r"""
|
| 622 |
+
attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)` or `dict`, *optional*):
|
| 623 |
+
Mask to avoid performing attention on padding token indices (1 keeps a token, 0 masks it), or the mask
|
| 624 |
+
dict of `create_masks_for_generate`. A mask dict is used as given: `generate` refuses packed
|
| 625 |
+
`position_ids` before its prefill.
|
| 626 |
+
"""
|
| 627 |
+
if (input_ids is None) ^ (inputs_embeds is not None):
|
| 628 |
+
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
| 629 |
+
# `seq_idx` is a shared kwarg (`DataCollatorWithFlattening(return_seq_idx=True)`, read by Lfm2); LFM3 derives
|
| 630 |
+
# the documents from the restarting `position_ids` that come with it
|
| 631 |
+
if kwargs.pop("seq_idx", None) is not None and position_ids is None:
|
| 632 |
+
raise ValueError("Pass packed documents as `position_ids` restarting at 0; LFM3 does not read `seq_idx`")
|
| 633 |
+
|
| 634 |
+
if inputs_embeds is None:
|
| 635 |
+
inputs_embeds = self.embed_tokens(input_ids)
|
| 636 |
+
|
| 637 |
+
if use_cache and past_key_values is None:
|
| 638 |
+
past_key_values = DynamicCache(config=self.config)
|
| 639 |
+
filled = past_key_values is not None and (
|
| 640 |
+
past_key_values.get_seq_length() > 0
|
| 641 |
+
if any(t != "conv" for t in self.config.layer_types)
|
| 642 |
+
else past_key_values.has_previous_state()
|
| 643 |
+
)
|
| 644 |
+
packed_cache = getattr(past_key_values, "lfm3_packed", None)
|
| 645 |
+
if isinstance(packed_cache, torch.Tensor) or packed_cache:
|
| 646 |
+
# Reset layers have no history to mix. Ordinary caches carry a Python `False` and skip this, so the decode
|
| 647 |
+
# steps do no tensor work here; only a compiled packed prefill leaves a tensor flag, and `filled` stays
|
| 648 |
+
# tensor-valued for static caches.
|
| 649 |
+
packed_message = (
|
| 650 |
+
"A cache filled by a packed-document forward cannot be continued: its attention spans the documents"
|
| 651 |
+
)
|
| 652 |
+
packed_cache = packed_cache & filled if isinstance(packed_cache, torch.Tensor) else filled
|
| 653 |
+
if isinstance(packed_cache, torch.Tensor):
|
| 654 |
+
torch_compilable_check(~packed_cache, packed_message)
|
| 655 |
+
if not is_tracing(packed_cache):
|
| 656 |
+
past_key_values.lfm3_packed = bool(packed_cache)
|
| 657 |
+
elif packed_cache:
|
| 658 |
+
raise ValueError(packed_message)
|
| 659 |
+
else:
|
| 660 |
+
past_key_values.lfm3_packed = False
|
| 661 |
+
|
| 662 |
+
# The mask utilities detect packing only with `attention_mask=None`, so a multi-token forward with explicit
|
| 663 |
+
# positions drops the all-ones mask tokenizers emit (a data-dependent check, hence not when traced).
|
| 664 |
+
explicit_positions = position_ids is not None and inputs_embeds.shape[1] > 1
|
| 665 |
+
if (
|
| 666 |
+
explicit_positions
|
| 667 |
+
and isinstance(attention_mask, torch.Tensor)
|
| 668 |
+
and attention_mask.ndim == 2
|
| 669 |
+
and not is_tracing(position_ids)
|
| 670 |
+
and bool(attention_mask.all())
|
| 671 |
+
):
|
| 672 |
+
attention_mask = None
|
| 673 |
+
seq_idx = None
|
| 674 |
+
if position_ids is None:
|
| 675 |
+
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
| 676 |
+
position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
|
| 677 |
+
position_ids = position_ids.unsqueeze(0)
|
| 678 |
+
elif explicit_positions and not isinstance(attention_mask, dict):
|
| 679 |
+
# a mask dict comes from `generate`, which refused packing before its prefill
|
| 680 |
+
if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 2:
|
| 681 |
+
# data-dependent: a traced forward skips it (`generate` refuses packing eagerly before its prefill)
|
| 682 |
+
if not is_tracing(position_ids):
|
| 683 |
+
torch_compilable_check(
|
| 684 |
+
~packed_under_padding(position_ids, attention_mask, filled),
|
| 685 |
+
"Packed documents (`position_ids` restarting) cannot be combined with an attention mask: "
|
| 686 |
+
"attention would not isolate the documents. Pass `attention_mask=None` without a static cache.",
|
| 687 |
+
)
|
| 688 |
+
else:
|
| 689 |
+
restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
|
| 690 |
+
document_ids = restarts.cumsum(-1)
|
| 691 |
+
if is_tracing(position_ids) or bool((document_ids[..., -1] > 0).any()):
|
| 692 |
+
seq_idx = document_ids.to(torch.int32)
|
| 693 |
+
# Packed documents are prefilled in one forward on an empty cache (cached tokens would join the first
|
| 694 |
+
# document), and their cache is never continued
|
| 695 |
+
if filled:
|
| 696 |
+
torch_compilable_check(
|
| 697 |
+
(~restarts[:, 1:]).all() & (position_ids[:, 0] != 0).all(),
|
| 698 |
+
"Packed documents must be prefilled in one forward; the cache already holds tokens",
|
| 699 |
+
)
|
| 700 |
+
# Ordinary continuation must retain the convolution window from the cached document.
|
| 701 |
+
seq_idx = None
|
| 702 |
+
if attention_mask is None:
|
| 703 |
+
# without a mask only the cacheless packed-sequence mask isolates repeated zero positions (left
|
| 704 |
+
# padding): a cache-sized mask and every later decode step would attend them, so they count too
|
| 705 |
+
packed = (document_ids[..., -1] > 0).any()
|
| 706 |
+
else:
|
| 707 |
+
# a prepared 4D mask is the caller's and carries the padding (repeated zero positions), but may
|
| 708 |
+
# or may not isolate documents, so the cache records the packing and is never continued
|
| 709 |
+
steps = torch.diff(position_ids, dim=-1)
|
| 710 |
+
packed = ((steps != 1) & ~((steps == 0) & (position_ids[:, :-1] == 0))).any()
|
| 711 |
+
if past_key_values is not None:
|
| 712 |
+
# resolved here, where the prefill already syncs, so the decode steps carry a Python flag: no
|
| 713 |
+
# per-step host sync, data-dependent op or CUDA-graph output on the cache
|
| 714 |
+
past_key_values.lfm3_packed = packed if is_tracing(position_ids) else bool(packed)
|
| 715 |
+
|
| 716 |
+
if not isinstance(causal_mask_mapping := attention_mask, dict):
|
| 717 |
+
mask_kwargs = {
|
| 718 |
+
"config": self.config,
|
| 719 |
+
"inputs_embeds": inputs_embeds,
|
| 720 |
+
"attention_mask": attention_mask,
|
| 721 |
+
"past_key_values": past_key_values,
|
| 722 |
+
"position_ids": position_ids,
|
| 723 |
+
}
|
| 724 |
+
# The mask utilities only isolate packed documents without a cache, and an empty cache adds nothing.
|
| 725 |
+
layer_types = self.config.layer_types
|
| 726 |
+
if seq_idx is not None and past_key_values is not None:
|
| 727 |
+
attention_layer = next((i for i, t in enumerate(layer_types) if t != "conv"), None)
|
| 728 |
+
if attention_layer is not None and past_key_values.get_seq_length(attention_layer) == 0:
|
| 729 |
+
if past_key_values.is_compileable:
|
| 730 |
+
# static layers always return their full length, so the masks are sized by the cache, which
|
| 731 |
+
# isolates neither documents nor positional-only left padding
|
| 732 |
+
torch_compilable_check(
|
| 733 |
+
~packed,
|
| 734 |
+
"Packed documents, or left padding without an attention mask, cannot be prefilled into a "
|
| 735 |
+
"static cache: its masks are sized by the cache. Pass the padding as `attention_mask`, or "
|
| 736 |
+
"use a `DynamicCache` for packed documents",
|
| 737 |
+
)
|
| 738 |
+
else:
|
| 739 |
+
mask_kwargs["past_key_values"] = None
|
| 740 |
+
causal_mask_mapping = create_masks_by_layer_type(**mask_kwargs)
|
| 741 |
+
|
| 742 |
+
hidden_states = inputs_embeds
|
| 743 |
+
position_embeddings = {
|
| 744 |
+
layer_type: self.rotary_emb(hidden_states, position_ids=position_ids, layer_type=layer_type)
|
| 745 |
+
for layer_type in self.rotary_emb.layer_types
|
| 746 |
+
}
|
| 747 |
+
|
| 748 |
+
for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
|
| 749 |
+
layer_type = self.config.layer_types[i]
|
| 750 |
+
hidden_states = decoder_layer(
|
| 751 |
+
hidden_states,
|
| 752 |
+
attention_mask=causal_mask_mapping[layer_type],
|
| 753 |
+
position_embeddings=position_embeddings.get(layer_type),
|
| 754 |
+
position_ids=position_ids,
|
| 755 |
+
past_key_values=past_key_values,
|
| 756 |
+
seq_idx=seq_idx,
|
| 757 |
+
**kwargs,
|
| 758 |
+
)
|
| 759 |
+
|
| 760 |
+
hidden_states = self.embedding_norm(hidden_states)
|
| 761 |
+
|
| 762 |
+
return BaseModelOutputWithPast(
|
| 763 |
+
last_hidden_state=hidden_states,
|
| 764 |
+
past_key_values=past_key_values,
|
| 765 |
+
)
|
| 766 |
+
|
| 767 |
+
|
| 768 |
+
def reject_packed_generation(config: PreTrainedConfig, input_ids, generation_config, model_kwargs, is_first_iteration):
|
| 769 |
+
"""`generate` decodes from the cache its prefill fills, and no mask it builds isolates packed documents
|
| 770 |
+
(`position_ids` restarting): the prompt is refused before the prefill, one forward or every chunk, writes to it."""
|
| 771 |
+
position_ids, attention_mask = model_kwargs.get("position_ids"), model_kwargs.get("attention_mask")
|
| 772 |
+
cache = model_kwargs.get("past_key_values")
|
| 773 |
+
if cache is None or position_ids is None or position_ids.ndim != 2:
|
| 774 |
+
return
|
| 775 |
+
if attention_mask is not None and attention_mask.ndim != 2:
|
| 776 |
+
return
|
| 777 |
+
has_attention = any(t != "conv" for t in config.get_text_config().layer_types)
|
| 778 |
+
filled = bool(cache.get_seq_length() > 0 if has_attention else cache.has_previous_state())
|
| 779 |
+
inputs_embeds = model_kwargs.get("inputs_embeds") if is_first_iteration else None
|
| 780 |
+
length = input_ids.shape[1] if inputs_embeds is None else inputs_embeds.shape[1]
|
| 781 |
+
if generation_config.prefill_chunk_size is not None:
|
| 782 |
+
# the chunks run from the start of the prompt
|
| 783 |
+
position_ids = position_ids[:, :length]
|
| 784 |
+
else:
|
| 785 |
+
if filled and has_attention:
|
| 786 |
+
# the slice `_prefill` feeds: a prompt that spans the cache is cut to its uncached part
|
| 787 |
+
mask_length = getattr(generation_config, "_mask_length", -1)
|
| 788 |
+
if attention_mask is not None:
|
| 789 |
+
mask_length = attention_mask.shape[1]
|
| 790 |
+
if inputs_embeds is not None or mask_length == length:
|
| 791 |
+
length -= int(cache.get_seq_length())
|
| 792 |
+
position_ids = position_ids[:, position_ids.shape[1] - length :]
|
| 793 |
+
if attention_mask is None:
|
| 794 |
+
attention_mask = torch.ones_like(position_ids)
|
| 795 |
+
if bool(packed_under_padding(position_ids, attention_mask, filled).any()):
|
| 796 |
+
raise ValueError(
|
| 797 |
+
"`generate` cannot decode packed documents (`position_ids` restarting, or starting at 0 on a cache that "
|
| 798 |
+
"holds tokens): attention would span them. Generate each document in its own row of a left-padded batch."
|
| 799 |
+
)
|
| 800 |
+
|
| 801 |
+
|
| 802 |
+
@auto_docstring
|
| 803 |
+
class Lfm3ForCausalLM(Lfm3PreTrainedModel, GenerationMixin):
|
| 804 |
+
_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
|
| 805 |
+
_tp_plan = {"lm_head": "colwise_gather_output"}
|
| 806 |
+
_pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
|
| 807 |
+
_fsdp_plan = {"lm_head": "keep_full_weight"}
|
| 808 |
+
|
| 809 |
+
def __init__(self, config):
|
| 810 |
+
super().__init__(config)
|
| 811 |
+
self.model = Lfm3Model(config)
|
| 812 |
+
self.vocab_size = config.vocab_size
|
| 813 |
+
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
| 814 |
+
|
| 815 |
+
# Initialize weights and apply final processing
|
| 816 |
+
self.post_init()
|
| 817 |
+
|
| 818 |
+
@can_return_tuple
|
| 819 |
+
@auto_docstring
|
| 820 |
+
def forward(
|
| 821 |
+
self,
|
| 822 |
+
input_ids: torch.LongTensor | None = None,
|
| 823 |
+
attention_mask: torch.Tensor | None = None,
|
| 824 |
+
position_ids: torch.LongTensor | None = None,
|
| 825 |
+
past_key_values: Cache | None = None,
|
| 826 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 827 |
+
labels: torch.LongTensor | None = None,
|
| 828 |
+
use_cache: bool | None = None,
|
| 829 |
+
logits_to_keep: int | torch.Tensor = 0,
|
| 830 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 831 |
+
) -> CausalLMOutputWithPast:
|
| 832 |
+
r"""
|
| 833 |
+
Example:
|
| 834 |
+
|
| 835 |
+
```python
|
| 836 |
+
>>> from transformers import AutoTokenizer, Lfm3ForCausalLM
|
| 837 |
+
|
| 838 |
+
>>> model = Lfm3ForCausalLM.from_pretrained("LiquidAI/LFM3-1.7B-Preview")
|
| 839 |
+
>>> tokenizer = AutoTokenizer.from_pretrained("LiquidAI/LFM3-1.7B-Preview")
|
| 840 |
+
|
| 841 |
+
>>> inputs = tokenizer.apply_chat_template(
|
| 842 |
+
... [{"role": "user", "content": "What is C. elegans?"}], add_generation_prompt=True, return_tensors="pt"
|
| 843 |
+
... )
|
| 844 |
+
>>> generate_ids = model.generate(**inputs, max_new_tokens=64)
|
| 845 |
+
>>> tokenizer.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True)
|
| 846 |
+
```"""
|
| 847 |
+
outputs: BaseModelOutputWithPast = self.model(
|
| 848 |
+
input_ids=input_ids,
|
| 849 |
+
attention_mask=attention_mask,
|
| 850 |
+
position_ids=position_ids,
|
| 851 |
+
past_key_values=past_key_values,
|
| 852 |
+
inputs_embeds=inputs_embeds,
|
| 853 |
+
use_cache=use_cache,
|
| 854 |
+
**kwargs,
|
| 855 |
+
)
|
| 856 |
+
|
| 857 |
+
hidden_states = outputs.last_hidden_state
|
| 858 |
+
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
|
| 859 |
+
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
| 860 |
+
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
| 861 |
+
|
| 862 |
+
loss = None
|
| 863 |
+
if labels is not None:
|
| 864 |
+
loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
|
| 865 |
+
|
| 866 |
+
return CausalLMOutputWithPast(
|
| 867 |
+
loss=loss,
|
| 868 |
+
logits=logits,
|
| 869 |
+
past_key_values=outputs.past_key_values,
|
| 870 |
+
hidden_states=outputs.hidden_states,
|
| 871 |
+
attentions=outputs.attentions,
|
| 872 |
+
)
|
| 873 |
+
|
| 874 |
+
def _prefill(self, input_ids, generation_config, model_kwargs, is_first_iteration=True):
|
| 875 |
+
# every prefill passes here once before its first forward: a check per chunk would leave the earlier chunks
|
| 876 |
+
# in the cache, and a dynamic cache would only refuse packing at the first decode step
|
| 877 |
+
reject_packed_generation(self.config, input_ids, generation_config, model_kwargs, is_first_iteration)
|
| 878 |
+
return super()._prefill(input_ids, generation_config, model_kwargs, is_first_iteration=is_first_iteration)
|
| 879 |
+
|
| 880 |
+
|
| 881 |
+
__all__ = ["Lfm3ForCausalLM", "Lfm3Model", "Lfm3PreTrainedModel"]
|
modeling_lfm3_omni.py
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
modeling_lfm3_vl.py
ADDED
|
@@ -0,0 +1,965 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
from collections.abc import Callable
|
| 15 |
+
from dataclasses import dataclass
|
| 16 |
+
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn.functional as F
|
| 19 |
+
from torch import nn
|
| 20 |
+
|
| 21 |
+
from transformers.activations import ACT2FN
|
| 22 |
+
from transformers.cache_utils import Cache
|
| 23 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 24 |
+
from transformers.generation import GenerationMixin
|
| 25 |
+
from transformers.integrations import use_kernel_forward_from_hub
|
| 26 |
+
from transformers.masking_utils import (
|
| 27 |
+
create_causal_mask,
|
| 28 |
+
create_recurrent_attention_mask,
|
| 29 |
+
create_sliding_window_causal_mask,
|
| 30 |
+
)
|
| 31 |
+
from transformers.modeling_layers import GradientCheckpointingLayer
|
| 32 |
+
from transformers.modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast
|
| 33 |
+
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
|
| 34 |
+
from transformers.processing_utils import Unpack
|
| 35 |
+
from transformers.utils import (
|
| 36 |
+
TransformersKwargs,
|
| 37 |
+
auto_docstring,
|
| 38 |
+
can_return_tuple,
|
| 39 |
+
torch_compilable_check,
|
| 40 |
+
)
|
| 41 |
+
from transformers.utils.generic import (
|
| 42 |
+
accepts_precomputed_kwargs,
|
| 43 |
+
is_flash_attention_requested,
|
| 44 |
+
maybe_autocast,
|
| 45 |
+
merge_with_config_defaults,
|
| 46 |
+
)
|
| 47 |
+
from transformers.utils.output_capturing import OutputRecorder, capture_outputs
|
| 48 |
+
from transformers.vision_utils import get_max_seqlen, get_vision_attention_seqlens, get_vision_position_ids
|
| 49 |
+
|
| 50 |
+
from .configuration_lfm3_vl import Lfm3VlConfig, Lfm3VlVisionConfig
|
| 51 |
+
from .modeling_lfm3 import Lfm3Model
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
@use_kernel_forward_from_hub("RMSNorm")
|
| 55 |
+
class Lfm3VlRMSNorm(nn.Module):
|
| 56 |
+
def __init__(self, hidden_size, eps: float = 1e-6) -> None:
|
| 57 |
+
"""
|
| 58 |
+
Lfm3VlRMSNorm is equivalent to T5LayerNorm
|
| 59 |
+
"""
|
| 60 |
+
super().__init__()
|
| 61 |
+
self.weight = nn.Parameter(torch.ones(hidden_size))
|
| 62 |
+
self.variance_epsilon = eps
|
| 63 |
+
|
| 64 |
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 65 |
+
input_dtype = hidden_states.dtype
|
| 66 |
+
hidden_states = hidden_states.to(torch.float32)
|
| 67 |
+
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
| 68 |
+
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
| 69 |
+
return self.weight * hidden_states.to(input_dtype)
|
| 70 |
+
|
| 71 |
+
def extra_repr(self):
|
| 72 |
+
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
class Lfm3VlMLP(nn.Module):
|
| 76 |
+
def __init__(self, config: Lfm3VlConfig):
|
| 77 |
+
super().__init__()
|
| 78 |
+
self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 79 |
+
self.w3 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 80 |
+
self.w2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
| 81 |
+
|
| 82 |
+
def forward(self, x):
|
| 83 |
+
return self.w2(F.silu(self.w1(x)) * self.w3(x))
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
class Lfm3VlVisionRotaryEmbedding(nn.Module):
|
| 87 |
+
"""
|
| 88 |
+
Simple axial 2D rope with same freqs used for H and W grids. The freqs are
|
| 89 |
+
pre-computed using `head-dim//4` which is later used to concat H and W positions.
|
| 90 |
+
The final angles rotate over the whole head dim, no partial rotation involved.
|
| 91 |
+
"""
|
| 92 |
+
|
| 93 |
+
def __init__(self, config: Lfm3VlVisionConfig):
|
| 94 |
+
super().__init__()
|
| 95 |
+
self.config = config
|
| 96 |
+
|
| 97 |
+
self.rope_type = self.config.rope_parameters["rope_type"]
|
| 98 |
+
rope_init_fn: Callable = self.compute_axial_rope_parameters
|
| 99 |
+
if self.rope_type != "axial":
|
| 100 |
+
raise ValueError(f"{self.__class__.__name__} supports only axial rope, but requested {self.rope_type}")
|
| 101 |
+
inv_freq, self.attention_scaling = rope_init_fn(self.config)
|
| 102 |
+
|
| 103 |
+
self.inv_freq = nn.Buffer(inv_freq, persistent=False)
|
| 104 |
+
self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False)
|
| 105 |
+
|
| 106 |
+
@staticmethod
|
| 107 |
+
def compute_axial_rope_parameters(config: Lfm3VlVisionConfig, **kwargs) -> tuple[torch.Tensor, float]:
|
| 108 |
+
"""Inverse frequencies of the axial RoPE over half the head dimension per axis, with an attention scaling of 1."""
|
| 109 |
+
base = config.rope_parameters["rope_theta"]
|
| 110 |
+
spatial_dim = config.hidden_size // config.num_attention_heads // 2
|
| 111 |
+
inv_freq = 1.0 / (base ** (torch.arange(0, spatial_dim, 2, dtype=torch.float) / spatial_dim))
|
| 112 |
+
return inv_freq, 1.0
|
| 113 |
+
|
| 114 |
+
@torch.no_grad()
|
| 115 |
+
def forward(self, x, position_ids):
|
| 116 |
+
# position_ids: (2, N) — row 0 = h coords, row 1 = w coords
|
| 117 |
+
position_ids_expanded = position_ids[..., None].float()
|
| 118 |
+
device_type = x.device.type if isinstance(x.device.type, str) else "cpu"
|
| 119 |
+
with maybe_autocast(device_type=device_type, enabled=False):
|
| 120 |
+
freqs = position_ids_expanded * self.inv_freq.float()
|
| 121 |
+
cos = freqs.cos() * self.attention_scaling
|
| 122 |
+
sin = freqs.sin() * self.attention_scaling
|
| 123 |
+
|
| 124 |
+
cos = self.recomposition_frequencies(cos)
|
| 125 |
+
sin = self.recomposition_frequencies(sin)
|
| 126 |
+
return cos, sin
|
| 127 |
+
|
| 128 |
+
def recomposition_frequencies(self, freq):
|
| 129 |
+
"""
|
| 130 |
+
Recompose the frequencies into the final spatial layout used per each grid.
|
| 131 |
+
"""
|
| 132 |
+
freq_h, freq_w = freq[:, 0], freq[:, 1]
|
| 133 |
+
freq_hw = torch.cat([freq_h, freq_w], dim=-1)
|
| 134 |
+
return torch.cat([freq_hw, freq_hw], dim=-1)
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
def rotate_half(x):
|
| 138 |
+
"""Rotates half the hidden dims of the input."""
|
| 139 |
+
x1 = x[..., : x.shape[-1] // 2]
|
| 140 |
+
x2 = x[..., x.shape[-1] // 2 :]
|
| 141 |
+
return torch.cat((-x2, x1), dim=-1)
|
| 142 |
+
|
| 143 |
+
|
| 144 |
+
def apply_rotary_pos_emb_vision(
|
| 145 |
+
q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
|
| 146 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 147 |
+
orig_q_dtype = q.dtype
|
| 148 |
+
orig_k_dtype = k.dtype
|
| 149 |
+
q, k = q.float(), k.float()
|
| 150 |
+
cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
|
| 151 |
+
q_embed = (q * cos) + (rotate_half(q) * sin)
|
| 152 |
+
k_embed = (k * cos) + (rotate_half(k) * sin)
|
| 153 |
+
q_embed = q_embed.to(orig_q_dtype)
|
| 154 |
+
k_embed = k_embed.to(orig_k_dtype)
|
| 155 |
+
return q_embed, k_embed
|
| 156 |
+
|
| 157 |
+
|
| 158 |
+
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
| 159 |
+
"""
|
| 160 |
+
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
| 161 |
+
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
| 162 |
+
"""
|
| 163 |
+
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
| 164 |
+
if n_rep == 1:
|
| 165 |
+
return hidden_states
|
| 166 |
+
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
| 167 |
+
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
def eager_attention_forward(
|
| 171 |
+
module: nn.Module,
|
| 172 |
+
query: torch.Tensor,
|
| 173 |
+
key: torch.Tensor,
|
| 174 |
+
value: torch.Tensor,
|
| 175 |
+
attention_mask: torch.Tensor | None,
|
| 176 |
+
scaling: float,
|
| 177 |
+
dropout: float = 0.0,
|
| 178 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 179 |
+
):
|
| 180 |
+
key_states = repeat_kv(key, module.num_key_value_groups)
|
| 181 |
+
value_states = repeat_kv(value, module.num_key_value_groups)
|
| 182 |
+
|
| 183 |
+
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
|
| 184 |
+
if attention_mask is not None:
|
| 185 |
+
attn_weights = attn_weights + attention_mask
|
| 186 |
+
|
| 187 |
+
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
|
| 188 |
+
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
|
| 189 |
+
attn_output = torch.matmul(attn_weights, value_states)
|
| 190 |
+
attn_output = attn_output.transpose(1, 2).contiguous()
|
| 191 |
+
|
| 192 |
+
return attn_output, attn_weights
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
class Lfm3VlVisionAttention(nn.Module):
|
| 196 |
+
"""Bidirectional grouped-query attention over the patches of one frame, with RMSNorm on the query and key heads."""
|
| 197 |
+
|
| 198 |
+
def __init__(self, config: Lfm3VlVisionConfig):
|
| 199 |
+
super().__init__()
|
| 200 |
+
self.config = config
|
| 201 |
+
self.num_heads = config.num_attention_heads
|
| 202 |
+
self.num_key_value_heads = config.num_key_value_heads
|
| 203 |
+
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
|
| 204 |
+
self.head_dim = config.hidden_size // config.num_attention_heads
|
| 205 |
+
self.scaling = self.head_dim**-0.5
|
| 206 |
+
self.attention_dropout = 0.0
|
| 207 |
+
self.is_causal = False
|
| 208 |
+
self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
|
| 209 |
+
self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
| 210 |
+
self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
| 211 |
+
self.out_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
|
| 212 |
+
self.q_layernorm = Lfm3VlRMSNorm(self.head_dim, eps=config.norm_eps)
|
| 213 |
+
self.k_layernorm = Lfm3VlRMSNorm(self.head_dim, eps=config.norm_eps)
|
| 214 |
+
|
| 215 |
+
def forward(
|
| 216 |
+
self,
|
| 217 |
+
hidden_states: torch.Tensor,
|
| 218 |
+
cu_seqlens: torch.Tensor,
|
| 219 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
| 220 |
+
max_seqlen: int | None = None,
|
| 221 |
+
frame_shapes: list[list[int]] | None = None,
|
| 222 |
+
**kwargs,
|
| 223 |
+
) -> torch.Tensor:
|
| 224 |
+
seq_length = hidden_states.shape[0]
|
| 225 |
+
query_states = self.q_layernorm(self.q_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim))
|
| 226 |
+
key_states = self.k_layernorm(
|
| 227 |
+
self.k_proj(hidden_states).view(seq_length, self.num_key_value_heads, self.head_dim)
|
| 228 |
+
)
|
| 229 |
+
value_states = self.v_proj(hidden_states).view(seq_length, self.num_key_value_heads, self.head_dim)
|
| 230 |
+
|
| 231 |
+
cos, sin = position_embeddings
|
| 232 |
+
query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
|
| 233 |
+
|
| 234 |
+
query_states = query_states.transpose(0, 1).unsqueeze(0)
|
| 235 |
+
key_states = key_states.transpose(0, 1).unsqueeze(0)
|
| 236 |
+
value_states = value_states.transpose(0, 1).unsqueeze(0)
|
| 237 |
+
|
| 238 |
+
attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
|
| 239 |
+
self.config._attn_implementation, eager_attention_forward
|
| 240 |
+
)
|
| 241 |
+
|
| 242 |
+
if is_flash_attention_requested(self.config):
|
| 243 |
+
max_seqlen = get_max_seqlen(cu_seqlens, self.config, kwargs={"max_seqlen": max_seqlen})
|
| 244 |
+
attn_output, attn_weights = attention_interface(
|
| 245 |
+
self,
|
| 246 |
+
query_states,
|
| 247 |
+
key_states,
|
| 248 |
+
value_states,
|
| 249 |
+
attention_mask=None,
|
| 250 |
+
scaling=self.scaling,
|
| 251 |
+
dropout=0.0,
|
| 252 |
+
cu_seq_lens_q=cu_seqlens,
|
| 253 |
+
cu_seq_lens_k=cu_seqlens,
|
| 254 |
+
max_length_q=max_seqlen,
|
| 255 |
+
max_length_k=max_seqlen,
|
| 256 |
+
is_causal=False,
|
| 257 |
+
**kwargs,
|
| 258 |
+
)
|
| 259 |
+
else:
|
| 260 |
+
# Other implementations: frames of one shape (every video, equal-size images) attend as one batch, mixed
|
| 261 |
+
# shapes frame by frame. `frame_shapes` gives the lengths without a device sync.
|
| 262 |
+
if frame_shapes is not None:
|
| 263 |
+
lengths = [height * width for height, width in frame_shapes]
|
| 264 |
+
else:
|
| 265 |
+
lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
|
| 266 |
+
if len(set(lengths)) == 1:
|
| 267 |
+
query_states, key_states, value_states = (
|
| 268 |
+
tensor.reshape(-1, len(lengths), lengths[0], self.head_dim).transpose(0, 1)
|
| 269 |
+
for tensor in (query_states, key_states, value_states)
|
| 270 |
+
)
|
| 271 |
+
attn_output, attn_weights = attention_interface(
|
| 272 |
+
self,
|
| 273 |
+
query_states,
|
| 274 |
+
key_states,
|
| 275 |
+
value_states,
|
| 276 |
+
attention_mask=None,
|
| 277 |
+
scaling=self.scaling,
|
| 278 |
+
dropout=0.0,
|
| 279 |
+
is_causal=False,
|
| 280 |
+
**kwargs,
|
| 281 |
+
)
|
| 282 |
+
else:
|
| 283 |
+
splits = [torch.split(tensor, lengths, dim=2) for tensor in (query_states, key_states, value_states)]
|
| 284 |
+
attn_outputs, attn_weights = [], []
|
| 285 |
+
for q, k, v in zip(*splits):
|
| 286 |
+
output, weights = attention_interface(
|
| 287 |
+
self,
|
| 288 |
+
q,
|
| 289 |
+
k,
|
| 290 |
+
v,
|
| 291 |
+
attention_mask=None,
|
| 292 |
+
scaling=self.scaling,
|
| 293 |
+
dropout=0.0,
|
| 294 |
+
is_causal=False,
|
| 295 |
+
**kwargs,
|
| 296 |
+
)
|
| 297 |
+
attn_outputs.append(output)
|
| 298 |
+
attn_weights.append(weights)
|
| 299 |
+
attn_output = torch.cat(attn_outputs, dim=1)
|
| 300 |
+
# frames of different shapes have differently sized maps: one entry per frame, or None under sdpa
|
| 301 |
+
attn_weights = None if attn_weights[0] is None else tuple(attn_weights)
|
| 302 |
+
|
| 303 |
+
attn_output = attn_output.reshape(seq_length, -1).contiguous()
|
| 304 |
+
attn_output = self.out_proj(attn_output)
|
| 305 |
+
return attn_output, attn_weights
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
class Lfm3VlVisionConv2d(nn.Module):
|
| 309 |
+
"""Gated depthwise 2D convolution over the patch grid of each frame."""
|
| 310 |
+
|
| 311 |
+
def __init__(self, config: Lfm3VlVisionConfig):
|
| 312 |
+
super().__init__()
|
| 313 |
+
self.spatial_merge_size = config.spatial_merge_size
|
| 314 |
+
self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False)
|
| 315 |
+
self.spatial_conv = nn.Conv2d(
|
| 316 |
+
config.hidden_size,
|
| 317 |
+
config.hidden_size,
|
| 318 |
+
kernel_size=config.conv_kernel_size,
|
| 319 |
+
padding=config.conv_kernel_size // 2,
|
| 320 |
+
groups=config.hidden_size,
|
| 321 |
+
bias=False,
|
| 322 |
+
)
|
| 323 |
+
self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
|
| 324 |
+
|
| 325 |
+
def forward(self, hidden_states: torch.Tensor, frame_shapes: list[list[int]]):
|
| 326 |
+
B, C, x = self.in_proj(hidden_states).chunk(3, dim=-1)
|
| 327 |
+
gated = B * x
|
| 328 |
+
merge = self.spatial_merge_size
|
| 329 |
+
num_frames = len(frame_shapes)
|
| 330 |
+
if len(set(map(tuple, frame_shapes))) == 1:
|
| 331 |
+
# frames of one shape (every video, equal-size images) convolve as one batch: merge-block order -> raster
|
| 332 |
+
# grid, convolve, back
|
| 333 |
+
height, width = frame_shapes[0]
|
| 334 |
+
frames = gated.view(num_frames, height // merge, width // merge, merge, merge, -1)
|
| 335 |
+
frames = frames.permute(0, 5, 1, 3, 2, 4).reshape(num_frames, -1, height, width)
|
| 336 |
+
frames = self.spatial_conv(frames)
|
| 337 |
+
frames = frames.view(num_frames, -1, height // merge, merge, width // merge, merge).permute(
|
| 338 |
+
0, 2, 4, 3, 5, 1
|
| 339 |
+
)
|
| 340 |
+
conv_output = frames.reshape(gated.shape[0], -1)
|
| 341 |
+
else:
|
| 342 |
+
conv_outputs = []
|
| 343 |
+
lengths = [height * width for height, width in frame_shapes]
|
| 344 |
+
for frame, (height, width) in zip(torch.split(gated, lengths, dim=0), frame_shapes):
|
| 345 |
+
frame = frame.view(height // merge, width // merge, merge, merge, -1)
|
| 346 |
+
frame = frame.permute(4, 0, 2, 1, 3).reshape(1, -1, height, width)
|
| 347 |
+
frame = self.spatial_conv(frame)
|
| 348 |
+
frame = frame.view(-1, height // merge, merge, width // merge, merge).permute(1, 3, 2, 4, 0)
|
| 349 |
+
conv_outputs.append(frame.reshape(height * width, -1))
|
| 350 |
+
conv_output = torch.cat(conv_outputs, dim=0)
|
| 351 |
+
return self.out_proj(C * conv_output)
|
| 352 |
+
|
| 353 |
+
|
| 354 |
+
class Lfm3VlVisionLayer(GradientCheckpointingLayer):
|
| 355 |
+
def __init__(self, config: Lfm3VlVisionConfig, layer_idx: int):
|
| 356 |
+
super().__init__()
|
| 357 |
+
self.is_attention_layer = config.layer_types[layer_idx] == "full_attention"
|
| 358 |
+
if self.is_attention_layer:
|
| 359 |
+
self.self_attn = Lfm3VlVisionAttention(config)
|
| 360 |
+
else:
|
| 361 |
+
self.conv = Lfm3VlVisionConv2d(config)
|
| 362 |
+
self.feed_forward = Lfm3VlMLP(config)
|
| 363 |
+
self.operator_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 364 |
+
self.ffn_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 365 |
+
|
| 366 |
+
def forward(
|
| 367 |
+
self,
|
| 368 |
+
hidden_states: torch.Tensor,
|
| 369 |
+
cu_seqlens: torch.Tensor,
|
| 370 |
+
frame_shapes: list[list[int]],
|
| 371 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
| 372 |
+
**kwargs,
|
| 373 |
+
) -> torch.Tensor:
|
| 374 |
+
if self.is_attention_layer:
|
| 375 |
+
operator_output, _ = self.self_attn(
|
| 376 |
+
self.operator_norm(hidden_states),
|
| 377 |
+
cu_seqlens=cu_seqlens,
|
| 378 |
+
position_embeddings=position_embeddings,
|
| 379 |
+
frame_shapes=frame_shapes,
|
| 380 |
+
**kwargs,
|
| 381 |
+
)
|
| 382 |
+
else:
|
| 383 |
+
operator_output = self.conv(self.operator_norm(hidden_states), frame_shapes)
|
| 384 |
+
hidden_states = hidden_states + operator_output
|
| 385 |
+
hidden_states = hidden_states + self.feed_forward(self.ffn_norm(hidden_states))
|
| 386 |
+
return hidden_states
|
| 387 |
+
|
| 388 |
+
|
| 389 |
+
def create_masks_by_layer_type(
|
| 390 |
+
config: PreTrainedConfig,
|
| 391 |
+
inputs_embeds: torch.Tensor,
|
| 392 |
+
attention_mask: torch.Tensor | None,
|
| 393 |
+
past_key_values: Cache | None,
|
| 394 |
+
position_ids: torch.Tensor | None = None,
|
| 395 |
+
) -> dict:
|
| 396 |
+
"""The cache sizes each mask from the layer it is asked about, and the generic mask utilities ask about the first
|
| 397 |
+
non-sliding layer, which for a leading `conv` layer redirects to the first attention layer of any type: name the
|
| 398 |
+
first layer of each mask's own type instead."""
|
| 399 |
+
config = config.get_text_config()
|
| 400 |
+
mask_kwargs = {
|
| 401 |
+
"config": config,
|
| 402 |
+
"inputs_embeds": inputs_embeds,
|
| 403 |
+
"attention_mask": attention_mask,
|
| 404 |
+
"past_key_values": past_key_values,
|
| 405 |
+
"position_ids": position_ids,
|
| 406 |
+
}
|
| 407 |
+
layer_types, causal_mask_mapping = config.layer_types, {}
|
| 408 |
+
if "full_attention" in layer_types:
|
| 409 |
+
layer_idx = layer_types.index("full_attention")
|
| 410 |
+
causal_mask_mapping["full_attention"] = create_causal_mask(**mask_kwargs, layer_idx=layer_idx)
|
| 411 |
+
if "sliding_attention" in layer_types:
|
| 412 |
+
layer_idx = layer_types.index("sliding_attention")
|
| 413 |
+
causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(
|
| 414 |
+
**mask_kwargs, layer_idx=layer_idx
|
| 415 |
+
)
|
| 416 |
+
if "conv" in layer_types:
|
| 417 |
+
causal_mask_mapping["conv"] = create_recurrent_attention_mask(**mask_kwargs)
|
| 418 |
+
return causal_mask_mapping
|
| 419 |
+
|
| 420 |
+
|
| 421 |
+
@auto_docstring
|
| 422 |
+
class Lfm3VlPreTrainedModel(PreTrainedModel):
|
| 423 |
+
config: Lfm3VlConfig
|
| 424 |
+
base_model_prefix = "model"
|
| 425 |
+
input_modalities = ("image", "video", "text")
|
| 426 |
+
supports_gradient_checkpointing = True
|
| 427 |
+
_skip_keys_device_placement = ["past_key_values"]
|
| 428 |
+
|
| 429 |
+
_supports_flash_attn = True
|
| 430 |
+
_supports_sdpa = True
|
| 431 |
+
# the cached decode step traces without graph breaks, like the text model
|
| 432 |
+
_can_compile_fullgraph = True
|
| 433 |
+
_supports_flex_attn = True
|
| 434 |
+
_supports_attention_backend = True
|
| 435 |
+
_no_split_modules = ["Lfm3VlVisionLayer"]
|
| 436 |
+
|
| 437 |
+
@staticmethod
|
| 438 |
+
def create_masks_for_generate(
|
| 439 |
+
config: PreTrainedConfig,
|
| 440 |
+
inputs_embeds: torch.Tensor,
|
| 441 |
+
attention_mask: torch.Tensor | None,
|
| 442 |
+
past_key_values: Cache | None,
|
| 443 |
+
position_ids: torch.Tensor | None = None,
|
| 444 |
+
**kwargs,
|
| 445 |
+
) -> dict:
|
| 446 |
+
# the static-cache masks of the text model are sized by layer type, not by the generic first attention layer
|
| 447 |
+
return create_masks_by_layer_type(config, inputs_embeds, attention_mask, past_key_values, position_ids)
|
| 448 |
+
|
| 449 |
+
|
| 450 |
+
def get_vision_frame_shapes(grid_thw: torch.Tensor) -> list[list[int]]:
|
| 451 |
+
return grid_thw[:, 1:].repeat_interleave(grid_thw[:, 0], dim=0).tolist()
|
| 452 |
+
|
| 453 |
+
|
| 454 |
+
@auto_docstring(
|
| 455 |
+
custom_intro="""
|
| 456 |
+
The LFM3-VL vision encoder: a stack of gated convolution and bidirectional attention layers over the patches of
|
| 457 |
+
each image or video frame, embedded by a linear patchifier.
|
| 458 |
+
"""
|
| 459 |
+
)
|
| 460 |
+
class Lfm3VlVisionModel(Lfm3VlPreTrainedModel):
|
| 461 |
+
config: Lfm3VlVisionConfig
|
| 462 |
+
input_modalities = ("image", "video")
|
| 463 |
+
_input_embed_layer = "patchifier"
|
| 464 |
+
_can_record_outputs = {
|
| 465 |
+
"hidden_states": Lfm3VlVisionLayer,
|
| 466 |
+
"attentions": OutputRecorder(Lfm3VlVisionAttention, index=1),
|
| 467 |
+
}
|
| 468 |
+
|
| 469 |
+
def __init__(self, config: Lfm3VlVisionConfig):
|
| 470 |
+
super().__init__(config)
|
| 471 |
+
self.spatial_merge_size = config.spatial_merge_size
|
| 472 |
+
patch_dim = config.in_channels * config.temporal_patch_size * config.patch_size**2
|
| 473 |
+
self.patchifier = nn.Linear(patch_dim, config.hidden_size, bias=False)
|
| 474 |
+
self.embedding_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 475 |
+
self.rotary_pos_emb = Lfm3VlVisionRotaryEmbedding(config)
|
| 476 |
+
self.layers = nn.ModuleList(
|
| 477 |
+
[Lfm3VlVisionLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
| 478 |
+
)
|
| 479 |
+
self.gradient_checkpointing = False
|
| 480 |
+
self.post_init()
|
| 481 |
+
|
| 482 |
+
@merge_with_config_defaults
|
| 483 |
+
@capture_outputs
|
| 484 |
+
@auto_docstring
|
| 485 |
+
def forward(
|
| 486 |
+
self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
|
| 487 |
+
) -> BaseModelOutputWithPooling:
|
| 488 |
+
r"""
|
| 489 |
+
pixel_values (`torch.Tensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
|
| 490 |
+
The flattened patches of every image or video, grouped by merge block.
|
| 491 |
+
grid_thw (`torch.LongTensor` of shape `(num_images_or_videos, 3)`):
|
| 492 |
+
The temporal, height and width dimensions of the patch grid of each image or video.
|
| 493 |
+
"""
|
| 494 |
+
position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size, kwargs=kwargs)
|
| 495 |
+
cu_seqlens, max_seqlen = get_vision_attention_seqlens(grid_thw, self.config, kwargs=kwargs)
|
| 496 |
+
frame_shapes = get_vision_frame_shapes(grid_thw)
|
| 497 |
+
|
| 498 |
+
hidden_states = self.patchifier(pixel_values.to(self.patchifier.weight.dtype))
|
| 499 |
+
hidden_states = self.embedding_norm(hidden_states)
|
| 500 |
+
position_embeddings = self.rotary_pos_emb(hidden_states, position_ids)
|
| 501 |
+
|
| 502 |
+
for layer in self.layers:
|
| 503 |
+
hidden_states = layer(
|
| 504 |
+
hidden_states,
|
| 505 |
+
cu_seqlens=cu_seqlens,
|
| 506 |
+
frame_shapes=frame_shapes,
|
| 507 |
+
max_seqlen=max_seqlen,
|
| 508 |
+
position_embeddings=position_embeddings,
|
| 509 |
+
**kwargs,
|
| 510 |
+
)
|
| 511 |
+
|
| 512 |
+
return BaseModelOutputWithPooling(last_hidden_state=hidden_states)
|
| 513 |
+
|
| 514 |
+
|
| 515 |
+
class Lfm3VlMultiModalProjector(nn.Module):
|
| 516 |
+
"""Merges each block of `downsample_factor**2` vision tokens into one language token."""
|
| 517 |
+
|
| 518 |
+
def __init__(self, config: Lfm3VlConfig):
|
| 519 |
+
super().__init__()
|
| 520 |
+
self.downsample_factor = config.downsample_factor
|
| 521 |
+
self.layer_norm = nn.LayerNorm(config.vision_config.hidden_size, eps=config.projector_layer_norm_eps)
|
| 522 |
+
self.linear_1 = nn.Linear(
|
| 523 |
+
config.vision_config.hidden_size * config.downsample_factor**2,
|
| 524 |
+
config.projector_hidden_size,
|
| 525 |
+
bias=config.projector_bias,
|
| 526 |
+
)
|
| 527 |
+
self.act = ACT2FN[config.projector_hidden_act]
|
| 528 |
+
self.linear_2 = nn.Linear(
|
| 529 |
+
config.projector_hidden_size, config.text_config.hidden_size, bias=config.projector_bias
|
| 530 |
+
)
|
| 531 |
+
|
| 532 |
+
def forward(self, image_features: torch.Tensor):
|
| 533 |
+
hidden_states = self.layer_norm(image_features).reshape(-1, self.linear_1.in_features)
|
| 534 |
+
hidden_states = self.linear_1(hidden_states)
|
| 535 |
+
hidden_states = self.act(hidden_states)
|
| 536 |
+
hidden_states = self.linear_2(hidden_states)
|
| 537 |
+
return hidden_states
|
| 538 |
+
|
| 539 |
+
|
| 540 |
+
@auto_docstring(
|
| 541 |
+
custom_intro="""
|
| 542 |
+
Base class for Lfm3Vl outputs, with hidden states and attentions.
|
| 543 |
+
"""
|
| 544 |
+
)
|
| 545 |
+
@dataclass
|
| 546 |
+
class Lfm3VlModelOutputWithPast(BaseModelOutputWithPast):
|
| 547 |
+
r"""
|
| 548 |
+
past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
|
| 549 |
+
It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
|
| 550 |
+
|
| 551 |
+
Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
|
| 552 |
+
`past_key_values` input) to speed up sequential decoding.
|
| 553 |
+
image_hidden_states (`torch.FloatTensor`, *optional*):
|
| 554 |
+
A `torch.FloatTensor` of size `(num_image_tokens, hidden_size)`.
|
| 555 |
+
image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
|
| 556 |
+
video_hidden_states (`torch.FloatTensor`, *optional*):
|
| 557 |
+
A `torch.FloatTensor` of size `(num_video_tokens, hidden_size)`.
|
| 558 |
+
video_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
|
| 559 |
+
"""
|
| 560 |
+
|
| 561 |
+
image_hidden_states: torch.FloatTensor | None = None
|
| 562 |
+
video_hidden_states: torch.FloatTensor | None = None
|
| 563 |
+
|
| 564 |
+
|
| 565 |
+
@auto_docstring(
|
| 566 |
+
custom_intro="""
|
| 567 |
+
Base class for Lfm3Vl causal language model (or autoregressive) outputs.
|
| 568 |
+
"""
|
| 569 |
+
)
|
| 570 |
+
@dataclass
|
| 571 |
+
class Lfm3VlCausalLMOutputWithPast(CausalLMOutputWithPast):
|
| 572 |
+
r"""
|
| 573 |
+
image_hidden_states (`torch.FloatTensor`, *optional*):
|
| 574 |
+
A `torch.FloatTensor` of size `(num_image_tokens, hidden_size)`.
|
| 575 |
+
image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
|
| 576 |
+
video_hidden_states (`torch.FloatTensor`, *optional*):
|
| 577 |
+
A `torch.FloatTensor` of size `(num_video_tokens, hidden_size)`.
|
| 578 |
+
video_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
|
| 579 |
+
"""
|
| 580 |
+
|
| 581 |
+
image_hidden_states: torch.FloatTensor | None = None
|
| 582 |
+
video_hidden_states: torch.FloatTensor | None = None
|
| 583 |
+
|
| 584 |
+
|
| 585 |
+
@auto_docstring(
|
| 586 |
+
custom_intro="""
|
| 587 |
+
The Lfm3Vl model which consists of a vision backbone and a language model, without a language modeling head.
|
| 588 |
+
"""
|
| 589 |
+
)
|
| 590 |
+
class Lfm3VlModel(Lfm3VlPreTrainedModel):
|
| 591 |
+
def __init__(self, config: Lfm3VlConfig):
|
| 592 |
+
super().__init__(config)
|
| 593 |
+
self.vision_tower = Lfm3VlVisionModel(config.vision_config)
|
| 594 |
+
self.multi_modal_projector = Lfm3VlMultiModalProjector(config)
|
| 595 |
+
self.language_model = Lfm3Model(config.text_config)
|
| 596 |
+
self.post_init()
|
| 597 |
+
|
| 598 |
+
@accepts_precomputed_kwargs(modality="image")
|
| 599 |
+
@can_return_tuple
|
| 600 |
+
@auto_docstring(
|
| 601 |
+
custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection."
|
| 602 |
+
)
|
| 603 |
+
def get_image_features(
|
| 604 |
+
self,
|
| 605 |
+
pixel_values: torch.FloatTensor,
|
| 606 |
+
image_grid_thw: torch.LongTensor,
|
| 607 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 608 |
+
) -> tuple | BaseModelOutputWithPooling:
|
| 609 |
+
r"""
|
| 610 |
+
pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
|
| 611 |
+
The flattened patches of every image, grouped by merge block.
|
| 612 |
+
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`):
|
| 613 |
+
The temporal, height and width dimensions of the patch grid of each image.
|
| 614 |
+
"""
|
| 615 |
+
image_outputs = self.vision_tower(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs)
|
| 616 |
+
image_features = self.multi_modal_projector(image_outputs.last_hidden_state)
|
| 617 |
+
split_sizes = (image_grid_thw.prod(-1) // self.vision_tower.spatial_merge_size**2).tolist()
|
| 618 |
+
image_outputs.pooler_output = torch.split(image_features, split_sizes)
|
| 619 |
+
return image_outputs
|
| 620 |
+
|
| 621 |
+
def get_placeholder_mask(
|
| 622 |
+
self,
|
| 623 |
+
input_ids: torch.LongTensor,
|
| 624 |
+
inputs_embeds: torch.FloatTensor,
|
| 625 |
+
image_features: torch.FloatTensor | None = None,
|
| 626 |
+
video_features: torch.FloatTensor | None = None,
|
| 627 |
+
):
|
| 628 |
+
"""Placeholder masks from `input_ids` or `inputs_embeds`, checked against the feature counts."""
|
| 629 |
+
if input_ids is None:
|
| 630 |
+
special_image_mask = inputs_embeds == self.get_input_embeddings()(
|
| 631 |
+
torch.full((), self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
|
| 632 |
+
)
|
| 633 |
+
special_image_mask = special_image_mask.all(-1)
|
| 634 |
+
special_video_mask = inputs_embeds == self.get_input_embeddings()(
|
| 635 |
+
torch.full((), self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
|
| 636 |
+
)
|
| 637 |
+
special_video_mask = special_video_mask.all(-1)
|
| 638 |
+
else:
|
| 639 |
+
special_image_mask = input_ids == self.config.image_token_id
|
| 640 |
+
special_video_mask = input_ids == self.config.video_token_id
|
| 641 |
+
|
| 642 |
+
n_image_tokens = special_image_mask.sum()
|
| 643 |
+
special_image_mask = special_image_mask.unsqueeze(-1).to(inputs_embeds.device)
|
| 644 |
+
if image_features is not None:
|
| 645 |
+
torch_compilable_check(
|
| 646 |
+
n_image_tokens * inputs_embeds.shape[-1] == image_features.numel(),
|
| 647 |
+
f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}",
|
| 648 |
+
)
|
| 649 |
+
|
| 650 |
+
n_video_tokens = special_video_mask.sum()
|
| 651 |
+
special_video_mask = special_video_mask.unsqueeze(-1).to(inputs_embeds.device)
|
| 652 |
+
if video_features is not None:
|
| 653 |
+
torch_compilable_check(
|
| 654 |
+
n_video_tokens * inputs_embeds.shape[-1] == video_features.numel(),
|
| 655 |
+
f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}",
|
| 656 |
+
)
|
| 657 |
+
return special_image_mask, special_video_mask
|
| 658 |
+
|
| 659 |
+
@can_return_tuple
|
| 660 |
+
@auto_docstring
|
| 661 |
+
def forward(
|
| 662 |
+
self,
|
| 663 |
+
input_ids: torch.LongTensor | None = None,
|
| 664 |
+
attention_mask: torch.Tensor | None = None,
|
| 665 |
+
position_ids: torch.LongTensor | None = None,
|
| 666 |
+
past_key_values: Cache | None = None,
|
| 667 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 668 |
+
pixel_values: torch.FloatTensor | None = None,
|
| 669 |
+
image_grid_thw: torch.LongTensor | None = None,
|
| 670 |
+
pixel_values_videos: torch.FloatTensor | None = None,
|
| 671 |
+
video_grid_thw: torch.LongTensor | None = None,
|
| 672 |
+
use_cache: bool | None = None,
|
| 673 |
+
mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None,
|
| 674 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 675 |
+
) -> tuple | Lfm3VlModelOutputWithPast:
|
| 676 |
+
r"""
|
| 677 |
+
pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
|
| 678 |
+
The flattened patches of every image, grouped by merge block.
|
| 679 |
+
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
|
| 680 |
+
The temporal, height and width dimensions of the patch grid of each image.
|
| 681 |
+
pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
|
| 682 |
+
The flattened patches of every video, grouped by merge block.
|
| 683 |
+
video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
|
| 684 |
+
The temporal, height and width dimensions of the patch grid of each video.
|
| 685 |
+
"""
|
| 686 |
+
if (input_ids is None) ^ (inputs_embeds is not None):
|
| 687 |
+
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
| 688 |
+
|
| 689 |
+
if (pixel_values is not None or pixel_values_videos is not None) and mm_encoder_outputs is not None:
|
| 690 |
+
raise ValueError(
|
| 691 |
+
"You cannot specify both pixel_values/pixel_values_videos and mm_encoder_outputs at the same time"
|
| 692 |
+
)
|
| 693 |
+
|
| 694 |
+
if inputs_embeds is None:
|
| 695 |
+
inputs_embeds = self.get_input_embeddings()(input_ids)
|
| 696 |
+
|
| 697 |
+
mm_encoder_outputs = mm_encoder_outputs if mm_encoder_outputs is not None else {}
|
| 698 |
+
if (pixel_values is None) != (image_grid_thw is None):
|
| 699 |
+
raise ValueError("pixel_values and image_grid_thw must be supplied together")
|
| 700 |
+
if (pixel_values_videos is None) != (video_grid_thw is None):
|
| 701 |
+
raise ValueError("pixel_values_videos and video_grid_thw must be supplied together")
|
| 702 |
+
if mm_encoder_outputs.get("image") is None and pixel_values is not None:
|
| 703 |
+
mm_encoder_outputs["image"] = self.get_image_features(
|
| 704 |
+
pixel_values, image_grid_thw, return_dict=True, **kwargs
|
| 705 |
+
)
|
| 706 |
+
if mm_encoder_outputs.get("video") is None and pixel_values_videos is not None:
|
| 707 |
+
mm_encoder_outputs["video"] = self.get_video_features(
|
| 708 |
+
pixel_values_videos, video_grid_thw, return_dict=True, **kwargs
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
image_features = None
|
| 712 |
+
if mm_encoder_outputs.get("image") is not None:
|
| 713 |
+
image_features = torch.cat(mm_encoder_outputs["image"].pooler_output, dim=0).to(
|
| 714 |
+
inputs_embeds.device, inputs_embeds.dtype
|
| 715 |
+
)
|
| 716 |
+
image_mask, _ = self.get_placeholder_mask(
|
| 717 |
+
input_ids, inputs_embeds=inputs_embeds, image_features=image_features
|
| 718 |
+
)
|
| 719 |
+
inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_features)
|
| 720 |
+
|
| 721 |
+
video_features = None
|
| 722 |
+
if mm_encoder_outputs.get("video") is not None:
|
| 723 |
+
video_features = torch.cat(mm_encoder_outputs["video"].pooler_output, dim=0).to(
|
| 724 |
+
inputs_embeds.device, inputs_embeds.dtype
|
| 725 |
+
)
|
| 726 |
+
_, video_mask = self.get_placeholder_mask(
|
| 727 |
+
input_ids, inputs_embeds=inputs_embeds, video_features=video_features
|
| 728 |
+
)
|
| 729 |
+
inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_features)
|
| 730 |
+
|
| 731 |
+
outputs = self.language_model(
|
| 732 |
+
attention_mask=attention_mask,
|
| 733 |
+
position_ids=position_ids,
|
| 734 |
+
past_key_values=past_key_values,
|
| 735 |
+
inputs_embeds=inputs_embeds,
|
| 736 |
+
use_cache=use_cache,
|
| 737 |
+
**kwargs,
|
| 738 |
+
)
|
| 739 |
+
|
| 740 |
+
return Lfm3VlModelOutputWithPast(
|
| 741 |
+
last_hidden_state=outputs.last_hidden_state,
|
| 742 |
+
past_key_values=outputs.past_key_values,
|
| 743 |
+
hidden_states=outputs.hidden_states,
|
| 744 |
+
attentions=outputs.attentions,
|
| 745 |
+
image_hidden_states=image_features,
|
| 746 |
+
video_hidden_states=video_features,
|
| 747 |
+
)
|
| 748 |
+
|
| 749 |
+
@accepts_precomputed_kwargs(modality="video")
|
| 750 |
+
@can_return_tuple
|
| 751 |
+
@auto_docstring(
|
| 752 |
+
custom_intro="Obtains video last hidden states from the vision tower and apply multimodal projection."
|
| 753 |
+
)
|
| 754 |
+
def get_video_features(
|
| 755 |
+
self,
|
| 756 |
+
pixel_values_videos: torch.FloatTensor,
|
| 757 |
+
video_grid_thw: torch.LongTensor,
|
| 758 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 759 |
+
) -> tuple | BaseModelOutputWithPooling:
|
| 760 |
+
r"""
|
| 761 |
+
pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
|
| 762 |
+
The flattened patches of every video, grouped by merge block.
|
| 763 |
+
video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`):
|
| 764 |
+
The temporal, height and width dimensions of the patch grid of each video.
|
| 765 |
+
"""
|
| 766 |
+
return self.get_image_features(pixel_values_videos, video_grid_thw, **kwargs)
|
| 767 |
+
|
| 768 |
+
|
| 769 |
+
def reject_packed_generation(config: PreTrainedConfig, input_ids, generation_config, model_kwargs, is_first_iteration):
|
| 770 |
+
"""`generate` decodes from the cache its prefill fills, and no mask it builds isolates packed documents
|
| 771 |
+
(`position_ids` restarting): the prompt is refused before the prefill, one forward or every chunk, writes to it."""
|
| 772 |
+
position_ids, attention_mask = model_kwargs.get("position_ids"), model_kwargs.get("attention_mask")
|
| 773 |
+
cache = model_kwargs.get("past_key_values")
|
| 774 |
+
if cache is None or position_ids is None or position_ids.ndim != 2:
|
| 775 |
+
return
|
| 776 |
+
if attention_mask is not None and attention_mask.ndim != 2:
|
| 777 |
+
return
|
| 778 |
+
has_attention = any(t != "conv" for t in config.get_text_config().layer_types)
|
| 779 |
+
filled = bool(cache.get_seq_length() > 0 if has_attention else cache.has_previous_state())
|
| 780 |
+
inputs_embeds = model_kwargs.get("inputs_embeds") if is_first_iteration else None
|
| 781 |
+
length = input_ids.shape[1] if inputs_embeds is None else inputs_embeds.shape[1]
|
| 782 |
+
if generation_config.prefill_chunk_size is not None:
|
| 783 |
+
# the chunks run from the start of the prompt
|
| 784 |
+
position_ids = position_ids[:, :length]
|
| 785 |
+
else:
|
| 786 |
+
if filled and has_attention:
|
| 787 |
+
# the slice `_prefill` feeds: a prompt that spans the cache is cut to its uncached part
|
| 788 |
+
mask_length = getattr(generation_config, "_mask_length", -1)
|
| 789 |
+
if attention_mask is not None:
|
| 790 |
+
mask_length = attention_mask.shape[1]
|
| 791 |
+
if inputs_embeds is not None or mask_length == length:
|
| 792 |
+
length -= int(cache.get_seq_length())
|
| 793 |
+
position_ids = position_ids[:, position_ids.shape[1] - length :]
|
| 794 |
+
if attention_mask is None:
|
| 795 |
+
attention_mask = torch.ones_like(position_ids)
|
| 796 |
+
if bool(packed_under_padding(position_ids, attention_mask, filled).any()):
|
| 797 |
+
raise ValueError(
|
| 798 |
+
"`generate` cannot decode packed documents (`position_ids` restarting, or starting at 0 on a cache that "
|
| 799 |
+
"holds tokens): attention would span them. Generate each document in its own row of a left-padded batch."
|
| 800 |
+
)
|
| 801 |
+
|
| 802 |
+
|
| 803 |
+
def packed_under_padding(position_ids: torch.Tensor, padding_mask: torch.Tensor, filled) -> torch.Tensor:
|
| 804 |
+
"""Padding repeats position 0, so packing is a restart between two kept tokens, or a kept position 0 on a filled
|
| 805 |
+
cache whose row holds a kept token; attention under a padding mask cannot isolate it."""
|
| 806 |
+
restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
|
| 807 |
+
valid = padding_mask[:, -restarts.shape[-1] :].bool()
|
| 808 |
+
packed = restarts[:, 1:] & valid[:, 1:] & valid[:, :-1]
|
| 809 |
+
if filled:
|
| 810 |
+
# a row whose cached tokens are all padding (chunked prefill of a left-padded batch) has no document to pack
|
| 811 |
+
# against; a mask covering only the chunk keeps the strict rule
|
| 812 |
+
cached = padding_mask[:, : padding_mask.shape[-1] - restarts.shape[-1]]
|
| 813 |
+
cached_valid = cached.bool().any(-1, keepdim=True) if cached.shape[-1] else torch.ones_like(valid[:, :1])
|
| 814 |
+
packed = packed | (valid[:, :1] & (position_ids[:, :1] == 0) & cached_valid)
|
| 815 |
+
return packed
|
| 816 |
+
|
| 817 |
+
|
| 818 |
+
@auto_docstring(
|
| 819 |
+
custom_intro="""
|
| 820 |
+
The LFM3_VL model which consists of a vision backbone and a language model.
|
| 821 |
+
"""
|
| 822 |
+
)
|
| 823 |
+
class Lfm3VlForConditionalGeneration(Lfm3VlPreTrainedModel, GenerationMixin):
|
| 824 |
+
_tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
|
| 825 |
+
|
| 826 |
+
def __init__(self, config: Lfm3VlConfig):
|
| 827 |
+
super().__init__(config)
|
| 828 |
+
self.model = Lfm3VlModel(config)
|
| 829 |
+
self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
|
| 830 |
+
self.post_init()
|
| 831 |
+
|
| 832 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 833 |
+
return self.lm_head
|
| 834 |
+
|
| 835 |
+
@auto_docstring
|
| 836 |
+
def get_image_features(
|
| 837 |
+
self,
|
| 838 |
+
pixel_values: torch.FloatTensor,
|
| 839 |
+
image_grid_thw: torch.LongTensor,
|
| 840 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 841 |
+
) -> tuple | BaseModelOutputWithPooling:
|
| 842 |
+
r"""
|
| 843 |
+
pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
|
| 844 |
+
The flattened patches of every image, grouped by merge block.
|
| 845 |
+
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`):
|
| 846 |
+
The temporal, height and width dimensions of the patch grid of each image.
|
| 847 |
+
"""
|
| 848 |
+
return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs)
|
| 849 |
+
|
| 850 |
+
@can_return_tuple
|
| 851 |
+
@auto_docstring
|
| 852 |
+
def forward(
|
| 853 |
+
self,
|
| 854 |
+
input_ids: torch.LongTensor | None = None,
|
| 855 |
+
attention_mask: torch.Tensor | None = None,
|
| 856 |
+
position_ids: torch.LongTensor | None = None,
|
| 857 |
+
past_key_values: Cache | None = None,
|
| 858 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 859 |
+
pixel_values: torch.FloatTensor | None = None,
|
| 860 |
+
image_grid_thw: torch.LongTensor | None = None,
|
| 861 |
+
pixel_values_videos: torch.FloatTensor | None = None,
|
| 862 |
+
video_grid_thw: torch.LongTensor | None = None,
|
| 863 |
+
labels: torch.LongTensor | None = None,
|
| 864 |
+
use_cache: bool | None = None,
|
| 865 |
+
logits_to_keep: int | torch.Tensor = 0,
|
| 866 |
+
mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None,
|
| 867 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 868 |
+
) -> tuple | Lfm3VlCausalLMOutputWithPast:
|
| 869 |
+
r"""
|
| 870 |
+
pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
|
| 871 |
+
The flattened patches of every image, grouped by merge block.
|
| 872 |
+
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
|
| 873 |
+
The temporal, height and width dimensions of the patch grid of each image.
|
| 874 |
+
pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
|
| 875 |
+
The flattened patches of every video, grouped by merge block.
|
| 876 |
+
video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
|
| 877 |
+
The temporal, height and width dimensions of the patch grid of each video.
|
| 878 |
+
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
| 879 |
+
Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
|
| 880 |
+
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
|
| 881 |
+
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
|
| 882 |
+
|
| 883 |
+
Example:
|
| 884 |
+
|
| 885 |
+
```python
|
| 886 |
+
>>> from transformers import AutoProcessor, AutoModelForImageTextToText
|
| 887 |
+
|
| 888 |
+
>>> model = AutoModelForImageTextToText.from_pretrained("LiquidAI/LFM3-VL-1.9B-Preview")
|
| 889 |
+
>>> processor = AutoProcessor.from_pretrained("LiquidAI/LFM3-VL-1.9B-Preview")
|
| 890 |
+
|
| 891 |
+
>>> conversation = [
|
| 892 |
+
... {
|
| 893 |
+
... "role": "user",
|
| 894 |
+
... "content": [
|
| 895 |
+
... {"type": "image", "url": "https://www.ilankelman.org/stopsigns/australia.jpg"},
|
| 896 |
+
... {"type": "text", "text": "What is in this image?"},
|
| 897 |
+
... ],
|
| 898 |
+
... },
|
| 899 |
+
... ]
|
| 900 |
+
|
| 901 |
+
>>> inputs = processor.apply_chat_template(
|
| 902 |
+
... conversation, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt"
|
| 903 |
+
... )
|
| 904 |
+
|
| 905 |
+
>>> outputs = model.generate(**inputs, max_new_tokens=45)
|
| 906 |
+
>>> processor.batch_decode(outputs, skip_special_tokens=True)[0]
|
| 907 |
+
```"""
|
| 908 |
+
outputs = self.model(
|
| 909 |
+
input_ids=input_ids,
|
| 910 |
+
attention_mask=attention_mask,
|
| 911 |
+
position_ids=position_ids,
|
| 912 |
+
past_key_values=past_key_values,
|
| 913 |
+
inputs_embeds=inputs_embeds,
|
| 914 |
+
pixel_values=pixel_values,
|
| 915 |
+
image_grid_thw=image_grid_thw,
|
| 916 |
+
pixel_values_videos=pixel_values_videos,
|
| 917 |
+
video_grid_thw=video_grid_thw,
|
| 918 |
+
use_cache=use_cache,
|
| 919 |
+
mm_encoder_outputs=mm_encoder_outputs,
|
| 920 |
+
**kwargs,
|
| 921 |
+
)
|
| 922 |
+
|
| 923 |
+
hidden_states = outputs[0]
|
| 924 |
+
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
|
| 925 |
+
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
| 926 |
+
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
| 927 |
+
|
| 928 |
+
loss = None
|
| 929 |
+
if labels is not None:
|
| 930 |
+
loss = self.loss_function(
|
| 931 |
+
logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
|
| 932 |
+
)
|
| 933 |
+
|
| 934 |
+
return Lfm3VlCausalLMOutputWithPast(
|
| 935 |
+
loss=loss,
|
| 936 |
+
logits=logits,
|
| 937 |
+
past_key_values=outputs.past_key_values,
|
| 938 |
+
hidden_states=outputs.hidden_states,
|
| 939 |
+
attentions=outputs.attentions,
|
| 940 |
+
image_hidden_states=outputs.image_hidden_states,
|
| 941 |
+
video_hidden_states=outputs.video_hidden_states,
|
| 942 |
+
)
|
| 943 |
+
|
| 944 |
+
def _prefill(self, input_ids, generation_config, model_kwargs, is_first_iteration=True):
|
| 945 |
+
# the text model's refusal, before the prefill writes to the cache
|
| 946 |
+
reject_packed_generation(self.config, input_ids, generation_config, model_kwargs, is_first_iteration)
|
| 947 |
+
return super()._prefill(input_ids, generation_config, model_kwargs, is_first_iteration=is_first_iteration)
|
| 948 |
+
|
| 949 |
+
@auto_docstring
|
| 950 |
+
def get_video_features(
|
| 951 |
+
self,
|
| 952 |
+
pixel_values_videos: torch.FloatTensor,
|
| 953 |
+
video_grid_thw: torch.LongTensor,
|
| 954 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 955 |
+
) -> tuple | BaseModelOutputWithPooling:
|
| 956 |
+
r"""
|
| 957 |
+
pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
|
| 958 |
+
The flattened patches of every video, grouped by merge block.
|
| 959 |
+
video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`):
|
| 960 |
+
The temporal, height and width dimensions of the patch grid of each video.
|
| 961 |
+
"""
|
| 962 |
+
return self.model.get_video_features(pixel_values_videos, video_grid_thw, **kwargs)
|
| 963 |
+
|
| 964 |
+
|
| 965 |
+
__all__ = ["Lfm3VlForConditionalGeneration", "Lfm3VlModel", "Lfm3VlPreTrainedModel", "Lfm3VlVisionModel"]
|
modeling_liquid_codec.py
ADDED
|
@@ -0,0 +1,381 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
import math
|
| 16 |
+
from dataclasses import dataclass
|
| 17 |
+
|
| 18 |
+
import torch
|
| 19 |
+
import torch.nn.functional as F
|
| 20 |
+
from torch import nn
|
| 21 |
+
|
| 22 |
+
from transformers import initialization as init
|
| 23 |
+
from transformers.cache_utils import DynamicCache
|
| 24 |
+
from transformers.conversion_mapping import register_checkpoint_conversion_mapping
|
| 25 |
+
from transformers.core_model_loading import WeightRenaming
|
| 26 |
+
from transformers.modeling_utils import PreTrainedAudioTokenizerBase
|
| 27 |
+
from transformers.utils import ModelOutput, auto_docstring, can_return_tuple, torch_compilable_check
|
| 28 |
+
from transformers.utils.import_utils import is_tracing
|
| 29 |
+
|
| 30 |
+
from .configuration_liquid_codec import LiquidCodecConfig
|
| 31 |
+
from .modeling_lfm3 import Lfm3Model
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
@auto_docstring
|
| 35 |
+
@dataclass
|
| 36 |
+
class LiquidCodecEncoderOutput(ModelOutput):
|
| 37 |
+
r"""
|
| 38 |
+
audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, num_frames)`):
|
| 39 |
+
The code of every codebook of every frame.
|
| 40 |
+
"""
|
| 41 |
+
|
| 42 |
+
audio_codes: torch.LongTensor = None
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
@auto_docstring
|
| 46 |
+
@dataclass
|
| 47 |
+
class LiquidCodecDecoderOutput(ModelOutput):
|
| 48 |
+
r"""
|
| 49 |
+
audio_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`):
|
| 50 |
+
The decoded 24 kHz audio in float32.
|
| 51 |
+
"""
|
| 52 |
+
|
| 53 |
+
audio_values: torch.FloatTensor = None
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
@auto_docstring
|
| 57 |
+
@dataclass
|
| 58 |
+
class LiquidCodecOutput(ModelOutput):
|
| 59 |
+
r"""
|
| 60 |
+
audio_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`):
|
| 61 |
+
The reconstructed audio, cut to the input length.
|
| 62 |
+
audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, num_frames)`):
|
| 63 |
+
The code of every codebook of every frame.
|
| 64 |
+
"""
|
| 65 |
+
|
| 66 |
+
audio_values: torch.FloatTensor = None
|
| 67 |
+
audio_codes: torch.LongTensor = None
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
class LiquidCodecStreamingCache:
|
| 71 |
+
"""Carries the decoder state between chunks: the LFM3 cache and the unfinished overlap-add tail."""
|
| 72 |
+
|
| 73 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 74 |
+
self.past_key_values = DynamicCache(config=config.lfm_config(config.decoder_window))
|
| 75 |
+
self.audio: torch.Tensor | None = None
|
| 76 |
+
self.envelope: torch.Tensor | None = None
|
| 77 |
+
self.started = False
|
| 78 |
+
self.finished = False
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
class LiquidCodecSTFT(nn.Module):
|
| 82 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 83 |
+
super().__init__()
|
| 84 |
+
self.n_fft = config.stft_n_fft
|
| 85 |
+
self.hop_length = config.stft_hop_length
|
| 86 |
+
self.win_length = config.stft_win_length
|
| 87 |
+
self.window = nn.Buffer(torch.hann_window(self.win_length, dtype=torch.float32), persistent=False)
|
| 88 |
+
self.num_features = 3 * (self.n_fft // 2 + 1)
|
| 89 |
+
|
| 90 |
+
def forward(self, audio: torch.Tensor) -> torch.Tensor:
|
| 91 |
+
total_pad = self.win_length - self.hop_length
|
| 92 |
+
audio = F.pad(audio.float(), (total_pad // 2, total_pad - total_pad // 2), mode="reflect")
|
| 93 |
+
spec = torch.stft(
|
| 94 |
+
audio,
|
| 95 |
+
n_fft=self.n_fft,
|
| 96 |
+
hop_length=self.hop_length,
|
| 97 |
+
win_length=self.win_length,
|
| 98 |
+
window=self.window.float(),
|
| 99 |
+
center=False,
|
| 100 |
+
return_complex=True,
|
| 101 |
+
)
|
| 102 |
+
magnitude = spec.abs()
|
| 103 |
+
phase_denominator = magnitude.clamp_min(1e-5)
|
| 104 |
+
return torch.cat((torch.log1p(magnitude), spec.real / phase_denominator, spec.imag / phase_denominator), dim=1)
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
class LiquidCodecFrontEnd(nn.Module):
|
| 108 |
+
def __init__(self, config: LiquidCodecConfig, num_features: int):
|
| 109 |
+
super().__init__()
|
| 110 |
+
self.frame_stack = config.stft_frame_stack
|
| 111 |
+
self.in_proj = nn.Conv1d(num_features, config.encoder_conv_dim, kernel_size=1)
|
| 112 |
+
self.conv = nn.Conv1d(config.encoder_conv_dim, config.encoder_conv_dim, kernel_size=3, padding=2)
|
| 113 |
+
self.out = nn.Conv1d(
|
| 114 |
+
config.encoder_conv_dim, config.encoder_dim, kernel_size=self.frame_stack, stride=self.frame_stack
|
| 115 |
+
)
|
| 116 |
+
self.act = nn.SiLU()
|
| 117 |
+
|
| 118 |
+
def forward(self, features: torch.Tensor) -> torch.Tensor:
|
| 119 |
+
hidden_states = self.act(self.in_proj(features.to(self.in_proj.weight.dtype)))
|
| 120 |
+
hidden_states = self.act(self.conv(hidden_states)[:, :, : hidden_states.shape[-1]])
|
| 121 |
+
return self.out(hidden_states)
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
class LiquidCodecLfm(nn.Module):
|
| 125 |
+
def __init__(self, config: LiquidCodecConfig, sliding_window: int | None = None):
|
| 126 |
+
super().__init__()
|
| 127 |
+
self.lfm = Lfm3Model(config.lfm_config(sliding_window))
|
| 128 |
+
self.lfm.register_module("embed_tokens", None)
|
| 129 |
+
# the projections run in the model dtype but the rotary tables stay in float32; a named hook keeps the
|
| 130 |
+
# module picklable
|
| 131 |
+
self.lfm.rotary_emb.register_forward_pre_hook(self._rotary_in_float32)
|
| 132 |
+
|
| 133 |
+
@staticmethod
|
| 134 |
+
def _rotary_in_float32(module, args):
|
| 135 |
+
return (args[0].float(), *args[1:])
|
| 136 |
+
|
| 137 |
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 138 |
+
return self.lfm(inputs_embeds=hidden_states, use_cache=False).last_hidden_state
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
class LiquidCodecQuantizer(nn.Module):
|
| 142 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 143 |
+
super().__init__()
|
| 144 |
+
levels = list(config.fsq_levels)
|
| 145 |
+
self.num_codebooks = config.n_q
|
| 146 |
+
self.codebook_dim = len(levels)
|
| 147 |
+
self.levels = nn.Buffer(torch.tensor(levels, dtype=torch.int32), persistent=False)
|
| 148 |
+
self.basis = nn.Buffer(
|
| 149 |
+
torch.tensor([math.prod(levels[:index]) for index in range(self.codebook_dim)], dtype=torch.int32),
|
| 150 |
+
persistent=False,
|
| 151 |
+
)
|
| 152 |
+
effective_dim = self.codebook_dim * self.num_codebooks
|
| 153 |
+
if config.quantizer_dim == effective_dim:
|
| 154 |
+
self.project_in = nn.Identity()
|
| 155 |
+
self.project_out = nn.Identity()
|
| 156 |
+
else:
|
| 157 |
+
self.project_in = nn.Linear(config.quantizer_dim, effective_dim)
|
| 158 |
+
self.project_out = nn.Linear(effective_dim, config.quantizer_dim)
|
| 159 |
+
|
| 160 |
+
def _bound(self, latents: torch.Tensor, eps: float = 1e-3) -> torch.Tensor:
|
| 161 |
+
half_levels = (self.levels - 1) * (1 + eps) / 2
|
| 162 |
+
offset = torch.where(self.levels % 2 == 0, 0.5, 0.0)
|
| 163 |
+
shift = torch.atanh(offset / half_levels)
|
| 164 |
+
return torch.tanh(latents + shift) * half_levels - offset
|
| 165 |
+
|
| 166 |
+
def encode(self, latents: torch.Tensor) -> torch.Tensor:
|
| 167 |
+
# the quantizer runs outside `forward`, so no `device_map` hook moves its inputs (here and in `decode`) to the
|
| 168 |
+
# device of its buffers
|
| 169 |
+
latents = self.project_in(latents.movedim(1, -1).contiguous()).float().to(self.levels.device)
|
| 170 |
+
if not is_tracing(latents):
|
| 171 |
+
# finite audio far above full scale overflows the float32 STFT, and NaN latents round to INT64_MIN codes
|
| 172 |
+
torch_compilable_check(
|
| 173 |
+
torch.isfinite(latents),
|
| 174 |
+
"The audio overflows the codec (non-finite latents); expected samples in [-1, 1]",
|
| 175 |
+
)
|
| 176 |
+
latents = latents.reshape(*latents.shape[:-1], self.num_codebooks, self.codebook_dim)
|
| 177 |
+
half_width = self.levels // 2
|
| 178 |
+
codes = self._bound(latents).round() / half_width
|
| 179 |
+
indices = ((codes * half_width + half_width) * self.basis).sum(-1).round().long()
|
| 180 |
+
return indices.transpose(1, 2)
|
| 181 |
+
|
| 182 |
+
def decode(self, codes: torch.Tensor) -> torch.Tensor:
|
| 183 |
+
indices = codes.to(self.basis.device).transpose(1, 2).unsqueeze(-1)
|
| 184 |
+
half_width = self.levels // 2
|
| 185 |
+
latents = ((indices // self.basis) % self.levels - half_width) / half_width
|
| 186 |
+
latents = latents.reshape(*latents.shape[:-2], -1)
|
| 187 |
+
if isinstance(self.project_out, nn.Linear):
|
| 188 |
+
latents = latents.to(self.project_out.weight.dtype)
|
| 189 |
+
return self.project_out(latents).movedim(-1, 1)
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
class LiquidCodecISTFT(nn.Module):
|
| 193 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 194 |
+
super().__init__()
|
| 195 |
+
self.n_fft = config.stft_n_fft
|
| 196 |
+
self.hop_length = config.stft_hop_length
|
| 197 |
+
self.win_length = config.stft_win_length
|
| 198 |
+
self.window = nn.Buffer(torch.hann_window(self.win_length, dtype=torch.float32))
|
| 199 |
+
|
| 200 |
+
def forward(self, spec: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
| 201 |
+
num_frames = spec.shape[-1]
|
| 202 |
+
frames = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") * self.window[None, :, None]
|
| 203 |
+
output_size = (num_frames - 1) * self.hop_length + self.win_length
|
| 204 |
+
fold = {"output_size": (1, output_size), "kernel_size": (1, self.win_length), "stride": (1, self.hop_length)}
|
| 205 |
+
audio = F.fold(frames, **fold)[:, 0, 0]
|
| 206 |
+
envelope = F.fold(self.window.square().expand(1, num_frames, -1).transpose(1, 2), **fold)[0, 0, 0]
|
| 207 |
+
return audio, envelope
|
| 208 |
+
|
| 209 |
+
|
| 210 |
+
class LiquidCodecDecoder(nn.Module):
|
| 211 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 212 |
+
super().__init__()
|
| 213 |
+
self.frame_stack = config.stft_frame_stack
|
| 214 |
+
self.lfm = LiquidCodecLfm(config, sliding_window=config.decoder_window).lfm
|
| 215 |
+
self.lin = nn.Linear(config.encoder_dim, 2 * (config.stft_n_fft // 2 + 1))
|
| 216 |
+
self.istft = LiquidCodecISTFT(config)
|
| 217 |
+
|
| 218 |
+
def forward(
|
| 219 |
+
self, embeddings: torch.Tensor, past_key_values: DynamicCache | None = None
|
| 220 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 221 |
+
hidden_states = self.lfm(
|
| 222 |
+
inputs_embeds=F.interpolate(
|
| 223 |
+
embeddings.mT, embeddings.shape[1] * self.frame_stack, mode="nearest-exact"
|
| 224 |
+
).mT,
|
| 225 |
+
past_key_values=past_key_values,
|
| 226 |
+
use_cache=past_key_values is not None,
|
| 227 |
+
).last_hidden_state
|
| 228 |
+
log_magnitude, angle = torch.chunk(self.lin(hidden_states).mT.contiguous(), 2, 1)
|
| 229 |
+
return self.istft(torch.polar(log_magnitude.float().exp(), angle.float()))
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
@auto_docstring(
|
| 233 |
+
custom_intro="The Liquid Codec neural audio codec: 24 kHz audio to 12 codebooks at 12.5 frames per second and back."
|
| 234 |
+
)
|
| 235 |
+
class LiquidCodecModel(PreTrainedAudioTokenizerBase):
|
| 236 |
+
config: LiquidCodecConfig
|
| 237 |
+
base_model_prefix = "codec"
|
| 238 |
+
main_input_name = "input_values"
|
| 239 |
+
# the analysis window is a non-persistent buffer cast in the forward; only the synthesis window is a weight
|
| 240 |
+
_keep_in_fp32_modules_strict = ["istft"]
|
| 241 |
+
_supports_flash_attn = False
|
| 242 |
+
_supports_sdpa = True
|
| 243 |
+
|
| 244 |
+
def __init__(self, config: LiquidCodecConfig):
|
| 245 |
+
super().__init__(config)
|
| 246 |
+
self.stft = LiquidCodecSTFT(config)
|
| 247 |
+
self.frontend = LiquidCodecFrontEnd(config, self.stft.num_features)
|
| 248 |
+
self.encoder = LiquidCodecLfm(config)
|
| 249 |
+
self.to_quantizer = nn.Linear(config.encoder_dim, config.quantizer_dim)
|
| 250 |
+
self.quantizer = LiquidCodecQuantizer(config)
|
| 251 |
+
self.decoder = LiquidCodecDecoder(config)
|
| 252 |
+
self.from_quantizer = nn.Linear(config.quantizer_dim, config.encoder_dim)
|
| 253 |
+
self.post_init()
|
| 254 |
+
|
| 255 |
+
@torch.no_grad()
|
| 256 |
+
def _init_weights(self, module):
|
| 257 |
+
super()._init_weights(module)
|
| 258 |
+
if isinstance(module, (LiquidCodecSTFT, LiquidCodecISTFT)):
|
| 259 |
+
init.copy_(module.window, torch.hann_window(module.win_length, dtype=torch.float32))
|
| 260 |
+
elif isinstance(module, LiquidCodecQuantizer):
|
| 261 |
+
levels = list(self.config.fsq_levels)
|
| 262 |
+
init.copy_(module.levels, torch.tensor(levels, dtype=torch.int32))
|
| 263 |
+
init.copy_(
|
| 264 |
+
module.basis,
|
| 265 |
+
torch.tensor([math.prod(levels[:index]) for index in range(len(levels))], dtype=torch.int32),
|
| 266 |
+
)
|
| 267 |
+
|
| 268 |
+
@can_return_tuple
|
| 269 |
+
def encode(self, input_values: torch.FloatTensor, *, return_dict: bool | None = None) -> LiquidCodecEncoderOutput:
|
| 270 |
+
r"""
|
| 271 |
+
Non-finite audio, or audio that overflows the codec, is refused eagerly; a compiled `encode` skips the check.
|
| 272 |
+
|
| 273 |
+
input_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`):
|
| 274 |
+
24 kHz audio. The tail is zero-padded to a whole codec frame of 1920 samples.
|
| 275 |
+
return_dict (`bool`, *optional*):
|
| 276 |
+
Whether to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
| 277 |
+
"""
|
| 278 |
+
if input_values.ndim != 2 or input_values.shape[-1] == 0:
|
| 279 |
+
raise ValueError(
|
| 280 |
+
f"Expected non-empty audio of shape (batch_size, num_samples), got {tuple(input_values.shape)}"
|
| 281 |
+
)
|
| 282 |
+
if not is_tracing(input_values):
|
| 283 |
+
# a non-finite sample would round to an INT64_MIN code that only fails far downstream; the tensor-valued
|
| 284 |
+
# check would break a `fullgraph` compile or a CUDA-graph capture
|
| 285 |
+
torch_compilable_check(torch.isfinite(input_values), "`input_values` must be finite")
|
| 286 |
+
frame_size = self.config.frame_size
|
| 287 |
+
num_samples = -(input_values.shape[-1] // -frame_size) * frame_size
|
| 288 |
+
input_values = F.pad(input_values, (0, num_samples - input_values.shape[-1]))
|
| 289 |
+
hidden_states = self.frontend(self.stft(input_values)).mT
|
| 290 |
+
latents = self.to_quantizer(self.encoder(hidden_states)).mT
|
| 291 |
+
return LiquidCodecEncoderOutput(audio_codes=self.quantizer.encode(latents))
|
| 292 |
+
|
| 293 |
+
@can_return_tuple
|
| 294 |
+
def decode(
|
| 295 |
+
self,
|
| 296 |
+
audio_codes: torch.LongTensor,
|
| 297 |
+
streaming_cache: LiquidCodecStreamingCache | None = None,
|
| 298 |
+
finalize: bool = True,
|
| 299 |
+
*,
|
| 300 |
+
return_dict: bool | None = None,
|
| 301 |
+
) -> LiquidCodecDecoderOutput:
|
| 302 |
+
r"""
|
| 303 |
+
audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, num_frames)`):
|
| 304 |
+
The codes to decode; every frame becomes 1920 samples.
|
| 305 |
+
streaming_cache ([`LiquidCodecStreamingCache`], *optional*):
|
| 306 |
+
The state carried between consecutive chunks of one stream. The chunks decode to the same audio as the
|
| 307 |
+
concatenated codes.
|
| 308 |
+
finalize (`bool`, *optional*, defaults to `True`):
|
| 309 |
+
Whether this is the last chunk of the stream. Until then the trailing samples that later frames still
|
| 310 |
+
contribute to are held back in the cache.
|
| 311 |
+
return_dict (`bool`, *optional*):
|
| 312 |
+
Whether to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
| 313 |
+
"""
|
| 314 |
+
if audio_codes.ndim != 3 or audio_codes.shape[1] != self.config.n_q:
|
| 315 |
+
raise ValueError(
|
| 316 |
+
f"Expected codes of shape (batch_size, {self.config.n_q}, num_frames), got {tuple(audio_codes.shape)}"
|
| 317 |
+
)
|
| 318 |
+
torch_compilable_check(
|
| 319 |
+
(audio_codes >= 0) & (audio_codes < self.config.codebook_size),
|
| 320 |
+
f"Codes must be in [0, {self.config.codebook_size})",
|
| 321 |
+
)
|
| 322 |
+
if streaming_cache is not None and streaming_cache.finished:
|
| 323 |
+
raise ValueError("This stream is finalized; start with a new cache")
|
| 324 |
+
trim = (self.config.stft_win_length - self.config.stft_hop_length) // 2
|
| 325 |
+
tail = 2 * trim
|
| 326 |
+
batch_size = audio_codes.shape[0]
|
| 327 |
+
if not audio_codes.shape[-1] and not (finalize and streaming_cache is not None and streaming_cache.started):
|
| 328 |
+
# no frames and no held-back tail to flush: the overlap-add state must stay untouched
|
| 329 |
+
if streaming_cache is not None and finalize:
|
| 330 |
+
streaming_cache.finished = True
|
| 331 |
+
return LiquidCodecDecoderOutput(audio_values=audio_codes.new_zeros((batch_size, 0), dtype=torch.float32))
|
| 332 |
+
if audio_codes.shape[-1]:
|
| 333 |
+
# `quantizer.decode` returns float32 latents when its projections are `nn.Identity`
|
| 334 |
+
latents = self.quantizer.decode(audio_codes).mT.to(self.from_quantizer.weight.dtype)
|
| 335 |
+
embeddings = self.from_quantizer(latents)
|
| 336 |
+
past_key_values = None if streaming_cache is None else streaming_cache.past_key_values
|
| 337 |
+
audio, envelope = self.decoder(embeddings, past_key_values)
|
| 338 |
+
else:
|
| 339 |
+
audio = self.decoder.istft.window.new_zeros(batch_size, tail)
|
| 340 |
+
envelope = self.decoder.istft.window.new_zeros(tail)
|
| 341 |
+
if streaming_cache is None:
|
| 342 |
+
audio, envelope = audio[:, trim:-trim], envelope[trim:-trim]
|
| 343 |
+
else:
|
| 344 |
+
if streaming_cache.audio is not None:
|
| 345 |
+
audio = torch.cat([audio[:, :tail] + streaming_cache.audio, audio[:, tail:]], dim=1)
|
| 346 |
+
envelope = torch.cat([envelope[:tail] + streaming_cache.envelope, envelope[tail:]])
|
| 347 |
+
if not streaming_cache.started:
|
| 348 |
+
audio, envelope = audio[:, trim:], envelope[trim:]
|
| 349 |
+
streaming_cache.started = True
|
| 350 |
+
if finalize:
|
| 351 |
+
audio, envelope = audio[:, :-trim], envelope[:-trim]
|
| 352 |
+
streaming_cache.audio, streaming_cache.envelope = None, None
|
| 353 |
+
streaming_cache.finished = True
|
| 354 |
+
else:
|
| 355 |
+
streaming_cache.audio, streaming_cache.envelope = audio[:, -tail:], envelope[-tail:]
|
| 356 |
+
audio, envelope = audio[:, :-tail], envelope[:-tail]
|
| 357 |
+
return LiquidCodecDecoderOutput(audio_values=(audio / envelope).float())
|
| 358 |
+
|
| 359 |
+
@can_return_tuple
|
| 360 |
+
@auto_docstring
|
| 361 |
+
def forward(self, input_values: torch.FloatTensor) -> LiquidCodecOutput:
|
| 362 |
+
r"""
|
| 363 |
+
input_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`):
|
| 364 |
+
24 kHz audio.
|
| 365 |
+
"""
|
| 366 |
+
audio_codes = self.encode(input_values, return_dict=True).audio_codes
|
| 367 |
+
audio = self.decode(audio_codes, return_dict=True).audio_values
|
| 368 |
+
return LiquidCodecOutput(audio_values=audio[..., : input_values.shape[-1]], audio_codes=audio_codes)
|
| 369 |
+
|
| 370 |
+
|
| 371 |
+
__all__ = ["LiquidCodecModel", "LiquidCodecStreamingCache"]
|
| 372 |
+
|
| 373 |
+
|
| 374 |
+
register_checkpoint_conversion_mapping(
|
| 375 |
+
"liquid_codec",
|
| 376 |
+
[
|
| 377 |
+
WeightRenaming(source_patterns=r"^_frontend\.", target_patterns="frontend."),
|
| 378 |
+
WeightRenaming(source_patterns=r"^quantizer\.fsq\.", target_patterns="quantizer."),
|
| 379 |
+
],
|
| 380 |
+
overwrite=True,
|
| 381 |
+
)
|
processing_lfm3_omni.py
ADDED
|
@@ -0,0 +1,769 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
|
| 15 |
+
import copy
|
| 16 |
+
import math
|
| 17 |
+
import os
|
| 18 |
+
import re
|
| 19 |
+
|
| 20 |
+
import numpy as np
|
| 21 |
+
import torch
|
| 22 |
+
|
| 23 |
+
from transformers.audio_utils import AudioInput, make_list_of_audio
|
| 24 |
+
from transformers.feature_extraction_utils import BatchFeature
|
| 25 |
+
from transformers.image_utils import is_pil_image, is_valid_image, make_flat_list_of_images
|
| 26 |
+
from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, VideosKwargs
|
| 27 |
+
from transformers.utils import auto_docstring, is_numpy_array, is_torch_tensor
|
| 28 |
+
from transformers.video_utils import is_valid_video, make_batched_metadata, make_batched_videos
|
| 29 |
+
|
| 30 |
+
from .modeling_lfm3_omni import Lfm3OmniFeatureExtractorKwargs
|
| 31 |
+
from .modeling_liquid_codec import LiquidCodecModel
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class Lfm3OmniVideosKwargs(VideosKwargs, total=False):
|
| 35 |
+
"""
|
| 36 |
+
use_audio_in_video (`bool`, *optional*, defaults to `False`):
|
| 37 |
+
Whether each video comes with its soundtrack: the first audio clip after the video in the text (the one
|
| 38 |
+
`apply_chat_template(load_audio_from_video=True)` puts right behind the video) is interleaved after the
|
| 39 |
+
video's frame groups instead of standing in its own media span.
|
| 40 |
+
"""
|
| 41 |
+
|
| 42 |
+
use_audio_in_video: bool
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
class Lfm3OmniProcessorKwargs(ProcessingKwargs, total=False):
|
| 46 |
+
audio_kwargs: Lfm3OmniFeatureExtractorKwargs
|
| 47 |
+
videos_kwargs: Lfm3OmniVideosKwargs
|
| 48 |
+
_defaults = {"videos_kwargs": {"return_metadata": True, "use_audio_in_video": False}}
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
# The prompt format of the training data: ChatML turns, one placeholder per clip, no thinking block
|
| 52 |
+
DEFAULT_CHAT_TEMPLATE = (
|
| 53 |
+
"{{- bos_token -}}"
|
| 54 |
+
"{%- for message in messages -%}"
|
| 55 |
+
"{{- '<|im_start|>' + message['role'] + '\\n' -}}"
|
| 56 |
+
"{%- if message['content'] is string -%}{{- message['content'] -}}"
|
| 57 |
+
"{%- else -%}{%- for content in message['content'] -%}"
|
| 58 |
+
"{%- if content['type'] == 'audio' -%}{{- '<|audio|>' -}}"
|
| 59 |
+
"{%- elif content['type'] == 'audio_out' -%}{{- '<|audio_out|>' -}}"
|
| 60 |
+
"{%- elif content['type'] == 'text' -%}{{- content['text'] -}}"
|
| 61 |
+
"{%- endif -%}{%- endfor -%}{%- endif -%}"
|
| 62 |
+
"{{- '<|im_end|>\\n' -}}"
|
| 63 |
+
"{%- endfor -%}"
|
| 64 |
+
"{%- if add_generation_prompt -%}{{- '<|im_start|>assistant\\n' -}}{%- endif -%}"
|
| 65 |
+
)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def _media_per_prompt(media, is_item) -> list[int] | None:
|
| 69 |
+
# rows name the prompt that owns them while the expansion consumes media in flat order; a row is a list of whole
|
| 70 |
+
# items or one stacked batch, and a list of single frames is one video, not a row
|
| 71 |
+
if not isinstance(media, (list, tuple)):
|
| 72 |
+
return None
|
| 73 |
+
counts = []
|
| 74 |
+
for row in media:
|
| 75 |
+
if isinstance(row, (list, tuple)) and all(is_item(item) for item in row):
|
| 76 |
+
counts.append(len(row))
|
| 77 |
+
elif (is_numpy_array(row) or is_torch_tensor(row)) and row.ndim == 4 and (len(row) == 0 or is_item(row[0])):
|
| 78 |
+
counts.append(len(row))
|
| 79 |
+
else:
|
| 80 |
+
return None
|
| 81 |
+
return counts
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def _check_vision_placeholder_counts(processor, prompts, images, videos) -> None:
|
| 85 |
+
for name, token, media, flatten, is_item in (
|
| 86 |
+
(
|
| 87 |
+
"images",
|
| 88 |
+
processor.image_token,
|
| 89 |
+
images,
|
| 90 |
+
make_flat_list_of_images,
|
| 91 |
+
lambda item: is_valid_image(item) and (is_pil_image(item) or item.ndim == 3),
|
| 92 |
+
),
|
| 93 |
+
(
|
| 94 |
+
"videos",
|
| 95 |
+
processor.video_token,
|
| 96 |
+
videos,
|
| 97 |
+
make_batched_videos,
|
| 98 |
+
lambda item: isinstance(item, str) or (not is_pil_image(item) and is_valid_video(item)),
|
| 99 |
+
),
|
| 100 |
+
):
|
| 101 |
+
num_placeholders = sum(prompt.count(token) for prompt in prompts)
|
| 102 |
+
num_media = 0 if media is None else len(flatten(media))
|
| 103 |
+
if num_placeholders != num_media:
|
| 104 |
+
raise ValueError(f"Found {num_placeholders} {token} tokens in the text but received {num_media} {name}")
|
| 105 |
+
# one prompt cannot take another's media
|
| 106 |
+
per_prompt = _media_per_prompt(media, is_item) if len(prompts) > 1 else None
|
| 107 |
+
in_text = [prompt.count(token) for prompt in prompts]
|
| 108 |
+
if per_prompt is not None and per_prompt != in_text:
|
| 109 |
+
raise ValueError(f"Found {in_text} {token} tokens per prompt but received {per_prompt} {name} per prompt")
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def _check_video_metadata_count(videos, video_metadata) -> None:
|
| 113 |
+
# the shared sampler pairs clips with metadata by `zip`, which drops the clips past the last entry
|
| 114 |
+
if video_metadata is None:
|
| 115 |
+
return
|
| 116 |
+
videos = make_batched_videos(videos)
|
| 117 |
+
entries = make_batched_metadata(videos, video_metadata) if video_metadata != [] else []
|
| 118 |
+
if len(entries) != len(videos):
|
| 119 |
+
raise ValueError(
|
| 120 |
+
f"Received {len(videos)} videos but {len(entries)} `video_metadata` entries; pass one entry per video"
|
| 121 |
+
)
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
DEFAULT_OMNI_CHAT_TEMPLATE = DEFAULT_CHAT_TEMPLATE.replace(
|
| 125 |
+
"{%- if content['type'] == 'audio' -%}",
|
| 126 |
+
"{%- if content['type'] == 'image' -%}{{- '<|image|>' -}}"
|
| 127 |
+
"{%- elif content['type'] == 'video' -%}{{- '<|video|>' -}}"
|
| 128 |
+
"{%- elif content['type'] == 'audio' -%}",
|
| 129 |
+
)
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
@auto_docstring
|
| 133 |
+
class Lfm3OmniProcessor(ProcessorMixin):
|
| 134 |
+
"""The LFM3-Audio processor extended with the LFM3-VL image and video processors: every image, video frame group
|
| 135 |
+
and input clip is expanded in place, wrapped in the media span tokens the Omni SFT data uses."""
|
| 136 |
+
|
| 137 |
+
attributes = ["image_processor", "feature_extractor", "tokenizer", "video_processor"]
|
| 138 |
+
tokenizer_class = "AutoTokenizer"
|
| 139 |
+
audio_tokenizer_class = "LiquidCodecModel"
|
| 140 |
+
valid_processor_kwargs = Lfm3OmniProcessorKwargs
|
| 141 |
+
|
| 142 |
+
def __init__(
|
| 143 |
+
self,
|
| 144 |
+
image_processor=None,
|
| 145 |
+
feature_extractor=None,
|
| 146 |
+
tokenizer=None,
|
| 147 |
+
video_processor=None,
|
| 148 |
+
audio_tokenizer=None,
|
| 149 |
+
chat_template=None,
|
| 150 |
+
audio_token="<|audio|>",
|
| 151 |
+
audio_out_token="<|audio_out|>",
|
| 152 |
+
image_token="<|image|>",
|
| 153 |
+
video_token="<|video|>",
|
| 154 |
+
vision_start_token="<|media_start|>",
|
| 155 |
+
vision_end_token="<|media_end|>",
|
| 156 |
+
vision_markers=True,
|
| 157 |
+
audio_markers=True,
|
| 158 |
+
**kwargs,
|
| 159 |
+
):
|
| 160 |
+
r"""
|
| 161 |
+
image_processor ([`Lfm3VlImageProcessor`]):
|
| 162 |
+
The image patch processor.
|
| 163 |
+
feature_extractor ([`Lfm3OmniFeatureExtractor`]):
|
| 164 |
+
The log-mel feature extractor.
|
| 165 |
+
tokenizer ([`PreTrainedTokenizerFast`]):
|
| 166 |
+
The LFM3 tokenizer.
|
| 167 |
+
video_processor ([`Lfm3VlVideoProcessor`], *optional*):
|
| 168 |
+
The video frame sampler and patch processor.
|
| 169 |
+
audio_tokenizer ([`LiquidCodecModel`], *optional*):
|
| 170 |
+
The codec for 24 kHz output audio.
|
| 171 |
+
chat_template (`str`, *optional*):
|
| 172 |
+
The conversation template.
|
| 173 |
+
audio_token (`str`, *optional*, defaults to `"<|audio|>"`):
|
| 174 |
+
The input audio placeholder.
|
| 175 |
+
audio_out_token (`str`, *optional*, defaults to `"<|audio_out|>"`):
|
| 176 |
+
The output audio placeholder.
|
| 177 |
+
image_token (`str`, *optional*, defaults to `"<|image|>"`):
|
| 178 |
+
The image placeholder.
|
| 179 |
+
video_token (`str`, *optional*, defaults to `"<|video|>"`):
|
| 180 |
+
The video placeholder.
|
| 181 |
+
vision_start_token (`str`, *optional*, defaults to `"<|media_start|>"`):
|
| 182 |
+
The token opening an image, a video and each of its timestamped frame groups.
|
| 183 |
+
vision_end_token (`str`, *optional*, defaults to `"<|media_end|>"`):
|
| 184 |
+
The token closing them.
|
| 185 |
+
vision_markers (`bool`, *optional*, defaults to `True`):
|
| 186 |
+
Whether images and videos are wrapped in the vision span tokens, as the Omni SFT data renders them.
|
| 187 |
+
audio_markers (`bool`, *optional*, defaults to `True`):
|
| 188 |
+
Whether input audio clips are wrapped in the same span tokens, as the Omni SFT data renders them.
|
| 189 |
+
"""
|
| 190 |
+
self.audio_token = audio_token
|
| 191 |
+
self.audio_out_token = audio_out_token
|
| 192 |
+
self.image_token = image_token
|
| 193 |
+
self.video_token = video_token
|
| 194 |
+
self.vision_start_token = vision_start_token
|
| 195 |
+
self.vision_end_token = vision_end_token
|
| 196 |
+
self.vision_markers = vision_markers
|
| 197 |
+
self.audio_markers = audio_markers
|
| 198 |
+
for token in (audio_token, audio_out_token, image_token, video_token):
|
| 199 |
+
if tokenizer.convert_tokens_to_ids(token) in (None, tokenizer.unk_token_id):
|
| 200 |
+
tokenizer.add_tokens([token], special_tokens=True)
|
| 201 |
+
# the span tokens exist in every LFM3 tokenizer under one of two names (`<|vision_*|>` in liquid7.1,
|
| 202 |
+
# `<|media_*|>` in liquid7.1_media) at the same ids; never mint new ones
|
| 203 |
+
for attribute, token, fallback in (
|
| 204 |
+
("vision_start_token", vision_start_token, "<|vision_start|>"),
|
| 205 |
+
("vision_end_token", vision_end_token, "<|vision_end|>"),
|
| 206 |
+
):
|
| 207 |
+
if tokenizer.convert_tokens_to_ids(token) in (None, tokenizer.unk_token_id):
|
| 208 |
+
if tokenizer.convert_tokens_to_ids(fallback) not in (None, tokenizer.unk_token_id):
|
| 209 |
+
setattr(self, attribute, fallback)
|
| 210 |
+
else:
|
| 211 |
+
tokenizer.add_tokens([token], special_tokens=True)
|
| 212 |
+
self.audio_token_id = tokenizer.convert_tokens_to_ids(audio_token)
|
| 213 |
+
self.audio_out_token_id = tokenizer.convert_tokens_to_ids(audio_out_token)
|
| 214 |
+
self.image_token_id = tokenizer.convert_tokens_to_ids(image_token)
|
| 215 |
+
self.video_token_id = tokenizer.convert_tokens_to_ids(video_token)
|
| 216 |
+
super().__init__(
|
| 217 |
+
image_processor,
|
| 218 |
+
feature_extractor,
|
| 219 |
+
tokenizer,
|
| 220 |
+
video_processor,
|
| 221 |
+
chat_template=chat_template if chat_template is not None else DEFAULT_OMNI_CHAT_TEMPLATE,
|
| 222 |
+
)
|
| 223 |
+
# a released transformers only accepts its own codec classes as `audio_tokenizer`
|
| 224 |
+
self.audio_tokenizer = audio_tokenizer
|
| 225 |
+
|
| 226 |
+
@classmethod
|
| 227 |
+
def from_pretrained(cls, pretrained_model_name_or_path, *args, audio_tokenizer=None, **kwargs):
|
| 228 |
+
# `ProcessorMixin` forwards unknown kwargs to the tokenizer loader, which would keep the codec object in the
|
| 229 |
+
# tokenizer config and break `save_pretrained`
|
| 230 |
+
processor = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
|
| 231 |
+
if audio_tokenizer is not None:
|
| 232 |
+
processor.audio_tokenizer = audio_tokenizer
|
| 233 |
+
return processor
|
| 234 |
+
|
| 235 |
+
def to_dict(self):
|
| 236 |
+
# `ProcessorMixin.to_dict` deep-copies the attributes; the codec weights are saved on their own. A shallow
|
| 237 |
+
# clone drops the codec, so concurrent calls on `self` keep it.
|
| 238 |
+
clone = copy.copy(self)
|
| 239 |
+
codec = clone.__dict__.pop("audio_tokenizer", None)
|
| 240 |
+
output = ProcessorMixin.to_dict(clone)
|
| 241 |
+
if codec is not None:
|
| 242 |
+
output["audio_tokenizer_subfolder"] = "audio_tokenizer"
|
| 243 |
+
return output
|
| 244 |
+
|
| 245 |
+
def save_pretrained(self, save_directory, push_to_hub=False, **kwargs):
|
| 246 |
+
files = super().save_pretrained(save_directory, push_to_hub=False, **kwargs)
|
| 247 |
+
if getattr(self, "audio_tokenizer", None) is not None:
|
| 248 |
+
self.audio_tokenizer.save_pretrained(os.path.join(save_directory, "audio_tokenizer"))
|
| 249 |
+
if push_to_hub:
|
| 250 |
+
repo_id = kwargs.pop("repo_id", os.path.basename(os.path.abspath(save_directory)))
|
| 251 |
+
self.push_to_hub(repo_id, **kwargs)
|
| 252 |
+
return files
|
| 253 |
+
|
| 254 |
+
@classmethod
|
| 255 |
+
def get_processor_dict(cls, pretrained_model_name_or_path, **kwargs):
|
| 256 |
+
codec_kwargs = {
|
| 257 |
+
key: value
|
| 258 |
+
for key, value in kwargs.items()
|
| 259 |
+
if key in ("cache_dir", "force_download", "local_files_only", "token", "revision")
|
| 260 |
+
}
|
| 261 |
+
parent = kwargs.get("subfolder") or ""
|
| 262 |
+
processor_dict, kwargs = super().get_processor_dict(pretrained_model_name_or_path, **kwargs)
|
| 263 |
+
subfolder = processor_dict.pop("audio_tokenizer_subfolder", None)
|
| 264 |
+
if subfolder is not None:
|
| 265 |
+
if subfolder != "audio_tokenizer":
|
| 266 |
+
raise ValueError("Expected the bundled codec in the audio_tokenizer subfolder")
|
| 267 |
+
codec_class = LiquidCodecModel
|
| 268 |
+
processor_dict["audio_tokenizer"] = codec_class.from_pretrained(
|
| 269 |
+
pretrained_model_name_or_path, subfolder=os.path.join(parent, subfolder), **codec_kwargs
|
| 270 |
+
)
|
| 271 |
+
return processor_dict, kwargs
|
| 272 |
+
|
| 273 |
+
@auto_docstring
|
| 274 |
+
def __call__(
|
| 275 |
+
self,
|
| 276 |
+
text=None,
|
| 277 |
+
audio: AudioInput | None = None,
|
| 278 |
+
audio_out: AudioInput | None = None,
|
| 279 |
+
images=None,
|
| 280 |
+
videos=None,
|
| 281 |
+
**kwargs,
|
| 282 |
+
):
|
| 283 |
+
r"""
|
| 284 |
+
audio_out (`np.ndarray`, `torch.Tensor` or a list of them, *optional*):
|
| 285 |
+
24 kHz clips, one per `audio_out_token` in the text, encoded with the codec into `audio_out_codes`; the
|
| 286 |
+
placeholder is repeated once per frame plus once for the end-of-speech frame that closes every clip.
|
| 287 |
+
videos (`VideoInput`, *optional*):
|
| 288 |
+
One clip per `video_token`. With `use_audio_in_video=True` each video takes the first audio clip after it
|
| 289 |
+
as its soundtrack, interleaved after each timestamped frame group like the Omni SFT data; otherwise every
|
| 290 |
+
audio clip is its own media span.
|
| 291 |
+
"""
|
| 292 |
+
for name, media in (("audio", audio), ("audio_out", audio_out), ("images", images), ("videos", videos)):
|
| 293 |
+
if isinstance(media, (list, tuple)) and len(media) == 0:
|
| 294 |
+
raise ValueError(f"`{name}` must hold at least one item; pass `{name}=None` for prompts without one")
|
| 295 |
+
if text is None:
|
| 296 |
+
# media-only calls return the features alone, like the audio processor
|
| 297 |
+
if audio_out is not None:
|
| 298 |
+
raise ValueError("`audio_out` needs the text holding its placeholders")
|
| 299 |
+
return super().__call__(audio=audio, images=images, videos=videos, **kwargs)
|
| 300 |
+
single = isinstance(text, str)
|
| 301 |
+
caller = [text] if single else list(text)
|
| 302 |
+
stripped = [self._strip_spans(prompt) for prompt in caller]
|
| 303 |
+
prompts, offsets = [prompt for prompt, _ in stripped], [row for _, row in stripped]
|
| 304 |
+
audio_out_codes = None
|
| 305 |
+
if audio_out is not None:
|
| 306 |
+
prompts, audio_out_codes, audio_out_offsets = self._expand_audio_out(prompts, audio_out)
|
| 307 |
+
offsets = self._compose_replacement_offsets(caller, offsets, audio_out_offsets)
|
| 308 |
+
merged = self._merge_kwargs(
|
| 309 |
+
self.valid_processor_kwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs
|
| 310 |
+
)
|
| 311 |
+
return_tensors = merged["text_kwargs"].get("return_tensors")
|
| 312 |
+
if videos is not None:
|
| 313 |
+
outputs = self._call_with_videos(prompts, audio, images, videos, single, **kwargs)
|
| 314 |
+
else:
|
| 315 |
+
outputs = super().__call__(
|
| 316 |
+
text=prompts[0] if single else prompts, audio=audio, images=images, videos=videos, **kwargs
|
| 317 |
+
)
|
| 318 |
+
if "text_replacement_offsets" in outputs:
|
| 319 |
+
# the expansion saw the prompts after the span pairs were stripped and the speech placeholders repeated
|
| 320 |
+
outputs["text_replacement_offsets"] = self._compose_replacement_offsets(
|
| 321 |
+
caller, offsets, outputs["text_replacement_offsets"]
|
| 322 |
+
)
|
| 323 |
+
if audio_out_codes is not None:
|
| 324 |
+
self._check_audio_out_placeholders(outputs["input_ids"], audio_out_codes.shape[0])
|
| 325 |
+
outputs["audio_out_codes"] = self._as_return_tensors(audio_out_codes, return_tensors)
|
| 326 |
+
return outputs
|
| 327 |
+
|
| 328 |
+
def _check_audio_out_placeholders(self, input_ids, num_frames: int) -> None:
|
| 329 |
+
# truncation may have dropped placeholders the codes still cover
|
| 330 |
+
rows = input_ids.tolist() if hasattr(input_ids, "tolist") else input_ids
|
| 331 |
+
rows = [rows] if rows and isinstance(rows[0], int) else rows
|
| 332 |
+
num_placeholders = sum(row.count(self.audio_out_token_id) for row in rows)
|
| 333 |
+
if num_placeholders != num_frames:
|
| 334 |
+
raise ValueError(
|
| 335 |
+
f"Found {num_placeholders} {self.audio_out_token} tokens after tokenization for {num_frames} "
|
| 336 |
+
"speech frames; disable truncation or shorten the prompt"
|
| 337 |
+
)
|
| 338 |
+
|
| 339 |
+
@staticmethod
|
| 340 |
+
def _as_return_tensors(codes: torch.Tensor, return_tensors):
|
| 341 |
+
return BatchFeature({"codes": codes}, tensor_type=return_tensors)["codes"]
|
| 342 |
+
|
| 343 |
+
@staticmethod
|
| 344 |
+
def _rewrite_with_offsets(prompt: str, pattern: str, replace, kind: str) -> tuple[str, list[dict]]:
|
| 345 |
+
"""Rewrites `prompt` like `ProcessorMixin.get_text_with_replacements`, with the same offset records."""
|
| 346 |
+
pieces, row, last, offset = [], [], 0, 0
|
| 347 |
+
for match in re.finditer(pattern, prompt):
|
| 348 |
+
start, end = match.span()
|
| 349 |
+
replacement = replace(match)
|
| 350 |
+
pieces.extend((prompt[last:start], replacement))
|
| 351 |
+
row.append(
|
| 352 |
+
{
|
| 353 |
+
"type": kind,
|
| 354 |
+
"span": (start, end),
|
| 355 |
+
"new_span": (start + offset, start + offset + len(replacement)),
|
| 356 |
+
"text": match.group(),
|
| 357 |
+
"replacement": replacement,
|
| 358 |
+
}
|
| 359 |
+
)
|
| 360 |
+
offset += len(replacement) - (end - start)
|
| 361 |
+
last = end
|
| 362 |
+
pieces.append(prompt[last:])
|
| 363 |
+
return "".join(pieces), row
|
| 364 |
+
|
| 365 |
+
@staticmethod
|
| 366 |
+
def _compose_replacement_offsets(prompts: list[str], first: list[list[dict]], second) -> list[list[dict]]:
|
| 367 |
+
"""The records of two successive rewrites (the caller's text A to B, then B to C) as one rewrite of A: an entry
|
| 368 |
+
of the second that covers entries of the first absorbs them, so every entry spans the caller's text."""
|
| 369 |
+
second = second or [[] for _ in prompts]
|
| 370 |
+
composed = []
|
| 371 |
+
for prompt, row_first, row_second in zip(prompts, first, second):
|
| 372 |
+
|
| 373 |
+
def to_first(position, row_first=row_first):
|
| 374 |
+
return position - sum(
|
| 375 |
+
len(r["replacement"]) - (r["span"][1] - r["span"][0])
|
| 376 |
+
for r in row_first
|
| 377 |
+
if r["new_span"][1] <= position
|
| 378 |
+
)
|
| 379 |
+
|
| 380 |
+
mapped = []
|
| 381 |
+
for entry in row_second:
|
| 382 |
+
start, end = to_first(entry["span"][0]), to_first(entry["span"][1])
|
| 383 |
+
mapped.append({**entry, "span": (start, end), "text": prompt[start:end]})
|
| 384 |
+
kept = [
|
| 385 |
+
r
|
| 386 |
+
for r in row_first
|
| 387 |
+
if not any(m["span"][0] <= r["span"][0] and r["span"][1] <= m["span"][1] for m in mapped)
|
| 388 |
+
]
|
| 389 |
+
row, offset = [], 0
|
| 390 |
+
for entry in sorted(kept + mapped, key=lambda r: r["span"]):
|
| 391 |
+
start, end = entry["span"]
|
| 392 |
+
row.append({**entry, "new_span": (start + offset, start + offset + len(entry["replacement"]))})
|
| 393 |
+
offset += len(entry["replacement"]) - (end - start)
|
| 394 |
+
composed.append(row)
|
| 395 |
+
return composed
|
| 396 |
+
|
| 397 |
+
def _expand_audio_out(
|
| 398 |
+
self, prompts: list[str], audio_out: AudioInput
|
| 399 |
+
) -> tuple[list[str], torch.Tensor, list[list[dict]]]:
|
| 400 |
+
if getattr(self, "audio_tokenizer", None) is None:
|
| 401 |
+
raise ValueError("`audio_out` requires the processor to hold the codec as `audio_tokenizer`")
|
| 402 |
+
clips = make_list_of_audio(audio_out)
|
| 403 |
+
if sum(prompt.count(self.audio_out_token) for prompt in prompts) != len(clips):
|
| 404 |
+
raise ValueError(f"The number of {self.audio_out_token} tokens in the text must match the number of clips")
|
| 405 |
+
codes = [self.encode_audio(clip) for clip in clips]
|
| 406 |
+
# the model ends speech with a frame of end-of-speech codes before the next text token; the prompt and the
|
| 407 |
+
# training targets carry it the same way `generate(return_audio_eos=True)` does
|
| 408 |
+
eos_frame = torch.full((1, codes[0].shape[1]), self.audio_eos_code, dtype=codes[0].dtype)
|
| 409 |
+
codes = [torch.cat((frames, eos_frame)) for frames in codes]
|
| 410 |
+
frame_counts = iter(code.shape[0] for code in codes)
|
| 411 |
+
expanded, offsets = [], []
|
| 412 |
+
for prompt in prompts:
|
| 413 |
+
prompt, row = self._rewrite_with_offsets(
|
| 414 |
+
prompt,
|
| 415 |
+
re.escape(self.audio_out_token),
|
| 416 |
+
lambda match: self.audio_out_token * next(frame_counts),
|
| 417 |
+
"audio_out",
|
| 418 |
+
)
|
| 419 |
+
expanded.append(prompt)
|
| 420 |
+
offsets.append(row)
|
| 421 |
+
return expanded, torch.cat(codes), offsets
|
| 422 |
+
|
| 423 |
+
@property
|
| 424 |
+
def audio_eos_code(self) -> int:
|
| 425 |
+
"""The end-of-speech code: the first id after the codec's vocabulary."""
|
| 426 |
+
return math.prod(self.audio_tokenizer.config.fsq_levels)
|
| 427 |
+
|
| 428 |
+
@torch.no_grad()
|
| 429 |
+
def encode_audio(self, audio) -> torch.Tensor:
|
| 430 |
+
codec = self.audio_tokenizer
|
| 431 |
+
waveform = torch.as_tensor(audio if torch.is_tensor(audio) else np.asarray(audio), dtype=torch.float32)
|
| 432 |
+
# the quantizer rounds non-finite values to INT64_MIN, which would surface as codes
|
| 433 |
+
if waveform.ndim != 1 or not torch.isfinite(waveform).all():
|
| 434 |
+
raise ValueError("`audio_out` clips must be mono waveforms of finite samples")
|
| 435 |
+
# an offloaded codec keeps some weights on `meta`: send the clip to a device that holds weights, and the hooks
|
| 436 |
+
# move it on from there
|
| 437 |
+
device = next(
|
| 438 |
+
(param.device for param in codec.parameters() if param.device.type != "meta"), torch.device("cpu")
|
| 439 |
+
)
|
| 440 |
+
# the codec's return type follows its config's `return_dict`
|
| 441 |
+
codes = codec.encode(waveform.to(device)[None], return_dict=True).audio_codes[0].T.cpu()
|
| 442 |
+
# a compiled codec skips its latent check, and a non-finite latent rounds to INT64_MIN
|
| 443 |
+
if codes.numel() and not 0 <= int(codes.min()) <= int(codes.max()) < self.audio_eos_code:
|
| 444 |
+
raise ValueError("The audio overflows the codec; expected `audio_out` samples in [-1, 1]")
|
| 445 |
+
return codes
|
| 446 |
+
|
| 447 |
+
def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwargs):
|
| 448 |
+
super().validate_inputs(images=images, text=text, videos=videos, audio=audio, **kwargs)
|
| 449 |
+
if text is None:
|
| 450 |
+
return
|
| 451 |
+
prompts = [text] if isinstance(text, str) else text
|
| 452 |
+
num_placeholders = sum(prompt.count(self.audio_token) for prompt in prompts)
|
| 453 |
+
num_clips = 0 if audio is None else len(make_list_of_audio(audio))
|
| 454 |
+
if num_placeholders != num_clips:
|
| 455 |
+
raise ValueError(
|
| 456 |
+
f"Found {num_placeholders} {self.audio_token} tokens in the text but received {num_clips} clips"
|
| 457 |
+
)
|
| 458 |
+
if text is not None:
|
| 459 |
+
_check_vision_placeholder_counts(self, [text] if isinstance(text, str) else text, images, videos)
|
| 460 |
+
|
| 461 |
+
def _process_audio(self, audio: AudioInput, **kwargs):
|
| 462 |
+
# the backend only serves `apply_chat_template`'s audio loading
|
| 463 |
+
kwargs.pop("load_audio_backend", None)
|
| 464 |
+
audio_inputs = self.feature_extractor(audio, **kwargs)
|
| 465 |
+
lengths = np.asarray(audio_inputs["feature_attention_mask"]).sum(-1).tolist()
|
| 466 |
+
audio_inputs["num_audio_tokens"] = [
|
| 467 |
+
self.feature_extractor.num_encoder_tokens(int(length)) for length in lengths
|
| 468 |
+
]
|
| 469 |
+
audio_replacements = [self.replace_audio_token(audio_inputs, audio_idx=idx) for idx in range(len(lengths))]
|
| 470 |
+
return audio_inputs, audio_replacements
|
| 471 |
+
|
| 472 |
+
def replace_audio_token(self, audio_inputs: dict, audio_idx: int, **kwargs) -> str:
|
| 473 |
+
return self._wrap(self.audio_token * audio_inputs["num_audio_tokens"][audio_idx], self.audio_markers)
|
| 474 |
+
|
| 475 |
+
@property
|
| 476 |
+
def unused_input_names(self) -> list[str]:
|
| 477 |
+
return ["num_audio_tokens"]
|
| 478 |
+
|
| 479 |
+
@property
|
| 480 |
+
def model_input_names(self):
|
| 481 |
+
return list(
|
| 482 |
+
dict.fromkeys(
|
| 483 |
+
self.tokenizer.model_input_names
|
| 484 |
+
+ self.feature_extractor.model_input_names
|
| 485 |
+
+ self.image_processor.model_input_names
|
| 486 |
+
+ (self.video_processor.model_input_names if self.video_processor is not None else [])
|
| 487 |
+
+ ["audio_out_codes"]
|
| 488 |
+
)
|
| 489 |
+
)
|
| 490 |
+
|
| 491 |
+
def _wrap(self, placeholders: str, markers: bool) -> str:
|
| 492 |
+
return self.vision_start_token + placeholders + self.vision_end_token if markers else placeholders
|
| 493 |
+
|
| 494 |
+
def _strip_spans(self, prompt: str) -> tuple[str, list[dict]]:
|
| 495 |
+
# a chat template may already wrap a placeholder in the span pair; the expansion adds exactly one pair
|
| 496 |
+
wrapped = {
|
| 497 |
+
self.vision_start_token + token + self.vision_end_token: token
|
| 498 |
+
for token in (self.image_token, self.video_token, self.audio_token)
|
| 499 |
+
}
|
| 500 |
+
return self._rewrite_with_offsets(
|
| 501 |
+
prompt, "|".join(map(re.escape, wrapped)), lambda match: wrapped[match.group()], "span"
|
| 502 |
+
)
|
| 503 |
+
|
| 504 |
+
def apply_chat_template(
|
| 505 |
+
self,
|
| 506 |
+
conversation,
|
| 507 |
+
chat_template=None,
|
| 508 |
+
tools=None,
|
| 509 |
+
documents=None,
|
| 510 |
+
add_generation_prompt=False,
|
| 511 |
+
continue_final_message=False,
|
| 512 |
+
return_assistant_tokens_mask=False,
|
| 513 |
+
tokenize=False,
|
| 514 |
+
return_tensors=None,
|
| 515 |
+
return_dict=False,
|
| 516 |
+
load_audio_from_video=False,
|
| 517 |
+
processor_kwargs=None,
|
| 518 |
+
**kwargs,
|
| 519 |
+
):
|
| 520 |
+
"""With `load_audio_from_video=True`, each video's soundtrack becomes the audio part right behind it, which
|
| 521 |
+
`use_audio_in_video` pairs with the video; standalone clips are kept (the shared implementation drops them)."""
|
| 522 |
+
if load_audio_from_video and tokenize:
|
| 523 |
+
load_audio_from_video = False
|
| 524 |
+
conversation = self._with_soundtrack_parts(conversation)
|
| 525 |
+
return super().apply_chat_template(
|
| 526 |
+
conversation,
|
| 527 |
+
chat_template,
|
| 528 |
+
tools=tools,
|
| 529 |
+
documents=documents,
|
| 530 |
+
add_generation_prompt=add_generation_prompt,
|
| 531 |
+
continue_final_message=continue_final_message,
|
| 532 |
+
return_assistant_tokens_mask=return_assistant_tokens_mask,
|
| 533 |
+
tokenize=tokenize,
|
| 534 |
+
return_tensors=return_tensors,
|
| 535 |
+
return_dict=return_dict,
|
| 536 |
+
load_audio_from_video=load_audio_from_video,
|
| 537 |
+
processor_kwargs=processor_kwargs,
|
| 538 |
+
**kwargs,
|
| 539 |
+
)
|
| 540 |
+
|
| 541 |
+
@staticmethod
|
| 542 |
+
def _with_soundtrack_parts(conversation):
|
| 543 |
+
if not conversation:
|
| 544 |
+
return conversation
|
| 545 |
+
batched = isinstance(conversation[0], (list, tuple))
|
| 546 |
+
rewritten = []
|
| 547 |
+
for messages in conversation if batched else [conversation]:
|
| 548 |
+
rewritten.append([])
|
| 549 |
+
for message in messages:
|
| 550 |
+
content = message.get("content")
|
| 551 |
+
if isinstance(content, list):
|
| 552 |
+
parts = []
|
| 553 |
+
for part in content:
|
| 554 |
+
parts.append(part)
|
| 555 |
+
source = next((part[key] for key in ("video", "url", "path") if key in part), None)
|
| 556 |
+
if part.get("type") == "video" and source is not None:
|
| 557 |
+
parts.append({"type": "audio", "audio": source})
|
| 558 |
+
message = {**message, "content": parts}
|
| 559 |
+
rewritten[-1].append(message)
|
| 560 |
+
return rewritten if batched else rewritten[0]
|
| 561 |
+
|
| 562 |
+
def _attach_soundtracks(self, prompt: str) -> tuple[str, list[dict]]:
|
| 563 |
+
"""`use_audio_in_video`: each video takes the first audio placeholder after it as its soundtrack, moved right
|
| 564 |
+
behind the video; returns the prompt and the offset records of the moves."""
|
| 565 |
+
row, taken = [], set()
|
| 566 |
+
audio = [match.start() for match in re.finditer(re.escape(self.audio_token), prompt)]
|
| 567 |
+
for video in re.finditer(re.escape(self.video_token), prompt):
|
| 568 |
+
position = next((p for p in audio if p >= video.end() and p not in taken), None)
|
| 569 |
+
if position is None:
|
| 570 |
+
raise ValueError("`use_audio_in_video` needs an audio clip after each video")
|
| 571 |
+
taken.add(position)
|
| 572 |
+
row.append(
|
| 573 |
+
{"type": "audio", "span": (video.end(), video.end()), "text": "", "replacement": self.audio_token}
|
| 574 |
+
)
|
| 575 |
+
row.append(
|
| 576 |
+
{
|
| 577 |
+
"type": "audio",
|
| 578 |
+
"span": (position, position + len(self.audio_token)),
|
| 579 |
+
"text": self.audio_token,
|
| 580 |
+
"replacement": "",
|
| 581 |
+
}
|
| 582 |
+
)
|
| 583 |
+
row.sort(key=lambda r: r["span"])
|
| 584 |
+
pieces, last, offset = [], 0, 0
|
| 585 |
+
for entry in row:
|
| 586 |
+
start, end = entry["span"]
|
| 587 |
+
pieces.extend((prompt[last:start], entry["replacement"]))
|
| 588 |
+
entry["new_span"] = (start + offset, start + offset + len(entry["replacement"]))
|
| 589 |
+
offset += len(entry["replacement"]) - (end - start)
|
| 590 |
+
last = end
|
| 591 |
+
pieces.append(prompt[last:])
|
| 592 |
+
return "".join(pieces), row
|
| 593 |
+
|
| 594 |
+
def _call_with_videos(self, prompts, audio, images, videos, single, **kwargs):
|
| 595 |
+
"""Expands the placeholders itself: a soundtrack's encoder tokens are split across its video's frame groups,
|
| 596 |
+
which the generic one-clip-one-run expansion cannot do."""
|
| 597 |
+
images, prompts, videos, audio = self.prepare_inputs_layout(
|
| 598 |
+
images=images, text=prompts, videos=videos, audio=audio, **kwargs
|
| 599 |
+
)
|
| 600 |
+
self.validate_inputs(images=images, text=prompts, videos=videos, audio=audio, **kwargs)
|
| 601 |
+
merged = self._merge_kwargs(
|
| 602 |
+
self.valid_processor_kwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs
|
| 603 |
+
)
|
| 604 |
+
return_tensors = merged["text_kwargs"].pop("return_tensors", None)
|
| 605 |
+
return_mm_token_type_ids = merged["text_kwargs"].pop("return_mm_token_type_ids", False)
|
| 606 |
+
return_text_replacement_offsets = merged["text_kwargs"].pop("return_text_replacement_offsets", False)
|
| 607 |
+
# Timestamps always need metadata internally, independently of whether the caller wants it returned.
|
| 608 |
+
return_metadata = kwargs.get("return_metadata", kwargs.get("videos_kwargs", {}).get("return_metadata", False))
|
| 609 |
+
merged["videos_kwargs"]["return_metadata"] = True
|
| 610 |
+
use_audio_in_video = merged["videos_kwargs"].pop("use_audio_in_video")
|
| 611 |
+
unattached, attach_offsets = prompts, [[] for _ in prompts]
|
| 612 |
+
if use_audio_in_video:
|
| 613 |
+
attached = [self._attach_soundtracks(prompt) for prompt in prompts]
|
| 614 |
+
prompts, attach_offsets = [prompt for prompt, _ in attached], [row for _, row in attached]
|
| 615 |
+
outputs = {}
|
| 616 |
+
clips = make_list_of_audio(audio) if audio is not None else []
|
| 617 |
+
videos = make_batched_videos(videos)
|
| 618 |
+
self._check_video_clock(videos, merged["videos_kwargs"].get("video_metadata"))
|
| 619 |
+
image_inputs = self.image_processor(images=images, **merged["images_kwargs"]) if images is not None else None
|
| 620 |
+
if image_inputs is not None:
|
| 621 |
+
outputs.update(image_inputs)
|
| 622 |
+
if clips:
|
| 623 |
+
# the decoding backend is a loader option, not a feature-extractor one (the shared path drops it too)
|
| 624 |
+
merged["audio_kwargs"].pop("load_audio_backend", None)
|
| 625 |
+
audio_inputs = self.feature_extractor(clips, **merged["audio_kwargs"])
|
| 626 |
+
lengths = np.asarray(audio_inputs["feature_attention_mask"]).sum(-1).tolist()
|
| 627 |
+
audio_tokens = [self.feature_extractor.num_encoder_tokens(int(length)) for length in lengths]
|
| 628 |
+
outputs.update(audio_inputs)
|
| 629 |
+
video_inputs = self.video_processor(videos=videos, **merged["videos_kwargs"])
|
| 630 |
+
outputs.update(
|
| 631 |
+
{key: value for key, value in video_inputs.items() if key != "video_metadata" or return_metadata}
|
| 632 |
+
)
|
| 633 |
+
video = self.video_token + (self.audio_token if use_audio_in_video else "")
|
| 634 |
+
pattern = re.compile("|".join(map(re.escape, (video, self.image_token, self.audio_token))))
|
| 635 |
+
counters = {"image": 0, "video": 0, "audio": 0}
|
| 636 |
+
|
| 637 |
+
def expand(match):
|
| 638 |
+
token = match.group(0)
|
| 639 |
+
if token == self.image_token:
|
| 640 |
+
text = self.replace_image_token(image_inputs, counters["image"])
|
| 641 |
+
counters["image"] += 1
|
| 642 |
+
return text
|
| 643 |
+
if token == self.audio_token:
|
| 644 |
+
text = self._wrap(self.audio_token * audio_tokens[counters["audio"]], self.audio_markers)
|
| 645 |
+
counters["audio"] += 1
|
| 646 |
+
return text
|
| 647 |
+
soundtrack = None
|
| 648 |
+
if token != self.video_token: # the pair `_attach_soundtracks` made
|
| 649 |
+
soundtrack = audio_tokens[counters["audio"]]
|
| 650 |
+
counters["audio"] += 1
|
| 651 |
+
text = self._video_expansion(video_inputs, counters["video"], soundtrack)
|
| 652 |
+
counters["video"] += 1
|
| 653 |
+
return text
|
| 654 |
+
|
| 655 |
+
expanded, batch_offsets = [], []
|
| 656 |
+
for prompt in prompts:
|
| 657 |
+
pieces, offsets = [], []
|
| 658 |
+
last, offset = 0, 0
|
| 659 |
+
for match in pattern.finditer(prompt):
|
| 660 |
+
start, end = match.span()
|
| 661 |
+
replacement = expand(match)
|
| 662 |
+
pieces.extend((prompt[last:start], replacement))
|
| 663 |
+
token = match.group(0)
|
| 664 |
+
offsets.append(
|
| 665 |
+
{
|
| 666 |
+
"type": "video"
|
| 667 |
+
if token.startswith(self.video_token)
|
| 668 |
+
else "image"
|
| 669 |
+
if token == self.image_token
|
| 670 |
+
else "audio",
|
| 671 |
+
"span": (start, end),
|
| 672 |
+
"new_span": (start + offset, start + offset + len(replacement)),
|
| 673 |
+
"text": token,
|
| 674 |
+
"replacement": replacement,
|
| 675 |
+
}
|
| 676 |
+
)
|
| 677 |
+
offset += len(replacement) - (end - start)
|
| 678 |
+
last = end
|
| 679 |
+
pieces.append(prompt[last:])
|
| 680 |
+
expanded.append("".join(pieces))
|
| 681 |
+
batch_offsets.append(offsets)
|
| 682 |
+
text_inputs = self.tokenizer(expanded, **merged["text_kwargs"])
|
| 683 |
+
self._check_special_mm_tokens(expanded, text_inputs, modalities=["image", "video", "audio"])
|
| 684 |
+
if return_text_replacement_offsets:
|
| 685 |
+
text_inputs["text_replacement_offsets"] = self._compose_replacement_offsets(
|
| 686 |
+
unattached, attach_offsets, batch_offsets
|
| 687 |
+
)
|
| 688 |
+
if return_mm_token_type_ids:
|
| 689 |
+
text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"])
|
| 690 |
+
outputs.update(text_inputs)
|
| 691 |
+
return BatchFeature(outputs, tensor_type=return_tensors, skip_tensor_conversion=self.skip_tensor_conversion)
|
| 692 |
+
|
| 693 |
+
def _video_expansion(self, video_inputs: dict, video_idx: int, soundtrack_tokens: int | None = None) -> str:
|
| 694 |
+
grid = video_inputs["video_grid_thw"][video_idx]
|
| 695 |
+
patch_tokens = int(grid[1:].prod() // self.video_processor.merge_size**2)
|
| 696 |
+
metadata = video_inputs["video_metadata"][video_idx]
|
| 697 |
+
if metadata.frames_indices is None or not (
|
| 698 |
+
metadata.fps is not None and math.isfinite(metadata.fps) and metadata.fps > 0
|
| 699 |
+
):
|
| 700 |
+
raise ValueError(
|
| 701 |
+
"LFM3-Omni prompts carry frame timestamps: pass videos with `video_metadata` holding the sampled frame "
|
| 702 |
+
"indices and a positive fps"
|
| 703 |
+
)
|
| 704 |
+
# a released transformers reads torchcodec's scanned average fps, which carries float noise
|
| 705 |
+
# (30.000000000000004) where the fork reads the container header
|
| 706 |
+
metadata.fps = round(metadata.fps, 6)
|
| 707 |
+
temporal = self.video_processor.temporal_patch_size
|
| 708 |
+
timestamps = self._calculate_timestamps(metadata.frames_indices, metadata.fps, temporal)[: int(grid[0])]
|
| 709 |
+
bounds = None
|
| 710 |
+
if soundtrack_tokens is not None:
|
| 711 |
+
# the soundtrack's encoder tokens are split at each frame group's first frame
|
| 712 |
+
starts = [index / metadata.fps for index in list(metadata.frames_indices)[::temporal]][: int(grid[0])]
|
| 713 |
+
seconds_per_token = (
|
| 714 |
+
self.feature_extractor.hop_length
|
| 715 |
+
* self.feature_extractor.encoder_stride
|
| 716 |
+
/ self.feature_extractor.sampling_rate
|
| 717 |
+
)
|
| 718 |
+
bounds = (
|
| 719 |
+
[0]
|
| 720 |
+
+ [min(soundtrack_tokens, round(start / seconds_per_token)) for start in starts[1:]]
|
| 721 |
+
+ [soundtrack_tokens]
|
| 722 |
+
)
|
| 723 |
+
bounds = [max(bounds[: i + 1]) for i in range(len(bounds))]
|
| 724 |
+
groups = ""
|
| 725 |
+
for index, timestamp in enumerate(timestamps):
|
| 726 |
+
groups += f"<{timestamp:.1f} seconds>" + self._wrap(self.video_token * patch_tokens, self.vision_markers)
|
| 727 |
+
if bounds is not None:
|
| 728 |
+
groups += self.audio_token * (bounds[index + 1] - bounds[index])
|
| 729 |
+
return self._wrap(groups, self.vision_markers)
|
| 730 |
+
|
| 731 |
+
def _check_video_clock(self, videos, video_metadata) -> None:
|
| 732 |
+
_check_video_metadata_count(videos, video_metadata)
|
| 733 |
+
if video_metadata is None:
|
| 734 |
+
if any(not isinstance(video, str) for video in make_batched_videos(videos)):
|
| 735 |
+
# decoded frames carry no clock: without metadata the timestamps would silently assume 24 fps
|
| 736 |
+
raise ValueError(
|
| 737 |
+
"LFM3-Omni prompts carry frame timestamps: pass decoded frames with `video_metadata` holding the "
|
| 738 |
+
"clip's fps and total_num_frames (a path or URL is decoded with its own metadata)"
|
| 739 |
+
)
|
| 740 |
+
# the shared sampler writes a 24 fps default into metadata lacking a clock before the expansion runs
|
| 741 |
+
elif any(metadata.fps is None for metadata in make_batched_metadata(videos, video_metadata)):
|
| 742 |
+
raise ValueError(
|
| 743 |
+
"LFM3-Omni prompts carry frame timestamps: `video_metadata` must hold the clip's fps "
|
| 744 |
+
"(a path or URL is decoded with its own metadata)"
|
| 745 |
+
)
|
| 746 |
+
|
| 747 |
+
def _process_videos(self, videos, **kwargs):
|
| 748 |
+
# only media-only calls reach the shared path (prompts with videos take `_call_with_videos`): the pairing
|
| 749 |
+
# flag is not a video-processor argument and no placeholder text is needed
|
| 750 |
+
kwargs.pop("use_audio_in_video", None)
|
| 751 |
+
self._check_video_clock(videos, kwargs.get("video_metadata"))
|
| 752 |
+
return self.video_processor(videos, **kwargs), []
|
| 753 |
+
|
| 754 |
+
def replace_image_token(self, image_inputs, image_idx, **kwargs):
|
| 755 |
+
count = image_inputs["image_grid_thw"][image_idx].prod() // self.image_processor.merge_size**2
|
| 756 |
+
return self._wrap(self.image_token * int(count), self.vision_markers)
|
| 757 |
+
|
| 758 |
+
@staticmethod
|
| 759 |
+
def _calculate_timestamps(indices, video_fps: float, temporal_patch_size: int) -> list[float]:
|
| 760 |
+
indices = list(indices)
|
| 761 |
+
if pad := -len(indices) % temporal_patch_size:
|
| 762 |
+
indices.extend([indices[-1]] * pad)
|
| 763 |
+
return [
|
| 764 |
+
(indices[i] + indices[i + temporal_patch_size - 1]) / (2 * video_fps)
|
| 765 |
+
for i in range(0, len(indices), temporal_patch_size)
|
| 766 |
+
]
|
| 767 |
+
|
| 768 |
+
|
| 769 |
+
__all__ = ["Lfm3OmniProcessor"]
|
processing_lfm3_vl.py
ADDED
|
@@ -0,0 +1,248 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
import math
|
| 15 |
+
|
| 16 |
+
import numpy as np
|
| 17 |
+
|
| 18 |
+
from transformers.image_utils import is_pil_image, is_valid_image, make_flat_list_of_images
|
| 19 |
+
from transformers.processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin
|
| 20 |
+
from transformers.utils import auto_docstring, is_numpy_array, is_torch_tensor, logging
|
| 21 |
+
from transformers.video_utils import VideoMetadata, is_valid_video, make_batched_videos
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
logger = logging.get_logger(__name__)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
class Lfm3VlProcessorKwargs(ProcessingKwargs, total=False):
|
| 28 |
+
_defaults = {
|
| 29 |
+
"text_kwargs": {
|
| 30 |
+
"padding": False,
|
| 31 |
+
"return_token_type_ids": False,
|
| 32 |
+
"return_mm_token_type_ids": False,
|
| 33 |
+
},
|
| 34 |
+
"videos_kwargs": {"return_metadata": True},
|
| 35 |
+
}
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def _check_video_fps(metadata: VideoMetadata) -> None:
|
| 39 |
+
if metadata.fps is not None and not (math.isfinite(metadata.fps) and metadata.fps > 0):
|
| 40 |
+
raise ValueError(f"`video_metadata.fps` must be a positive number, got {metadata.fps}")
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def _media_per_prompt(media, is_item) -> list[int] | None:
|
| 44 |
+
# rows name the prompt that owns them while the expansion consumes media in flat order; a row is a list of whole
|
| 45 |
+
# items or one stacked batch, and a list of single frames is one video, not a row
|
| 46 |
+
if not isinstance(media, (list, tuple)):
|
| 47 |
+
return None
|
| 48 |
+
counts = []
|
| 49 |
+
for row in media:
|
| 50 |
+
if isinstance(row, (list, tuple)) and all(is_item(item) for item in row):
|
| 51 |
+
counts.append(len(row))
|
| 52 |
+
elif (is_numpy_array(row) or is_torch_tensor(row)) and row.ndim == 4 and (len(row) == 0 or is_item(row[0])):
|
| 53 |
+
counts.append(len(row))
|
| 54 |
+
else:
|
| 55 |
+
return None
|
| 56 |
+
return counts
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def _check_vision_placeholder_counts(processor, prompts, images, videos) -> None:
|
| 60 |
+
for name, token, media, flatten, is_item in (
|
| 61 |
+
(
|
| 62 |
+
"images",
|
| 63 |
+
processor.image_token,
|
| 64 |
+
images,
|
| 65 |
+
make_flat_list_of_images,
|
| 66 |
+
lambda item: is_valid_image(item) and (is_pil_image(item) or item.ndim == 3),
|
| 67 |
+
),
|
| 68 |
+
(
|
| 69 |
+
"videos",
|
| 70 |
+
processor.video_token,
|
| 71 |
+
videos,
|
| 72 |
+
make_batched_videos,
|
| 73 |
+
lambda item: isinstance(item, str) or (not is_pil_image(item) and is_valid_video(item)),
|
| 74 |
+
),
|
| 75 |
+
):
|
| 76 |
+
num_placeholders = sum(prompt.count(token) for prompt in prompts)
|
| 77 |
+
num_media = 0 if media is None else len(flatten(media))
|
| 78 |
+
if num_placeholders != num_media:
|
| 79 |
+
raise ValueError(f"Found {num_placeholders} {token} tokens in the text but received {num_media} {name}")
|
| 80 |
+
# one prompt cannot take another's media
|
| 81 |
+
per_prompt = _media_per_prompt(media, is_item) if len(prompts) > 1 else None
|
| 82 |
+
in_text = [prompt.count(token) for prompt in prompts]
|
| 83 |
+
if per_prompt is not None and per_prompt != in_text:
|
| 84 |
+
raise ValueError(f"Found {in_text} {token} tokens per prompt but received {per_prompt} {name} per prompt")
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
@auto_docstring
|
| 88 |
+
class Lfm3VlProcessor(ProcessorMixin):
|
| 89 |
+
valid_processor_kwargs = Lfm3VlProcessorKwargs
|
| 90 |
+
|
| 91 |
+
def __init__(self, image_processor=None, tokenizer=None, video_processor=None, chat_template=None, **kwargs):
|
| 92 |
+
self.image_token = getattr(tokenizer, "image_token", "<|image|>")
|
| 93 |
+
self.video_token = getattr(tokenizer, "video_token", "<|video|>")
|
| 94 |
+
self.vision_start_token = getattr(tokenizer, "vision_start_token", "<|vision_start|>")
|
| 95 |
+
self.vision_end_token = getattr(tokenizer, "vision_end_token", "<|vision_end|>")
|
| 96 |
+
self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token)
|
| 97 |
+
self.video_token_id = tokenizer.convert_tokens_to_ids(self.video_token)
|
| 98 |
+
self.vision_start_token_id = tokenizer.convert_tokens_to_ids(self.vision_start_token)
|
| 99 |
+
self.vision_end_token_id = tokenizer.convert_tokens_to_ids(self.vision_end_token)
|
| 100 |
+
super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template)
|
| 101 |
+
|
| 102 |
+
def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> str:
|
| 103 |
+
merge_length = self.image_processor.merge_size**2
|
| 104 |
+
num_image_tokens = image_inputs["image_grid_thw"][image_idx].prod() // merge_length
|
| 105 |
+
return self.image_token * num_image_tokens
|
| 106 |
+
|
| 107 |
+
def replace_video_token(self, video_inputs: dict, video_idx: int, **kwargs) -> str:
|
| 108 |
+
metadata = video_inputs["video_metadata"][video_idx]
|
| 109 |
+
if metadata.frames_indices is None:
|
| 110 |
+
raise ValueError("Frame timestamps need the sampled frame indices in `video_metadata`")
|
| 111 |
+
_check_video_fps(metadata)
|
| 112 |
+
merge_length = self.video_processor.merge_size**2
|
| 113 |
+
num_frames = video_inputs["video_grid_thw"][video_idx][0]
|
| 114 |
+
frame_seqlen = video_inputs["video_grid_thw"][video_idx][1:].prod() // merge_length
|
| 115 |
+
metadata = video_inputs["video_metadata"][video_idx]
|
| 116 |
+
video_placeholder = ""
|
| 117 |
+
|
| 118 |
+
if metadata.fps is None:
|
| 119 |
+
logger.warning_once(
|
| 120 |
+
"Lfm3Vl requires frame timestamps to construct prompts, but the `fps` of the input video could not be inferred. "
|
| 121 |
+
"Probably `video_metadata` was missing from inputs and you passed pre-sampled frames. "
|
| 122 |
+
"Defaulting to `fps=24`. Please provide `video_metadata` for more accurate results."
|
| 123 |
+
)
|
| 124 |
+
# a released transformers reads torchcodec's scanned average fps, which carries float noise
|
| 125 |
+
# (30.000000000000004) where the fork reads the container header
|
| 126 |
+
metadata.fps = 24 if metadata.fps is None else round(metadata.fps, 6)
|
| 127 |
+
|
| 128 |
+
# if timestamps are not provided, calculate them
|
| 129 |
+
curr_timestamp = self._calculate_timestamps(
|
| 130 |
+
metadata.frames_indices,
|
| 131 |
+
metadata.fps,
|
| 132 |
+
self.video_processor.temporal_patch_size,
|
| 133 |
+
)
|
| 134 |
+
|
| 135 |
+
for frame_idx in range(num_frames):
|
| 136 |
+
curr_time = curr_timestamp[frame_idx]
|
| 137 |
+
video_placeholder += f"<{curr_time:.1f} seconds>"
|
| 138 |
+
video_placeholder += self.vision_start_token + self.video_token * frame_seqlen + self.vision_end_token
|
| 139 |
+
return video_placeholder
|
| 140 |
+
|
| 141 |
+
def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs):
|
| 142 |
+
"""
|
| 143 |
+
Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
|
| 144 |
+
|
| 145 |
+
Args:
|
| 146 |
+
image_sizes (`list[list[int]]`, *optional*):
|
| 147 |
+
The input sizes formatted as (height, width) per each image.
|
| 148 |
+
video_sizes (`list[list[int]]`, *optional*):
|
| 149 |
+
The input sizes formatted as (num_frames, height, width) per each video.
|
| 150 |
+
|
| 151 |
+
Returns:
|
| 152 |
+
`MultiModalData`: the number of tokens per provided modality.
|
| 153 |
+
"""
|
| 154 |
+
vision_data = {}
|
| 155 |
+
if image_sizes is not None:
|
| 156 |
+
# copy: `update` on the class-level defaults would leak kwargs into every later call
|
| 157 |
+
images_kwargs = dict(Lfm3VlProcessorKwargs._defaults.get("images_kwargs", {}))
|
| 158 |
+
images_kwargs.update(kwargs)
|
| 159 |
+
merge_size = images_kwargs.get("merge_size") or self.image_processor.merge_size
|
| 160 |
+
|
| 161 |
+
num_image_patches = [
|
| 162 |
+
self.image_processor.get_number_of_image_patches(*image_size, images_kwargs)
|
| 163 |
+
for image_size in image_sizes
|
| 164 |
+
]
|
| 165 |
+
num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches]
|
| 166 |
+
vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
|
| 167 |
+
|
| 168 |
+
if video_sizes is not None:
|
| 169 |
+
videos_kwargs = dict(Lfm3VlProcessorKwargs._defaults.get("videos_kwargs", {}))
|
| 170 |
+
videos_kwargs.update(kwargs)
|
| 171 |
+
merge_size = videos_kwargs.get("merge_size") or self.video_processor.merge_size
|
| 172 |
+
num_video_patches = [
|
| 173 |
+
self.video_processor.get_num_of_video_patches(*video_size, videos_kwargs) for video_size in video_sizes
|
| 174 |
+
]
|
| 175 |
+
num_video_tokens = [(num_patches // merge_size**2) for num_patches in num_video_patches]
|
| 176 |
+
vision_data["num_video_tokens"] = num_video_tokens
|
| 177 |
+
|
| 178 |
+
return MultiModalData(**vision_data)
|
| 179 |
+
|
| 180 |
+
def post_process_image_text_to_text(
|
| 181 |
+
self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs
|
| 182 |
+
):
|
| 183 |
+
"""
|
| 184 |
+
Post-process the output of the model to decode the text.
|
| 185 |
+
|
| 186 |
+
Args:
|
| 187 |
+
generated_outputs (`torch.Tensor` or `np.ndarray`):
|
| 188 |
+
The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
|
| 189 |
+
or `(sequence_length,)`.
|
| 190 |
+
skip_special_tokens (`bool`, *optional*, defaults to `True`):
|
| 191 |
+
Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
|
| 192 |
+
clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
|
| 193 |
+
Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.
|
| 194 |
+
**kwargs:
|
| 195 |
+
Additional arguments to be passed to the tokenizer's `batch_decode method`.
|
| 196 |
+
|
| 197 |
+
Returns:
|
| 198 |
+
`list[str]`: The decoded text.
|
| 199 |
+
"""
|
| 200 |
+
return self.tokenizer.batch_decode(
|
| 201 |
+
generated_outputs,
|
| 202 |
+
skip_special_tokens=skip_special_tokens,
|
| 203 |
+
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
|
| 204 |
+
**kwargs,
|
| 205 |
+
)
|
| 206 |
+
|
| 207 |
+
@property
|
| 208 |
+
def model_input_names(self):
|
| 209 |
+
return list(
|
| 210 |
+
dict.fromkeys(
|
| 211 |
+
self.tokenizer.model_input_names
|
| 212 |
+
+ self.image_processor.model_input_names
|
| 213 |
+
+ self.video_processor.model_input_names
|
| 214 |
+
)
|
| 215 |
+
)
|
| 216 |
+
|
| 217 |
+
def _calculate_timestamps(self, indices: list[int] | np.ndarray, video_fps: float, merge_size: int = 2):
|
| 218 |
+
if not isinstance(indices, list):
|
| 219 |
+
indices = indices.tolist()
|
| 220 |
+
if len(indices) % merge_size != 0:
|
| 221 |
+
indices.extend(indices[-1] for _ in range(merge_size - len(indices) % merge_size))
|
| 222 |
+
timestamps = [idx / video_fps for idx in indices]
|
| 223 |
+
# @JJJYmmm frames are merged by self.merge_size, \
|
| 224 |
+
# so we need to average the timestamps between the first/last frame within the temporal patch
|
| 225 |
+
timestamps = [
|
| 226 |
+
(timestamps[i] + timestamps[i + merge_size - 1]) / 2 for i in range(0, len(timestamps), merge_size)
|
| 227 |
+
]
|
| 228 |
+
return timestamps
|
| 229 |
+
|
| 230 |
+
def prepare_inputs_layout(self, images=None, text=None, videos=None, audio=None, **kwargs):
|
| 231 |
+
for name, media in (("images", images), ("videos", videos)):
|
| 232 |
+
if isinstance(media, (list, tuple)) and len(media) == 0:
|
| 233 |
+
raise ValueError(f"`{name}` must hold at least one item; pass `{name}=None` for prompts without one")
|
| 234 |
+
return super().prepare_inputs_layout(images=images, text=text, videos=videos, audio=audio, **kwargs)
|
| 235 |
+
|
| 236 |
+
def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwargs):
|
| 237 |
+
super().validate_inputs(images=images, text=text, videos=videos, audio=audio, **kwargs)
|
| 238 |
+
if text is not None:
|
| 239 |
+
_check_vision_placeholder_counts(self, [text] if isinstance(text, str) else text, images, videos)
|
| 240 |
+
|
| 241 |
+
def _process_videos(self, videos, **kwargs):
|
| 242 |
+
# the frame timestamps need the metadata whatever the caller's `return_metadata`; `__call__` drops it from
|
| 243 |
+
# the output unless it was asked for
|
| 244 |
+
kwargs["return_metadata"] = True
|
| 245 |
+
return super()._process_videos(videos, **kwargs)
|
| 246 |
+
|
| 247 |
+
|
| 248 |
+
__all__ = ["Lfm3VlProcessor"]
|
processor_config.json
ADDED
|
@@ -0,0 +1,98 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"audio_markers": true,
|
| 3 |
+
"audio_out_token": "<|audio_out|>",
|
| 4 |
+
"audio_token": "<|audio|>",
|
| 5 |
+
"audio_tokenizer_subfolder": "audio_tokenizer",
|
| 6 |
+
"feature_extractor": {
|
| 7 |
+
"encoder_stride": 8,
|
| 8 |
+
"feature_extractor_type": "Lfm3OmniFeatureExtractor",
|
| 9 |
+
"feature_size": 128,
|
| 10 |
+
"hop_length": 160,
|
| 11 |
+
"log_floor": -12.0,
|
| 12 |
+
"n_fft": 400,
|
| 13 |
+
"normalization_divisor": 4.0,
|
| 14 |
+
"normalization_offset": 4.0,
|
| 15 |
+
"padding_side": "right",
|
| 16 |
+
"padding_value": 0.0,
|
| 17 |
+
"return_attention_mask": true,
|
| 18 |
+
"sampling_rate": 16000,
|
| 19 |
+
"win_length": 400,
|
| 20 |
+
"auto_map": {
|
| 21 |
+
"AutoFeatureExtractor": "feature_extraction_lfm3_omni.Lfm3OmniFeatureExtractor"
|
| 22 |
+
}
|
| 23 |
+
},
|
| 24 |
+
"image_processor": {
|
| 25 |
+
"do_convert_rgb": true,
|
| 26 |
+
"do_normalize": true,
|
| 27 |
+
"do_rescale": true,
|
| 28 |
+
"do_resize": true,
|
| 29 |
+
"image_mean": [
|
| 30 |
+
0.5,
|
| 31 |
+
0.5,
|
| 32 |
+
0.5
|
| 33 |
+
],
|
| 34 |
+
"image_processor_type": "Lfm3VlImageProcessor",
|
| 35 |
+
"image_std": [
|
| 36 |
+
0.5,
|
| 37 |
+
0.5,
|
| 38 |
+
0.5
|
| 39 |
+
],
|
| 40 |
+
"merge_size": 2,
|
| 41 |
+
"patch_size": 16,
|
| 42 |
+
"resample": 3,
|
| 43 |
+
"rescale_factor": 0.00392156862745098,
|
| 44 |
+
"size": {
|
| 45 |
+
"longest_edge": 2097152,
|
| 46 |
+
"shortest_edge": 65536
|
| 47 |
+
},
|
| 48 |
+
"temporal_patch_size": 2,
|
| 49 |
+
"auto_map": {
|
| 50 |
+
"AutoImageProcessor": "image_processing_lfm3_vl.Lfm3VlImageProcessor"
|
| 51 |
+
}
|
| 52 |
+
},
|
| 53 |
+
"image_token": "<|image|>",
|
| 54 |
+
"processor_class": "Lfm3OmniProcessor",
|
| 55 |
+
"video_processor": {
|
| 56 |
+
"cap_pixels_per_frame": false,
|
| 57 |
+
"do_convert_rgb": true,
|
| 58 |
+
"do_normalize": true,
|
| 59 |
+
"do_rescale": true,
|
| 60 |
+
"do_resize": true,
|
| 61 |
+
"do_sample_frames": true,
|
| 62 |
+
"fps": 2.0,
|
| 63 |
+
"image_mean": [
|
| 64 |
+
0.5,
|
| 65 |
+
0.5,
|
| 66 |
+
0.5
|
| 67 |
+
],
|
| 68 |
+
"image_std": [
|
| 69 |
+
0.5,
|
| 70 |
+
0.5,
|
| 71 |
+
0.5
|
| 72 |
+
],
|
| 73 |
+
"max_frames": 384,
|
| 74 |
+
"max_video_tokens": 768,
|
| 75 |
+
"merge_size": 2,
|
| 76 |
+
"min_frames": 4,
|
| 77 |
+
"patch_size": 16,
|
| 78 |
+
"resample": 3,
|
| 79 |
+
"rescale_factor": 0.00392156862745098,
|
| 80 |
+
"return_metadata": false,
|
| 81 |
+
"size": {
|
| 82 |
+
"longest_edge": 25165824,
|
| 83 |
+
"shortest_edge": 4096
|
| 84 |
+
},
|
| 85 |
+
"temporal_patch_size": 2,
|
| 86 |
+
"video_processor_type": "Lfm3VlVideoProcessor",
|
| 87 |
+
"auto_map": {
|
| 88 |
+
"AutoVideoProcessor": "video_processing_lfm3_vl.Lfm3VlVideoProcessor"
|
| 89 |
+
}
|
| 90 |
+
},
|
| 91 |
+
"video_token": "<|video|>",
|
| 92 |
+
"vision_end_token": "<|vision_end|>",
|
| 93 |
+
"vision_markers": true,
|
| 94 |
+
"vision_start_token": "<|vision_start|>",
|
| 95 |
+
"auto_map": {
|
| 96 |
+
"AutoProcessor": "processing_lfm3_omni.Lfm3OmniProcessor"
|
| 97 |
+
}
|
| 98 |
+
}
|
tokenizer.json
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:41927bd314d17d7d7fadb44e10cd0919f7b71528d7741596d8cbd15c6945ec44
|
| 3 |
+
size 17908634
|
tokenizer_config.json
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"backend": "tokenizers",
|
| 3 |
+
"bos_token": "<|startoftext|>",
|
| 4 |
+
"eos_token": "<|im_end|>",
|
| 5 |
+
"extra_special_tokens": [],
|
| 6 |
+
"image_token": "<|image|>",
|
| 7 |
+
"is_local": true,
|
| 8 |
+
"local_files_only": false,
|
| 9 |
+
"model_max_length": 1000000000000000019884624838656,
|
| 10 |
+
"model_specific_special_tokens": {
|
| 11 |
+
"image_token": "<|image|>",
|
| 12 |
+
"video_token": "<|video|>",
|
| 13 |
+
"vision_end_token": "<|vision_end|>",
|
| 14 |
+
"vision_start_token": "<|vision_start|>"
|
| 15 |
+
},
|
| 16 |
+
"pad_token": "<|pad|>",
|
| 17 |
+
"processor_class": "Lfm3OmniProcessor",
|
| 18 |
+
"tokenizer_class": "TokenizersBackend",
|
| 19 |
+
"video_token": "<|video|>",
|
| 20 |
+
"vision_end_token": "<|vision_end|>",
|
| 21 |
+
"vision_start_token": "<|vision_start|>"
|
| 22 |
+
}
|
video_preprocessor_config.json
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"cap_pixels_per_frame": false,
|
| 3 |
+
"do_convert_rgb": true,
|
| 4 |
+
"do_normalize": true,
|
| 5 |
+
"do_rescale": true,
|
| 6 |
+
"do_resize": true,
|
| 7 |
+
"do_sample_frames": true,
|
| 8 |
+
"fps": 2.0,
|
| 9 |
+
"image_mean": [
|
| 10 |
+
0.5,
|
| 11 |
+
0.5,
|
| 12 |
+
0.5
|
| 13 |
+
],
|
| 14 |
+
"image_std": [
|
| 15 |
+
0.5,
|
| 16 |
+
0.5,
|
| 17 |
+
0.5
|
| 18 |
+
],
|
| 19 |
+
"max_frames": 384,
|
| 20 |
+
"max_video_tokens": 768,
|
| 21 |
+
"merge_size": 2,
|
| 22 |
+
"min_frames": 4,
|
| 23 |
+
"patch_size": 16,
|
| 24 |
+
"resample": 3,
|
| 25 |
+
"rescale_factor": 0.00392156862745098,
|
| 26 |
+
"return_metadata": false,
|
| 27 |
+
"size": {
|
| 28 |
+
"longest_edge": 25165824,
|
| 29 |
+
"shortest_edge": 4096
|
| 30 |
+
},
|
| 31 |
+
"temporal_patch_size": 2,
|
| 32 |
+
"video_processor_type": "Lfm3VlVideoProcessor",
|
| 33 |
+
"auto_map": {
|
| 34 |
+
"AutoVideoProcessor": "video_processing_lfm3_vl.Lfm3VlVideoProcessor"
|
| 35 |
+
}
|
| 36 |
+
}
|
video_processing_lfm3_vl.py
ADDED
|
@@ -0,0 +1,396 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
import math
|
| 15 |
+
|
| 16 |
+
import numpy as np
|
| 17 |
+
import torch
|
| 18 |
+
from torchvision.transforms.v2 import functional as tvF
|
| 19 |
+
|
| 20 |
+
from transformers.image_processing_utils import BatchFeature
|
| 21 |
+
from transformers.image_utils import PILImageResampling, SizeDict
|
| 22 |
+
from transformers.processing_utils import Unpack, VideosKwargs
|
| 23 |
+
from transformers.utils import TensorType, auto_docstring, logging
|
| 24 |
+
from transformers.video_processing_utils import BaseVideoProcessor
|
| 25 |
+
from transformers.video_utils import (
|
| 26 |
+
VideoMetadata,
|
| 27 |
+
group_videos_by_shape,
|
| 28 |
+
make_batched_metadata,
|
| 29 |
+
make_batched_videos,
|
| 30 |
+
reorder_videos,
|
| 31 |
+
)
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
logger = logging.get_logger(__name__)
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
class Lfm3VlVideoProcessorInitKwargs(VideosKwargs, total=False):
|
| 38 |
+
r"""
|
| 39 |
+
patch_size (`int`, *optional*, defaults to 16):
|
| 40 |
+
The spatial patch size of the vision encoder.
|
| 41 |
+
temporal_patch_size (`int`, *optional*, defaults to 2):
|
| 42 |
+
The temporal patch size of the vision encoder.
|
| 43 |
+
merge_size (`int`, *optional*, defaults to 2):
|
| 44 |
+
The merge size of the vision encoder to llm encoder.
|
| 45 |
+
min_frames (`int`, *optional*, defaults to 4):
|
| 46 |
+
The minimum number of frames to sample from video.
|
| 47 |
+
max_frames (`int`, *optional*, defaults to 768):
|
| 48 |
+
The maximum number of frames to sample from video.
|
| 49 |
+
cap_pixels_per_frame (`bool`, *optional*):
|
| 50 |
+
Whether to cap each frame's pixel cost the way the reference implementation (qwen-vl-utils) does:
|
| 51 |
+
per-frame pixels are limited to `min(max_video_tokens * factor**2, size["longest_edge"] / num_frames)`
|
| 52 |
+
and floored at `1.05 * size["shortest_edge"]`, so token cost scales with clip duration. Without the
|
| 53 |
+
cap, videos that sample few frames spend the whole `size["longest_edge"]` budget on those frames and
|
| 54 |
+
keep near-native per-frame resolution, so a short clip can cost almost as many tokens as a long
|
| 55 |
+
video. If unset, the current uncapped behavior is kept and a warning is emitted: the default will
|
| 56 |
+
change to `True` in v5.22, after which the argument will be removed.
|
| 57 |
+
max_video_tokens (`int`, *optional*, defaults to 768):
|
| 58 |
+
The per-frame token ceiling applied by `cap_pixels_per_frame`, in vision tokens per frame
|
| 59 |
+
(qwen-vl-utils' `VIDEO_MAX_TOKEN_NUM`).
|
| 60 |
+
"""
|
| 61 |
+
|
| 62 |
+
patch_size: int
|
| 63 |
+
temporal_patch_size: int
|
| 64 |
+
merge_size: int
|
| 65 |
+
min_frames: int
|
| 66 |
+
max_frames: int
|
| 67 |
+
cap_pixels_per_frame: bool
|
| 68 |
+
max_video_tokens: int
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
def smart_resize(
|
| 72 |
+
num_frames: int,
|
| 73 |
+
height: int,
|
| 74 |
+
width: int,
|
| 75 |
+
temporal_factor: int = 2,
|
| 76 |
+
factor: int = 28,
|
| 77 |
+
min_pixels: int = 112 * 112,
|
| 78 |
+
max_pixels: int = 14 * 14 * 2 * 2 * 2 * 6144,
|
| 79 |
+
):
|
| 80 |
+
if num_frames < temporal_factor:
|
| 81 |
+
raise ValueError(f"t:{num_frames} must be larger than temporal_factor:{temporal_factor}")
|
| 82 |
+
if height < factor or width < factor:
|
| 83 |
+
scale = max(factor / height, factor / width)
|
| 84 |
+
height = int(height * scale)
|
| 85 |
+
width = int(width * scale)
|
| 86 |
+
|
| 87 |
+
if max(height, width) / min(height, width) > 200:
|
| 88 |
+
raise ValueError(
|
| 89 |
+
f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
|
| 90 |
+
)
|
| 91 |
+
h_bar = round(height / factor) * factor
|
| 92 |
+
w_bar = round(width / factor) * factor
|
| 93 |
+
t_bar = round(num_frames / temporal_factor) * temporal_factor
|
| 94 |
+
|
| 95 |
+
if t_bar * h_bar * w_bar > max_pixels:
|
| 96 |
+
beta = math.sqrt((num_frames * height * width) / max_pixels)
|
| 97 |
+
h_bar = max(factor, math.floor(height / beta / factor) * factor)
|
| 98 |
+
w_bar = max(factor, math.floor(width / beta / factor) * factor)
|
| 99 |
+
elif t_bar * h_bar * w_bar < min_pixels:
|
| 100 |
+
beta = math.sqrt(min_pixels / (num_frames * height * width))
|
| 101 |
+
h_bar = math.ceil(height * beta / factor) * factor
|
| 102 |
+
w_bar = math.ceil(width * beta / factor) * factor
|
| 103 |
+
|
| 104 |
+
return h_bar, w_bar
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
def _check_video_fps(metadata: VideoMetadata) -> None:
|
| 108 |
+
if metadata.fps is not None and not (math.isfinite(metadata.fps) and metadata.fps > 0):
|
| 109 |
+
raise ValueError(f"`video_metadata.fps` must be a positive number, got {metadata.fps}")
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def _check_video_metadata_count(videos, video_metadata) -> None:
|
| 113 |
+
# the shared sampler pairs clips with metadata by `zip`, which drops the clips past the last entry
|
| 114 |
+
if video_metadata is None:
|
| 115 |
+
return
|
| 116 |
+
videos = make_batched_videos(videos)
|
| 117 |
+
entries = make_batched_metadata(videos, video_metadata) if video_metadata != [] else []
|
| 118 |
+
if len(entries) != len(videos):
|
| 119 |
+
raise ValueError(
|
| 120 |
+
f"Received {len(videos)} videos but {len(entries)} `video_metadata` entries; pass one entry per video"
|
| 121 |
+
)
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
@auto_docstring
|
| 125 |
+
class Lfm3VlVideoProcessor(BaseVideoProcessor):
|
| 126 |
+
resample = PILImageResampling.BICUBIC
|
| 127 |
+
size = {"shortest_edge": 4_096, "longest_edge": 25_165_824}
|
| 128 |
+
image_mean = [0.5, 0.5, 0.5]
|
| 129 |
+
image_std = [0.5, 0.5, 0.5]
|
| 130 |
+
do_resize = True
|
| 131 |
+
do_rescale = True
|
| 132 |
+
do_normalize = True
|
| 133 |
+
do_convert_rgb = True
|
| 134 |
+
patch_size = 16
|
| 135 |
+
temporal_patch_size = 2
|
| 136 |
+
merge_size = 2
|
| 137 |
+
min_frames = 4
|
| 138 |
+
max_frames = 768
|
| 139 |
+
do_sample_frames = True
|
| 140 |
+
cap_pixels_per_frame = False
|
| 141 |
+
max_video_tokens = 768
|
| 142 |
+
valid_kwargs = Lfm3VlVideoProcessorInitKwargs
|
| 143 |
+
model_input_names = ["pixel_values_videos", "video_grid_thw"]
|
| 144 |
+
max_image_size = {"longest_edge": 28 * 28 * 2 * 30000}
|
| 145 |
+
max_duration = None
|
| 146 |
+
num_frames = None
|
| 147 |
+
fps = 2.0
|
| 148 |
+
|
| 149 |
+
def __init__(self, **kwargs: Unpack[Lfm3VlVideoProcessorInitKwargs]):
|
| 150 |
+
super().__init__(**kwargs)
|
| 151 |
+
|
| 152 |
+
def sample_frames(
|
| 153 |
+
self, metadata: VideoMetadata, num_frames: int | None = None, fps: int | float | None = None, **kwargs
|
| 154 |
+
):
|
| 155 |
+
"""Decoded frames whose metadata carries no fps are sampled and timestamped at the 24 fps the shared sampler
|
| 156 |
+
assumes (with a warning); pass `video_metadata` with the clip's fps for accurate timestamps."""
|
| 157 |
+
_check_video_fps(metadata)
|
| 158 |
+
if fps is not None and not (math.isfinite(fps) and fps > 0):
|
| 159 |
+
raise ValueError(f"`fps` must be a positive number, got {fps}")
|
| 160 |
+
if fps is not None and num_frames is not None:
|
| 161 |
+
raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!")
|
| 162 |
+
|
| 163 |
+
total_num_frames = metadata.total_num_frames
|
| 164 |
+
fps = fps if fps is not None else self.fps
|
| 165 |
+
|
| 166 |
+
# If num_frames is not given but fps is, calculate num_frames from fps
|
| 167 |
+
if num_frames is None and fps is not None:
|
| 168 |
+
if metadata.fps is None:
|
| 169 |
+
metadata.fps = 24
|
| 170 |
+
logger.warning_once(
|
| 171 |
+
"Asked to sample `fps` frames per second but no video metadata was provided which is required when sampling with `fps`. "
|
| 172 |
+
"Defaulting to `fps=24`. Please provide `video_metadata` for more accurate results."
|
| 173 |
+
)
|
| 174 |
+
num_frames = int(total_num_frames / metadata.fps * fps)
|
| 175 |
+
num_frames = min(max(num_frames, self.min_frames), self.max_frames, total_num_frames)
|
| 176 |
+
|
| 177 |
+
if num_frames is None:
|
| 178 |
+
num_frames = min(max(total_num_frames, self.min_frames), self.max_frames)
|
| 179 |
+
|
| 180 |
+
indices = np.linspace(0, total_num_frames - 1, num_frames).round().astype(int)
|
| 181 |
+
|
| 182 |
+
return indices
|
| 183 |
+
|
| 184 |
+
def resize(
|
| 185 |
+
self,
|
| 186 |
+
videos: "torch.Tensor",
|
| 187 |
+
size: SizeDict,
|
| 188 |
+
resample: "PILImageResampling | tvF.InterpolationMode | int | None",
|
| 189 |
+
factor: int,
|
| 190 |
+
temporal_factor: int,
|
| 191 |
+
cap_pixels_per_frame: bool | None = None,
|
| 192 |
+
**kwargs,
|
| 193 |
+
) -> "torch.Tensor":
|
| 194 |
+
"""Resize dynamically based on input video aspect ratio."""
|
| 195 |
+
if not size.shortest_edge or not size.longest_edge:
|
| 196 |
+
raise ValueError(f"`size` dict must contain 'shortest_edge' and 'longest_edge' keys but got {size}.")
|
| 197 |
+
|
| 198 |
+
num_frames = videos.shape[1]
|
| 199 |
+
max_pixels = size.longest_edge
|
| 200 |
+
if cap_pixels_per_frame:
|
| 201 |
+
# per-frame pixels are capped at `max_video_tokens` patches or the budget's even share per frame
|
| 202 |
+
frame_cap = self.max_video_tokens * factor * factor
|
| 203 |
+
pixels_per_frame = max(min(frame_cap, size.longest_edge // num_frames), int(size.shortest_edge * 1.05))
|
| 204 |
+
max_pixels = pixels_per_frame * num_frames
|
| 205 |
+
|
| 206 |
+
height, width = videos.shape[-2:]
|
| 207 |
+
resized_height, resized_width = smart_resize(
|
| 208 |
+
height=height,
|
| 209 |
+
width=width,
|
| 210 |
+
num_frames=num_frames,
|
| 211 |
+
factor=factor,
|
| 212 |
+
temporal_factor=temporal_factor,
|
| 213 |
+
min_pixels=size.shortest_edge,
|
| 214 |
+
max_pixels=max_pixels,
|
| 215 |
+
)
|
| 216 |
+
return super().resize(
|
| 217 |
+
image=videos,
|
| 218 |
+
size=SizeDict(height=resized_height, width=resized_width),
|
| 219 |
+
resample=resample,
|
| 220 |
+
)
|
| 221 |
+
|
| 222 |
+
def patchify(
|
| 223 |
+
self,
|
| 224 |
+
videos: "torch.Tensor",
|
| 225 |
+
patch_size: int,
|
| 226 |
+
merge_size: int,
|
| 227 |
+
temporal_patch_size: int,
|
| 228 |
+
) -> tuple["torch.Tensor", int, int]:
|
| 229 |
+
"Patchifies each video into flat layout of shape (`seq_len`, `patch_dim`) so we can concat dynamically shaped pixels."
|
| 230 |
+
batch_size, num_frames, channel, resized_height, resized_width = videos.shape
|
| 231 |
+
|
| 232 |
+
# Check that videos have `num_frames` divisible by `temporal_patch_size`
|
| 233 |
+
if pad := -num_frames % temporal_patch_size:
|
| 234 |
+
repeats = videos[:, -1:].expand(-1, pad, -1, -1, -1)
|
| 235 |
+
videos = torch.cat((videos, repeats), dim=1)
|
| 236 |
+
num_frames += pad
|
| 237 |
+
|
| 238 |
+
grid_t = num_frames // temporal_patch_size
|
| 239 |
+
grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
|
| 240 |
+
|
| 241 |
+
patches = videos.view(
|
| 242 |
+
batch_size,
|
| 243 |
+
grid_t,
|
| 244 |
+
temporal_patch_size,
|
| 245 |
+
channel,
|
| 246 |
+
grid_h // merge_size,
|
| 247 |
+
merge_size,
|
| 248 |
+
patch_size,
|
| 249 |
+
grid_w // merge_size,
|
| 250 |
+
merge_size,
|
| 251 |
+
patch_size,
|
| 252 |
+
)
|
| 253 |
+
patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
|
| 254 |
+
flatten_patches = patches.reshape(
|
| 255 |
+
batch_size,
|
| 256 |
+
grid_t * grid_h * grid_w,
|
| 257 |
+
channel * temporal_patch_size * patch_size * patch_size,
|
| 258 |
+
)
|
| 259 |
+
|
| 260 |
+
return flatten_patches, grid_t, grid_h, grid_w
|
| 261 |
+
|
| 262 |
+
def _preprocess(
|
| 263 |
+
self,
|
| 264 |
+
videos: list["torch.Tensor"],
|
| 265 |
+
do_resize: bool,
|
| 266 |
+
size: SizeDict,
|
| 267 |
+
resample: "PILImageResampling | tvF.InterpolationMode | int | None",
|
| 268 |
+
do_rescale: bool,
|
| 269 |
+
rescale_factor: float,
|
| 270 |
+
do_normalize: bool,
|
| 271 |
+
image_mean: float | list[float] | None,
|
| 272 |
+
image_std: float | list[float] | None,
|
| 273 |
+
patch_size: int | None = None,
|
| 274 |
+
temporal_patch_size: int | None = None,
|
| 275 |
+
merge_size: int | None = None,
|
| 276 |
+
cap_pixels_per_frame: bool | None = None,
|
| 277 |
+
return_tensors: str | TensorType | None = None,
|
| 278 |
+
**kwargs,
|
| 279 |
+
):
|
| 280 |
+
if cap_pixels_per_frame is None:
|
| 281 |
+
logger.warning_once(
|
| 282 |
+
"Lfm3Vl video processing does not apply the per-frame pixel cap the reference "
|
| 283 |
+
"implementation (qwen-vl-utils) applies, so some videos cost far more tokens than they "
|
| 284 |
+
"would there. In v5.22 the capped behavior will become the default and "
|
| 285 |
+
"`cap_pixels_per_frame` will be removed. Pass `cap_pixels_per_frame=True` to adopt the "
|
| 286 |
+
"reference behavior now, or `False` to keep the current behavior and silence this "
|
| 287 |
+
"warning."
|
| 288 |
+
)
|
| 289 |
+
cap_pixels_per_frame = False
|
| 290 |
+
# Group videos by size for batched resizing
|
| 291 |
+
grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
|
| 292 |
+
resized_videos_grouped = {}
|
| 293 |
+
for shape, stacked_videos in grouped_videos.items():
|
| 294 |
+
if do_resize:
|
| 295 |
+
stacked_videos = self.resize(
|
| 296 |
+
videos=stacked_videos,
|
| 297 |
+
size=size,
|
| 298 |
+
resample=resample,
|
| 299 |
+
factor=patch_size * merge_size,
|
| 300 |
+
temporal_factor=temporal_patch_size,
|
| 301 |
+
cap_pixels_per_frame=cap_pixels_per_frame,
|
| 302 |
+
)
|
| 303 |
+
resized_videos_grouped[shape] = stacked_videos
|
| 304 |
+
resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
|
| 305 |
+
|
| 306 |
+
# Group videos by size for further processing
|
| 307 |
+
# Needed in case do_resize is False, or resize returns videos with different sizes
|
| 308 |
+
grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
|
| 309 |
+
processed_videos_grouped = {}
|
| 310 |
+
processed_grids = {}
|
| 311 |
+
for shape, stacked_videos in grouped_videos.items():
|
| 312 |
+
# Fused rescale and normalize
|
| 313 |
+
stacked_videos = self.rescale_and_normalize(
|
| 314 |
+
stacked_videos, do_rescale, rescale_factor, do_normalize, image_mean, image_std
|
| 315 |
+
)
|
| 316 |
+
patches, grid_t, grid_h, grid_w = self.patchify(
|
| 317 |
+
stacked_videos,
|
| 318 |
+
patch_size=patch_size,
|
| 319 |
+
merge_size=merge_size,
|
| 320 |
+
temporal_patch_size=temporal_patch_size,
|
| 321 |
+
)
|
| 322 |
+
|
| 323 |
+
processed_videos_grouped[shape] = patches
|
| 324 |
+
processed_grids[shape] = [[grid_t, grid_h, grid_w]] * len(stacked_videos)
|
| 325 |
+
|
| 326 |
+
processed_videos = reorder_videos(processed_videos_grouped, grouped_videos_index)
|
| 327 |
+
processed_grids = reorder_videos(processed_grids, grouped_videos_index)
|
| 328 |
+
pixel_values_videos = torch.cat(processed_videos, dim=0)
|
| 329 |
+
video_grid_thw = torch.tensor(processed_grids)
|
| 330 |
+
|
| 331 |
+
return BatchFeature(
|
| 332 |
+
data={"pixel_values_videos": pixel_values_videos, "video_grid_thw": video_grid_thw},
|
| 333 |
+
tensor_type=return_tensors,
|
| 334 |
+
)
|
| 335 |
+
|
| 336 |
+
def get_num_of_video_patches(self, num_frames: int, height: int, width: int, videos_kwargs=None):
|
| 337 |
+
"""
|
| 338 |
+
A utility that returns number of video patches a given video size.
|
| 339 |
+
|
| 340 |
+
Args:
|
| 341 |
+
num_frames (`int`):
|
| 342 |
+
Number of frames in the input video.
|
| 343 |
+
height (`int`):
|
| 344 |
+
Height of the input video.
|
| 345 |
+
width (`int`):
|
| 346 |
+
Width of the input video.
|
| 347 |
+
videos_kwargs (`dict`, *optional*)
|
| 348 |
+
Any kwargs to override defaults of the video processor.
|
| 349 |
+
Returns:
|
| 350 |
+
`int`: Number of video patches per video.
|
| 351 |
+
"""
|
| 352 |
+
videos_kwargs = videos_kwargs if videos_kwargs is not None else {}
|
| 353 |
+
size = videos_kwargs.get("size", None) or self.size
|
| 354 |
+
min_pixels = size["shortest_edge"]
|
| 355 |
+
max_pixels = size["longest_edge"]
|
| 356 |
+
patch_size = videos_kwargs.get("patch_size", None) or self.patch_size
|
| 357 |
+
merge_size = videos_kwargs.get("merge_size", None) or self.merge_size
|
| 358 |
+
temporal_patch_size = videos_kwargs.get("temporal_patch_size", None) or self.temporal_patch_size
|
| 359 |
+
cap_pixels_per_frame = videos_kwargs.get("cap_pixels_per_frame", None) or self.cap_pixels_per_frame
|
| 360 |
+
|
| 361 |
+
factor = patch_size * merge_size
|
| 362 |
+
if cap_pixels_per_frame:
|
| 363 |
+
# Keep the count in sync with `resize` when the per-frame cap is active.
|
| 364 |
+
frame_cap = self.max_video_tokens * factor * factor
|
| 365 |
+
pixels_per_frame = max(min(frame_cap, max_pixels // num_frames), int(min_pixels * 1.05))
|
| 366 |
+
max_pixels = pixels_per_frame * num_frames
|
| 367 |
+
|
| 368 |
+
resized_height, resized_width = smart_resize(
|
| 369 |
+
num_frames,
|
| 370 |
+
height,
|
| 371 |
+
width,
|
| 372 |
+
temporal_factor=temporal_patch_size,
|
| 373 |
+
factor=factor,
|
| 374 |
+
min_pixels=min_pixels,
|
| 375 |
+
max_pixels=max_pixels,
|
| 376 |
+
)
|
| 377 |
+
grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
|
| 378 |
+
grid_t = (num_frames + -num_frames % temporal_patch_size) // temporal_patch_size
|
| 379 |
+
return grid_t * grid_h * grid_w
|
| 380 |
+
|
| 381 |
+
def preprocess(self, videos, **kwargs):
|
| 382 |
+
_check_video_metadata_count(videos, kwargs.get("video_metadata"))
|
| 383 |
+
# the shared `preprocess` fills every unset option from the instance before `sample_frames` runs, so the
|
| 384 |
+
# caller's options are told apart from the instance ones here: a caller option beats an instance one, and a
|
| 385 |
+
# configured `num_frames` beats the default fps
|
| 386 |
+
num_frames, fps = kwargs.get("num_frames"), kwargs.get("fps")
|
| 387 |
+
if num_frames is not None and fps is not None:
|
| 388 |
+
raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!")
|
| 389 |
+
if num_frames is not None or (fps is None and self.num_frames is not None):
|
| 390 |
+
kwargs["fps"] = None
|
| 391 |
+
elif fps is not None:
|
| 392 |
+
kwargs["num_frames"] = None
|
| 393 |
+
return super().preprocess(videos, **kwargs)
|
| 394 |
+
|
| 395 |
+
|
| 396 |
+
__all__ = ["Lfm3VlVideoProcessor"]
|