Update README

README.md

@@ -1,2 +1,67 @@
-# stable-diffusion-burn
+# Stable-Diffusion-Burn
-Stable Diffusion v1.4 ported to Rust's burn framework
+
+Stable-Diffusion-Burn is a Rust-based project which ports the V1 stable diffusion model into the deep learning framework, Burn. This repository is licensed under the MIT Licence.
+
+## How To Use
+
+### Step 1: Download the Model and Set Environment Variables
+
+Start by downloading the SDv1-4.bin model provided on HuggingFace.
+
+```bash
+wget https://huggingface.co/Gadersd/Stable-Diffusion-Burn/resolve/main/V1/SDv1-4.bin
+```
+
+Next, set the appropriate CUDA version.
+
+```bash
+export TORCH_CUDA_VERSION=cu113
+```
+### Step 2: Run the Sample Binary
+
+Invoke the sample binary provided in the rust code, as shown below:
+
+```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
+```
+
+This command will generate an image according to the provided prompt, which will be saved as 'apple.png'.
+
+### Optional: Extract and Convert a Fine-Tuned Model
+
+If users are interested in using a fine-tuned version of stable diffusion, the Python scripts provided in this project can be used to transform a weight dump into a Burn model file.
+
+```bash
+# Step into the Python directory
+cd python
+
+# Download the model, this is just the base v1.4 model as an example
+wget https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt
+
+# Extract the weights
+python3 dump.py sd-v1-4.ckpt
+
+# Move the extracted weight folder out
+mv params ..
+
+# Step out of the Python directory
+cd ..
+
+# Convert the weights into a usable form
+cargo run --release --bin convert params SDv1-4
+```
+
+The binaries 'convert' and 'sample' are contained in Rust. Convert works on CPU whereas sample needs CUDA.
+
+Remember, `convert` should be used if you're planning on using the fine-tuned version of the stable diffusion. 
+
+## License
+
+This project is licensed under MIT license.
+
+## Example Inference
+
+INSER IMAGE HERE
+
+We wish you a productive time using this project. Enjoy!

python/dump.py

@@ -589,6 +589,7 @@ class StableDiffusion:
 # this is sd-v1-4.ckpt
 FILENAME = Path(__file__).parent.parent / "weights/sd-v1-4.ckpt"
 
+import sys
 import clip as clipsave
 import autoencoder as autoencodersave
 import unet as unetsave
@@ -631,14 +632,20 @@ if __name__ == "__main__":
   output = unet(input, timesteps, context)
   #print(output.numpy())'''
 
+  if len(sys.argv) != 2:
+        print(f"Wrong command line parameters, Usage: python3 {sys.argv[0]} <model_filename>")
+        sys.exit()
+
+  FILENAME = sys.argv[1]
 
   Tensor.no_grad = True
   model = StableDiffusion()
 
   # load in weights
-  download_file('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', FILENAME)
+  #download_file('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', FILENAME)
   load_state_dict(model, torch_load(FILENAME)['state_dict'], strict=False)
 
-  print('Saving model...')
+  print('Dumping model...')
   sdsave.save_stable_diffusion(model, "params")
-  print('Model saved.')
+  print('Model weights saved in params.')
+