Support loading original stable diffusion checkpoints

.gitignore

@@ -201,4 +201,6 @@ cython_debug/
 data/
 ddata/
 migrators/
-data.*.tar
+data.*.tar
+sd-v1-5.safetensors
+out/

README.md

@@ -95,6 +95,24 @@ prompts = ["a photograph of an astronaut riding a horse"]
 images = pipeline.generate(prompts, models=models, device='cuda', idle_device='cpu')
 ```
 
+Load an original stable diffusion model in safetensors format:
+```py
+import safetensors.torch
+from stable_diffusion_pytorch import pipeline, convert_from_sdmodel, model_loader
+
+# load the checkpoint file
+state_dict = safetensors.torch.load_file('/PATH/TO/sd-v1-5.safetensors')
+
+# split the state_dict into the 4 state_dicts for our models
+state_dicts = convert_from_sdmodel.split_state_dict(state_dict)
+
+# create the model objects, and apply the weights in state_dicts
+models = model_loader.load_models(state_dicts, 'cpu')
+
+prompts = ["a photograph of an astronaut riding a horse"]
+images = pipeline.generate(prompts, models=models, device='cuda', idle_device='cpu')
+```
+
 Image-to-image generation:
 ```py
 from PIL import Image

convert_model.py

@@ -0,0 +1,28 @@
+import argparse
+import safetensors.torch
+from stable_diffusion_pytorch import convert_from_sdmodel
+
+parser = argparse.ArgumentParser(description='Convert an original stable diffusion checkpoint file in safetensor format to 4 separate safetensors for this library.')
+parser.add_argument(
+    "--sd_model", type=str, required=True,
+    help="Stable Diffusion model to load in safetensors format"
+)
+parser.add_argument(
+    "--save_to", type=str, required=True,
+    help="The prefix of the path to save to, for example \"./mymodel_\" " \
+        "will save 4 files, mymodel_clip.safetensors, mymodel_decoder.safetensors, " \
+        "mymodel_encoder.safetensors, and mymodel_diffusion.safetensors"
+)
+
+args = parser.parse_args()
+
+# load the checkpoint file
+state_dict = safetensors.torch.load_file(args.sd_model)
+
+# convert to the state_dicts format that our library expects
+state_dicts = convert_from_sdmodel.split_state_dict(state_dict)
+
+for key in state_dicts:
+    outPath = args.save_to + key + '.safetensors'
+    print('Writing', outPath)
+    safetensors.torch.save_file(state_dicts[key], outPath)

demo_sdmodel_generate_image.py

@@ -0,0 +1,46 @@
+# python3
+# A demo of loading a third party model (everythingV5) meant for the original stable diffusion code, and generating images.
+import os
+import safetensors.torch
+import torch
+from stable_diffusion_pytorch import pipeline, convert_from_sdmodel, model_loader
+
+# Runs on intel gpu by default, to run on cuda remove the following import and change 'xpu' to 'cuda'.
+import intel_extension_for_pytorch as ipex
+device = 'xpu'
+
+
+# where to store generated images
+outDir = './out'
+
+# path to SD 1.4/1.5 based model safetensors file
+modelPath = './sd-v1-5.safetensors' # download from https://huggingface.co/tabtap/sd-v1-5.safetensor/tree/main
+
+# either float16 or float32
+dtype = torch.float16
+
+
+
+os.makedirs(outDir, exist_ok=True)
+
+# load the checkpoint file
+state_dict = safetensors.torch.load_file(modelPath)
+
+# convert to the state_dicts format that our library expects
+state_dicts = convert_from_sdmodel.split_state_dict(state_dict)
+
+# create the model objects, and apply the weights in state_dicts
+models = model_loader.load_models(state_dicts, device, dtype)
+
+steps = 40
+seed = 12345
+prompt = '1girl,cirno,anime style'
+negativePrompt = 'bad anatomy,bad hands,missing fingers,extra fingers'
+
+fileName = prompt.replace(' ', '_').replace('\\', '＼').replace(':', '⦂').replace(',', '-') + '.' + str(seed) + '.png'
+
+images = pipeline.generate([prompt], uncond_prompts=[negativePrompt],
+            models=models, n_inference_steps=steps, seed=seed, device=device,
+            height=768, width=512)
+
+images[0].save(outDir + '/' + fileName)