h3-rs 0.2.0

//! End-to-end generation: prompt (+ optional first / last keyframes) -> MP4 with stereo audio.
//! Stages run sequentially (text encoder -> VAE encode -> DiT -> VAE decodes) and, with
//! `MemoryConfig::sequential`, each model is dropped before the next one loads.

use std::path::{Path, PathBuf};
use std::time::Instant;

use anyhow::{bail, Context, Result};
use candle_core::{DType, Device, Tensor};

use crate::config::{vram_info, MemoryConfig};
use crate::dit::{log, Dit, DitCond, KeyframeCond, RefCond, AUDIO_COND_TIMESTEP, FRAME_PER_TOKEN, VISUAL_COND_TIMESTEP};
use crate::lora::Lora;
use crate::media;
use crate::progress::{self, Stage};
use crate::sampling::{self, Sampler};
use crate::text::{RefItem, TextEncoder, Tokenizer};
use crate::upscaler::{self, LatentUpscaler};
use crate::vae_audio::{AudioEncoder, AudioVae};
use crate::vae_video::VideoVae;

pub const FPS: usize = 24;
/// DiT activation memory of a forward over `tokens` (hidden 5376): ~135 KB per token with the fused qkv, ~110 KB per
/// token once attention runs in head groups (long sequences, see `Dit::attn_groups`; measured on a 24 GB card with
/// 153k tokens: 13 resident blocks ran out of memory, 9 fit), plus the streamed weights' buffers. Used to size the
/// VRAM an upscaled request's second pass needs; too low is caught by the out-of-memory retry.
fn dit_activation_bytes(tokens: usize) -> usize {
    if 3 * tokens * 5376 * 2 <= 1_500_000_000 {
        tokens * 135_000 + (3 << 29)
    } else {
        tokens * 110_000 + (2 << 30)
    }
}
/// Upscaler activation memory per output latent voxel of a 32-frame segment (512 channels, f16, ~8 live tensors).
const UPSCALER_BYTES_PER_VOXEL: usize = 8_500;
const CHUNK_FRAMES: usize = 42;
pub const AUDIO_LATENT_FPS: usize = 40;

/// Model files, ComfyUI layout (`models/{diffusion_models,text_encoders,vae,loras}`).
#[derive(Debug, Clone)]
pub struct ModelPaths {
    pub dit: PathBuf,
    pub text_encoder: PathBuf,
    pub video_vae: PathBuf,
    pub audio_vae: PathBuf,
    /// Latent upscaler checkpoint (`models/latent_upscale_models/`), when one was found.
    pub upscaler: Option<PathBuf>,
}

/// Latent upscaler looked up by default.
pub const DEFAULT_UPSCALER: &str = "minimax_h3_latent_upscaler_3d_conv_v1_fp16.safetensors";

impl ModelPaths {
    /// Resolve each name against `<models_dir>/<subdir>/` unless it is already a path to a file.
    pub fn resolve(models_dir: &Path, dit: &str, te: &str, vvae: &str, avae: &str) -> Result<Self> {
        let r = |sub: &str, n: &str| -> Result<PathBuf> {
            let p = Path::new(n);
            if p.is_file() {
                return Ok(p.to_path_buf());
            }
            let c = models_dir.join(sub).join(n);
            if c.is_file() {
                return Ok(c);
            }
            bail!("model file {n} not found (looked in {})", models_dir.join(sub).display())
        };
        Ok(Self { dit: r("diffusion_models", dit)?, text_encoder: r("text_encoders", te)?, video_vae: r("vae", vvae)?, audio_vae: r("vae", avae)?, upscaler: None })
    }

    /// Look for the latent upscaler `name` (a path, or a file in `<models_dir>/latent_upscale_models/`).
    /// Without `required`, a missing file just leaves upscaling unavailable.
    pub fn with_upscaler(mut self, models_dir: &Path, name: &str, required: bool) -> Result<Self> {
        let direct = Path::new(name);
        let in_dir = models_dir.join("latent_upscale_models").join(name);
        self.upscaler = if direct.is_file() {
            Some(direct.to_path_buf())
        } else if in_dir.is_file() {
            Some(in_dir)
        } else if required {
            bail!("latent upscaler {name} not found (looked in {}); see the README for where to get it", models_dir.join("latent_upscale_models").display())
        } else {
            None
        };
        Ok(self)
    }
}

