cortiq_engine/
ltxdur.rs

1//! The LTX-2.5 duration head: how long the shot the prompt describes should
2//! be, read off the connector outputs before a single denoising step runs.
3//!
4//! Both connector streams are projected into a shared 256-wide space, tagged
5//! with a learnable per-modality embedding so the pooler can tell them
6//! apart, and cross-attended by one learnable query. A two-layer MLP turns
7//! the pooled vector into a *log*-duration — the head was trained in log
8//! seconds so its loss spreads evenly across orders of magnitude — and the
9//! exponential of that is the answer.
10
11use crate::ltxdit::{Lin, gelu_tanh, softmax};
12use cortiq_core::CmfModel;
13use std::sync::Arc;
14
15fn vecf(model: &Arc<CmfModel>, name: &str) -> Result<Vec<f32>, String> {
16    crate::dit::cmf_f32(model, name)
17}
18
19pub struct DurationHead {
20    v_proj: Lin,
21    a_proj: Lin,
22    v_emb: Vec<f32>,
23    a_emb: Vec<f32>,
24    query: Vec<f32>,
25    in_proj_w: Vec<f32>,
26    in_proj_b: Vec<f32>,
27    out_proj: Lin,
28    mlp_hidden: Lin,
29    mlp_out: Lin,
30    dim: usize,
31    heads: usize,
32}
33
34impl DurationHead {
35    pub fn from_cmf(model: &Arc<CmfModel>) -> Result<DurationHead, String> {
36        let p = "dhead.duration_head";
37        let query = vecf(model, &format!("{p}.attention_pooler.query_tokens"))?;
38        let dim = query.len();
39        Ok(DurationHead {
40            v_proj: Lin::load(model, &format!("{p}.video_input_proj"), true)?,
41            a_proj: Lin::load(model, &format!("{p}.audio_input_proj"), true)?,
42            v_emb: vecf(model, &format!("{p}.video_modality_emb"))?,
43            a_emb: vecf(model, &format!("{p}.audio_modality_emb"))?,
44            query,
45            in_proj_w: vecf(model, &format!("{p}.attention_pooler.cross_attn.in_proj_weight"))?,
46            in_proj_b: vecf(model, &format!("{p}.attention_pooler.cross_attn.in_proj_bias"))?,
47            out_proj: Lin::load(model, &format!("{p}.attention_pooler.cross_attn.out_proj"), true)?,
48            mlp_hidden: Lin::load(model, &format!("{p}.mlp_hidden"), true)?,
49            mlp_out: Lin::load(model, &format!("{p}.mlp_out"), true)?,
50            dim,
51            heads: 4,
52        })
53    }
54
55    /// Seconds, from the two connector outputs.
56    pub fn seconds(
57        &self,
58        video: &[f32],
59        audio: &[f32],
60        ctx_len: usize,
61        pool: Option<&crate::pool::Pool>,
62    ) -> f32 {
63        let d = self.dim;
64        let mut tokens = self.v_proj.apply(video, ctx_len, pool);
65        for (i, v) in tokens.iter_mut().enumerate() {
66            *v += self.v_emb[i % d];
67        }
68        let mut at = self.a_proj.apply(audio, ctx_len, pool);
69        for (i, v) in at.iter_mut().enumerate() {
70            *v += self.a_emb[i % d];
71        }
72        tokens.extend_from_slice(&at);
73        let n = 2 * ctx_len;
74
75        // one query, cross-attending every token: q from the learnable
76        // token, k and v from the stream. torch packs the three input
77        // projections into one matrix, in that order.
78        let proj = |x: &[f32], off: usize| -> Vec<f32> {
79            (0..d)
80                .map(|o| {
81                    let row = &self.in_proj_w[(off + o) * d..(off + o) * d + d];
82                    self.in_proj_b[off + o] + row.iter().zip(x).map(|(&a, &b)| a * b).sum::<f32>()
83                })
84                .collect()
85        };
86        let q = proj(&self.query, 0);
87        let hd = d / self.heads;
88        let mut ctx = vec![0f32; d];
89        let mut k = vec![0f32; d];
90        let mut v = vec![0f32; d];
91        let mut scores = vec![vec![0f32; n]; self.heads];
92        let mut vs = vec![0f32; n * d];
93        for t in 0..n {
94            let row = &tokens[t * d..(t + 1) * d];
95            k.copy_from_slice(&proj(row, d));
96            v.copy_from_slice(&proj(row, 2 * d));
97            vs[t * d..(t + 1) * d].copy_from_slice(&v);
98            for h in 0..self.heads {
99                let s: f32 = (0..hd).map(|i| q[h * hd + i] * k[h * hd + i]).sum();
100                scores[h][t] = s / (hd as f32).sqrt();
101            }
102        }
103        for h in 0..self.heads {
104            softmax(&mut scores[h]);
105            for (t, &p) in scores[h].iter().enumerate() {
106                for i in 0..hd {
107                    ctx[h * hd + i] += p * vs[t * d + h * hd + i];
108                }
109            }
110        }
111        let pooled = self.out_proj.apply(&ctx, 1, pool);
112        let mut hidden = self.mlp_hidden.apply(&pooled, 1, pool);
113        hidden.iter_mut().for_each(|v| *v = gelu_tanh(*v));
114        self.mlp_out.apply(&hidden, 1, pool)[0].exp()
115    }
116}
117
118/// A duration in seconds → a frame count on the VAE's `8k + 1` grid,
119/// clamped so a misbehaving prediction cannot ask for a degenerate or
120/// enormous generation.
121pub fn frames_for(seconds: f32, fps: f64, min_seconds: f64, max_seconds: f64) -> usize {
122    let scale = crate::ltxpipe::SCALE_TIME;
123    let min_frames = (min_seconds * fps).round().max(1.0) as usize;
124    let max_frames = (max_seconds * fps).round() as usize;
125    let raw = ((seconds as f64 * fps).round() as usize).clamp(min_frames, max_frames);
126    // snap down to the grid, then up if that fell under the floor
127    let frames = ((raw.saturating_sub(1)) / scale) * scale + 1;
128    if frames < min_frames {
129        (min_frames.saturating_sub(1)).div_ceil(scale) * scale + 1
130    } else {
131        frames
132    }
133    .min(max_frames)
134}
cortiq_engine/ltxdur.rs

cortiq_engine/
ltxdur.rs
