Use buffer to increase model loading/saving speed tremendously

src/bin/convert/main.rs

@@ -16,7 +16,8 @@ use burn::{
 
 use burn_tch::{TchBackend, TchDevice};
 
-use burn::record::{self, Recorder, BinFileRecorder, FullPrecisionSettings};
+use burn::record::{self, Recorder, FullPrecisionSettings};
+use stablediffusion::binrecorder::{BinFileRecorderBuffered};
 
 fn convert_dump_to_model<B: Backend>(dump_path: &str, model_name: &str, device: &B::Device) -> Result<(), Box<dyn Error>> {
     println!("Loading dump...");
@@ -29,7 +30,7 @@ fn convert_dump_to_model<B: Backend>(dump_path: &str, model_name: &str, device:
 }
 
 fn save_model_file<B: Backend>(model: StableDiffusion<B>, name: &str) -> Result<(), record::RecorderError> {
-    BinFileRecorder::<FullPrecisionSettings>::new()
+    BinFileRecorderBuffered::<FullPrecisionSettings>::new()
     .record(
         model.into_record(),
         name.into(),

src/bin/sample/main.rs

@@ -15,10 +15,11 @@ use std::env;
 use std::io;
 use std::process;
 
-use burn::record::{self, Recorder, BinFileRecorder, FullPrecisionSettings};
+use burn::record::{self, Recorder, FullPrecisionSettings};
+use stablediffusion::binrecorder::{BinFileRecorderBuffered};
 
 fn load_stable_diffusion_model_file<B: Backend>(filename: &str) -> Result<StableDiffusion<B>, record::RecorderError> {
-    BinFileRecorder::<FullPrecisionSettings>::new()
+    BinFileRecorderBuffered::<FullPrecisionSettings>::new()
     .load(filename.into())
     .map(|record| StableDiffusionConfig::new().init().load_record(record))
 }