1use crate::ltxdit::{Attn, Lin, Rope, gelu_tanh, rms_plain, rows};
26use crate::pool::Pool;
27use crate::qtensor::QTensor;
28use cortiq_core::CmfModel;
29use std::sync::Arc;
30
31const EPS: f64 = 1e-6;
32
33fn rms(x: &[f32], w: &[f32], dst: &mut [f32]) {
36 let ss = x.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / x.len() as f64;
37 let inv = 1.0 / (ss + EPS).sqrt();
38 for ((d, &v), &g) in dst.iter_mut().zip(x).zip(w) {
39 *d = (v as f64 * inv) as f32 * g;
40 }
41}
42
43fn rms_nw(x: &mut [f32]) {
45 let ss = x.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / x.len() as f64;
46 let inv = 1.0 / (ss + EPS).sqrt();
47 for v in x.iter_mut() {
48 *v = (*v as f64 * inv) as f32;
49 }
50}
51
52struct Rot {
55 cos: Vec<f32>,
56 sin: Vec<f32>,
57 half: usize,
58}
59
60impl Rot {
61 fn build(seq: usize, head_dim: usize, base: f64, rotary: f64) -> Rot {
66 let half = head_dim / 2;
67 let rope_angles = (rotary * head_dim as f64 / 2.0) as usize;
68 let inv: Vec<f64> = (0..half)
69 .map(|j| {
70 if j < rope_angles {
71 1.0 / base.powf((2 * j) as f64 / head_dim as f64)
72 } else {
73 0.0
74 }
75 })
76 .collect();
77 let mut cos = vec![0f32; seq * half];
78 let mut sin = vec![0f32; seq * half];
79 for p in 0..seq {
80 for (j, &f) in inv.iter().enumerate() {
81 let a = p as f64 * f;
82 cos[p * half + j] = a.cos() as f32;
83 sin[p * half + j] = a.sin() as f32;
84 }
85 }
86 Rot { cos, sin, half }
87 }
88
89 fn apply(&self, t: usize, row: &mut [f32]) {
90 let h = self.half;
91 let (c, s) = (&self.cos[t * h..(t + 1) * h], &self.sin[t * h..(t + 1) * h]);
92 for i in 0..h {
93 let (a, b) = (row[i], row[i + h]);
94 row[i] = a * c[i] - b * s[i];
95 row[i + h] = b * c[i] + a * s[i];
96 }
97 }
98}
99
100struct GemmaLayer {
101 q: Lin,
102 k: Lin,
103 v: Option<Lin>,
104 o: Lin,
105 q_norm: Vec<f32>,
106 k_norm: Vec<f32>,
107 in_norm: Vec<f32>,
108 post_attn_norm: Vec<f32>,
109 pre_ff_norm: Vec<f32>,
110 post_ff_norm: Vec<f32>,
111 gate: Lin,
112 up: Lin,
113 down: Lin,
114 scalar: f32,
115 head_dim: usize,
116 q_heads: usize,
117 kv_heads: usize,
118 sliding: bool,
119}
120
121pub struct LtxTextEncoder {
123 embed: QTensor,
124 mapped_bytes: u64,
125 layers: Vec<GemmaLayer>,
126 norm: Vec<f32>,
127 video_agg: Lin,
128 audio_agg: Lin,
129 v_conn: Connector,
130 a_conn: Connector,
131 hidden: usize,
132 embed_scale: f32,
133 pub max_len: usize,
134 pub bos: u32,
135 pub pad: u32,
136}
137
138struct Connector {
139 blocks: Vec<(Attn, Lin, Lin)>,
140 registers: Vec<f32>,
141 dim: usize,
142 heads: usize,
143 dh: usize,
144 max_pos: f64,
145}
146
147fn vecf(model: &Arc<CmfModel>, name: &str) -> Result<Vec<f32>, String> {
148 crate::dit::cmf_f32(model, name)
149}
150
151impl LtxTextEncoder {
152 pub fn from_cmf(model: &Arc<CmfModel>) -> Result<LtxTextEncoder, String> {
153 let cfg_bytes = ["te.gemma_config_json", "ltx.gemma_config_json"]
154 .iter()
155 .find_map(|n| model.tensor(n).map(|e| model.entry_bytes(e)));
156 let cfg: serde_json::Value = match cfg_bytes {
157 Some(b) => serde_json::from_slice(b).map_err(|e| format!("gemma config: {e}"))?,
158 None => serde_json::Value::Null,
159 };
160 let tc = cfg.get("text_config").cloned().unwrap_or(serde_json::Value::Null);
161 let g = |k: &str, d: f64| tc.get(k).and_then(|v| v.as_f64()).unwrap_or(d);
162 let hidden = g("hidden_size", 3840.0) as usize;
163 let n_layers = g("num_hidden_layers", 48.0) as usize;
164 let head_dim = g("head_dim", 256.0) as usize;
165 let global_head_dim = g("global_head_dim", 512.0) as usize;
166 let types: Vec<String> = tc
167 .get("layer_types")
168 .and_then(|v| v.as_array())
169 .map(|a| a.iter().map(|s| s.as_str().unwrap_or("sliding_attention").to_string()).collect())
170 .unwrap_or_else(|| {
171 (0..n_layers)
173 .map(|i| {
174 if i % 6 == 5 { "full_attention".into() } else { "sliding_attention".into() }
175 })
176 .collect()
177 });
178
179 let mut layers = Vec::with_capacity(n_layers);
180 for i in 0..n_layers {
181 let p = format!("te.model.layers.{i}");
182 let sliding = types[i] != "full_attention";
183 let hd = if sliding { head_dim } else { global_head_dim };
184 let q = Lin::load(model, &format!("{p}.self_attn.q_proj"), false)?;
185 let k = Lin::load(model, &format!("{p}.self_attn.k_proj"), false)?;
186 let v = match model.tensor(&format!("{p}.self_attn.v_proj.weight")) {
187 Some(_) => Some(Lin::load(model, &format!("{p}.self_attn.v_proj"), false)?),
188 None => None,
189 };
190 let q_rows = model
191 .tensor(&format!("{p}.self_attn.q_proj.weight"))
192 .ok_or_else(|| format!("missing {p}.self_attn.q_proj.weight"))?
193 .shape[0];
194 let k_rows = model
195 .tensor(&format!("{p}.self_attn.k_proj.weight"))
196 .ok_or("missing k_proj")?
197 .shape[0];
198 layers.push(GemmaLayer {
199 q,
200 k,
201 v,
202 o: Lin::load(model, &format!("{p}.self_attn.o_proj"), false)?,
203 q_norm: vecf(model, &format!("{p}.self_attn.q_norm.weight"))?,
204 k_norm: vecf(model, &format!("{p}.self_attn.k_norm.weight"))?,
205 in_norm: vecf(model, &format!("{p}.input_layernorm.weight"))?,
206 post_attn_norm: vecf(model, &format!("{p}.post_attention_layernorm.weight"))?,
207 pre_ff_norm: vecf(model, &format!("{p}.pre_feedforward_layernorm.weight"))?,
208 post_ff_norm: vecf(model, &format!("{p}.post_feedforward_layernorm.weight"))?,
209 gate: Lin::load(model, &format!("{p}.mlp.gate_proj"), false)?,
210 up: Lin::load(model, &format!("{p}.mlp.up_proj"), false)?,
211 down: Lin::load(model, &format!("{p}.mlp.down_proj"), false)?,
212 scalar: vecf(model, &format!("{p}.layer_scalar"))?[0],
213 head_dim: hd,
214 q_heads: q_rows / hd,
215 kv_heads: k_rows / hd,
216 sliding,
217 });
218 }
219
220 let conn = |prefix: &str, dim: usize, heads: usize, dh: usize| -> Result<Connector, String> {
221 let mut blocks = Vec::new();
222 let mut i = 0usize;
223 while model
224 .tensor(&format!("{prefix}.transformer_1d_blocks.{i}.attn1.to_q.weight"))
225 .is_some()
226 {
227 let p = format!("{prefix}.transformer_1d_blocks.{i}");
228 blocks.push((
229 Attn::load(model, &format!("{p}.attn1"), heads, dh)?,
230 Lin::load(model, &format!("{p}.ff.net.0.proj"), true)?,
231 Lin::load(model, &format!("{p}.ff.net.2"), true)?,
232 ));
233 i += 1;
234 }
235 Ok(Connector {
236 blocks,
237 registers: vecf(model, &format!("{prefix}.learnable_registers"))?,
238 dim,
239 heads,
240 dh,
241 max_pos: 4096.0,
242 })
243 };
244
245 Ok(LtxTextEncoder {
246 embed: QTensor::from_model(model, "te.model.embed_tokens.weight")?,
247 mapped_bytes: model.primary_bytes().len() as u64,
248 layers,
249 norm: vecf(model, "te.model.norm.weight")?,
250 video_agg: Lin::load(model, "te.text_embedding_projection.video_aggregate_embed", true)?,
251 audio_agg: Lin::load(model, "te.text_embedding_projection.audio_aggregate_embed", true)?,
252 v_conn: conn("dit.video_embeddings_connector", 4096, 32, 128)?,
253 a_conn: conn("dit.audio_embeddings_connector", 2048, 32, 64)?,
254 hidden,
255 embed_scale: (hidden as f64).sqrt() as f32,
256 max_len: 1024,
257 bos: tc.get("bos_token_id").and_then(|v| v.as_u64()).unwrap_or(2) as u32,
258 pad: tc.get("pad_token_id").and_then(|v| v.as_u64()).unwrap_or(0) as u32,
259 })
260 }
261
262 pub fn pad_ids(&self, ids: &[u32]) -> (Vec<u32>, Vec<f32>) {
265 let mut body = Vec::with_capacity(self.max_len);
266 if ids.first() != Some(&self.bos) {
267 body.push(self.bos);
268 }
269 body.extend_from_slice(ids);
270 body.truncate(self.max_len);
271 let padlen = self.max_len - body.len();
272 let mut out = vec![self.pad; padlen];
273 out.extend_from_slice(&body);
274 let mut mask = vec![0f32; padlen];
275 mask.extend(std::iter::repeat_n(1f32, body.len()));
276 (out, mask)
277 }
278
279 pub fn hidden_states(&self, ids: &[u32], mask: &[f32], pool: Option<&Pool>) -> Vec<Vec<f32>> {
282 let _pause = self.crowds_the_device().then(crate::gpu::pause_gpu);
289 self.hidden_states_inner(ids, mask, pool)
290 }
291
292 #[cfg(target_os = "macos")]
293 fn crowds_the_device(&self) -> bool {
294 let one_window = crate::gpu_metal::max_buffer_bytes();
300 let wired = crate::gpu_metal::working_set_bytes();
301 (one_window > 0 && self.mapped_bytes > one_window)
302 || (wired > 0 && self.mapped_bytes > wired)
303 }
304
305 #[cfg(not(target_os = "macos"))]
306 fn crowds_the_device(&self) -> bool {
307 false
308 }
309
310 fn hidden_states_inner(&self, ids: &[u32], mask: &[f32], pool: Option<&Pool>) -> Vec<Vec<f32>> {
311 let t = ids.len();
312 let d = self.hidden;
313 let mut x = vec![0f32; t * d];
314 for (i, &id) in ids.iter().enumerate() {
315 self.embed.row_f32(id as usize, &mut x[i * d..(i + 1) * d]);
316 for v in x[i * d..(i + 1) * d].iter_mut() {
317 *v *= self.embed_scale;
318 }
319 }
320 let mut out = vec![x.clone()];
321 let rot_slide = Rot::build(t, self.layers[0].head_dim, 10000.0, 1.0);
323 let full = self.layers.iter().find(|l| !l.sliding);
324 let rot_full = full.map(|l| Rot::build(t, l.head_dim, 1_000_000.0, 0.25));
325
326 for layer in &self.layers {
327 let rot = if layer.sliding { &rot_slide } else { rot_full.as_ref().unwrap() };
328 let mut h = vec![0f32; t * d];
329 for i in 0..t {
330 rms(&x[i * d..(i + 1) * d], &layer.in_norm, &mut h[i * d..(i + 1) * d]);
331 }
332 let attn = self.attention(layer, &h, t, mask, rot, pool);
333 for i in 0..t {
334 let mut n = vec![0f32; d];
335 rms(&attn[i * d..(i + 1) * d], &layer.post_attn_norm, &mut n);
336 for (v, &a) in x[i * d..(i + 1) * d].iter_mut().zip(&n) {
337 *v += a;
338 }
339 }
340 let mut h2 = vec![0f32; t * d];
341 for i in 0..t {
342 rms(&x[i * d..(i + 1) * d], &layer.pre_ff_norm, &mut h2[i * d..(i + 1) * d]);
343 }
344 let mut g = layer.gate.apply(&h2, t, pool);
345 let u = layer.up.apply(&h2, t, pool);
346 for (a, &b) in g.iter_mut().zip(&u) {
347 *a = gelu_tanh(*a) * b;
348 }
349 let ff = layer.down.apply(&g, t, pool);
350 for i in 0..t {
351 let mut n = vec![0f32; d];
352 rms(&ff[i * d..(i + 1) * d], &layer.post_ff_norm, &mut n);
353 for (v, &a) in x[i * d..(i + 1) * d].iter_mut().zip(&n) {
354 *v += a;
355 }
356 }
357 for v in x.iter_mut() {
358 *v *= layer.scalar;
359 }
360 out.push(x.clone());
361 }
362 if let Some(last) = out.last_mut() {
368 let mut n = vec![0f32; t * d];
369 for i in 0..t {
370 rms(&last[i * d..(i + 1) * d], &self.norm, &mut n[i * d..(i + 1) * d]);
371 }
372 *last = n;
373 }
374 out
375 }
376
377 fn attention(
378 &self,
379 l: &GemmaLayer,
380 h: &[f32],
381 t: usize,
382 mask: &[f32],
383 rot: &Rot,
384 pool: Option<&Pool>,
385 ) -> Vec<f32> {
386 let hd = l.head_dim;
387 let qi = l.q_heads * hd;
388 let ki = l.kv_heads * hd;
389 let mut q = l.q.apply(h, t, pool);
390 let mut k = l.k.apply(h, t, pool);
391 let mut v = match &l.v {
392 Some(p) => p.apply(h, t, pool),
393 None => k.clone(),
394 };
395 for i in 0..t {
396 for hh in 0..l.q_heads {
397 let r = &mut q[i * qi + hh * hd..i * qi + (hh + 1) * hd];
398 let mut n = vec![0f32; hd];
399 rms(r, &l.q_norm, &mut n);
400 r.copy_from_slice(&n);
401 rot.apply(i, r);
402 }
403 for hh in 0..l.kv_heads {
404 let r = &mut k[i * ki + hh * hd..i * ki + (hh + 1) * hd];
405 let mut n = vec![0f32; hd];
406 rms(r, &l.k_norm, &mut n);
407 r.copy_from_slice(&n);
408 rot.apply(i, r);
409 rms_nw(&mut v[i * ki + hh * hd..i * ki + (hh + 1) * hd]);
410 }
411 }
412 let window = if l.sliding { 1024usize } else { usize::MAX };
416 let mut bias = vec![0f32; t * t];
417 for i in 0..t {
418 for j in 0..t {
419 let blocked =
420 j > i || (window != usize::MAX && i - j >= window) || mask[j] == 0.0;
421 bias[i * t + j] = if blocked { f32::NEG_INFINITY } else { 0.0 };
422 }
423 }
424 let dead: Vec<bool> = (0..t).map(|i| mask[i] == 0.0).collect();
430 let mut out = vec![0f32; t * qi];
431 let mut qh = vec![0f32; t * hd];
432 let mut kh = vec![0f32; t * hd];
433 let mut vh = vec![0f32; t * hd];
434 let mut sc = vec![0f32; t * t];
435 let mut oh = vec![0f32; t * hd];
436 for hh in 0..l.q_heads {
437 let kv = hh * l.kv_heads / l.q_heads;
438 for i in 0..t {
439 qh[i * hd..(i + 1) * hd].copy_from_slice(&q[i * qi + hh * hd..][..hd]);
440 kh[i * hd..(i + 1) * hd].copy_from_slice(&k[i * ki + kv * hd..][..hd]);
441 vh[i * hd..(i + 1) * hd].copy_from_slice(&v[i * ki + kv * hd..][..hd]);
442 }
443 crate::fcd_ops::gemm_nt(&qh, &kh, &mut sc, t, hd, t, pool);
444 let sp = crate::ltxdit::Shared(sc.as_mut_ptr());
445 rows(pool, t, &|s, e| {
446 let r = unsafe { sp.at(s * t, (e - s) * t) };
447 for (row, i) in r.chunks_exact_mut(t).zip(s..e) {
448 if dead[i] {
449 row.iter_mut().for_each(|x| *x = 0.0);
450 continue;
451 }
452 for (x, b) in row.iter_mut().zip(&bias[i * t..(i + 1) * t]) {
453 *x += *b;
454 }
455 crate::ltxdit::softmax(row);
456 }
457 });
458 oh.iter_mut().for_each(|x| *x = 0.0);
459 crate::fcd_ops::gemm_dx(&sc, &vh, &mut oh, t, hd, t, pool);
460 for i in 0..t {
461 out[i * qi + hh * hd..i * qi + (hh + 1) * hd]
462 .copy_from_slice(&oh[i * hd..(i + 1) * hd]);
463 }
464 }
465 l.o.apply(&out, t, pool)
466 }
467
468 pub fn encode_ids(
470 &self,
471 ids: &[u32],
472 mask: &[f32],
473 pool: Option<&Pool>,
474 ) -> (Vec<f32>, Vec<f32>, usize) {
475 let _pause = self.crowds_the_device().then(crate::gpu::pause_gpu);
480 let hs = self.hidden_states_inner(ids, mask, pool);
481 let (t, d, l) = (ids.len(), self.hidden, hs.len());
482 let mut feats = vec![0f32; t * d * l];
485 for (li, layer) in hs.iter().enumerate() {
486 for i in 0..t {
487 let row = &layer[i * d..(i + 1) * d];
488 let var = row.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / d as f64;
489 let inv = 1.0 / (var + 1e-6).sqrt();
490 let keep = mask[i] != 0.0;
491 for j in 0..d {
492 feats[i * d * l + j * l + li] = if keep { (row[j] as f64 * inv) as f32 } else { 0.0 };
493 }
494 }
495 }
496 let project = |agg: &Lin, out_dim: usize| -> Vec<f32> {
501 let scale = ((out_dim as f64) / (d as f64)).sqrt() as f32;
502 let scaled: Vec<f32> = feats.iter().map(|&v| v * scale).collect();
503 crate::gpu::cpu_scope(|| agg.apply(&scaled, t, pool))
504 };
505 let vfeat = project(&self.video_agg, 4096);
506 let afeat = project(&self.audio_agg, 2048);
507 let order: Vec<usize> = (0..t)
509 .filter(|&i| mask[i] != 0.0)
510 .chain((0..t).filter(|&i| mask[i] == 0.0))
511 .collect();
512 let valid = mask.iter().filter(|&&m| m != 0.0).count();
513 let reorder = |x: &[f32], dim: usize| -> Vec<f32> {
514 let mut o = vec![0f32; t * dim];
515 for (new, &old) in order.iter().enumerate() {
516 o[new * dim..(new + 1) * dim].copy_from_slice(&x[old * dim..(old + 1) * dim]);
517 }
518 o
519 };
520 let v = self.v_conn.run(&reorder(&vfeat, 4096), t, valid, pool);
521 let a = self.a_conn.run(&reorder(&afeat, 2048), t, valid, pool);
522 (v, a, t)
523 }
524}
525
526impl Connector {
527 fn run(&self, x: &[f32], t: usize, valid: usize, pool: Option<&Pool>) -> Vec<f32> {
532 let d = self.dim;
533 let regs = self.registers.len() / d;
534 let mut h = x.to_vec();
535 for i in valid..t {
536 let r = i % regs;
537 h[i * d..(i + 1) * d].copy_from_slice(&self.registers[r * d..(r + 1) * d]);
538 }
539 let pos: Vec<Vec<f64>> = (0..t).map(|i| vec![i as f64]).collect();
540 let pe = Rope::build(&pos, &[self.max_pos], d, self.heads, 10000.0);
541 for (attn, ff_in, ff_out) in &self.blocks {
542 let mut n = vec![0f32; t * d];
543 for i in 0..t {
544 rms_plain(&h[i * d..(i + 1) * d], &mut n[i * d..(i + 1) * d]);
545 }
546 let a = attn.forward(&n, t, &n, t, Some(&pe), Some(&pe), None, pool);
547 for (v, &y) in h.iter_mut().zip(&a) {
548 *v += y;
549 }
550 let mut n2 = vec![0f32; t * d];
551 for i in 0..t {
552 rms_plain(&h[i * d..(i + 1) * d], &mut n2[i * d..(i + 1) * d]);
553 }
554 let mut g = ff_in.apply(&n2, t, pool);
555 for v in g.iter_mut() {
556 *v = gelu_tanh(*v);
557 }
558 let f = ff_out.apply(&g, t, pool);
559 for (v, &y) in h.iter_mut().zip(&f) {
560 *v += y;
561 }
562 }
563 let mut out = vec![0f32; t * d];
564 for i in 0..t {
565 rms_plain(&h[i * d..(i + 1) * d], &mut out[i * d..(i + 1) * d]);
566 }
567 let _ = self.dh;
568 out
569 }
570}