Update README.md

README.md

@@ -2,7 +2,7 @@
 
 This is a collection of scripts that can be useful for various tasks related to the [diffusers library](https://github.com/huggingface/diffusers)
 
-## Test against original checkpoints
+## 1. Test against original checkpoints
 
 **It's very important to have visually the exact same results as the original code bases.!**
 
@@ -37,4 +37,62 @@ image.save("/home/patrick_huggingface_co/images/aa_comp.png")
 
 Both commands should give the following image on a V100:
 
-![image](https://huggingface.co/diffusers/tools/resolve/main/aa_orig_comp%20(6).png)
+
+## 2. Test against [k-diffusion](https://github.com/crowsonkb/k-diffusion):
+
+You can run the following script to compare against k-diffusion.
+
+See results [here](https://huggingface.co/datasets/patrickvonplaten/images)
+
+```python
+from diffusers import StableDiffusionKDiffusionPipeline, HeunDiscreteScheduler, StableDiffusionPipeline, DPMSolverMultistepScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler
+import torch
+import os
+
+seed = 13
+inference_steps = 25
+#checkpoint = "CompVis/stable-diffusion-v1-4"
+checkpoint = "stabilityai/stable-diffusion-2-1"
+prompts = ["astronaut riding horse", "whale falling from sky", "magical forest", "highly photorealistic picture of johnny depp"]
+prompts = 8 * ["highly photorealistic picture of johnny depp"]
+#prompts = prompts[:1]
+samplers = ["sample_dpmpp_2m", "sample_euler", "sample_heun", "sample_dpm_2", "sample_lms"]
+#samplers = samplers[:1]
+
+pipe = StableDiffusionKDiffusionPipeline.from_pretrained(checkpoint, torch_dtype=torch.float16, safety_checker=None)
+pipe = pipe.to("cuda")
+
+for i, prompt in enumerate(prompts):
+    prompt_f = f"{'_'.join(prompt.split())}_{i}"
+    for sampler in samplers:
+        pipe.set_scheduler(sampler)
+        torch.manual_seed(seed + i)
+        image = pipe(prompt, num_inference_steps=inference_steps).images[0]
+        checkpoint_f = f"{'--'.join(checkpoint.split('/'))}"
+        os.makedirs(f"/home/patrick_huggingface_co/images/{checkpoint_f}", exist_ok=True)
+        os.makedirs(f"/home/patrick_huggingface_co/images/{checkpoint_f}/{sampler}", exist_ok=True)
+        image.save(f"/home/patrick_huggingface_co/images/{checkpoint_f}/{sampler}/{prompt_f}.png")
+
+
+pipe = StableDiffusionPipeline(**pipe.components)
+pipe = pipe.to("cuda")
+
+for i, prompt in enumerate(prompts):
+    prompt_f = f"{'_'.join(prompt.split())}_{i}"
+    for sampler in samplers:
+        if sampler == "sample_euler":
+            pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
+        elif sampler == "sample_heun":
+            pipe.scheduler = HeunDiscreteScheduler.from_config(pipe.scheduler.config)
+        elif sampler == "sample_dpmpp_2m":
+            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+        elif sampler == "sample_lms":
+            pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)
+ 
+        torch.manual_seed(seed + i)
+        image = pipe(prompt, num_inference_steps=inference_steps).images[0]
+        checkpoint_f = f"{'--'.join(checkpoint.split('/'))}"
+        os.makedirs("/home/patrick_huggingface_co/images/{checkpoint_f}", exist_ok=True)
+        os.makedirs(f"/home/patrick_huggingface_co/images/{checkpoint_f}/{sampler}", exist_ok=True)
+        image.save(f"/home/patrick_huggingface_co/images/{checkpoint_f}/{sampler}/{prompt_f}_hf.png")
+```