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

Update README

Browse Source
This commit is contained in:
Gadersd
2023-08-05 10:39:51 -04:00
parent 9e247777fa
commit c408de50b7
2 changed files with 77 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
README.md
View File

@@ -1,2 +1,67 @@
# stable-diffusion-burn
Stable Diffusion v1.4 ported to Rust's burn framework
# Stable-Diffusion-Burn
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!

13
python/dump.py
View File

@@ -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.')