Skip to main content

mold_server/
chain_limits.rs

1//! Chain-limits computation for the `/api/capabilities/chain-limits` route.
2//!
3//! The model's hardcoded per-clip cap is the primary constraint; the
4//! hardware-derived recommended value is `min(cap, free_vram_adjusted)` and
5//! is inert for distilled LTX-2 today because 97 is model-capped.
6
7use serde::{Deserialize, Serialize};
8use utoipa::ToSchema;
9
10#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
11pub struct ChainLimits {
12    pub model: String,
13    pub frames_per_clip_cap: u32,
14    pub frames_per_clip_recommended: u32,
15    pub max_stages: u32,
16    pub max_total_frames: u32,
17    pub fade_frames_max: u32,
18    pub transition_modes: Vec<String>,
19    pub quantization_family: String,
20    /// Whether this model's family has an audio decode path. The SPA reads
21    /// this to decide whether to show the chain-level "Generate audio"
22    /// toggle; the chain endpoint refuses `enable_audio: true` upstream
23    /// when this is false. Single source of truth: `mold_inference::chain::capability_for_family`.
24    pub supports_audio: bool,
25}
26
27/// Per-model-family hardcoded caps. Keyed by the family string returned by
28/// `mold_core::manifest::resolve_family`.
29pub fn family_cap(family: &str) -> Option<u32> {
30    mold_inference::chain::capability_for_family(family).map(|c| c.frames_per_clip_cap)
31}
32
33/// Whether a chain-capable family also has an audio path. The chain handler
34/// rejects requests with `enable_audio: true` when this returns false, so
35/// users get a clear upfront error instead of silently-dropped audio.
36pub fn family_supports_audio(family: &str) -> bool {
37    mold_inference::chain::capability_for_family(family).is_some_and(|c| c.supports_audio)
38}
39
40/// Compute the chain-limits response for a resolved model name.
41///
42/// `family` is the canonical family string (e.g. "ltx2").
43/// `quant` is the quantization slug ("fp8", "fp16", "q8", ...).
44/// `free_vram_bytes` is the current free VRAM on the primary GPU.
45pub fn compute_limits(model: &str, family: &str, quant: &str, free_vram_bytes: u64) -> ChainLimits {
46    let cap = family_cap(family).unwrap_or(97);
47    // Hardware-derived recommended: for distilled LTX-2, 97 is already
48    // the binding constraint. Reserve the derivation scaffolding for
49    // future non-distilled models.
50    let _ = free_vram_bytes; // suppress unused for now; D wires this up
51    let recommended = cap;
52
53    const MAX_STAGES: u32 = 16;
54    ChainLimits {
55        model: model.to_string(),
56        frames_per_clip_cap: cap,
57        frames_per_clip_recommended: recommended,
58        max_stages: MAX_STAGES,
59        max_total_frames: cap * MAX_STAGES,
60        fade_frames_max: 32,
61        transition_modes: vec!["smooth".into(), "cut".into(), "fade".into()],
62        quantization_family: quant.to_string(),
63        supports_audio: family_supports_audio(family),
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[test]
72    fn ltx2_cap_is_97() {
73        // ltx2 family covers both v2 19B and v2.3 22B (dev and distilled);
74        // both resolve to family="ltx2" via `resolve_family`.
75        assert_eq!(family_cap("ltx2"), Some(97));
76    }
77
78    #[test]
79    fn ltx_video_cap_is_97() {
80        // LTX-Video uses the img2vid-less fallback; same per-clip cap as
81        // ltx2 because the chain endpoint stitches independent clips at the
82        // pixel level once the cap is hit.
83        assert_eq!(family_cap("ltx-video"), Some(97));
84    }
85
86    #[test]
87    fn audio_capability_is_ltx2_only() {
88        // Only the LTX-2 / LTX-2.3 AV transformer has an audio decode path.
89        // LTX-Video is video-only; FLUX/SDXL aren't in chain support at all.
90        assert!(family_supports_audio("ltx2"));
91        assert!(!family_supports_audio("ltx-video"));
92        assert!(!family_supports_audio("flux"));
93        assert!(!family_supports_audio(""));
94    }
95
96    #[test]
97    fn unknown_family_has_no_cap() {
98        assert_eq!(family_cap("flux"), None);
99        assert_eq!(family_cap("sdxl"), None);
100    }
101
102    #[test]
103    fn compute_limits_for_distilled() {
104        let lim = compute_limits("ltx-2-19b-distilled:fp8", "ltx2", "fp8", 8_000_000_000);
105        assert_eq!(lim.frames_per_clip_cap, 97);
106        assert_eq!(lim.frames_per_clip_recommended, 97);
107        assert_eq!(lim.max_stages, 16);
108        assert_eq!(lim.max_total_frames, 97 * 16);
109        assert_eq!(
110            lim.transition_modes,
111            vec!["smooth".to_string(), "cut".into(), "fade".into()]
112        );
113        assert!(
114            lim.supports_audio,
115            "ltx2 family has the AV transformer + audio VAE / vocoder path",
116        );
117    }
118
119    #[test]
120    fn compute_limits_for_ltx_video_has_no_audio() {
121        // LTX-Video is video-only; the SPA must hide the audio toggle and the
122        // chain endpoint will reject `enable_audio: true` upstream regardless.
123        let lim = compute_limits("ltx-video-0.9.7-distilled:fp8", "ltx-video", "fp8", 0);
124        assert!(
125            !lim.supports_audio,
126            "ltx-video has no audio path — toggle must stay off",
127        );
128    }
129}