Use buffer to increase model loading/saving speed tremendously

This commit is contained in:
Gadersd
2023-08-05 16:48:26 -04:00
parent ce91c8838f
commit 2f3d94f9ad
5 changed files with 95 additions and 6 deletions

View File

@@ -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(),