Skip to main content

lumen_engine/node/processing/
color_grade.rs

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