Skip to main content

lumen_engine/node/
media_output.rs

1use crate::node::{NodeId, PortRef};
2
3use crate::gpu::{CompiledOutput, GpuCompileNode, RasterHandle, compiler};
4
5pub(crate) const SHADER: &str = include_str!("media_output.wgsl");
6pub(crate) const RENDER_SHADER: &str = include_str!("media_output_render.wgsl");
7
8/// Copies the compiled raster into the final composition output.
9#[derive(Debug, Clone, lumen_macros::Node)]
10#[node(kind = "media_output", name = "Media Output", category = "output")]
11pub struct MediaOutput {
12    pub id: NodeId,
13    #[input()]
14    pub source: PortRef,
15}
16
17impl Default for MediaOutput {
18    fn default() -> Self {
19        Self {
20            id: NodeId::new(0),
21            source: PortRef::empty(),
22        }
23    }
24}
25
26impl GpuCompileNode for MediaOutput {
27    fn compile_gpu(
28        &self,
29        ctx: &mut crate::gpu::CompileContext<'_>,
30        port: &PortRef,
31    ) -> crate::Result<CompiledOutput> {
32        if port.port != "output" {
33            return Err(ctx.missing_output(self.id, &port.port));
34        }
35
36        let source = ctx
37            .compile_port(&self.source)?
38            .into_raster(self.source.id, &self.source.port)?;
39        let size = lumen_gpu::Size::new(
40            ctx.composition().render_settings.width.max(1),
41            ctx.composition().render_settings.height.max(1),
42        );
43        let output_format = ctx.output_format();
44        if output_format == lumen_gpu::wgpu::TextureFormat::Rgba8Unorm
45            && source.domain == lumen_gpu::TextureDomain::full_frame(size)
46        {
47            return Ok(CompiledOutput::Raster(source));
48        }
49
50        let output = ctx.builder_mut().texture_for(
51            lumen_gpu::NodeKey(self.id.0),
52            Some("media-output:final".to_string()),
53            media_output_texture_desc(size, output_format),
54        );
55        if output_format == lumen_gpu::wgpu::TextureFormat::Rgba8Unorm {
56            let program = ctx.builder_mut().program_for(
57                lumen_gpu::NodeKey(self.id.0),
58                lumen_gpu::ProgramDesc::Compute(lumen_gpu::ComputeProgramDesc {
59                    label: Some("media-output".to_string()),
60                    shader: SHADER.to_string(),
61                    entry: "cs_main".to_string(),
62                    bind_groups: lumen_gpu::BindGroupLayoutSpec::single(vec![
63                        lumen_gpu::BindingLayoutEntry::texture(
64                            0,
65                            lumen_gpu::wgpu::ShaderStages::COMPUTE,
66                        ),
67                        lumen_gpu::BindingLayoutEntry::storage_texture(
68                            1,
69                            lumen_gpu::wgpu::ShaderStages::COMPUTE,
70                            lumen_gpu::wgpu::TextureFormat::Rgba8Unorm,
71                            lumen_gpu::wgpu::StorageTextureAccess::WriteOnly,
72                        ),
73                    ]),
74                }),
75            );
76            ctx.builder_mut().compute_pass(lumen_gpu::ComputePassDesc {
77                label: Some("media-output:copy".to_string()),
78                owner: Some(lumen_gpu::NodeKey(self.id.0)),
79                program,
80                bindings: vec![
81                    lumen_gpu::Binding::sampled_texture(0, 0, source.texture),
82                    lumen_gpu::Binding::storage_texture(0, 1, output),
83                ],
84                dispatch: compiler::dispatch_for(size).into(),
85            });
86        } else {
87            let sampler = ctx.builder_mut().sampler(
88                Some("media-output:sampler".to_string()),
89                lumen_gpu::wgpu::SamplerDescriptor {
90                    address_mode_u: lumen_gpu::wgpu::AddressMode::ClampToEdge,
91                    address_mode_v: lumen_gpu::wgpu::AddressMode::ClampToEdge,
92                    address_mode_w: lumen_gpu::wgpu::AddressMode::ClampToEdge,
93                    mag_filter: lumen_gpu::wgpu::FilterMode::Nearest,
94                    min_filter: lumen_gpu::wgpu::FilterMode::Nearest,
95                    mipmap_filter: lumen_gpu::wgpu::MipmapFilterMode::Nearest,
96                    ..Default::default()
97                },
98            );
99            let program = ctx.builder_mut().program_for(
100                lumen_gpu::NodeKey(self.id.0),
101                lumen_gpu::ProgramDesc::Render(lumen_gpu::RenderProgramDesc {
102                    label: Some("media-output".to_string()),
103                    shader: RENDER_SHADER.to_string(),
104                    vertex_entry: "vs_main".to_string(),
105                    fragment_entry: "fs_main".to_string(),
106                    bind_groups: lumen_gpu::BindGroupLayoutSpec::single(vec![
107                        lumen_gpu::BindingLayoutEntry::texture(
108                            0,
109                            lumen_gpu::wgpu::ShaderStages::FRAGMENT,
110                        ),
111                        lumen_gpu::BindingLayoutEntry::sampler(
112                            1,
113                            lumen_gpu::wgpu::ShaderStages::FRAGMENT,
114                        ),
115                    ]),
116                    targets: vec![Some(lumen_gpu::wgpu::ColorTargetState {
117                        format: output_format,
118                        blend: Some(lumen_gpu::wgpu::BlendState::REPLACE),
119                        write_mask: lumen_gpu::wgpu::ColorWrites::ALL,
120                    })],
121                    vertex_buffers: Vec::new(),
122                    primitive: lumen_gpu::wgpu::PrimitiveState::default(),
123                }),
124            );
125            ctx.builder_mut().render_pass(lumen_gpu::RenderPassDesc {
126                label: Some("media-output:render".to_string()),
127                owner: Some(lumen_gpu::NodeKey(self.id.0)),
128                program,
129                targets: vec![lumen_gpu::RenderTargetRef {
130                    texture: output,
131                    load: lumen_gpu::LoadOp::Clear(lumen_gpu::wgpu::Color::TRANSPARENT),
132                    store: lumen_gpu::wgpu::StoreOp::Store,
133                }],
134                bindings: vec![
135                    lumen_gpu::Binding::sampled_texture(0, 0, source.texture),
136                    lumen_gpu::Binding::sampler(0, 1, sampler),
137                ],
138                vertex_buffers: Vec::new(),
139                index_buffer: None,
140                draw: lumen_gpu::DrawCommand::Draw(lumen_gpu::Draw {
141                    vertices: 0..3,
142                    instances: 0..1,
143                }),
144                scissor: None,
145            });
146        }
147
148        Ok(CompiledOutput::Raster(RasterHandle {
149            texture: output,
150            domain: lumen_gpu::TextureDomain::full_frame(size),
151            metadata: source.metadata,
152        }))
153    }
154}
155
156fn media_output_texture_desc(
157    size: lumen_gpu::Size,
158    format: lumen_gpu::wgpu::TextureFormat,
159) -> lumen_gpu::TextureDesc {
160    if format == lumen_gpu::wgpu::TextureFormat::Rgba8Unorm {
161        compiler::copyable_texture_desc(size)
162    } else {
163        lumen_gpu::TextureDesc {
164            domain: lumen_gpu::TextureDomain::full_frame(size),
165            format,
166            usage: lumen_gpu::wgpu::TextureUsages::COPY_SRC
167                | lumen_gpu::wgpu::TextureUsages::TEXTURE_BINDING
168                | lumen_gpu::wgpu::TextureUsages::RENDER_ATTACHMENT,
169        }
170    }
171}