Add batching to generate two images at a time

2023-08-05 19:41:11 -04:00
parent 57b446c08d
commit 6e065035cd
2 changed files with 3 additions and 3 deletions
--- a/src/bin/sample/main.rs
+++ b/src/bin/sample/main.rs
@@ -66,7 +66,7 @@ fn main() {
    let sd = sd.to_device(&device);

    let unconditional_context = sd.unconditional_context(&tokenizer);
-    let context = sd.context(&tokenizer, prompt).unsqueeze();
+    let context = sd.context(&tokenizer, prompt).unsqueeze().repeat(0, 2); // generate 2 samples

    println!("Sampling image...");
    let images = sd.sample_image(context, unconditional_context, unconditional_guidance_scale, n_steps);