/// Two-pass generation through the latent upscaler (the ComfyUI "Minimax H3 Latent Upscaler" workflow):
/// the first `split` of the sampling steps run at the requested canvas, the clean estimate is upscaled
/// `scale` times in latent space (no VAE round trip), re-noised to where the schedule stands, and the remaining
/// steps refine it at the larger canvas, which is what gets delivered.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Upscale {
    /// Factor between the requested canvas (sampled first) and the delivered one, in (1, 4].
    pub scale: f64,
    /// Fraction of the steps sampled at the low resolution, in (0, 1).
    pub split: f64,
}

impl Default for Upscale {
    fn default() -> Self {
        Self { scale: 1.5, split: 0.5 }
    }
}

#[derive(Debug, Clone)]
pub struct Request {
    pub prompt: String,
    /// Explicit canvas; when unset it follows the first keyframe's (else `aspect`'s) aspect ratio
    /// at `megapixels`, or defaults to 1344x768.
    pub width: Option<usize>,
    pub height: Option<usize>,
    pub megapixels: f64,
    /// Aspect ratio `(w, h)` for an unset canvas without a keyframe.
    pub aspect: Option<(f64, f64)>,
    /// How keyframes are fitted to the canvas (default: first stretched, last center-cropped,
    /// like ComfyUI).
    pub fit: Option<Fit>,
    /// ref2va references (images, videos, audio), in tag order per type.
    pub refs: Vec<RefInput>,
    pub ref_image_size: RefImageSize,
    /// Frame count at 24 fps; snapped up to the 17k+5 grid.
    pub length: usize,
    pub steps: usize,
    pub seed: u64,
    pub sampler: Sampler,
    pub shift_video: f64,
    pub shift_audio: f64,
    pub first_frame: Option<String>,
    pub last_frame: Option<String>,
    pub lora: Option<(String, f64)>,
    /// Audio anchored at frame 0 (ComfyUI `MiniMaxH3AddGuide` with audio): the generated video and its
    /// soundtrack follow it, e.g. speech for a talking avatar. Cropped to the video's duration.
    pub audio_guide: Option<String>,
    /// Sample at the requested size, upscale the latent, refine at `scale` times that size (see [`Upscale`]).
    pub upscale: Option<Upscale>,
    pub output: String,
    pub save_latents: Option<String>,
    pub crf: u32,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
pub enum Fit {
    Stretch,
    Crop,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
pub enum RefImageSize {
    /// Scale each reference (down only) to the generation's pixel area.
    Match,
    /// 2048 px short edge (best identity fidelity, slower).
    Max,
}

#[derive(Debug, Clone)]
pub enum RefInput {
    /// `<Picture i>`
    Image(String),
    /// `<Video k>`; `with_audio` also uses its soundtrack as `<Audio j>` (or a separate file).
    Video { path: String, audio: Option<String> },
    /// `<Audio j>`
    Audio(String),
}

/// Round to a multiple of 32 (at least 32).
fn r32(v: f64) -> usize {
    ((v / 32.0).round() as usize * 32).max(32)
}

/// ComfyUI `ImageScaleToTotalPixels`: keep the aspect ratio at `megapixels` (x 1024^2), step 32.
pub fn canvas_for_aspect(w: f64, h: f64, megapixels: f64) -> (usize, usize) {
    let s = (megapixels * 1024.0 * 1024.0 / (w * h)).sqrt();
    (r32(w * s), r32(h * s))
}

/// The reference pipeline's video canvas: 768 short edge, 768*1344 area cap, step 32.
fn adapt_canvas(width: usize, height: usize) -> (usize, usize) {
    let ratio = width as f64 / height as f64;
    let (mut w, mut h) = if ratio >= 1.0 { (768.0 * ratio, 768.0) } else { (768.0, 768.0 / ratio) };
    let max = 768.0 * 1344.0;
    if w * h > max {
        let s = (max / (w * h)).sqrt();
        w *= s;
        h *= s;
    }
    (r32(w), r32(h))
}

impl Default for Request {
    fn default() -> Self {
        Self {
            prompt: String::new(),
            width: None,
            height: None,
            megapixels: 1.0,
            aspect: None,
            fit: None,
            refs: vec![],
            ref_image_size: RefImageSize::Match,
            length: 124,
            steps: 20,
            seed: 1,
            sampler: Sampler::ResMultistep,
            shift_video: 12.0,
            shift_audio: 3.0,
            first_frame: None,
            last_frame: None,
            lora: None,
            audio_guide: None,
            upscale: None,
            output: "h3_output.mp4".into(),
            save_latents: None,
            crf: 18,
        }
    }
}

pub fn align_frame_count(mut n: usize) -> usize {
    while n % 17 != 5 {
        n += 1;
    }
    n
}

pub fn video_latent_t(frames: usize) -> usize {
    if frames <= 5 {
        2
    } else {
        ((frames - 5) / 17) * 5 + 2
    }
}

/// (frame_count, latent_t, audio_t)
pub fn temporal_shape(length: usize) -> (usize, usize, usize) {
    let fc = align_frame_count(length.max(5));
    (fc, video_latent_t(fc), (fc as f64 / FPS as f64 * AUDIO_LATENT_FPS as f64).round() as usize)
}

/// Loaded models, reused across requests when `keep` is set (the server); otherwise each model is
/// dropped as soon as its stage finishes (the CLI, `MemoryConfig::sequential`).
pub struct Engine {
    pub paths: ModelPaths,
    /// DiT used when a request has references.
    pub ref_dit: Option<PathBuf>,
    pub mem: MemoryConfig,
    pub dev: Device,
    pub keep: bool,
    tok: Option<Tokenizer>,
    te: Option<TextEncoder>,
    dit: Option<((PathBuf, Option<(String, u64)>), Dit)>,
    vvae: Option<VideoVae>,
    avae: Option<AudioVae>,
    aenc: Option<AudioEncoder>,
}

impl Engine {
    pub fn new(paths: ModelPaths, ref_dit: Option<PathBuf>, mem: MemoryConfig, dev: Device, keep: bool) -> Self {
        Self { paths, ref_dit, mem, dev, keep, tok: None, te: None, dit: None, vvae: None, avae: None, aenc: None }
    }

