Skip to main content

audiovis/
audiovis.rs

1use cuneus::compute::{ComputeShader, COMPUTE_TEXTURE_FORMAT_RGBA16};
2use cuneus::{
3    Core, ExportManager, RenderKit, ShaderApp, ShaderControls, ShaderManager, UniformProvider,
4};
5use log::error;
6use cuneus::WindowEvent;
7
8#[repr(C)]
9#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
10struct AudioVisParams {
11    red_power: f32,
12    green_power: f32,
13    blue_power: f32,
14    green_boost: f32,
15    contrast: f32,
16    gamma: f32,
17    glow: f32,
18    _padding: f32,
19}
20
21impl Default for AudioVisParams {
22    fn default() -> Self {
23        Self {
24            red_power: 0.98,
25            green_power: 0.85,
26            blue_power: 0.90,
27            green_boost: 1.62,
28            contrast: 1.0,
29            gamma: 1.0,
30            glow: 0.05,
31            _padding: 0.0,
32        }
33    }
34}
35
36impl UniformProvider for AudioVisParams {
37    fn as_bytes(&self) -> &[u8] {
38        bytemuck::bytes_of(self)
39    }
40}
41
42struct AudioVisCompute {
43    base: RenderKit,
44    compute_shader: ComputeShader,
45    current_params: AudioVisParams,
46}
47
48impl ShaderManager for AudioVisCompute {
49    fn init(core: &Core) -> Self {
50        let initial_params = AudioVisParams::default();
51        let base = RenderKit::new(core);
52
53        let config = ComputeShader::builder()
54            .with_entry_point("main")
55            .with_custom_uniforms::<AudioVisParams>()
56            .with_audio_spectrum(69) // 64 spectrum + BPM + 4 energy values
57            .with_workgroup_size([16, 16, 1])
58            .with_texture_format(COMPUTE_TEXTURE_FORMAT_RGBA16)
59            .with_label("Audio Visualizer Compute")
60            .build();
61
62        let compute_shader = cuneus::compute_shader!(core, "shaders/audiovis.wgsl", config);
63
64
65        // Initialize custom uniform with initial parameters
66        compute_shader.set_custom_params(initial_params, &core.queue);
67
68        Self {
69            base,
70            compute_shader,
71            current_params: initial_params,
72        }
73    }
74
75    fn update(&mut self, core: &Core) {
76        // Update time
77        let current_time = self.base.controls.get_time(&self.base.start_time);
78        let delta = 1.0 / 60.0;
79        self.compute_shader
80            .set_time(current_time, delta, &core.queue);
81
82        // Update audio spectrum - energy values are computed in spectrum.rs
83        // and included in the buffer at indices 65-68
84        self.base.update_audio_spectrum(&core.queue);
85        self.compute_shader
86            .update_audio_spectrum(&self.base.resolution_uniform.data, &core.queue);
87        // Handle export
88        self.compute_shader.handle_export(core, &mut self.base);
89    }
90
91    fn resize(&mut self, core: &Core) {
92        self.base.default_resize(core, &mut self.compute_shader);
93    }
94
95    fn render(&mut self, core: &Core) -> Result<(), cuneus::SurfaceError> {
96        let mut frame = self.base.begin_frame(core)?;
97
98        // Update video texture (this triggers spectrum data polling!)
99        let _video_updated = if self.base.using_video_texture {
100            self.base.update_video_texture(core, &core.queue)
101        } else {
102            false
103        };
104        let _webcam_updated = if self.base.using_webcam_texture {
105            self.base.update_webcam_texture(core, &core.queue)
106        } else {
107            false
108        };
109
110        let mut params = self.current_params;
111        let mut changed = false;
112        let mut should_start_export = false;
113        let mut export_request = self.base.export_manager.get_ui_request();
114        let mut controls_request = self
115            .base
116            .controls
117            .get_ui_request(&self.base.start_time, &core.size, self.base.fps_tracker.fps());
118
119        let using_video_texture = self.base.using_video_texture;
120        let using_hdri_texture = self.base.using_hdri_texture;
121        let using_webcam_texture = self.base.using_webcam_texture;
122        let video_info = self.base.get_video_info();
123        let hdri_info = self.base.get_hdri_info();
124        let webcam_info = self.base.get_webcam_info();
125
126        let full_output = if self.base.key_handler.show_ui {
127            self.base.render_ui(core, |ctx| {
128                RenderKit::apply_default_style(ctx);
129
130                egui::Window::new("Audio Visualizer")
131                    .collapsible(true)
132                    .resizable(true)
133                    .default_width(300.0)
134                    .show(ctx, |ui| {
135                        // Media controls
136                        ShaderControls::render_media_panel(
137                            ui,
138                            &mut controls_request,
139                            using_video_texture,
140                            video_info,
141                            using_hdri_texture,
142                            hdri_info,
143                            using_webcam_texture,
144                            webcam_info,
145                        );
146
147                        ui.separator();
148
149                        egui::CollapsingHeader::new("Visual Parameters")
150                            .default_open(true)
151                            .show(ui, |ui| {
152                                ui.label("Color Power:");
153                                changed |= ui
154                                    .add(
155                                        egui::Slider::new(&mut params.red_power, 0.1..=2.0)
156                                            .text("Red Power"),
157                                    )
158                                    .changed();
159                                changed |= ui
160                                    .add(
161                                        egui::Slider::new(&mut params.green_power, 0.1..=2.0)
162                                            .text("Green Power"),
163                                    )
164                                    .changed();
165                                changed |= ui
166                                    .add(
167                                        egui::Slider::new(&mut params.blue_power, 0.1..=2.0)
168                                            .text("Blue Power"),
169                                    )
170                                    .changed();
171
172                                ui.separator();
173                                changed |= ui
174                                    .add(
175                                        egui::Slider::new(&mut params.green_boost, 0.0..=3.0)
176                                            .text("Green Boost"),
177                                    )
178                                    .changed();
179                                changed |= ui
180                                    .add(
181                                        egui::Slider::new(&mut params.contrast, 0.1..=3.0)
182                                            .text("Contrast"),
183                                    )
184                                    .changed();
185                                changed |= ui
186                                    .add(
187                                        egui::Slider::new(&mut params.gamma, 0.1..=3.0)
188                                            .text("Gamma"),
189                                    )
190                                    .changed();
191                                changed |= ui
192                                    .add(
193                                        egui::Slider::new(&mut params.glow, 0.0..=1.0).text("Glow"),
194                                    )
195                                    .changed();
196                            });
197
198                        ui.separator();
199
200                        ShaderControls::render_controls_widget(ui, &mut controls_request);
201
202                        ui.separator();
203
204                        should_start_export =
205                            ExportManager::render_export_ui_widget(ui, &mut export_request);
206
207                        ui.separator();
208                        ui.label(format!("Frame: {}", self.compute_shader.current_frame));
209                    });
210            })
211        } else {
212            self.base.render_ui(core, |_ctx| {})
213        };
214
215        // Apply controls
216        self.base.export_manager.apply_ui_request(export_request);
217        self.base.apply_media_requests(core, &controls_request);
218
219        // Apply parameter changes
220        if changed {
221            self.current_params = params;
222            self.compute_shader.set_custom_params(params, &core.queue);
223        }
224
225        if should_start_export {
226            self.base.export_manager.start_export();
227        }
228
229        // Create command encoder
230
231        // Dispatch compute shader
232        self.compute_shader.dispatch(&mut frame.encoder, core);
233
234        self.base.renderer.render_to_view(&mut frame.encoder, &frame.view, &self.compute_shader.get_output_texture().bind_group);
235
236        self.base.end_frame(core, frame, full_output);
237
238        Ok(())
239    }
240
241    fn handle_input(&mut self, core: &Core, event: &WindowEvent) -> bool {
242        if self.base.default_handle_input(core, event) {
243            return true;
244        }
245        if let WindowEvent::DroppedFile(path) = event {
246            if let Err(e) = self.base.load_media(core, path) {
247                error!("Failed to load dropped file: {e:?}");
248            }
249            return true;
250        }
251        false
252    }
253}
254
255fn main() -> Result<(), Box<dyn std::error::Error>> {
256    env_logger::init();
257    cuneus::gst::init()?;
258    let (app, event_loop) = ShaderApp::new("Audio Visualizer", 800, 600);
259
260    app.run(event_loop, AudioVisCompute::init)
261}