Update to burn v0.14.0 and switch to .mpk model file

This commit is contained in:
Hermes
2024-10-05 14:19:49 -04:00
committed by Ben_Kosytorz
parent 3c49b0a151
commit 75f0cedd9f
19 changed files with 366 additions and 311 deletions

View File

@@ -11,9 +11,9 @@ use burn::{
tensor::{backend::Backend, Tensor},
};
use burn_ndarray::{NdArrayBackend, NdArrayDevice};
use burn_ndarray::{NdArray, NdArrayDevice};
use burn::record::{self, BinFileRecorder, FullPrecisionSettings, Recorder};
use burn::record::{self, NamedMpkFileRecorder, FullPrecisionSettings, Recorder};
fn convert_dump_to_model<B: Backend>(
dump_path: &str,
@@ -33,11 +33,11 @@ fn save_model_file<B: Backend>(
model: StableDiffusion<B>,
name: &str,
) -> Result<(), record::RecorderError> {
BinFileRecorder::<FullPrecisionSettings>::new().record(model.into_record(), name.into())
NamedMpkFileRecorder::<FullPrecisionSettings>::new().record(model.into_record(), name.into())
}
fn main() {
type Backend = NdArrayBackend<f32>;
type Backend = NdArray<f32>;
let device = NdArrayDevice::Cpu;
let args: Vec<String> = env::args().collect();