    /// Load everything a plain t2va/fl2va request needs (server warm-up). The DiT loads last so its
    /// automatic GPU residency sees what the other models left free.
    pub fn preload_with(&mut self, lora: Option<&(String, f64)>) -> Result<()> {
        self.tokenizer()?;
        self.text_encoder()?;
        self.video_vae()?;
        self.audio_vae()?;
        self.audio_encoder()?;
        let p = self.paths.dit.clone();
        self.dit(&p, lora)?;
        Ok(())
    }

    fn tokenizer(&mut self) -> Result<&Tokenizer> {
        if self.tok.is_none() {
            self.tok = Some(Tokenizer::load_default()?);
        }
        Ok(self.tok.as_ref().unwrap())
    }

    fn text_encoder(&mut self) -> Result<&TextEncoder> {
        if self.te.is_none() {
            self.te = Some(TextEncoder::load(&self.paths.text_encoder.display().to_string(), &self.dev, &self.mem)?);
            log::info(&format!("text encoder loaded ({})", vram()));
        }
        Ok(self.te.as_ref().unwrap())
    }

    pub fn video_vae(&mut self) -> Result<&VideoVae> {
        if self.vvae.is_none() {
            self.vvae = Some(VideoVae::load(&self.paths.video_vae.display().to_string(), &self.dev, &self.mem)?);
        }
        Ok(self.vvae.as_ref().unwrap())
    }

    fn audio_vae(&mut self) -> Result<&AudioVae> {
        if self.avae.is_none() {
            self.avae = Some(AudioVae::load(&self.paths.audio_vae.display().to_string(), &self.dev)?);
        }
        Ok(self.avae.as_ref().unwrap())
    }

    fn audio_encoder(&mut self) -> Result<&AudioEncoder> {
        if self.aenc.is_none() {
            self.aenc = Some(AudioEncoder::load(&self.paths.audio_vae.display().to_string(), &self.dev)?);
        }
        Ok(self.aenc.as_ref().unwrap())
    }

    fn dit(&mut self, path: &Path, lora: Option<&(String, f64)>) -> Result<&mut Dit> {
        let key = (path.to_path_buf(), lora.map(|(p, s)| (p.clone(), s.to_bits())));
        if self.dit.as_ref().map(|(k, _)| k != &key).unwrap_or(true) {
            self.dit = None; // free the old one first
            let t = Instant::now();
            let l = lora.map(|(p, s)| Lora::open(p).map(|l| (l, *s))).transpose()?;
            let d = Dit::load(&path.display().to_string(), &self.dev, &self.mem, l.as_ref().map(|(l, s)| (l, *s)))?;
            log::info(&format!("DiT loaded in {:.1}s ({})", t.elapsed().as_secs_f64(), vram()));
            self.dit = Some((key, d));
        }
        Ok(&mut self.dit.as_mut().unwrap().1)
    }

