lumen_engine/node/processing/
color_grade.rs1use 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#[derive(Debug, Clone, lumen_macros::Delegate)]
14pub struct ColorGradeParams {
15 #[meta(
17 name = "LUT source",
18 role = "lut_source",
19 multiline,
20 recommended_rows = 4
21 )]
22 pub lut_source: String,
23 #[meta(min = 0, max = 1, step = 0.01)]
25 pub strength: f64,
26 #[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#[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}