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(
46                model,
47                &format!("{p}.attention_pooler.cross_attn.in_proj_weight"),
48            )?,
49            in_proj_b: vecf(
50                model,
51                &format!("{p}.attention_pooler.cross_attn.in_proj_bias"),
52            )?,
53            out_proj: Lin::load(
54                model,
55                &format!("{p}.attention_pooler.cross_attn.out_proj"),
56                true,
57            )?,
58            mlp_hidden: Lin::load(model, &format!("{p}.mlp_hidden"), true)?,
59            mlp_out: Lin::load(model, &format!("{p}.mlp_out"), true)?,
60            dim,
61            heads: 4,
62        })
63    }
64
65    /// Seconds, from the two connector outputs.
66    pub fn seconds(
67        &self,
68        video: &[f32],
69        audio: &[f32],
70        ctx_len: usize,
71        pool: Option<&crate::pool::Pool>,
72    ) -> f32 {
73        let d = self.dim;
74        let mut tokens = self.v_proj.apply(video, ctx_len, pool);
75        for (i, v) in tokens.iter_mut().enumerate() {
76            *v += self.v_emb[i % d];
77        }
78        let mut at = self.a_proj.apply(audio, ctx_len, pool);
79        for (i, v) in at.iter_mut().enumerate() {
80            *v += self.a_emb[i % d];
81        }
82        tokens.extend_from_slice(&at);
83        let n = 2 * ctx_len;
84
85        // one query, cross-attending every token: q from the learnable
86        // token, k and v from the stream. torch packs the three input
87        // projections into one matrix, in that order.
88        let proj = |x: &[f32], off: usize| -> Vec<f32> {
89            (0..d)
90                .map(|o| {
91                    let row = &self.in_proj_w[(off + o) * d..(off + o) * d + d];
92                    self.in_proj_b[off + o] + row.iter().zip(x).map(|(&a, &b)| a * b).sum::<f32>()
93                })
94                .collect()
95        };
96        let q = proj(&self.query, 0);
97        let hd = d / self.heads;
98        let mut ctx = vec![0f32; d];
99        let mut k = vec![0f32; d];
100        let mut v = vec![0f32; d];
101        let mut scores = vec![vec![0f32; n]; self.heads];
102        let mut vs = vec![0f32; n * d];
103        for t in 0..n {
104            let row = &tokens[t * d..(t + 1) * d];
105            k.copy_from_slice(&proj(row, d));
106            v.copy_from_slice(&proj(row, 2 * d));
107            vs[t * d..(t + 1) * d].copy_from_slice(&v);
108            for h in 0..self.heads {
109                let s: f32 = (0..hd).map(|i| q[h * hd + i] * k[h * hd + i]).sum();
110                scores[h][t] = s / (hd as f32).sqrt();
111            }
112        }
113        for h in 0..self.heads {
114            softmax(&mut scores[h]);
115            for (t, &p) in scores[h].iter().enumerate() {
116                for i in 0..hd {
117                    ctx[h * hd + i] += p * vs[t * d + h * hd + i];
118                }
119            }
120        }
121        let pooled = self.out_proj.apply(&ctx, 1, pool);
122        let mut hidden = self.mlp_hidden.apply(&pooled, 1, pool);
123        hidden.iter_mut().for_each(|v| *v = gelu_tanh(*v));
124        self.mlp_out.apply(&hidden, 1, pool)[0].exp()
125    }
126}
127
128/// A duration in seconds → a frame count on the VAE's `8k + 1` grid,
129/// clamped so a misbehaving prediction cannot ask for a degenerate or
130/// enormous generation.
131pub fn frames_for(seconds: f32, fps: f64, min_seconds: f64, max_seconds: f64) -> usize {
132    let scale = crate::ltxpipe::SCALE_TIME;
133    let min_frames = (min_seconds * fps).round().max(1.0) as usize;
134    let max_frames = (max_seconds * fps).round() as usize;
135    let raw = ((seconds as f64 * fps).round() as usize).clamp(min_frames, max_frames);
136    // snap down to the grid, then up if that fell under the floor
137    let frames = ((raw.saturating_sub(1)) / scale) * scale + 1;
138    if frames < min_frames {
139        (min_frames.saturating_sub(1)).div_ceil(scale) * scale + 1
140    } else {
141        frames
142    }
143    .min(max_frames)
144}
cortiq_engine/ltxdur.rs

cortiq_engine/
ltxdur.rs