    /// Drop a stage's model unless the engine keeps models loaded.
    fn release(&mut self, what: &str) {
        if self.keep {
            return;
        }
        match what {
            "te" => self.te = None,
            "dit" => self.dit = None,
            "vvae" => self.vvae = None,
            "avae" => self.avae = None,
            "aenc" => self.aenc = None,
            _ => {}
        }
    }
}

fn vram() -> String {
    match vram_info() {
        Ok((f, t)) => format!("{:.1}/{:.1} GB VRAM used", (t - f) as f64 / 1e9, t as f64 / 1e9),
        Err(_) => String::new(),
    }
}

pub fn generate(req: &Request, paths: &ModelPaths, mem: &MemoryConfig, dev: &Device) -> Result<()> {
    let mut eng = Engine::new(paths.clone(), None, mem.clone(), dev.clone(), false);
    eng.generate(req)
}

impl Engine {
pub fn generate(&mut self, req: &Request) -> Result<()> {
    let dev = &self.dev.clone();
    let t_start = Instant::now();
    let first = req.first_frame.as_deref().map(media::load_image).transpose()?;
    let last = req.last_frame.as_deref().map(media::load_image).transpose()?;
    let src_aspect = first.as_ref().or(last.as_ref()).map(|i| (i.dim(1).unwrap() as f64, i.dim(0).unwrap() as f64)).or(req.aspect);
    let (w, h) = match (req.width, req.height, src_aspect) {
        (Some(w), Some(h), _) => (w / 32 * 32, h / 32 * 32),
        (Some(w), None, a) => {
            let (aw, ah) = a.unwrap_or((16.0, 9.0));
            (w / 32 * 32, r32(w as f64 * ah / aw))
        }
        (None, Some(h), a) => {
            let (aw, ah) = a.unwrap_or((16.0, 9.0));
            (r32(h as f64 * aw / ah), h / 32 * 32)
        }
        (None, None, Some((aw, ah))) => canvas_for_aspect(aw, ah, req.megapixels),
        (None, None, None) => (1344, 768),
    };
    if w == 0 || h == 0 {
        bail!("width/height must be at least 32");
    }
    // `(w, h)` is the requested canvas, sampled first; with an upscale `(fw, fh)` is the larger one delivered
    let (fw, fh) = match &req.upscale {
        Some(u) => {
            if !(u.scale > 1.0 && u.scale <= upscaler::MAX_SCALE) {
                bail!("upscale must be above 1 and at most {}", upscaler::MAX_SCALE);
            }
            if !(u.split > 0.0 && u.split < 1.0) || req.steps < 2 {
                bail!("upscale needs at least 2 steps and a split strictly between 0 and 1");
            }
            let (fw, fh) = (r32(w as f64 * u.scale), r32(h as f64 * u.scale));
            if fw <= w && fh <= h {
                bail!("upscale {} leaves a {w}x{h} canvas unchanged", u.scale);
            }
            (fw, fh)
        }
        None => (w, h),
    };
    let (frame_count, latent_t, audio_t) = temporal_shape(req.length);
    progress::stage(Stage::Encoding);
    log::info(&format!(
        "{fw}x{fh}{}, {frame_count} frames ({:.2} s) -> latent {latent_t}x{}x{}, audio {audio_t}",
        if req.upscale.is_some() { format!(" (sampled at {w}x{h}, then upscaled {:.2}x)", fw as f64 / w as f64) } else { String::new() },
        frame_count as f64 / FPS as f64,
        h / 16,
        w / 16
    ));

    // keyframe images on the canvas
    let mut images = vec![];
    let mut kf_index = vec![];
    if let Some(img) = &first {
        images.push(media::resize_lanczos(img, w, h, req.fit == Some(Fit::Crop))?);
        kf_index.push(0usize);
    }
    if let Some(img) = &last {
        images.push(media::resize_lanczos(img, w, h, req.fit != Some(Fit::Stretch))?);
        kf_index.push(frame_count - 1);
    }

    // the same keyframes on the delivered canvas, for the refining pass
    let mut images_final = vec![];
    if req.upscale.is_some() {
        if let Some(img) = &first {
            images_final.push(media::resize_lanczos(img, fw, fh, req.fit == Some(Fit::Crop))?);
        }
        if let Some(img) = &last {
            images_final.push(media::resize_lanczos(img, fw, fh, req.fit != Some(Fit::Stretch))?);
        }
    }

    // references (MiniMaxH3ReferenceToVideo): images, then videos (soundtrack label first), then audio
    let mut ref_items = vec![];
    let mut ref_imgs: Vec<Tensor> = vec![];
    let mut ref_vids: Vec<(Tensor, Option<Tensor>)> = vec![];
    let mut ref_auds: Vec<Tensor> = vec![];
    let ordered: Vec<&RefInput> = req
        .refs
        .iter()
        .filter(|r| matches!(r, RefInput::Image(_)))
        .chain(req.refs.iter().filter(|r| matches!(r, RefInput::Video { .. })))
        .chain(req.refs.iter().filter(|r| matches!(r, RefInput::Audio(_))))
        .collect();
    for r in ordered {
        match r {
            RefInput::Image(p) => {
                let img = media::load_image(p)?;
                let (ih, iw) = (img.dim(0)? as f64, img.dim(1)? as f64);
                let scale = match req.ref_image_size {
                    RefImageSize::Match => (((w * h) as f64) / (iw * ih)).sqrt().min(1.0),
                    RefImageSize::Max => (2048.0 / iw.min(ih)).min(1.0),
                };
                let img = media::resize_lanczos(&img, r32(iw * scale), r32(ih * scale), false)?;
                ref_items.push(RefItem::Image(img.clone()));
                ref_imgs.push(img);
            }
            RefInput::Video { path, audio } => {
                let (vw, vh) = media::video_size(path)?;
                let (mut cw, mut ch) = adapt_canvas(vw, vh);
                if vw * vh < cw * ch {
                    (cw, ch) = (r32(vw as f64), r32(vh as f64));
                }
                let raw = media::load_video(path, FPS as f64, frame_count)?;
                let mut n = raw.dim(0)?;
                if n < 5 {
                    bail!("reference video {path} needs at least 5 frames (~0.2 s at 24 fps)");
                }
                while n % 17 != 5 {
                    n -= 1;
                }
                let mut fr = vec![];
                for i in 0..n {
                    fr.push(media::resize_lanczos(&raw.get(i)?, cw, ch, false)?);
                }
                let frames = Tensor::stack(&fr, 0)?; // [n, ch, cw, 3]
                let sound = match audio {
                    Some(a) => Some(media::load_audio(a, 32000, None)?),
                    None => None,
                };
                if sound.is_some() {
                    ref_items.push(RefItem::Audio);
                }
                // Qwen sees the clip at 2 fps with timestamps
                let idx: Vec<usize> = (0..n).step_by(FPS / 2).collect();
                let q = Tensor::stack(&idx.iter().map(|&i| frames.get(i)).collect::<candle_core::Result<Vec<_>>>()?, 0)?;
                ref_items.push(RefItem::Video { frames: q, timestamps: (0..idx.len()).map(|i| i as f64 / 2.0).collect() });
                log::info(&format!("reference video {path}: {n} frames at {cw}x{ch}{}", if sound.is_some() { " + soundtrack" } else { "" }));
                ref_vids.push((frames, sound));
            }
            RefInput::Audio(p) => {
                ref_items.push(RefItem::Audio);
                ref_auds.push(media::load_audio(p, 32000, None)?);
            }
        }
    }
    let ref_mode = !req.refs.is_empty();

    // 1. text (+ vision) conditioning
    let t0 = Instant::now();
    let tok = self.tokenizer()?;
    let pieces = if ref_mode {
        if !images.is_empty() {
            bail!("keyframes (--image/--last-image) and references can't be combined");
        }
        tok.ref_presentation(&req.prompt, &ref_items)?
    } else {
        tok.presentation(&req.prompt, &images)?
    };
    let enc = self.text_encoder()?.encode(&pieces)?;
    self.release("te");
    log::info(&format!("conditioning: {} tokens in {:.1}s", enc.hidden.dim(0)?, t0.elapsed().as_secs_f64()));

    // 2. keyframe / reference latents
    let mut keyframes = vec![];
    let mut keyframes_final = vec![];
    let mut refs = vec![];
    if !images.is_empty() || !ref_imgs.is_empty() || !ref_vids.is_empty() {
        let vae = self.video_vae()?;
        for (img, idx) in images.iter().zip(&kf_index) {
            keyframes.push(KeyframeCond { frame_index: *idx, video: Some(vae.encode_image(img)?), audio: None });
        }
        for (img, idx) in images_final.iter().zip(&kf_index) {
            keyframes_final.push(KeyframeCond { frame_index: *idx, video: Some(vae.encode_image(img)?), audio: None });
        }
        for img in &ref_imgs {
            refs.push(RefCond::Image { latent: vae.encode_image(img)? });
        }
        for (frames, _) in &ref_vids {
            refs.push(RefCond::Video { latent: vae.encode_frames(frames)?, audio: None });
        }
        self.release("vvae");
    }
    if let Some(path) = &req.audio_guide {
        let wave = media::load_audio(path, 32000, None)?;
        let z = self.audio_encoder()?.encode(&wave)?;
        // the streams share one time axis; at frame 0 the guide may span the whole soundtrack
        let rt = z.dim(3)?.min(audio_t);
        let z = z.narrow(3, 0, rt)?.contiguous()?;
        log::info(&format!("audio guide {path}: {:.2}s", rt as f64 / AUDIO_LATENT_FPS as f64));
        for (pass, kfs) in [&mut keyframes, &mut keyframes_final].into_iter().enumerate() {
            if pass == 1 && req.upscale.is_none() {
                continue;
            }
            match kfs.iter_mut().find(|k| k.frame_index == 0) {
                Some(kf) => kf.audio = Some(z.clone()),
                None => kfs.insert(0, KeyframeCond { frame_index: 0, video: None, audio: Some(z.clone()) }),
            }
        }
        self.release("aenc");
    }
    if ref_vids.iter().any(|v| v.1.is_some()) || !ref_auds.is_empty() {
        let aenc = self.audio_encoder()?;
        let mut k = 0;
        for r in refs.iter_mut() {
            if let RefCond::Video { audio, .. } = r {
                if let Some(s) = &ref_vids[k].1 {
                    *audio = Some(aenc.encode(s)?);
                }
                k += 1;
            }
        }
        for a in &ref_auds {
            refs.push(RefCond::Audio { latent: aenc.encode(a)? });
        }
        self.release("aenc");
    }

    // 3. sample (with an upscale: first part small, latent upscale, rest at the delivered size)
    let dit_path = if ref_mode { self.ref_dit.clone().unwrap_or(self.paths.dit.clone()) } else { self.paths.dit.clone() };
    let upscaler_path = self.paths.upscaler.clone();
    if req.upscale.is_some() && upscaler_path.is_none() {
        bail!("no latent upscaler model: put {DEFAULT_UPSCALER} in models/latent_upscale_models/ (or pass --upscaler)");
    }
    let dev_for_upscaler = self.dev.clone();
    let keep_models = self.keep;
    let dit = self.dit(&dit_path, req.lora.as_ref())?;
    dit.cfg.shift_video = req.shift_video;
    dit.cfg.shift_audio = req.shift_audio;
    let sigmas = sampling::simple_sigmas(req.shift_video, req.steps);
    // steps sampled at the first size
    let split = match &req.upscale {
        Some(u) => ((req.steps as f64 * u.split).round() as usize).clamp(1, req.steps - 1),
        None => req.steps,
    };
    let cond = DitCond {
        text: enc.hidden,
        tags: Some(enc.tags),
        keyframes,
        refs,
        seed: req.seed,
        visual_cond_noise_aug: VISUAL_COND_TIMESTEP,
        audio_cond_noise_aug: AUDIO_COND_TIMESTEP,
    };
    let cond_final = req.upscale.map(|_| DitCond { keyframes: std::mem::take(&mut keyframes_final), ..cond.clone() });
    let vshape = [1usize, 24, latent_t, h / 16, w / 16];
    let run = dit.prepare(&cond, &vshape, audio_t, sigmas[..=split].to_vec())?;
    drop(cond);
    // noise for the packed latent (audio carried on the video schedule): sigma_max = 1
    // same stream as ComfyUI's RandomNoise: one torch.manual_seed(seed) generator, video then audio
    let mut gen = crate::rng::TorchGen::new(req.seed);
    let xv = gen.randn_tensor(&vshape, dev)?;
    let xa = gen.randn_tensor(&[1, 32, 2, audio_t], dev)?;
    let xv = (xv * sigmas[0])?;
    let xa = (xa * sigmas[0])?;
    let t2 = Instant::now();
    let mut last = Instant::now();
    progress::count(Stage::Sampling, 0, req.steps);
    let step_log = |done: usize, last: &mut Instant| {
        progress::count(Stage::Sampling, done, req.steps);
        log::info(&format!("step {done}/{}  {:.3}s/it  ({})", req.steps, last.elapsed().as_secs_f64(), vram()));
        *last = Instant::now();
    };
    let first_pass = sampling::sample_denoised(&*dit, &run, xv, xa, &sigmas[..=split], req.sampler, |i, _| step_log(i, &mut last))?;
    drop(run);
    let (zv, za) = match (&req.upscale, cond_final, upscaler_path) {
        (Some(_), Some(cond_final), Some(up_path)) => {
            drop(first_pass.xv);
            // The DiT was loaded with headroom for the first canvas. The upscaler's activations and then a pass with
            // ~scale^2 times the tokens (activations grow with them) need more: stream some resident blocks from
            // RAM to make room, and stream more and redo the stage if the estimate was short.
            let cells = (fh / 16, fw / 16);
            let tokens = latent_t * (fh / 32) * (fw / 32);
            let upscaler_need = (CHUNK_FRAMES.min(latent_t + 10) * cells.0 * cells.1) * UPSCALER_BYTES_PER_VOXEL + (6 << 28);
            let need = dit_activation_bytes(tokens).max(upscaler_need);
            crate::kernels::trim_pool();
            let free = vram_info().map(|(f, _)| f).unwrap_or(usize::MAX);
            let base_resident = dit.resident_blocks();
            let mut keep = base_resident;
            if free < need {
                keep = base_resident.saturating_sub((need - free).div_ceil(dit.block_bytes().max(1)));
                log::info(&format!("upscaling and a {tokens}-token second pass need ~{:.1} GB: {keep}/{base_resident} DiT blocks stay resident", need as f64 / 1e9));
                dit.set_resident(keep)?;
                crate::kernels::trim_pool();
            }
            let is_oom = |e: &anyhow::Error| format!("{e:?}").contains("OUT_OF_MEMORY");
            let mut attempt = 0;
            // the clean estimate of the first pass, upscaled in latent space
            let zv_up = loop {
                attempt += 1;
                match upscale_video_latent(&up_path, &dev_for_upscaler, &first_pass.denoised_v, cells) {
                    Ok(z) => break z,
                    Err(e) if keep > 0 && attempt < 6 && is_oom(&e) => {
                        keep = keep.saturating_sub(keep.div_ceil(4).max(4));
                        log::warn(&format!("out of memory upscaling, retrying with {keep} resident DiT blocks"));
                        dit.set_resident(keep)?;
                        crate::kernels::trim_pool();
                    }
                    Err(e) => return Err(e),
                }
            };
            let sig = &sigmas[split..];
            let vshape = [1usize, 24, latent_t, fh / 16, fw / 16];
            attempt = 0;
            let out = loop {
                attempt += 1;
                let r = (|| -> Result<(Tensor, Tensor)> {
                    // re-noise to where the schedule stands (CONST: x = sigma * noise + (1 - sigma) * latent), then finish
                    let run = dit.prepare(&cond_final, &vshape, audio_t, sig.to_vec())?;
                    let mut gen = crate::rng::TorchGen::new(req.seed);
                    let nv = gen.randn_tensor(&vshape, dev)?;
                    let na = gen.randn_tensor(&[1, 32, 2, audio_t], dev)?;
                    let xv = ((nv * sig[0])? + (&zv_up * (1.0 - sig[0]))?)?;
                    let xa = ((na * sig[0])? + (&first_pass.denoised_a * (1.0 - sig[0]))?)?;
                    progress::count(Stage::Sampling, split, req.steps);
                    last = Instant::now();
                    sampling::sample(&*dit, &run, xv, xa, sig, req.sampler, |i, _| step_log(split + i, &mut last))
                })();
                match r {
                    Ok(o) => break o,
                    Err(e) if keep > 0 && attempt < 6 && is_oom(&e) => {
                        keep = keep.saturating_sub(keep.div_ceil(4).max(4));
                        log::warn(&format!("out of memory in the second pass, retrying with {keep} resident DiT blocks"));
                        dit.set_resident(keep)?;
                        crate::kernels::trim_pool();
                    }
                    Err(e) => return Err(e),
                }
            };
            drop(cond_final);
            // the next request starts with the blocks it was loaded with
            if keep_models {
                crate::kernels::trim_pool();
                dit.set_resident(base_resident)?;
            }
            out
        }
        _ => (first_pass.xv, first_pass.xa),
    };
    let audio_scale = dit.audio_scale();
    log::info(&format!("sampling done in {:.1}s", t2.elapsed().as_secs_f64()));
    self.release("dit");
    let za = (za / audio_scale)?; // carried -> own audio latent

    if let Some(p) = &req.save_latents {
        save_latents(p, &zv, &za)?;
        log::info(&format!("latents saved to {p}"));
    }

    self.decode_and_write(&zv, &za, frame_count, &req.output, req.crf)?;
    log::info(&format!("done in {:.1}s -> {}", t_start.elapsed().as_secs_f64(), req.output));
    Ok(())
}

/// Decode a video latent to `[F, 3, H, W]` frames in [0, 1] (at most `frame_count` of them).
pub fn decode_frames(&mut self, zv: &Tensor, frame_count: usize) -> Result<Tensor> {
    let t = Instant::now();
    progress::stage(Stage::Decoding);
    let frames = self.video_vae()?.decode(zv)?;
    self.release("vvae");
    let nf = frames.dim(1)?.min(frame_count);
    let frames = frames.narrow(1, 0, nf)?.permute((1, 0, 2, 3))?.contiguous()?; // [F, 3, H, W]
    log::info(&format!("video decoded: {nf} frames in {:.1}s", t.elapsed().as_secs_f64()));
    Ok(frames)
}

/// Decode latents and mux them into an MP4.
pub fn decode_and_write(&mut self, zv: &Tensor, za: &Tensor, frame_count: usize, out: &str, crf: u32) -> Result<()> {
    let frames = self.decode_frames(zv, frame_count)?;
    let nf = frames.dim(0)?;
    let t = Instant::now();
    progress::stage(Stage::Muxing);
    let avae = self.audio_vae()?;
    let sr = avae.sample_rate;
    let wave = avae.decode(za)?;
    self.release("avae");
    let samples = ((nf as f64 / FPS as f64) * sr as f64).round() as usize;
    let wave = wave.narrow(1, 0, samples.min(wave.dim(1)?))?;
    log::info(&format!("audio decoded: {:.2}s in {:.1}s", wave.dim(1)? as f64 / sr as f64, t.elapsed().as_secs_f64()));
    media::write_mp4(out, &frames, FPS as f64, Some((&wave, sr)), crf)?;
    if out.ends_with(".mp4") {
        let _ = media::write_wav(&out.replace(".mp4", ".wav"), &wave, sr);
    }
    Ok(())
}
}

/// Decode latents and mux them into an MP4 (one-shot).
#[allow(clippy::too_many_arguments)]
pub fn decode_and_write(zv: &Tensor, za: &Tensor, frame_count: usize, paths: &ModelPaths, mem: &MemoryConfig, dev: &Device, out: &str, crf: u32) -> Result<()> {
    Engine::new(paths.clone(), None, mem.clone(), dev.clone(), false).decode_and_write(zv, za, frame_count, out, crf)
}

/// Frame count implied by a video latent with `t` frames (5k+2 -> 17k+5).
pub fn frames_for_latent(t: usize) -> usize {
    (0..t).map(|k| FRAME_PER_TOKEN[k % 5]).sum()
}

pub fn load_latents(path: &str, dev: &Device) -> Result<(Tensor, Tensor)> {
    let m = candle_core::safetensors::load(path, dev)?;
    Ok((m.get("video").context("video latent")?.to_dtype(DType::F32)?, m.get("audio").context("audio latent")?.to_dtype(DType::F32)?))
}

/// Write sampled latents (`video` `[1, 24, T, H, W]`, `audio` `[1, 32, 2, T]`) for `h3 decode` / `h3 upscale`.
pub fn save_latents(path: &str, zv: &Tensor, za: &Tensor) -> Result<()> {
    let m: std::collections::HashMap<String, Tensor> =
        [("video".to_string(), zv.to_device(&Device::Cpu)?), ("audio".to_string(), za.to_device(&Device::Cpu)?)].into_iter().collect();
    candle_core::safetensors::save(&m, path)?;
    Ok(())
}

/// Upscale a video latent `[1, 24, T, H, W]` to `cells = (height, width)` latent cells with the upscaler
/// checkpoint at `path`. The model is loaded for the call and dropped again, so it never competes with the DiT.
pub fn upscale_video_latent(path: &Path, dev: &Device, zv: &Tensor, cells: (usize, usize)) -> Result<Tensor> {
    let t = Instant::now();
    progress::stage(Stage::Upscaling);
    let up = LatentUpscaler::load(path, dev)?;
    let from = (zv.dim(3)?, zv.dim(4)?);
    let out = up.upscale(zv, cells.0, cells.1, None)?;
    dev.synchronize()?;
    log::info(&format!(
        "latent upscaled {}x{} -> {}x{} ({:.2}x) in {:.1}s ({})",
        from.1,
        from.0,
        cells.1,
        cells.0,
        LatentUpscaler::scale_for(from, cells),
        t.elapsed().as_secs_f64(),
        vram()
    ));
    Ok(out)
}