fix: k_dpm_fast on MPS platform

fixes #75

As discussed here: https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/2446#issuecomment-1279733739

Makefile

@@ -127,7 +127,9 @@ vendorize_kdiffusion:
 	sed -i '' -e 's#return (x - denoised) / utils.append_dims(sigma, x.ndim)#return (x - denoised) / sigma#g' imaginairy/vendored/k_diffusion/sampling.py
 	sed -i '' -e 's#x = x + torch.randn_like(x) \* sigma_up#x = x + torch.randn_like(x, device="cpu").to(x.device) \* sigma_up#g' imaginairy/vendored/k_diffusion/sampling.py
  	# https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/4558#issuecomment-1310387114
-	sed -i '' -e 's#t_fn = lambda sigma: sigma.log().neg()#t_fn = lambda sigma: sigma.to('cpu').log().neg().to(x.device)#g' imaginairy/vendored/k_diffusion/sampling.py
+	sed -i '' -e 's#t_fn = lambda sigma: sigma.log().neg()#t_fn = lambda sigma: sigma.to("cpu").log().neg().to(x.device)#g' imaginairy/vendored/k_diffusion/sampling.py
+	sed -i '' -e 's#return (x - denoised) / sigma#return (x - denoised) / sigma#g' imaginairy/vendored/k_diffusion/sampling.py
+	sed -i '' -e 's#return t.neg().exp()#return t.to("cpu").neg().exp().to(self.model.device)#g' imaginairy/vendored/k_diffusion/sampling.py
 	make af
 
 vendorize_noodle_soup:

README.md

@@ -225,9 +225,10 @@ docker run -it --gpus all -v $HOME/.cache/huggingface:/root/.cache/huggingface -
 ## ChangeLog
 
 - feature: added `DPM++ 2S a` and `DPM++ 2M` samplers.
+- feature: improve progress image logging
 - fix: fix bug with `--show-work`
 - fix: add workaround for pytorch bug affecting MacOS users using the new `DPM++ 2S a` and `DPM++ 2M` samplers.
-- feature: improve progress image logging
+- fix: add workaround for pytorch mps bug affecting `k_dpm_fast` sampler. fixes #75
 
 **5.0.0**
 - feature: 🎉 inpainting support using new inpainting model from RunwayML. It works really well! (Unfortunately it requires a HuggingFace token).

imaginairy/vendored/k_diffusion/sampling.py

@@ -37,7 +37,7 @@ def get_sigmas_vp(n, beta_d=19.9, beta_min=0.1, eps_s=1e-3, device="cpu"):
 
 def to_d(x, sigma, denoised):
     """Converts a denoiser output to a Karras ODE derivative."""
-    return (x - denoised) / sigma
+    return ((x - denoised) / sigma.to("cpu")).to(x.device)
 
 
 def get_ancestral_step(sigma_from, sigma_to, eta=1.0):
@@ -394,7 +394,7 @@ class DPMSolver(nn.Module):
         return -sigma.log()
 
     def sigma(self, t):
-        return t.neg().exp()
+        return t.to("cpu").neg().exp().to(self.model.device)
 
     def eps(self, eps_cache, key, x, t, *args, **kwargs):
         if key in eps_cache: