Skip to main content

lumen_server/render/
gpu.rs

1use std::{env, fs, sync::mpsc};
2
3use anyhow::Context;
4use lumen_engine::{composition::Composition, gpu::GpuCompositionRenderer, media::MediaStore};
5use lumen_ffmpeg::VideoCodec;
6
7use super::{
8    ProjectBundle, RenderError, RenderOptions, RenderProgress,
9    encoder::{
10        ENCODER_FRAME_QUEUE_CAPACITY, EncoderFrame, LumenFfmpegEncoder, has_audio,
11        write_composited_audio,
12    },
13    media::{LocalMediaStore, media_root},
14};
15
16#[cfg(all(target_os = "linux", feature = "cuda", feature = "vulkan"))]
17mod linux_cuda;
18#[cfg(not(all(target_os = "linux", feature = "cuda", feature = "vulkan")))]
19mod unsupported_cuda;
20
21#[cfg(all(target_os = "linux", feature = "cuda", feature = "vulkan"))]
22use linux_cuda::render_project_mp4_cuda;
23#[cfg(not(all(target_os = "linux", feature = "cuda", feature = "vulkan")))]
24use unsupported_cuda::render_project_mp4_cuda;
25
26pub(super) const MEDIA_PREFETCH_LOOKAHEAD_FRAMES: u32 = 30;
27
28pub fn render_project_mp4(
29    bundle: &ProjectBundle,
30    options: &RenderOptions,
31    on_progress: &mut dyn FnMut(RenderProgress),
32) -> Result<Vec<u8>, RenderError> {
33    let composition = &bundle.composition;
34    let width = composition.render_settings.width;
35    let height = composition.render_settings.height;
36    let fps = composition.timeline.fps;
37    let total_frames = composition.timeline.duration_frames;
38
39    if fps <= 0.0 {
40        return Err(RenderError {
41            code: "invalid_project_payload",
42            message: format!("invalid timeline fps: {fps}"),
43            retryable: false,
44        });
45    }
46
47    if total_frames == 0 {
48        return Err(RenderError {
49            code: "invalid_project_payload",
50            message: "composition duration_frames must be greater than zero".to_string(),
51            retryable: false,
52        });
53    }
54
55    let media_root = media_root(options.media_root.as_deref()).map_err(|err| RenderError {
56        code: "media_root_error",
57        message: err.to_string(),
58        retryable: false,
59    })?;
60    let media_store = LocalMediaStore::new(media_root);
61    let mut renderer = create_gpu_renderer().map_err(|err| RenderError {
62        code: "render_failed",
63        message: format!("failed to create GPU renderer: {err:#}"),
64        retryable: true,
65    })?;
66    renderer
67        .compile_with_media(
68            composition,
69            &media_store,
70            lumen_gpu::wgpu::TextureFormat::Rgba8Unorm,
71        )
72        .map_err(|err| RenderError {
73            code: "render_failed",
74            message: format!("failed to compile composition: {err}"),
75            retryable: false,
76        })?;
77
78    let encoder =
79        choose_video_encoder(options.video_encoder.as_deref()).map_err(|err| RenderError {
80            code: "invalid_render_profile",
81            message: err.to_string(),
82            retryable: false,
83        })?;
84    tracing::info!(
85        encoder = encoder.name,
86        codec = ?encoder.codec,
87        cuda_fast_path = cfg!(all(target_os = "linux", feature = "cuda", feature = "vulkan"))
88            && encoder.cuda_fast_path,
89            "selected video encoder"
90    );
91    if encoder.cuda_fast_path {
92        return render_project_mp4_cuda(
93            composition,
94            &media_store,
95            width,
96            height,
97            fps,
98            total_frames,
99            encoder.codec,
100            options.verbose_debug,
101            on_progress,
102        );
103    }
104
105    let tmp = tempfile::tempdir().map_err(|err| RenderError {
106        code: "encode_failed",
107        message: format!("failed to create temp dir: {err}"),
108        retryable: true,
109    })?;
110    let output_path = tmp.path().join("output.mp4");
111    let include_audio = has_audio(composition);
112    let encoder = LumenFfmpegEncoder::create(
113        &output_path,
114        width,
115        height,
116        fps,
117        &encoder.name,
118        encoder.codec,
119        include_audio,
120    )
121    .map_err(|err| RenderError {
122        code: "encode_failed",
123        message: err.to_string(),
124        retryable: true,
125    })?;
126    let (pixel_recycle_tx, _pixel_recycle_rx) =
127        mpsc::sync_channel::<Vec<u8>>(ENCODER_FRAME_QUEUE_CAPACITY + 2);
128
129    for frame in 0..total_frames {
130        prefetch_media_frames(composition, &media_store, frame);
131        let (raster, _submission) = renderer
132            .render_frame_submitted(composition, frame, &media_store)
133            .map_err(|err| RenderError {
134                code: "render_failed",
135                message: format!("render failed at frame {frame}: {err}"),
136                retryable: true,
137            })?;
138        let _ = renderer.precompile_frame_window(
139            composition,
140            frame.saturating_add(1),
141            MEDIA_PREFETCH_LOOKAHEAD_FRAMES,
142            &media_store,
143        );
144        renderer
145            .gpu_renderer()
146            .device
147            .poll(lumen_gpu::wgpu::PollType::Poll)
148            .map_err(|err| RenderError {
149                code: "render_failed",
150                message: format!("GPU poll failed at frame {frame}: {err}"),
151                retryable: true,
152            })?;
153        let size = raster.domain.storage_size;
154        if size.width != width || size.height != height {
155            return Err(RenderError {
156                code: "render_failed",
157                message: format!(
158                    "frame {frame} dimensions {}x{} do not match composition {}x{}",
159                    size.width, size.height, width, height
160                ),
161                retryable: true,
162            });
163        }
164        let pixels =
165            read_texture_rgba8(renderer.gpu_renderer(), raster.texture, size).map_err(|err| {
166                RenderError {
167                    code: "render_failed",
168                    message: format!("failed reading rendered pixels for frame {frame}: {err}"),
169                    retryable: true,
170                }
171            })?;
172        encoder
173            .send(EncoderFrame {
174                frame,
175                pixels,
176                recycle_tx: pixel_recycle_tx.clone(),
177            })
178            .map_err(|err| RenderError {
179                code: "encode_failed",
180                message: err.to_string(),
181                retryable: true,
182            })?;
183
184        let completed = frame.saturating_add(1);
185        let ratio = (completed as f32 / total_frames as f32).clamp(0.0, 1.0);
186        on_progress(RenderProgress {
187            stage: "rendering",
188            frame: completed,
189            total_frames,
190            ratio,
191        });
192    }
193
194    if include_audio {
195        write_composited_audio(composition, &media_store, &encoder).map_err(|err| RenderError {
196            code: "audio_render_failed",
197            message: err.to_string(),
198            retryable: true,
199        })?;
200    }
201
202    encoder.finish().map_err(|err| RenderError {
203        code: "encode_failed",
204        message: err.to_string(),
205        retryable: true,
206    })?;
207
208    fs::read(&output_path).map_err(|err| RenderError {
209        code: "encode_failed",
210        message: format!("failed to read encoded output: {err}"),
211        retryable: true,
212    })
213}
214
215pub(super) fn prefetch_media_frames(
216    composition: &Composition,
217    media_store: &LocalMediaStore,
218    frame: u32,
219) {
220    let Ok(requirements) =
221        lumen_engine::media::collect_frame_requirements(composition, media_store, frame)
222    else {
223        return;
224    };
225    for video in requirements.videos {
226        let Some(resolver) = media_store.get_video_resolver(&video.stream_id) else {
227            continue;
228        };
229        let mut frames = video.frames;
230        for required_frame in frames.clone() {
231            for offset in 1..=MEDIA_PREFETCH_LOOKAHEAD_FRAMES {
232                frames.push(required_frame.saturating_add(offset));
233            }
234        }
235        frames.sort_unstable();
236        frames.dedup();
237        for frame in &frames {
238            let _ = resolver.enqueue_frame(*frame);
239        }
240        resolver.retain_frames(&frames);
241    }
242}
243
244fn read_texture_rgba8(
245    renderer: &lumen_gpu::Renderer,
246    id: lumen_gpu::TextureId,
247    size: lumen_gpu::Size,
248) -> anyhow::Result<Vec<u8>> {
249    let bytes_per_pixel = 4;
250    let unpadded_bytes_per_row = size.width.saturating_mul(bytes_per_pixel);
251    let padded_bytes_per_row = align_to(
252        unpadded_bytes_per_row,
253        lumen_gpu::wgpu::COPY_BYTES_PER_ROW_ALIGNMENT,
254    );
255    let output_size = u64::from(padded_bytes_per_row).saturating_mul(u64::from(size.height));
256    let output = renderer
257        .device
258        .create_buffer(&lumen_gpu::wgpu::BufferDescriptor {
259            label: Some("lumen-server readback"),
260            size: output_size.max(1),
261            usage: lumen_gpu::wgpu::BufferUsages::COPY_DST
262                | lumen_gpu::wgpu::BufferUsages::MAP_READ,
263            mapped_at_creation: false,
264        });
265    let mut encoder =
266        renderer
267            .device
268            .create_command_encoder(&lumen_gpu::wgpu::CommandEncoderDescriptor {
269                label: Some("lumen-server readback encoder"),
270            });
271    let texture = renderer
272        .texture(id)
273        .ok_or_else(|| anyhow::anyhow!("render output texture {id:?} is unavailable"))?;
274    encoder.copy_texture_to_buffer(
275        texture.as_image_copy(),
276        lumen_gpu::wgpu::TexelCopyBufferInfo {
277            buffer: &output,
278            layout: lumen_gpu::wgpu::TexelCopyBufferLayout {
279                offset: 0,
280                bytes_per_row: Some(padded_bytes_per_row),
281                rows_per_image: Some(size.height),
282            },
283        },
284        lumen_gpu::wgpu::Extent3d {
285            width: size.width,
286            height: size.height,
287            depth_or_array_layers: 1,
288        },
289    );
290    renderer.queue.submit([encoder.finish()]);
291
292    let slice = output.slice(..);
293    let (tx, rx) = mpsc::channel();
294    slice.map_async(lumen_gpu::wgpu::MapMode::Read, move |result| {
295        let _ = tx.send(result);
296    });
297    renderer
298        .device
299        .poll(lumen_gpu::wgpu::PollType::wait_indefinitely())
300        .map_err(|error| anyhow::anyhow!("GPU readback poll failed: {error}"))?;
301    rx.recv()
302        .map_err(|_| anyhow::anyhow!("GPU readback channel closed"))?
303        .map_err(|error| anyhow::anyhow!("GPU readback map failed: {error}"))?;
304
305    let mapped = slice.get_mapped_range();
306    let mut pixels = vec![
307        0;
308        (size.width as usize)
309            .saturating_mul(size.height as usize)
310            .saturating_mul(bytes_per_pixel as usize)
311    ];
312    for row in 0..size.height as usize {
313        let src_start = row.saturating_mul(padded_bytes_per_row as usize);
314        let src_end = src_start.saturating_add(unpadded_bytes_per_row as usize);
315        let dst_start = row.saturating_mul(unpadded_bytes_per_row as usize);
316        let dst_end = dst_start.saturating_add(unpadded_bytes_per_row as usize);
317        pixels[dst_start..dst_end].copy_from_slice(&mapped[src_start..src_end]);
318    }
319    drop(mapped);
320    output.unmap();
321    Ok(pixels)
322}
323
324fn align_to(value: u32, alignment: u32) -> u32 {
325    value.div_ceil(alignment) * alignment
326}
327
328pub(super) fn create_gpu_renderer() -> anyhow::Result<GpuCompositionRenderer> {
329    let handle = tokio::runtime::Handle::try_current()
330        .context("GPU renderer creation requires a Tokio runtime")?;
331    handle
332        .block_on(GpuCompositionRenderer::new())
333        .context("failed to initialize wgpu renderer")
334}
335
336#[derive(Debug)]
337struct VideoEncoderSelection {
338    name: String,
339    codec: VideoCodec,
340    cuda_fast_path: bool,
341}
342
343fn choose_video_encoder(override_encoder: Option<&str>) -> anyhow::Result<VideoEncoderSelection> {
344    if let Some(encoder) = override_encoder {
345        let encoder = encoder.trim();
346        if !encoder.is_empty() {
347            return video_encoder_selection(encoder);
348        }
349    }
350
351    if let Ok(encoder) = env::var("LUMEN_VIDEO_ENCODER") {
352        let encoder = encoder.trim();
353        if !encoder.is_empty() {
354            return video_encoder_selection(encoder);
355        }
356    }
357
358    let name = if cfg!(target_os = "macos") {
359        "h264_videotoolbox"
360    } else {
361        "libx264"
362    };
363    video_encoder_selection(name)
364}
365
366fn video_encoder_selection(name: &str) -> anyhow::Result<VideoEncoderSelection> {
367    let codec = match name {
368        "h264" | "libx264" | "h264_videotoolbox" | "h264_vulkan" | "h264_nvenc" => VideoCodec::H264,
369        "hevc" | "h265" | "libx265" | "hevc_videotoolbox" | "hevc_vulkan" | "hevc_nvenc" => {
370            VideoCodec::Hevc
371        }
372        "av1" | "libaom-av1" | "av1_nvenc" => VideoCodec::Av1,
373        _ => anyhow::bail!(
374            "unsupported video encoder `{name}`; expected a known H.264, HEVC, or AV1 encoder"
375        ),
376    };
377    Ok(VideoEncoderSelection {
378        name: name.to_string(),
379        codec,
380        cuda_fast_path: matches!(name, "h264_nvenc" | "hevc_nvenc" | "av1_nvenc"),
381    })
382}
383
384#[cfg(test)]
385mod tests {
386    use super::*;
387
388    #[test]
389    fn nvenc_encoders_select_cuda_capable_codecs() {
390        let h264 = video_encoder_selection("h264_nvenc").unwrap();
391        assert_eq!(h264.name, "h264_nvenc");
392        assert_eq!(h264.codec, VideoCodec::H264);
393
394        let hevc = video_encoder_selection("hevc_nvenc").unwrap();
395        assert_eq!(hevc.name, "hevc_nvenc");
396        assert_eq!(hevc.codec, VideoCodec::Hevc);
397
398        let av1 = video_encoder_selection("av1_nvenc").unwrap();
399        assert_eq!(av1.name, "av1_nvenc");
400        assert_eq!(av1.codec, VideoCodec::Av1);
401    }
402
403    #[test]
404    fn unknown_encoder_is_rejected_before_rendering() {
405        let error = video_encoder_selection("definitely_not_an_encoder").unwrap_err();
406        assert!(
407            error
408                .to_string()
409                .contains("expected a known H.264, HEVC, or AV1 encoder")
410        );
411    }
412}