rasmus/burn-stablediffusion-vibecode
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add ability to load dump or burn model in sample binary

Browse Source
This commit is contained in:
Gadersd
2023-08-05 13:27:37 -04:00
parent c408de50b7
commit 08f7347ae5
1 changed files with 21 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/bin/sample/main.rs
View File

@@ -1,4 +1,4 @@
use stablediffusion::{tokenizer::SimpleTokenizer, model::stablediffusion::*};
use stablediffusion::{tokenizer::SimpleTokenizer, model::stablediffusion::{*, load::load_stable_diffusion}};
use burn::{
config::Config,
@@ -30,30 +30,39 @@ fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() != 6 {
eprintln!("Usage: {} <model_name> <unconditional_guidance_scale> <n_diffusion_steps> <prompt> <output_image_name>", args[0]);
eprintln!("Usage: {} <model_type(burn or dump)> <model_name> <unconditional_guidance_scale> <n_diffusion_steps> <prompt> <output_image_name>", args[0]);
process::exit(1);
}
let model_name = &args[1];
let unconditional_guidance_scale: f64 = args[2].parse().unwrap_or_else(|_| {
let model_type = &args[1];
let model_name = &args[2];
let unconditional_guidance_scale: f64 = args[3].parse().unwrap_or_else(|_| {
eprintln!("Error: Invalid unconditional guidance scale.");
process::exit(1);
});
let n_steps: usize = args[3].parse().unwrap_or_else(|_| {
let n_steps: usize = args[4].parse().unwrap_or_else(|_| {
eprintln!("Error: Invalid number of diffusion steps.");
process::exit(1);
});
let prompt = &args[4];
let output_image_name = &args[5];
let prompt = &args[5];
let output_image_name = &args[6];
println!("Loading tokenizer...");
let tokenizer = SimpleTokenizer::new().unwrap();
println!("Loading model...");
let sd: StableDiffusion<Backend> = load_stable_diffusion_model_file(model_name).unwrap_or_else(|err| {
let sd: StableDiffusion<Backend> = if model_type == "burn" {
load_stable_diffusion_model_file(model_name).unwrap_or_else(|err| {
eprintln!("Error loading model: {}", err);
process::exit(1);
});
})
} else {
load_stable_diffusion(model_name, &device).unwrap_or_else(|err| {
eprintln!("Error loading model dump: {}", err);
process::exit(1);
})
};
let sd = sd.to_device(&device);
let unconditional_context = sd.unconditional_context(&tokenizer);