lumen_engine/node/processing/
color_grade.rs1use 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#[derive(Debug, Clone, lumen_macros::Node)]
14#[node(kind = "color_grade", name = "Color Grade", category = "processing")]
15pub struct ColorGrade {
16 pub id: NodeId,
17 #[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 #[property(kind = "float", min = 0, max = 1, step = 0.01)]
28 pub strength: NodeProperty,
29 #[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(¶ms));
191 bound.write_buffer(*lut_buffer, 0, bytemuck::bytes_of(&lut));
192 Ok(())
193 }
194}