Skip to main content

lumen_engine/node/processing/
color_grade.rs

1use crate::node::{NodeId, NodeProperty, PortRef};
2
3use crate::gpu::{
4    BoundFrame, CompiledOutput, FrameBindContext, FrameBinding, GpuCompileNode, GpuFrameBindNode,
5    RasterHandle, compiler,
6};
7
8pub(crate) const SHADER: &str = include_str!("color_grade.wgsl");
9
10pub const IDENTITY_LUT: &str = "identity";
11
12/// Applies a LUT-driven color transform to a raster.
13#[derive(Debug, Clone, lumen_macros::Node)]
14#[node(kind = "color_grade", name = "Color Grade", category = "processing")]
15pub struct ColorGrade {
16    pub id: NodeId,
17    /// LUT data source or named LUT preset.
18    #[property(
19        kind = "string",
20        name = "LUT source",
21        role = "lut_source",
22        multiline,
23        recommended_rows = 4
24    )]
25    pub lut_source: NodeProperty,
26    /// Blend amount for the LUT transform.
27    #[property(kind = "float", min = 0, max = 1, step = 0.01)]
28    pub strength: NodeProperty,
29    /// Sampling filter used when reading LUT data.
30    #[property(kind = "int", format = "sampling_mode")]
31    pub interpolation: NodeProperty,
32    #[input()]
33    pub source: PortRef,
34}
35
36impl Default for ColorGrade {
37    fn default() -> Self {
38        Self {
39            id: NodeId::new(0),
40            lut_source: NodeProperty::String(IDENTITY_LUT.to_string()),
41            strength: NodeProperty::Float(1.0),
42            interpolation: NodeProperty::Int(1),
43            source: PortRef::empty(),
44        }
45    }
46}
47
48impl GpuCompileNode for ColorGrade {
49    fn compile_gpu(
50        &self,
51        ctx: &mut crate::gpu::CompileContext<'_>,
52        port: &PortRef,
53    ) -> crate::Result<CompiledOutput> {
54        if port.port != "output" {
55            return Err(ctx.missing_output(self.id, &port.port));
56        }
57
58        let source = ctx
59            .compile_port(&self.source)?
60            .into_raster(self.source.id, &self.source.port)?;
61        let size = source.domain.storage_size;
62        let texture = ctx.builder_mut().texture_for(
63            lumen_gpu::NodeKey(self.id.0),
64            Some(format!("color-grade:{}:output", self.id.0)),
65            lumen_gpu::TextureDesc::storage(size, lumen_gpu::wgpu::TextureFormat::Rgba8Unorm),
66        );
67        let params =
68            ctx.builder_mut().buffer_for(
69                lumen_gpu::NodeKey(self.id.0),
70                Some(format!("color-grade:{}:params", self.id.0)),
71                lumen_gpu::BufferDesc::uniform(
72                    std::mem::size_of::<compiler::ColorGradeParams>() as u64
73                ),
74            );
75        let lut = ctx.builder_mut().buffer_for(
76            lumen_gpu::NodeKey(self.id.0),
77            Some(format!("color-grade:{}:lut", self.id.0)),
78            lumen_gpu::BufferDesc::storage(std::mem::size_of::<compiler::ColorGradeLut>() as u64),
79        );
80        let program = ctx.builder_mut().program_for(
81            lumen_gpu::NodeKey(self.id.0),
82            lumen_gpu::ProgramDesc::Compute(lumen_gpu::ComputeProgramDesc {
83                label: Some("color-grade".to_string()),
84                shader: SHADER.to_string(),
85                entry: "cs_main".to_string(),
86                bind_groups: lumen_gpu::BindGroupLayoutSpec::single(vec![
87                    lumen_gpu::BindingLayoutEntry::texture(
88                        0,
89                        lumen_gpu::wgpu::ShaderStages::COMPUTE,
90                    ),
91                    lumen_gpu::BindingLayoutEntry::uniform(
92                        1,
93                        lumen_gpu::wgpu::ShaderStages::COMPUTE,
94                    ),
95                    lumen_gpu::BindingLayoutEntry::storage(
96                        2,
97                        lumen_gpu::wgpu::ShaderStages::COMPUTE,
98                        true,
99                    ),
100                    lumen_gpu::BindingLayoutEntry::storage_texture(
101                        3,
102                        lumen_gpu::wgpu::ShaderStages::COMPUTE,
103                        lumen_gpu::wgpu::TextureFormat::Rgba8Unorm,
104                        lumen_gpu::wgpu::StorageTextureAccess::WriteOnly,
105                    ),
106                ]),
107            }),
108        );
109        ctx.builder_mut().compute_pass(lumen_gpu::ComputePassDesc {
110            label: Some(format!("color-grade:{}:apply", self.id.0)),
111            owner: Some(lumen_gpu::NodeKey(self.id.0)),
112            program,
113            bindings: vec![
114                lumen_gpu::Binding::sampled_texture(0, 0, source.texture),
115                lumen_gpu::Binding::uniform(0, 1, params),
116                lumen_gpu::Binding::storage_buffer(0, 2, lut),
117                lumen_gpu::Binding::storage_texture(0, 3, texture),
118            ],
119            dispatch: compiler::dispatch_for(size).into(),
120        });
121        ctx.builder_mut().param(
122            lumen_gpu::ParamKey {
123                owner: lumen_gpu::NodeKey(self.id.0),
124                slot: 0,
125            },
126            lumen_gpu::ParamTarget::Buffer(params),
127        );
128        ctx.builder_mut().param(
129            lumen_gpu::ParamKey {
130                owner: lumen_gpu::NodeKey(self.id.0),
131                slot: 1,
132            },
133            lumen_gpu::ParamTarget::Buffer(lut),
134        );
135        ctx.push_frame_binding(FrameBinding::ColorGrade {
136            node_id: self.id,
137            lut_source: self.lut_source.clone(),
138            strength: self.strength.clone(),
139            interpolation: self.interpolation.clone(),
140            params_buffer: params,
141            lut_buffer: lut,
142        });
143
144        Ok(CompiledOutput::Raster(RasterHandle {
145            texture,
146            domain: source.domain,
147            metadata: source.metadata,
148        }))
149    }
150}
151
152impl GpuFrameBindNode for ColorGrade {
153    fn bind_gpu_frame(
154        &self,
155        ctx: &FrameBindContext<'_>,
156        binding: &FrameBinding,
157        bound: &mut BoundFrame,
158    ) -> crate::Result<()> {
159        let FrameBinding::ColorGrade {
160            node_id,
161            lut_source,
162            strength,
163            interpolation,
164            params_buffer,
165            lut_buffer,
166        } = binding
167        else {
168            return Ok(());
169        };
170        let lut_source = lut_source.resolve_string(
171            *node_id,
172            "lut_source",
173            &ctx.expr_context(*node_id, "lut_source"),
174        )?;
175        let interpolation = interpolation.resolve_int(
176            *node_id,
177            "interpolation",
178            &ctx.expr_context(*node_id, "interpolation"),
179        )?;
180        let params = compiler::ColorGradeParams {
181            strength: strength.resolve_float(
182                *node_id,
183                "strength",
184                &ctx.expr_context(*node_id, "strength"),
185            )? as f32,
186            interpolation: if interpolation == 0 { 0 } else { 1 },
187            _pad: [0; 2],
188        };
189        let lut = compiler::ColorGradeLut::parse(*node_id, ctx.frame(), &lut_source)?;
190        bound.write_buffer(*params_buffer, 0, bytemuck::bytes_of(&params));
191        bound.write_buffer(*lut_buffer, 0, bytemuck::bytes_of(&lut));
192        Ok(())
193    }
194}