add cumsum workaround for newer PyTorch pytorch/pytorch#89784

Author: Birch-san

scripts/play.py

@@ -2,12 +2,14 @@
 
 # monkey-patch _randn to use CPU random before k-diffusion uses it
 from helpers.brownian_tree_mps_fix import reassuring_message
+from helpers.cumsum_mps_fix import reassuring_message as reassuring_message_2
 from helpers.device import DeviceLiteral, get_device_type
 from helpers.diffusers_denoiser import DiffusersSDDenoiser, DiffusersSD2Denoiser
 from helpers.cfg_denoiser import Denoiser, DenoiserFactory
 from helpers.log_intermediates import LogIntermediates, make_log_intermediates
 from helpers.schedules import KarrasScheduleParams, KarrasScheduleTemplate, get_template_schedule
 print(reassuring_message) # avoid "unused" import :P
+print(reassuring_message_2)
 
 import torch
 from torch import Generator, Tensor, randn, no_grad, argmin, zeros

src/helpers/cumsum_mps_fix.py

@@ -0,0 +1,13 @@
+# monkey-patch cumsum to fallback to CPU
+# https://github.com/pytorch/pytorch/issues/89784
+import torch
+from torch import cumsum, Tensor
+torch.cumsum = lambda input, *args, **kwargs: (
+  cumsum(input.cpu() if input.device.type == 'mps' else input, *args, **kwargs).to(input.device)
+)
+orig_cumsum = torch.Tensor.cumsum
+def patched_cumsum(self: Tensor, *args, **kwargs):
+    return orig_cumsum(self.cpu() if self.device.type == 'mps' else self, *args, **kwargs).to(self.device)
+torch.Tensor.cumsum = patched_cumsum
+
+reassuring_message = "monkey-patched cumsum to fallback to CPU, for compatibility on MPS backend; see https://github.com/pytorch/pytorch/issues/89784"