Skip to main content

viser_ffmpeg/
encode.rs

1use std::path::{Path, PathBuf};
2use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
3use tokio::io::AsyncBufReadExt;
4use tokio::process::Command;
5
6use crate::{Codec, EncoderBackend, RateControlMode, Resolution, ffmpeg_path, probe};
7
8/// Parameters for a single encode.
9#[derive(Debug, Clone)]
10pub struct EncodeJob {
11    /// Source media file path.
12    pub input: String,
13    /// Destination file path for the encoded output.
14    pub output: String,
15    /// Optional target resolution; when set, scales with the lanczos filter.
16    pub resolution: Option<Resolution>,
17    /// Video codec to encode with.
18    pub codec: Codec,
19    /// Constant rate factor / quantizer value (interpretation depends on `rate_control`).
20    pub crf: i32,
21    /// Rate-control mode that determines how `crf`/bitrate fields are applied.
22    pub rate_control: RateControlMode,
23    /// Target bitrate in kbps; used for VBR mode.
24    pub target_bitrate: f64, // kbps, used for VBR mode
25    /// Maximum bitrate cap in kbps; used for capped CRF mode.
26    pub max_bitrate: f64, // kbps, used for capped CRF mode
27    /// VBV buffer size in kbps; used for capped CRF mode.
28    pub bufsize: f64, // kbps, used for capped CRF mode
29    /// Encoder speed preset (e.g. `"medium"`); empty leaves the encoder default.
30    pub preset: String,
31    /// Extra raw FFmpeg arguments appended verbatim before the output path.
32    pub extra_args: Vec<String>,
33}
34
35/// Output of a completed encode.
36#[derive(Debug, Clone)]
37pub struct EncodeResult {
38    /// The job that produced this result.
39    pub job: EncodeJob,
40    /// Average bitrate of the output in kbps, measured by probing it.
41    pub bitrate: f64, // kbps (average)
42    /// Output file size in bytes.
43    pub file_size: u64, // bytes
44    /// Wall-clock time taken to encode.
45    pub duration: Duration, // wall-clock encode time
46}
47
48/// Real-time encoding progress info parsed from FFmpeg.
49#[derive(Debug, Clone, Default)]
50pub struct Progress {
51    /// Number of frames encoded so far.
52    pub frame: i64,
53    /// Current encoding rate in frames per second.
54    pub fps: f64,
55    /// Current output bitrate in kbps.
56    pub bitrate: f64, // kbps
57    /// Encoding speed relative to real time (e.g. 2.5 means 2.5x).
58    pub speed: f64, // e.g. 2.5x
59    /// Output timestamp reached so far.
60    pub time: Duration,
61}
62
63/// Runs an FFmpeg encode job. Progress updates are sent on the channel if provided.
64pub async fn encode(
65    job: EncodeJob,
66    progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
67) -> anyhow::Result<EncodeResult> {
68    match job.rate_control {
69        RateControlMode::Vbr => encode_two_pass(job, progress_tx).await,
70        _ => encode_single_pass(job, progress_tx).await,
71    }
72}
73
74async fn encode_single_pass(
75    job: EncodeJob,
76    progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
77) -> anyhow::Result<EncodeResult> {
78    let args = build_encode_args(&job, EncodePass::Single)?;
79    run_encode(job, args, progress_tx).await
80}
81
82async fn encode_two_pass(
83    job: EncodeJob,
84    progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
85) -> anyhow::Result<EncodeResult> {
86    if job.target_bitrate <= 0.0 {
87        anyhow::bail!("target bitrate must be greater than zero for VBR mode");
88    }
89
90    let passlog_prefix = make_passlog_prefix(&job.output);
91    let cleanup = PasslogCleanup::new(passlog_prefix.clone());
92
93    let first_pass_args = build_encode_args(&job, EncodePass::First(&passlog_prefix))?;
94    run_ffmpeg(first_pass_args, None).await?;
95
96    let second_pass_args = build_encode_args(&job, EncodePass::Second(&passlog_prefix))?;
97    let result = run_encode(job, second_pass_args, progress_tx).await;
98
99    cleanup.run();
100    result
101}
102
103async fn run_encode(
104    job: EncodeJob,
105    args: Vec<String>,
106    progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
107) -> anyhow::Result<EncodeResult> {
108    let start = Instant::now();
109    run_ffmpeg(args, progress_tx).await?;
110
111    let elapsed = start.elapsed();
112
113    // Probe the output to get actual bitrate and file size
114    let meta = std::fs::metadata(&job.output)
115        .map_err(|e| anyhow::anyhow!("failed to stat output: {e}"))?;
116
117    let probe_result = probe(&job.output).await?;
118    let bitrate = probe_result.format.bit_rate as f64 / 1000.0;
119
120    Ok(EncodeResult { job, bitrate, file_size: meta.len(), duration: elapsed })
121}
122
123async fn run_ffmpeg(
124    args: Vec<String>,
125    progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
126) -> anyhow::Result<()> {
127    let mut cmd = Command::new(ffmpeg_path());
128    cmd.args(&args).stdout(std::process::Stdio::piped()).stderr(std::process::Stdio::piped());
129
130    let mut child = cmd.spawn().map_err(|e| anyhow::anyhow!("failed to start ffmpeg: {e}"))?;
131
132    // Parse progress from stdout
133    if let Some(stdout) = child.stdout.take() {
134        let tx = progress_tx.clone();
135        tokio::spawn(async move {
136            let reader = tokio::io::BufReader::new(stdout);
137            let mut lines = reader.lines();
138            let mut p = Progress::default();
139            while let Ok(Some(line)) = lines.next_line().await {
140                if parse_progress_line(&line, &mut p) {
141                    if let Some(ref tx) = tx {
142                        let _ = tx.try_send(p.clone());
143                    }
144                }
145            }
146        });
147    }
148
149    let output = child.wait_with_output().await?;
150    if !output.status.success() {
151        let stderr = String::from_utf8_lossy(&output.stderr);
152        anyhow::bail!("ffmpeg encode failed: {stderr}");
153    }
154
155    Ok(())
156}
157
158/// Copies a segment of a video file without re-encoding.
159pub async fn extract(input: &str, output: &str, start: f64, duration: f64) -> anyhow::Result<()> {
160    let args = vec![
161        "-y".to_string(),
162        "-ss".into(),
163        format!("{start:.6}"),
164        "-i".into(),
165        input.into(),
166        "-t".into(),
167        format!("{duration:.6}"),
168        "-c".into(),
169        "copy".into(),
170        "-avoid_negative_ts".into(),
171        "make_zero".into(),
172        output.into(),
173    ];
174
175    let output = Command::new(ffmpeg_path())
176        .args(&args)
177        .stderr(std::process::Stdio::piped())
178        .output()
179        .await?;
180
181    if !output.status.success() {
182        let stderr = String::from_utf8_lossy(&output.stderr);
183        anyhow::bail!("ffmpeg extract failed: {stderr}");
184    }
185    Ok(())
186}
187
188/// Concatenates multiple encoded chunks into a single output without re-encoding.
189pub async fn concat(inputs: &[String], output: &str) -> anyhow::Result<()> {
190    if inputs.is_empty() {
191        anyhow::bail!("cannot concat an empty input list");
192    }
193
194    let list_path = make_concat_list_path(output);
195    let list_body = inputs
196        .iter()
197        .map(|path| format!("file '{}'", path.replace('\'', "'\\''")))
198        .collect::<Vec<_>>()
199        .join("\n");
200    std::fs::write(&list_path, format!("{list_body}\n"))?;
201
202    let args = vec![
203        "-y".to_string(),
204        "-f".into(),
205        "concat".into(),
206        "-safe".into(),
207        "0".into(),
208        "-i".into(),
209        list_path.to_string_lossy().into_owned(),
210        "-c".into(),
211        "copy".into(),
212        output.into(),
213    ];
214
215    let result = run_ffmpeg(args, None).await;
216    let _ = std::fs::remove_file(&list_path);
217    result
218}
219
220enum EncodePass<'a> {
221    Single,
222    First(&'a Path),
223    Second(&'a Path),
224}
225
226fn build_encode_args(job: &EncodeJob, pass: EncodePass<'_>) -> anyhow::Result<Vec<String>> {
227    let mut args = vec!["-y".into(), "-i".into(), job.input.clone(), "-an".into()];
228
229    if !matches!(pass, EncodePass::First(_)) {
230        args.extend(["-progress".into(), "pipe:1".into(), "-nostats".into()]);
231    }
232
233    args.extend(["-c:v".into(), job.codec.as_str().into()]);
234
235    if job.codec.is_hardware() {
236        build_hw_args(&mut args, job, &pass)?;
237    } else {
238        build_sw_args(&mut args, job, &pass)?;
239    }
240
241    if !job.preset.is_empty() {
242        if job.codec.is_hardware() {
243            add_hw_preset(&mut args, job.codec, &job.preset);
244        } else {
245            args.extend(["-preset".into(), job.preset.clone()]);
246        }
247    }
248
249    if let Some(ref res) = job.resolution {
250        if res.width > 0 && res.height > 0 {
251            args.extend([
252                "-vf".into(),
253                format!("scale={}:{}:flags=lanczos", res.width, res.height),
254            ]);
255        }
256    }
257
258    args.extend(job.extra_args.iter().cloned());
259
260    match pass {
261        EncodePass::First(_) => {
262            args.extend(["-f".into(), "null".into()]);
263            args.push(null_output_path().into());
264        }
265        EncodePass::Single | EncodePass::Second(_) => args.push(job.output.clone()),
266    }
267
268    Ok(args)
269}
270
271fn build_sw_args(
272    args: &mut Vec<String>,
273    job: &EncodeJob,
274    pass: &EncodePass<'_>,
275) -> anyhow::Result<()> {
276    match job.rate_control {
277        RateControlMode::Qp => {
278            if job.codec == Codec::SvtAv1 {
279                args.extend(["-qp".into(), job.crf.to_string()]);
280                args.extend(["-svtav1-params".into(), "enable-adaptive-quantization=0".into()]);
281            } else {
282                args.extend(["-qp".into(), job.crf.to_string()]);
283            }
284        }
285        RateControlMode::CappedCrf => {
286            if job.max_bitrate <= 0.0 {
287                anyhow::bail!("max bitrate must be greater than zero for capped CRF mode");
288            }
289            let bufsize = if job.bufsize > 0.0 { job.bufsize } else { job.max_bitrate * 2.0 };
290            args.extend(["-crf".into(), job.crf.to_string()]);
291            args.extend(["-maxrate".into(), format!("{:.0}k", job.max_bitrate)]);
292            args.extend(["-bufsize".into(), format!("{bufsize:.0}k")]);
293        }
294        RateControlMode::Vbr => {
295            if job.target_bitrate <= 0.0 {
296                anyhow::bail!("target bitrate must be greater than zero for VBR mode");
297            }
298            args.extend(["-b:v".into(), format!("{:.0}k", job.target_bitrate)]);
299            args.extend(["-maxrate".into(), format!("{:.0}k", job.target_bitrate * 2.0)]);
300            args.extend(["-bufsize".into(), format!("{:.0}k", job.target_bitrate * 4.0)]);
301
302            let passlog = match pass {
303                EncodePass::First(path) => {
304                    args.extend(["-pass".into(), "1".into()]);
305                    path
306                }
307                EncodePass::Second(path) => {
308                    args.extend(["-pass".into(), "2".into()]);
309                    path
310                }
311                EncodePass::Single => {
312                    anyhow::bail!("VBR mode requires a two-pass encode flow");
313                }
314            };
315            args.extend(["-passlogfile".into(), passlog.to_string_lossy().into_owned()]);
316        }
317        RateControlMode::Crf => {
318            args.extend(["-crf".into(), job.crf.to_string()]);
319        }
320    }
321    Ok(())
322}
323
324fn build_hw_args(
325    args: &mut Vec<String>,
326    job: &EncodeJob,
327    _pass: &EncodePass<'_>,
328) -> anyhow::Result<()> {
329    let backend = job.codec.backend();
330    match job.rate_control {
331        RateControlMode::Crf | RateControlMode::CappedCrf => {
332            match backend {
333                EncoderBackend::Nvenc => {
334                    let cq = crf_to_nvenc_cq(job.crf);
335                    args.extend(["-cq".into(), cq.to_string()]);
336                    args.extend(["-rc".into(), "constqp".into()]);
337                }
338                EncoderBackend::Qsv => {
339                    let gq = crf_to_qsv_quality(job.crf);
340                    args.extend(["-global_quality".into(), gq.to_string()]);
341                }
342                EncoderBackend::VideoToolbox => {
343                    let qual = crf_to_vt_quality(job.crf);
344                    args.extend(["-quality".into(), qual.to_string()]);
345                }
346                EncoderBackend::Vaapi => {
347                    let gq = crf_to_qsv_quality(job.crf);
348                    args.extend(["-global_quality".into(), gq.to_string()]);
349                }
350                EncoderBackend::Amf => {
351                    args.extend(["-qp_i".into(), job.crf.to_string()]);
352                    args.extend(["-qp_p".into(), (job.crf + 2).to_string()]);
353                    args.extend(["-usage".into(), "transcoding".into()]);
354                }
355                EncoderBackend::Software => unreachable!(),
356            }
357            // VBV / maxrate for capped mode
358            if let RateControlMode::CappedCrf = job.rate_control {
359                if job.max_bitrate <= 0.0 {
360                    anyhow::bail!("max bitrate must be greater than zero for capped CRF mode");
361                }
362                let bufsize = if job.bufsize > 0.0 { job.bufsize } else { job.max_bitrate * 2.0 };
363                args.extend(["-maxrate".into(), format!("{:.0}k", job.max_bitrate)]);
364                args.extend(["-bufsize".into(), format!("{bufsize:.0}k")]);
365            }
366        }
367        RateControlMode::Qp => match backend {
368            EncoderBackend::VideoToolbox => {
369                anyhow::bail!("VideoToolbox does not support QP rate control mode");
370            }
371            _ => {
372                args.extend(["-qp".into(), job.crf.to_string()]);
373            }
374        },
375        RateControlMode::Vbr => {
376            if job.target_bitrate <= 0.0 {
377                anyhow::bail!("target bitrate must be greater than zero for VBR mode");
378            }
379            args.extend(["-b:v".into(), format!("{:.0}k", job.target_bitrate)]);
380            args.extend(["-maxrate".into(), format!("{:.0}k", job.target_bitrate * 2.0)]);
381            args.extend(["-bufsize".into(), format!("{:.0}k", job.target_bitrate * 4.0)]);
382
383            if backend == EncoderBackend::Nvenc {
384                args.extend(["-rc".into(), "vbr_hq".into()]);
385            }
386        }
387    }
388    Ok(())
389}
390
391fn crf_to_nvenc_cq(crf: i32) -> i32 {
392    let cq = (crf * 51) / 63;
393    cq.clamp(1, 51)
394}
395
396fn crf_to_qsv_quality(crf: i32) -> i32 {
397    let gq = 100 - ((crf * 100) / 51);
398    gq.clamp(1, 100)
399}
400
401fn crf_to_vt_quality(crf: i32) -> f64 {
402    let q = 1.0 - (crf as f64 / 51.0);
403    q.clamp(0.0, 1.0)
404}
405
406fn add_hw_preset(args: &mut Vec<String>, codec: Codec, preset: &str) {
407    match codec.backend() {
408        EncoderBackend::Nvenc => {
409            let p = map_nvenc_preset(preset);
410            args.extend(["-preset".into(), p.into()]);
411        }
412        EncoderBackend::Qsv => {
413            args.extend(["-preset".into(), preset.to_string()]);
414        }
415        EncoderBackend::Vaapi => {
416            args.extend(["-compression_level".into(), map_vaapi_preset(preset).into()]);
417        }
418        EncoderBackend::Amf => {
419            args.extend(["-quality".into(), map_amf_quality(preset).into()]);
420        }
421        EncoderBackend::VideoToolbox => {
422            if preset == "ultrafast" || preset == "superfast" || preset == "veryfast" {
423                args.extend(["-realtime".into(), "1".into()]);
424            }
425        }
426        EncoderBackend::Software => unreachable!(),
427    }
428}
429
430fn map_nvenc_preset(preset: &str) -> &str {
431    match preset {
432        "ultrafast" | "superfast" => "p1",
433        "veryfast" => "p2",
434        "faster" => "p3",
435        "fast" => "p4",
436        "medium" => "p5",
437        "slow" => "p6",
438        "slower" | "veryslow" => "p7",
439        other => other,
440    }
441}
442
443fn map_vaapi_preset(preset: &str) -> &str {
444    match preset {
445        "ultrafast" | "superfast" => "1",
446        "veryfast" | "faster" => "2",
447        "fast" | "medium" => "3",
448        "slow" => "4",
449        "slower" | "veryslow" => "5",
450        other => other,
451    }
452}
453
454fn map_amf_quality(preset: &str) -> &str {
455    match preset {
456        "ultrafast" | "superfast" => "speed",
457        "veryfast" | "faster" | "fast" => "balanced",
458        "medium" | "slow" | "slower" | "veryslow" => "quality",
459        other => other,
460    }
461}
462
463fn make_passlog_prefix(output: &str) -> PathBuf {
464    let output_path = Path::new(output);
465    let parent =
466        output_path.parent().filter(|p| !p.as_os_str().is_empty()).unwrap_or(Path::new("."));
467    let stem = output_path.file_stem().and_then(|s| s.to_str()).unwrap_or("viser");
468    let unique = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis()).unwrap_or(0);
469    parent.join(format!(".{stem}.viser-passlog-{unique}-{}", std::process::id()))
470}
471
472fn make_concat_list_path(output: &str) -> PathBuf {
473    let output_path = Path::new(output);
474    let parent =
475        output_path.parent().filter(|p| !p.as_os_str().is_empty()).unwrap_or(Path::new("."));
476    let stem = output_path.file_stem().and_then(|s| s.to_str()).unwrap_or("viser");
477    let unique = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis()).unwrap_or(0);
478    parent.join(format!(".{stem}.viser-concat-{unique}-{}.txt", std::process::id()))
479}
480
481fn null_output_path() -> &'static str {
482    if cfg!(windows) { "NUL" } else { "/dev/null" }
483}
484
485struct PasslogCleanup {
486    parent: PathBuf,
487    prefix: String,
488}
489
490impl PasslogCleanup {
491    fn new(path: PathBuf) -> Self {
492        let parent = path.parent().unwrap_or(Path::new(".")).to_path_buf();
493        let prefix = path.file_name().and_then(|name| name.to_str()).unwrap_or_default().to_owned();
494        Self { parent, prefix }
495    }
496
497    fn run(&self) {
498        let Ok(entries) = std::fs::read_dir(&self.parent) else {
499            return;
500        };
501
502        for entry in entries.flatten() {
503            let path = entry.path();
504            let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
505                continue;
506            };
507            if !name.starts_with(&self.prefix) {
508                continue;
509            }
510            if let Err(err) = std::fs::remove_file(&path) {
511                tracing::debug!(?path, ?err, "failed to remove ffmpeg two-pass log file");
512            }
513        }
514    }
515}
516
517/// Returns true when a complete progress block is ready.
518fn parse_progress_line(line: &str, p: &mut Progress) -> bool {
519    let Some((key, value)) = line.split_once('=') else {
520        return false;
521    };
522
523    match key {
524        "frame" => {
525            p.frame = value.parse().unwrap_or(0);
526        }
527        "fps" => {
528            p.fps = value.parse().unwrap_or(0.0);
529        }
530        "bitrate" => {
531            let v = value.trim_end_matches("kbits/s");
532            p.bitrate = v.parse().unwrap_or(0.0);
533        }
534        "speed" => {
535            let v = value.trim_end_matches('x');
536            p.speed = v.parse().unwrap_or(0.0);
537        }
538        "out_time_us" => {
539            let us: u64 = value.parse().unwrap_or(0);
540            p.time = Duration::from_micros(us);
541        }
542        "progress" => return true,
543        _ => {}
544    }
545    false
546}
547
548#[cfg(test)]
549mod tests {
550    use super::*;
551    use crate::Codec;
552
553    fn sample_job(mode: RateControlMode) -> EncodeJob {
554        EncodeJob {
555            input: "input.mp4".into(),
556            output: "output.mp4".into(),
557            resolution: Some(crate::Resolution::new(1280, 720)),
558            codec: Codec::X264,
559            crf: 23,
560            rate_control: mode,
561            target_bitrate: 2500.0,
562            max_bitrate: 3000.0,
563            bufsize: 6000.0,
564            preset: "medium".into(),
565            extra_args: vec![],
566        }
567    }
568
569    #[test]
570    fn test_build_encode_args_crf_single_pass() {
571        let args =
572            build_encode_args(&sample_job(RateControlMode::Crf), EncodePass::Single).unwrap();
573        assert!(args.windows(2).any(|w| w == ["-crf", "23"]));
574        assert_eq!(args.last().unwrap(), "output.mp4");
575    }
576
577    #[test]
578    fn test_build_encode_args_vbr_first_pass_uses_null_output() {
579        let job = sample_job(RateControlMode::Vbr);
580        let passlog = Path::new("/tmp/viser-passlog");
581        let args = build_encode_args(&job, EncodePass::First(passlog)).unwrap();
582        assert!(args.windows(2).any(|w| w == ["-pass", "1"]));
583        assert!(args.windows(2).any(|w| w == ["-f", "null"]));
584        assert_eq!(args.last().unwrap(), null_output_path());
585    }
586
587    #[test]
588    fn test_build_encode_args_vbr_second_pass_writes_output() {
589        let job = sample_job(RateControlMode::Vbr);
590        let passlog = Path::new("/tmp/viser-passlog");
591        let args = build_encode_args(&job, EncodePass::Second(passlog)).unwrap();
592        assert!(args.windows(2).any(|w| w == ["-pass", "2"]));
593        assert_eq!(args.last().unwrap(), "output.mp4");
594    }
595
596    #[test]
597    fn test_build_encode_args_capped_crf_sets_vbv() {
598        let args =
599            build_encode_args(&sample_job(RateControlMode::CappedCrf), EncodePass::Single).unwrap();
600        assert!(args.windows(2).any(|w| w == ["-crf", "23"]));
601        assert!(args.windows(2).any(|w| w == ["-maxrate", "3000k"]));
602        assert!(args.windows(2).any(|w| w == ["-bufsize", "6000k"]));
603    }
604}