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

@@ -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))
}