From 10ce5ac89d0a20ea11108144ae97a94d3eef1e72 Mon Sep 17 00:00:00 2001
From: Gadersd <gadersd@gmail.com>
Date: Sat, 5 Aug 2023 13:46:14 -0400
Subject: [PATCH] Correct a few trivial bugs

---
 README.md              | 8 +++++---
 src/bin/sample/main.rs | 2 +-
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index f72b02a..fb9e1f5 100644
--- a/README.md
+++ b/README.md
@@ -19,11 +19,13 @@ export TORCH_CUDA_VERSION=cu113
 ```
 ### Step 2: Run the Sample Binary
 
-Invoke the sample binary provided in the rust code, as shown below:
+Invoke the sample binary provided in the rust code, as shown below. Loading the burn model file
+is currently very slow, but hopefully that will be rectified soon. You can also dump a torch model's weights and load that
+which is currently much faster.
 
 ```bash
-# Arguments: model unconditional_guidance_scale n_diffusion_steps prompt output_image
-cargo run --release --bin sample SDv1-4 7.5 20 "A half-eaten apple sitting on a desk." apple.png
+# Arguments: <model_type(burn or dump)> <model> <unconditional_guidance_scale> <n_diffusion_steps> <prompt> <output_image>
+cargo run --release --bin sample burn SDv1-4 7.5 20 "A half-eaten apple sitting on a desk." apple
 ```
 
 This command will generate an image according to the provided prompt, which will be saved as 'apple.png'.
diff --git a/src/bin/sample/main.rs b/src/bin/sample/main.rs
index 14b4169..4d0179e 100644
--- a/src/bin/sample/main.rs
+++ b/src/bin/sample/main.rs
@@ -29,7 +29,7 @@ fn main() {
     let device = TchDevice::Cuda(0);
 
     let args: Vec<String> = std::env::args().collect();
-    if args.len() != 6 {
+    if args.len() != 7 {
         eprintln!("Usage: {} <model_type(burn or dump)> <model_name> <unconditional_guidance_scale> <n_diffusion_steps> <prompt> <output_image_name>", args[0]);
         process::exit(1);
     }