Reduce memory use for progressive lossless images

jxl/src/frame/modular/transforms/apply.rs

@@ -382,6 +382,20 @@ impl TransformStepChunk {
                 }
                 buffers[buf_in[0]].buffer_grid[in_grid].mark_used(is_final);
                 buffers[buf_in[1]].buffer_grid[res_grid].mark_used(is_final);
+                // Release the weak neighbor grids read above (next average and previous
+                // decoded), which are counted as uses in the transform graph.
+                let (gx, gy) = self.grid_pos;
+                if gx + 1 < buffers[*buf_out].grid_shape.0 {
+                    let next_avg_grid =
+                        buffers[buf_in[0]].get_grid_idx(out_grid_kind, (gx + 1, gy));
+                    if next_avg_grid != in_grid {
+                        buffers[buf_in[0]].buffer_grid[next_avg_grid].mark_used(is_final);
+                    }
+                }
+                if gx > 0 {
+                    let prev_out_grid = buffers[*buf_out].get_grid_idx(out_grid_kind, (gx - 1, gy));
+                    buffers[*buf_out].buffer_grid[prev_out_grid].mark_used(is_final);
+                }
             }
             TransformStep::VSqueeze {
                 buf_in,
@@ -491,6 +505,20 @@ impl TransformStepChunk {
                 }
                 buffers[buf_in[0]].buffer_grid[in_grid].mark_used(is_final);
                 buffers[buf_in[1]].buffer_grid[res_grid].mark_used(is_final);
+                // Release the weak neighbor grids read above (next average and previous
+                // decoded), which are counted as uses in the transform graph.
+                let (gx, gy) = self.grid_pos;
+                if gy + 1 < buffers[*buf_out].grid_shape.1 {
+                    let next_avg_grid =
+                        buffers[buf_in[0]].get_grid_idx(out_grid_kind, (gx, gy + 1));
+                    if next_avg_grid != in_grid {
+                        buffers[buf_in[0]].buffer_grid[next_avg_grid].mark_used(is_final);
+                    }
+                }
+                if gy > 0 {
+                    let prev_out_grid = buffers[*buf_out].get_grid_idx(out_grid_kind, (gx, gy - 1));
+                    buffers[*buf_out].buffer_grid[prev_out_grid].mark_used(is_final);
+                }
             }
         };
 

jxl_cli/benches/decode.rs

@@ -68,6 +68,7 @@ fn decode_benches(c: &mut Criterion) {
                         false,
                         None,
                         false,
+                        false,
                     )
                     .unwrap();
                 })

jxl_cli/src/dec/mod.rs

@@ -139,6 +139,7 @@ pub fn decode_frames<In: JxlBitstreamInputExt>(
     linear_output: bool,
     render_interval: Option<usize>,
     allow_partial_files: bool,
+    store_partial_renders: bool,
 ) -> Result<(DecodeOutput, Duration)> {
     let start = Instant::now();
 
@@ -282,7 +283,7 @@ pub fn decode_frames<In: JxlBitstreamInputExt>(
                     // render and retry.
                     if render_interval.is_some() && input.available_bytes()? > 0 {
                         has_rendered_data |= fallback.flush_pixels(&mut output_bufs)?;
-                        if has_rendered_data {
+                        if has_rendered_data && store_partial_renders {
                             partial_renders.push(
                                 outputs
                                     .iter()
@@ -332,7 +333,7 @@ pub fn decode_frames<In: JxlBitstreamInputExt>(
                     // render and retry.
                     if render_interval.is_some() && input.available_bytes()? > 0 {
                         has_rendered_data |= fallback.flush_pixels(&mut output_bufs)?;
-                        if has_rendered_data {
+                        if has_rendered_data && store_partial_renders {
                             partial_renders.push(
                                 outputs
                                     .iter()

jxl_cli/src/lib.rs

@@ -77,6 +77,7 @@ mod tests {
             false,
             None,
             false,
+            false,
         )
         .unwrap()
         .0
@@ -188,6 +189,7 @@ mod tests {
                 false,
                 None,
                 false,
+                false,
             )
             .unwrap();
         }

jxl_cli/src/main.rs

@@ -164,6 +164,7 @@ fn main() -> Result<()> {
             let linear_output = matches!(output_format, Some(OutputFormat::Exr));
             #[cfg(not(feature = "exr"))]
             let linear_output = false;
+            let store_partial_renders = output_format.is_some() && opt.render_interval.is_some();
             let (mut output, duration) = dec::decode_frames(
                 $input,
                 options(skip_preview),
@@ -176,6 +177,7 @@ fn main() -> Result<()> {
                 linear_output,
                 opt.render_interval,
                 opt.allow_partial_files,
+                store_partial_renders,
             )?;
             if opt.preview {
                 output.frames.truncate(1);