From 8d1ec94ac21f9552ecac7970f49bcdf5d2f8bc4c Mon Sep 17 00:00:00 2001 From: Rhys Lloyd Date: Tue, 10 Mar 2026 09:10:04 -0700 Subject: [PATCH] add video encoder --- Cargo.lock | 136 +++++++++- Cargo.toml | 1 + video-encoder/Cargo.toml | 14 + video-encoder/README.md | 4 + video-encoder/shaders/rgb_to_yuv.wgsl | 51 ++++ video-encoder/src/main.rs | 367 ++++++++++++++++++++++++++ 6 files changed, 571 insertions(+), 2 deletions(-) create mode 100644 video-encoder/Cargo.toml create mode 100644 video-encoder/README.md create mode 100644 video-encoder/shaders/rgb_to_yuv.wgsl create mode 100644 video-encoder/src/main.rs diff --git a/Cargo.lock b/Cargo.lock index 5b33c59..d126d4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -169,6 +169,12 @@ version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" +[[package]] +name = "bitstream-io" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6099cdc01846bc367c4e7dd630dc5966dccf36b652fae7a74e17b640411a91b2" + [[package]] name = "block" version = "0.1.6" @@ -405,6 +411,17 @@ dependencies = [ "num-traits", ] +[[package]] +name = "derivative" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fcc3dd5e9e9c0b295d6e1e4d811fb6f157d5ffd784b8d202fc62eac8035a770b" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "dispatch" version = "0.2.0" @@ -531,6 +548,12 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" +[[package]] +name = "four-cc" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "795cbfc56d419a7ce47ccbb7504dd9a5b7c484c083c356e797de08bd988d9629" + [[package]] name = "futures-core" version = "0.3.32" @@ -649,6 +672,19 @@ dependencies = [ "bitflags 2.11.0", ] +[[package]] +name = "h264-reader" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "036a78b2620d92f0ec57690bc792b3bb87348632ee5225302ba2e66a48021c6c" +dependencies = [ + "bitstream-io", + "hex-slice", + "log", + "memchr", + "rfc6381-codec", +] + [[package]] name = "half" version = "2.7.1" @@ -687,6 +723,12 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" +[[package]] +name = "hex-slice" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5491a308e0214554f07a81d8944abe45f552871c12e3c3c6e7e5d354039a6c4c" + [[package]] name = "hexf-parse" version = "0.2.1" @@ -899,6 +941,21 @@ dependencies = [ "paste", ] +[[package]] +name = "mp4ra-rust" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdbc3d3867085d66ac6270482e66f3dd2c5a18451a3dc9ad7269e94844a536b7" +dependencies = [ + "four-cc", +] + +[[package]] +name = "mpeg4-audio-const" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96a1fe2275b68991faded2c80aa4a33dba398b77d276038b8f50701a22e55918" + [[package]] name = "naga" version = "28.0.0" @@ -919,7 +976,7 @@ dependencies = [ "log", "num-traits", "once_cell", - "rustc-hash", + "rustc-hash 1.1.0", "spirv", "thiserror 2.0.18", "unicode-ident", @@ -1452,12 +1509,28 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19b30a45b0cd0bcca8037f3d0dc3421eaf95327a17cad11964fb8179b4fc4832" +[[package]] +name = "rfc6381-codec" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed54c20f5c3ec82eab6d998b313dc75ec5d5650d4f57675e61d72489040297fd" +dependencies = [ + "mp4ra-rust", + "mpeg4-audio-const", +] + [[package]] name = "rustc-hash" version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" +[[package]] +name = "rustc-hash" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" + [[package]] name = "rustix" version = "0.38.44" @@ -1862,14 +1935,29 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "pin-project-lite", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "tracing-core" version = "0.1.36" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] [[package]] name = "ttf-parser" @@ -1901,6 +1989,50 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "video-encoder" +version = "0.1.0" +dependencies = [ + "glam", + "strafesnet_common", + "strafesnet_graphics", + "strafesnet_roblox_bot_file", + "strafesnet_roblox_bot_player", + "strafesnet_snf", + "vk-video", + "wgpu", +] + +[[package]] +name = "vk-mem" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cb12b79bcec57a3334d0284f1364c1846f378bb47df9779c6dbfcfc245c9404" +dependencies = [ + "ash", + "bitflags 2.11.0", + "cc", +] + +[[package]] +name = "vk-video" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66b2031523b3ed32d99d650ace95b70606bba44fc7a1178ba7bbbe1c17fa0a2b" +dependencies = [ + "ash", + "bytes", + "cfg_aliases", + "derivative", + "h264-reader", + "memchr", + "rustc-hash 2.1.1", + "thiserror 1.0.69", + "tracing", + "vk-mem", + "wgpu", +] + [[package]] name = "walkdir" version = "2.5.0" @@ -2160,7 +2292,7 @@ dependencies = [ "portable-atomic", "profiling", "raw-window-handle", - "rustc-hash", + "rustc-hash 1.1.0", "smallvec", "thiserror 2.0.18", "wgpu-core-deps-apple", diff --git a/Cargo.toml b/Cargo.toml index a28b193..911efdd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,6 +2,7 @@ members = [ "lib", "native-player", + "video-encoder", "wasm-module" ] resolver = "3" diff --git a/video-encoder/Cargo.toml b/video-encoder/Cargo.toml new file mode 100644 index 0000000..53380c9 --- /dev/null +++ b/video-encoder/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "video-encoder" +version = "0.1.0" +edition = "2024" + +[dependencies] +glam.workspace = true +wgpu = "28.0.0" +strafesnet_roblox_bot_player.workspace = true +strafesnet_common.workspace = true +strafesnet_graphics.workspace = true +strafesnet_roblox_bot_file.workspace = true +strafesnet_snf.workspace = true +vk-video = "0.2.0" diff --git a/video-encoder/README.md b/video-encoder/README.md new file mode 100644 index 0000000..32dcb4a --- /dev/null +++ b/video-encoder/README.md @@ -0,0 +1,4 @@ +### How it works +- Render RGB to graphics_texture +- Convert RGB to YUV on video_texture +- Encode video frame diff --git a/video-encoder/shaders/rgb_to_yuv.wgsl b/video-encoder/shaders/rgb_to_yuv.wgsl new file mode 100644 index 0000000..4fee893 --- /dev/null +++ b/video-encoder/shaders/rgb_to_yuv.wgsl @@ -0,0 +1,51 @@ +struct VertexOutput { + @builtin(position) position: vec4, + @location(1) uv: vec2, +} + +@vertex +fn vs_main(@builtin(vertex_index) vertex_index: u32) -> VertexOutput { + // hacky way to draw a large triangle + let tmp1 = i32(vertex_index) / 2; + let tmp2 = i32(vertex_index) & 1; + var result:VertexOutput; + result.position=vec4( + f32(tmp1) * 4.0 - 1.0, + f32(tmp2) * 4.0 - 1.0, + 1.0, + 1.0 + ); + result.uv=vec2( + f32(tmp1) * 2.0, + 1.0 - f32(tmp2) * 2.0 + ); + return result; +} + +@group(0) +@binding(0) +var texture: texture_2d; +@group(0) +@binding(1) +var texture_sampler: sampler; + +@fragment +fn fs_main_y(input: VertexOutput) -> @location(0) f32 { + let conversion_weights = vec3(0.2126, 0.7152, 0.0722); + let color = textureSample(texture, texture_sampler, input.uv).rgb; + + return clamp(dot(color, conversion_weights), 0.0, 1.0); +} + +@fragment +fn fs_main_uv(input: VertexOutput) -> @location(0) vec2 { + let conversion_weights = mat3x2( + -0.1146, 0.5, + -0.3854, -0.4542, + 0.5, -0.0458, + ); + let conversion_bias = vec2(0.5, 0.5); + let color = textureSample(texture, texture_sampler, input.uv).rgb; + + return clamp(conversion_weights * color + conversion_bias, vec2(0.0, 0.0), vec2(1.0, 1.0)); +} diff --git a/video-encoder/src/main.rs b/video-encoder/src/main.rs new file mode 100644 index 0000000..6d71955 --- /dev/null +++ b/video-encoder/src/main.rs @@ -0,0 +1,367 @@ +use std::io::Write; +use strafesnet_common::session::Time as SessionTime; + +pub fn main(){ + let vulkan_instance = vk_video::VulkanInstance::new().unwrap(); + let vulkan_adapter = vulkan_instance.create_adapter(None).unwrap(); + let vulkan_device = vulkan_adapter + .create_device( + wgpu::Features::TEXTURE_COMPRESSION_BC, + wgpu::ExperimentalFeatures::disabled(), + wgpu::Limits::defaults(), + ) + .unwrap(); + + let size = glam::uvec2(1920,1080); + let target_framerate = 60; + let average_bitrate = 10_000_000; + let max_bitrate = 20_000_000; + + let bot_file=include_bytes!("../../web-demo/bhop_marble_7cf33a64-7120-4514-b9fa-4fe29d9523d.qbot"); + let map_file=include_bytes!("../../web-demo/bhop_marble_5692093612.snfm"); + + // decode + let timelines=strafesnet_roblox_bot_file::v0::read_all_to_block(std::io::Cursor::new(bot_file)).unwrap(); + let map=strafesnet_snf::read_map(std::io::Cursor::new(map_file)).unwrap().into_complete_map().unwrap(); + + // playback + let bot=strafesnet_roblox_bot_player::bot::CompleteBot::new(timelines); + let mut playback_head=strafesnet_roblox_bot_player::head::PlaybackHead::new(&bot,SessionTime::ZERO); + + let mut wgpu_state = WgpuState::new( + vulkan_device.wgpu_device(), + vulkan_device.wgpu_queue(), + size, + ); + + wgpu_state.change_map(&map); + + let mut encoder = vulkan_device + .create_wgpu_textures_encoder( + vulkan_device + .encoder_parameters_high_quality( + vk_video::parameters::VideoParameters { + width:size.x.try_into().unwrap(), + height:size.y.try_into().unwrap(), + target_framerate:target_framerate.into(), + }, + vk_video::parameters::RateControl::VariableBitrate { + average_bitrate, + max_bitrate, + virtual_buffer_size: std::time::Duration::from_secs(2), + }, + ) + .unwrap(), + ) + .unwrap(); + + let mut output_file = std::fs::File::create("output.h264").unwrap(); + + let duration = bot.duration(); + for i in 0..duration.get()*target_framerate as i64/SessionTime::ONE_SECOND.get() { + let time=SessionTime::raw(i*SessionTime::ONE_SECOND.get()/target_framerate as i64); + playback_head.advance_time(&bot,time); + let (pos,angles)=playback_head.get_position_angles(&bot,time); + wgpu_state.render(pos,angles); + + let res = unsafe { + encoder + .encode( + vk_video::Frame { + data: wgpu_state.video_texture.clone(), + pts: None, + }, + false, + ) + .unwrap() + }; + + output_file.write_all(&res.data).unwrap(); + } +} + +struct WgpuState { + device: wgpu::Device, + queue: wgpu::Queue, + // graphics output + graphics:strafesnet_roblox_bot_player::graphics::Graphics, + // not sure if this needs to stay bound to keep the TextureView valid + #[expect(unused)] + graphics_texture: wgpu::Texture, + graphics_texture_view: wgpu::TextureView, + // video output + video_texture: wgpu::Texture, + y_renderer: PlaneRenderer, + uv_renderer: PlaneRenderer, +} + +impl WgpuState { + fn new( + device: wgpu::Device, + queue: wgpu::Queue, + size: glam::UVec2, + ) -> WgpuState { + const FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Bgra8UnormSrgb; + let graphics = strafesnet_roblox_bot_player::graphics::Graphics::new(&device,&queue,size,FORMAT); + + let shader = wgpu::include_wgsl!("../shaders/rgb_to_yuv.wgsl"); + let shader = device.create_shader_module(shader); + + let graphics_texture_bind_group_layout=device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor{ + label:Some("RGB Bind Group Layout"), + entries:&[ + wgpu::BindGroupLayoutEntry{ + binding:0, + visibility:wgpu::ShaderStages::FRAGMENT, + ty:wgpu::BindingType::Texture{ + sample_type:wgpu::TextureSampleType::Float{filterable:true}, + multisampled:false, + view_dimension:wgpu::TextureViewDimension::D2, + }, + count:None, + }, + wgpu::BindGroupLayoutEntry{ + binding:1, + visibility:wgpu::ShaderStages::FRAGMENT, + ty:wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering), + count:None, + }, + ], + }); + + let graphics_texture=device.create_texture(&wgpu::TextureDescriptor{ + label:Some("RGB texture"), + format:FORMAT, + size:wgpu::Extent3d{ + width:size.x, + height:size.y, + depth_or_array_layers:1, + }, + mip_level_count:1, + sample_count:1, + dimension:wgpu::TextureDimension::D2, + usage:wgpu::TextureUsages::RENDER_ATTACHMENT|wgpu::TextureUsages::TEXTURE_BINDING, + view_formats:&[], + }); + let graphics_texture_view = graphics_texture.create_view(&wgpu::TextureViewDescriptor { + label: Some("RGB texture view"), + aspect: wgpu::TextureAspect::All, + usage: Some(wgpu::TextureUsages::RENDER_ATTACHMENT|wgpu::TextureUsages::TEXTURE_BINDING), + ..Default::default() + }); + let clamp_sampler=device.create_sampler(&wgpu::SamplerDescriptor{ + label:Some("Clamp Sampler"), + address_mode_u:wgpu::AddressMode::ClampToEdge, + address_mode_v:wgpu::AddressMode::ClampToEdge, + address_mode_w:wgpu::AddressMode::ClampToEdge, + mag_filter:wgpu::FilterMode::Linear, + min_filter:wgpu::FilterMode::Linear, + mipmap_filter:wgpu::MipmapFilterMode::Linear, + ..Default::default() + }); + let graphics_texture_bind_group=device.create_bind_group(&wgpu::BindGroupDescriptor{ + layout:&graphics_texture_bind_group_layout, + entries:&[ + wgpu::BindGroupEntry{ + binding:0, + resource:wgpu::BindingResource::TextureView(&graphics_texture_view), + }, + wgpu::BindGroupEntry{ + binding:1, + resource:wgpu::BindingResource::Sampler(&clamp_sampler), + }, + ], + label:Some("Graphics Texture"), + }); + + let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("wgpu pipeline layout"), + bind_group_layouts: &[ + &graphics_texture_bind_group_layout + ], + immediate_size: 0, + }); + + let video_texture = device.create_texture(&wgpu::TextureDescriptor { + label: Some("wgpu render target"), + format: wgpu::TextureFormat::NV12, + usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC, + dimension: wgpu::TextureDimension::D2, + sample_count: 1, + view_formats: &[], + mip_level_count: 1, + size: wgpu::Extent3d { + width: size.x, + height: size.y, + depth_or_array_layers: 1, + }, + }); + + let y_renderer = PlaneRenderer::new( + &device, + &pipeline_layout, + &shader, + "fs_main_y", + &video_texture, + wgpu::TextureAspect::Plane0, + graphics_texture_bind_group.clone(), + ); + let uv_renderer = PlaneRenderer::new( + &device, + &pipeline_layout, + &shader, + "fs_main_uv", + &video_texture, + wgpu::TextureAspect::Plane1, + graphics_texture_bind_group, + ); + + WgpuState { + device, + queue, + graphics, + graphics_texture, + graphics_texture_view, + video_texture, + y_renderer, + uv_renderer, + } + } + + fn change_map(&mut self,map:&strafesnet_common::map::CompleteMap){ + self.graphics.change_map(&self.device,&self.queue,map); + } + + fn render(&mut self,pos:glam::Vec3,angles:glam::Vec2) { + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("wgpu encoder"), + }); + + self.graphics.encode_commands(&mut encoder,&self.graphics_texture_view,pos,angles); + + self.y_renderer.render(&mut encoder); + self.uv_renderer.render(&mut encoder); + + encoder.transition_resources( + [].into_iter(), + [wgpu::TextureTransition { + texture: &self.video_texture, + state: wgpu::TextureUses::COPY_SRC, + selector: None, + }] + .into_iter(), + ); + + let buffer = encoder.finish(); + + self.queue.submit([buffer]); + } +} + +struct PlaneRenderer { + graphics_texture_bind_group: wgpu::BindGroup, + pipeline: wgpu::RenderPipeline, + plane: wgpu::TextureAspect, + plane_view: wgpu::TextureView, +} + +impl PlaneRenderer { + fn new( + device: &wgpu::Device, + pipeline_layout: &wgpu::PipelineLayout, + shader: &wgpu::ShaderModule, + fragment_entry_point: &str, + texture: &wgpu::Texture, + plane: wgpu::TextureAspect, + graphics_texture_bind_group: wgpu::BindGroup, + ) -> Self { + let format = texture.format().aspect_specific_format(plane).unwrap(); + let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor { + label: Some("wgpu pipeline"), + layout: Some(pipeline_layout), + cache: None, + vertex: wgpu::VertexState { + module: shader, + buffers: &[], + entry_point: None, + compilation_options: Default::default(), + }, + fragment: Some(wgpu::FragmentState { + module: shader, + entry_point: Some(fragment_entry_point), + compilation_options: Default::default(), + targets: &[Some(wgpu::ColorTargetState { + blend: None, + format, + write_mask: wgpu::ColorWrites::ALL, + })], + }), + primitive: wgpu::PrimitiveState { + topology: wgpu::PrimitiveTopology::TriangleList, + cull_mode: Some(wgpu::Face::Back), + polygon_mode: wgpu::PolygonMode::Fill, + front_face: wgpu::FrontFace::Cw, + conservative: false, + unclipped_depth: false, + strip_index_format: None, + }, + multiview_mask: None, + multisample: wgpu::MultisampleState { + count: 1, + mask: !0, + alpha_to_coverage_enabled: false, + }, + depth_stencil: None, + }); + + let plane_view = texture.create_view(&wgpu::TextureViewDescriptor { + label: Some("wgpu render target plane view"), + aspect: plane, + usage: Some(wgpu::TextureUsages::RENDER_ATTACHMENT), + ..Default::default() + }); + + Self { + graphics_texture_bind_group, + pipeline, + plane, + plane_view, + } + } + + fn render(&self, encoder: &mut wgpu::CommandEncoder) { + let clear_color = match self.plane { + wgpu::TextureAspect::Plane0 => wgpu::Color::BLACK, + wgpu::TextureAspect::Plane1 => wgpu::Color { + r: 0.5, + g: 0.5, + b: 0.0, + a: 1.0, + }, + _ => unreachable!(), + }; + + let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { + label: Some("wgpu render pass"), + timestamp_writes: None, + occlusion_query_set: None, + depth_stencil_attachment: None, + color_attachments: &[Some(wgpu::RenderPassColorAttachment { + view: &self.plane_view, + ops: wgpu::Operations { + load: wgpu::LoadOp::Clear(clear_color), + store: wgpu::StoreOp::Store, + }, + resolve_target: None, + depth_slice: None, + })], + multiview_mask: None, + }); + + render_pass.set_bind_group(0,&self.graphics_texture_bind_group,&[]); + render_pass.set_pipeline(&self.pipeline); + render_pass.draw(0..3, 0..1); + } +}