Repository: huggingface/diffusers
Pull request: Inference support for `mps` device
This addresses the issues identified when assessing inference on Apple Silicon, see https://github.com/huggingface/diffusers/issues/292#issuecomment-1236112207 for details.

### Current status

- Stable Diffusion pipeline works.
- Results in CPU / MPU are reproducible when using the same seeds. Generators do not work in the `mps` device, so we needed some minor adjustments.
- Some incompatible ops identified by failing tests were rewritten (or we fall back to CPU). But I have only verified `test_models_unet` and `test_models_vae`.
  - See f20a0dda6f1d2b46ac872852be4ad58188e63a5e
  - See 7f40f24297e4d3dc0279fdfe38ae4015de8b2747 and https://github.com/pytorch/pytorch/issues/84535
- Determinism tests pass, but the solution is a workaround hack.

### The hack

We perform a one-time "warmup" forward pass through `unet` and `vae` because the result from the first pass is different than subsequent results. I suspect something related to randomness might be at play, but I haven't identified the root cause. We have several options here:
- Remove the hack, it's probably overkill. We can recommend users to run a full pass through the pipeline (1 step is enough) after moving it to the device. Downside is that determinism tests fail in `test_models_unet` and `test_models_vae`.
- Find the cause and apply a proper solution.

I'd like to merge soon but don't like the hack. Can we remove it and live with some failing tests until we can find what's causing the issue?

Update: we only perform the hack during testing. Users are recommended to perform an initial pass if they care about reproducibility.

### To do:
- [x] Some `unet` and `vae` tests still fail. For example, `test_output_pretrained` in `AutoencoderKLTests`.
- [x] Translate some incompatible ops or fall back to CPU.
- [x] Ensure tests pass.
- [x] ~~Find a more principled workaround to perform the initial "warm up" pass~~. I tried to use forward hooks, but couldn't do it because they don't pass keyword arguments, which we use in many of our `forward` implementations.
- Create fail case for first-pass issue (results are different than in subsequent passes). Deferred to be tracked in #372.

Fixes #292.

Problem statement:
Add inference support for `mps` device (Apple Silicon)
There is a lot of community interest in running diffusers on Apple Silicon. A first step could be to introduce support for the `mps` device, which currently requires PyTorch nightly builds. Another step down the line could be to convert to Core ML and optimize to make sure that the use of the ANE (Neural Engine) is maximized.

This PR deals with the first approach.

**Describe the solution you'd like**
The following should work:

```Python
pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True)
pipe = pipe.to("mps")

# Rest of inference code remains the same
```

`pipe.to` would determine whether `mps` is available in the computer, and would raise an error if it's not.

One consideration is that perhaps not _all_ the components have to be moved to the `mps` device. Perhaps it's more efficient or practical to keep the text encoder in CPU, for instance. If so, `to()` would move the required components and the pipelines would need to be adapted to move the tensors transparently.

**Describe alternatives you've considered**
Conversion to Core ML, as mentioned above.

**Additional context**
I have tested the `unet` module in `mps` vs `cpu`, and these are some preliminary results in my computer (M1 Max with 64 GB of unified RAM), when iterating through the 51 default steps of Stable Diffusion with the default scheduler and no classifier-free guidance:

```
Device: cpu, torch: 1.12.1,  time: 92.09s
Device: cpu, torch: 1.13.0.dev20220830,  time: 92.12s
Device: mps:0, torch: 1.13.0.dev20220830,  time: 17.20s
```

Execution in `mps` complains about one operation having to be performed in CPU:

```
The operator 'aten::masked_select' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
```

We need to investigate the reason and whether an alternative would yield better performance.

Hints from issue discussion:
If it is of any help with this issue, here are some community-implemented or contributed Silicon compatibility with Stable Diffusion: 

Using `diffusers`
PR to this repo: https://github.com/huggingface/diffusers/pull/278/
Guide by wandb: https://wandb.ai/morgan/stable-diffusion/reports/Running-Stable-Diffusion-on-an-Apple-M1-Mac-with-Hugging-Face-Diffusers--VmlldzoyNTU2ODc2

Using the CompVis Stable Diffusion repo: 
- https://github.com/einanao/stable-diffusion/tree/apple-silicon
- https://github.com/magnusviri/stable-diffusion/tree/apple-mps-support
- https://github.com/lstein/stable-diffusion

Here's everything I've learned from the above linked resources as a Jupyter notebook that can be run in VSCode on MacOS: [StableDiffusion-MPS.ipynb](https://gist.github.com/mja/4cf7957cd6505820717c45b32dacd2ed). 

When running the pipeline I've confirmed using Activity Monitor that the Python process is using the GPU.
Regarding the alternative, i.e. converting the model to CoreML:

@madebyollin made [gist](https://gist.github.com/madebyollin/86b9596ffa4ab0fa7674a16ca2aeab3d) that can convert the U-net to a CoreML model, which can be injected into the pipeline, or used as a standalone model in a native app. However, it is not a complete model as this library (`diffusers`) is doing an extra work around the model.

From what I've seen so far this task entails the following:

- [x] Make the pipeline compatible when using classifier-free guidance. This is the `contiguous()` patch mentioned by `mja` in his notebook and addressed in #278. Interestingly, the `unet` works fine when not using CFG (i.e., input tensors are not concatenated).
- [x] Reproducibility. We need a workaround for seeds / generators not working properly, see https://github.com/pytorch/pytorch/issues/84288
- [ ] "Warmup". The first time you run a batch through the forward pass the results are different to those from subsequent passes (for the same input). This is most puzzling, and occurs in both the `unet` and `vae` models. This breaks determinism tests, as pointed out by @kashif. I'd like to do two things here:
    - [x] Create a temporary workaround so we can use `mps` in the codebase soon.
    - [ ] Identify the root cause and open an issue in PyTorch (if it's a PyTorch issue, which we don't know yet).
- [x] Disable training tests for the `mps` device. There are additional problems with gradients when using `mps`, so we'll work on training later.
- [ ] Attempt to replace the ops still not ported (`aten::masked_select`).

Regarding @mja's notebook, he mentions that it's very slow, but it's much faster than the CPU in my tests. These are my preliminary results on M1 Max using `StableDiffusionPipeline` with 50 inference steps for a single image:

```
[mps] -  30.81s
[cpu] - 213.46s
```

Is this similar to what you are seeing, @mja?
> ```
> [mps] -  30.81s
> [cpu] - 213.46s
> ```
> 
> Is this similar to what you are seeing, @mja?

On an M1 Mac Mini I'm getting 188s on MPS and 343s on CPU for 50 steps at 512x512.

Edit: Doing a bit more testing, I think it's the RAM that might be the limit on base M1 devices that is causing poor performance. I see the Python process grow to around 14GB and lots of red in the Activity Monitor Memory Pressure graph.

Commit message: Do not use float64 when using LDMScheduler.

Fixes #358.

Patch to review:
diff --git a/README.md b/README.md
index 28e1be164bb5..918b787ef500 100644
--- a/README.md
+++ b/README.md
@@ -39,6 +39,10 @@ pip install --upgrade diffusers  # should install diffusers 0.2.4
 conda install -c conda-forge diffusers
 ```
 
+**Apple Silicon (M1/M2) support**
+
+See [README_mps.md](README_mps.md).
+
 ## Contributing
 
 We ❤️  contributions from the open-source community! 
diff --git a/README_mps.md b/README_mps.md
new file mode 100644
index 000000000000..12f82f9711de
--- /dev/null
+++ b/README_mps.md
@@ -0,0 +1,46 @@
+## How to use Stable Diffusion in Apple Silicon (M1/M2)
+
+🤗 Diffusers is compatible with Apple silicon for Stable Diffusion inference, using the PyTorch `mps` device. These are the steps you need to follow to use your M1 or M2 computer with Stable Diffusion.
+
+### Requirements
+
+- Mac computer with Apple silicon (M1/M2) hardware.
+- macOS 12.3 or later.
+- arm64 version of Python.
+
+### Install PyTorch Nightly 
+
+Install the latest [Preview (Nightly) build of PyTorch](https://pytorch.org/get-started/locally/) on your Apple silicon Mac. Please, make sure it meets the requirements above.
+
+We have tested inference using PyTorch version `1.13.0.dev20220830`. After installing, please verify that your version is at least that one.
+
+### Inference Pipeline
+
+The snippet shown below demonstrates how to use the `mps` backend using the familiar `to()` interface to move the Stable Diffusion pipeline to your M1 or M2 device.
+
+We recommend to "prime" the pipeline using an additional one-time pass through it. This is a temporary workaround for a weird issue we have detected: the first inference pass produces slightly different results than subsequent ones. You only need to do this pass once, and it's ok to use just one inference step and discard the result.
+
+```python
+# make sure you're logged in with `huggingface-cli login`
+from diffusers import StableDiffusionPipeline
+
+pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=True)
+pipe = pipe.to("mps")
+
+prompt = "a photo of an astronaut riding a horse on mars"
+
+# First-time "warmup" pass (see explanation above)
+_ = pipe(prompt, num_inference_steps=1)
+
+# Results match those from the CPU device after the warmup pass.
+image = pipe(prompt)["sample"][0]  
+```
+
+### Performance
+
+These are the results we got on a M1 Max MacBook Pro with 64 GB of RAM, running macOS Ventura Version 13.0 Beta (22A5331f). We performed Stable Diffusion text-to-image generation of the same prompt for 50 inference steps, using a guidance scale of 7.5.
+
+| Device | Steps | Time    |
+|--------|-------|---------|
+| CPU    | 50    | 213.46s |
+| MPS    | 50    | 30.81s  |
diff --git a/src/diffusers/modeling_utils.py b/src/diffusers/modeling_utils.py
index ec501e2ae1f8..891fab1a8b3b 100644
--- a/src/diffusers/modeling_utils.py
+++ b/src/diffusers/modeling_utils.py
@@ -585,6 +585,37 @@ def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool
         else:
             return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
 
+    def _mps_warmup_inputs(self, batch_size=None) -> Optional[Tuple]:
+        r"""
+        Temporary procedure to run a one-time forward pass on some models, when using the `mps` device.
+
+        It has been observed that the output of some models (`unet`, `vae`) is different the first time they run than the
+        rest, for the same inputs. We are investigating the root cause of the problem, but meanwhile these methods cand be
+        used, if desired, to warmup those modules so their outputs are consistent.
+
+        Return inputs suitable for the forward pass of this model.
+        These will usually be a tuple of tensors that will be automatically moved to the `mps` device on warmup.
+
+        Return `None` if no warmup is required.
+        """
+        return None
+
+    def _mps_warmup(self, batch_size=None, **kwargs):
+        r"""
+        Temporary procedure to run a one-time forward pass on some models, when using the `mps` device.
+
+        Applies the warmup using `warmup_inputs`.
+        """
+        if self.device.type != "mps":
+            return
+
+        with torch.no_grad():
+            w_inputs = self._mps_warmup_inputs(batch_size)
+            if w_inputs is None:
+                return
+            w_inputs = [w.to("mps") for w in w_inputs]
+            self.__call__(*w_inputs, **kwargs)
+
 
 def unwrap_model(model: torch.nn.Module) -> torch.nn.Module:
     """
diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py
index 8d52ee9bde92..f87c4fa7bbb7 100644
--- a/src/diffusers/models/attention.py
+++ b/src/diffusers/models/attention.py
@@ -137,6 +137,7 @@ def __init__(self, dim, n_heads, d_head, dropout=0.0, context_dim=None, gated_ff
         self.checkpoint = checkpoint
 
     def forward(self, x, context=None):
+        x = x.contiguous() if x.device.type == "mps" else x
         x = self.attn1(self.norm1(x)) + x
         x = self.attn2(self.norm2(x), context=context) + x
         x = self.ff(self.norm3(x)) + x
diff --git a/src/diffusers/models/resnet.py b/src/diffusers/models/resnet.py
index 50382bcab37d..9ed8adf37192 100644
--- a/src/diffusers/models/resnet.py
+++ b/src/diffusers/models/resnet.py
@@ -446,10 +446,17 @@ def upfirdn2d_native(input, kernel, up=1, down=1, pad=(0, 0)):
     kernel_h, kernel_w = kernel.shape
 
     out = input.view(-1, in_h, 1, in_w, 1, minor)
+
+    # Temporary workaround for mps specific issue: https://github.com/pytorch/pytorch/issues/84535
+    using_mps = out.device.type == "mps"
+    if using_mps:
+        out = out.to("cpu")
     out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
     out = out.view(-1, in_h * up_y, in_w * up_x, minor)
 
     out = F.pad(out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
+    if using_mps:
+        out = out.to("mps")
     out = out[
         :,
         max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
diff --git a/src/diffusers/models/unet_2d.py b/src/diffusers/models/unet_2d.py
index 8634456346fe..ccb3b057066d 100644
--- a/src/diffusers/models/unet_2d.py
+++ b/src/diffusers/models/unet_2d.py
@@ -184,3 +184,9 @@ def forward(
         output = {"sample": sample}
 
         return output
+
+    def _mps_warmup_inputs(self, batch_size) -> Tuple:
+        batch_size = 1 if batch_size is None else batch_size
+        w_sample = torch.randn((batch_size, self.in_channels, 32, 32))
+        t = torch.tensor([10], dtype=torch.int32)
+        return (w_sample, t)
diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py
index a6f6eb791ee1..bb23350db965 100644
--- a/src/diffusers/models/unet_2d_condition.py
+++ b/src/diffusers/models/unet_2d_condition.py
@@ -135,7 +135,9 @@ def forward(
         if not torch.is_tensor(timesteps):
             timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device)
         elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:
-            timesteps = timesteps[None].to(sample.device)
+            if sample.device.type == "mps":
+                timesteps = timesteps.to(dtype=torch.float32)
+            timesteps = timesteps[None].to(device=sample.device)
 
         # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
         timesteps = timesteps * torch.ones(sample.shape[0], dtype=timesteps.dtype, device=timesteps.device)
@@ -186,3 +188,10 @@ def forward(
         output = {"sample": sample}
 
         return output
+
+    def _mps_warmup_inputs(self, batch_size) -> Tuple:
+        batch_size = 1 if batch_size is None else batch_size
+        w_sample = torch.randn((batch_size, self.in_channels, 64, 64))
+        t = torch.tensor([10], dtype=torch.int32)
+        w_encoded = torch.rand((batch_size, 77, 768))
+        return (w_sample, t, w_encoded)
diff --git a/src/diffusers/models/vae.py b/src/diffusers/models/vae.py
index d9fe89609803..b342d2b65690 100644
--- a/src/diffusers/models/vae.py
+++ b/src/diffusers/models/vae.py
@@ -296,7 +296,10 @@ def __init__(self, parameters, deterministic=False):
             self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
 
     def sample(self, generator: Optional[torch.Generator] = None) -> torch.FloatTensor:
-        x = self.mean + self.std * torch.randn(self.mean.shape, generator=generator, device=self.parameters.device)
+        device = self.parameters.device
+        sample_device = "cpu" if device.type == "mps" else device
+        sample = torch.randn(self.mean.shape, generator=generator, device=sample_device).to(device)
+        x = self.mean + self.std * sample
         return x
 
     def kl(self, other=None):
@@ -451,3 +454,8 @@ def forward(self, sample: torch.FloatTensor, sample_posterior: bool = False) ->
             z = posterior.mode()
         dec = self.decode(z)
         return dec
+
+    def _mps_warmup_inputs(self, batch_size) -> Tuple:
+        batch_size = 4 if batch_size is None else batch_size
+        w_sample = torch.randn((batch_size, self.in_channels, self.sample_size, self.sample_size))
+        return (w_sample,)
diff --git a/src/diffusers/pipeline_utils.py b/src/diffusers/pipeline_utils.py
index 214133bc5f17..58f571d709df 100644
--- a/src/diffusers/pipeline_utils.py
+++ b/src/diffusers/pipeline_utils.py
@@ -55,7 +55,6 @@
 
 
 class DiffusionPipeline(ConfigMixin):
-
     config_name = "model_index.json"
 
     def register_modules(self, **kwargs):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
index d4290da6f030..17a65aab1f07 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
@@ -104,12 +104,11 @@ def __call__(
             latents = torch.randn(
                 latents_shape,
                 generator=generator,
-                device=self.device,
             )
         else:
             if latents.shape != latents_shape:
                 raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}")
-            latents = latents.to(self.device)
+        latents = latents.to(self.device)
 
         # set timesteps
         accepts_offset = "offset" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())
diff --git a/src/diffusers/schedulers/scheduling_pndm.py b/src/diffusers/schedulers/scheduling_pndm.py
index 57ed4fb7e369..e2fef9022583 100644
--- a/src/diffusers/schedulers/scheduling_pndm.py
+++ b/src/diffusers/schedulers/scheduling_pndm.py
@@ -270,7 +270,8 @@ def add_noise(
         noise: Union[torch.FloatTensor, np.ndarray],
         timesteps: Union[torch.IntTensor, np.ndarray],
     ) -> torch.Tensor:
-
+        # mps requires indices to be in the same device, so we use cpu as is the default with cuda
+        timesteps = timesteps.to(self.alphas_cumprod.device)
         sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
         sqrt_alpha_prod = self.match_shape(sqrt_alpha_prod, original_samples)
         sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
diff --git a/src/diffusers/testing_utils.py b/src/diffusers/testing_utils.py
index 13f6332a9432..a1288b4edb3d 100644
--- a/src/diffusers/testing_utils.py
+++ b/src/diffusers/testing_utils.py
@@ -8,6 +8,7 @@
 
 global_rng = random.Random()
 torch_device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_device = "mps" if torch.backends.mps.is_available() else torch_device
 
 
 def parse_flag_from_env(key, default=False):
diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
index e19238050559..6cd4961e12ae 100644
--- a/tests/test_modeling_common.py
+++ b/tests/test_modeling_common.py
@@ -18,6 +18,7 @@
 
 import numpy as np
 import torch
+from diffusers.modeling_utils import ModelMixin
 
 from diffusers.testing_utils import torch_device
 from diffusers.training_utils import EMAModel
@@ -29,12 +30,16 @@ def test_from_pretrained_save_pretrained(self):
 
         model = self.model_class(**init_dict)
         model.to(torch_device)
+        if isinstance(model, ModelMixin):
+            model._mps_warmup()
         model.eval()
 
         with tempfile.TemporaryDirectory() as tmpdirname:
             model.save_pretrained(tmpdirname)
             new_model = self.model_class.from_pretrained(tmpdirname)
             new_model.to(torch_device)
+            if isinstance(new_model, ModelMixin):
+                model._mps_warmup(inputs_dict["sample"].shape[0])
 
         with torch.no_grad():
             image = model(**inputs_dict)
@@ -54,6 +59,9 @@ def test_determinism(self):
         model = self.model_class(**init_dict)
         model.to(torch_device)
         model.eval()
+        if isinstance(model, ModelMixin):
+            model._mps_warmup(inputs_dict["sample"].shape[0])
+
         with torch.no_grad():
             first = model(**inputs_dict)
             if isinstance(first, dict):
@@ -132,6 +140,9 @@ def test_model_from_config(self):
         self.assertEqual(output_1.shape, output_2.shape)
 
     def test_training(self):
+        if torch_device == "mps":
+            return
+
         init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
 
         model = self.model_class(**init_dict)
@@ -147,6 +158,9 @@ def test_training(self):
         loss.backward()
 
     def test_ema_training(self):
+        if torch_device == "mps":
+            return
+
         init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
 
         model = self.model_class(**init_dict)
diff --git a/tests/test_models_unet.py b/tests/test_models_unet.py
index 47f562dd5341..bf7a500d8109 100644
--- a/tests/test_models_unet.py
+++ b/tests/test_models_unet.py
@@ -191,7 +191,9 @@ def dummy_input(self, sizes=(32, 32)):
         num_channels = 3
 
         noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
-        time_step = torch.tensor(batch_size * [10]).to(torch_device)
+        time_step = torch.tensor(batch_size * [10]).to(device=torch_device)
+        if torch_device == "mps":
+            time_step = time_step.to(dtype=torch.int32)
 
         return {"sample": noise, "timestep": time_step}
 
diff --git a/tests/test_models_vae.py b/tests/test_models_vae.py
index 7df6b42bf796..852b0b08e13f 100644
--- a/tests/test_models_vae.py
+++ b/tests/test_models_vae.py
@@ -18,6 +18,7 @@
 import torch
 
 from diffusers import AutoencoderKL
+from diffusers.modeling_utils import ModelMixin
 from diffusers.testing_utils import floats_tensor, torch_device
 
 from .test_modeling_common import ModelTesterMixin
@@ -79,6 +80,8 @@ def test_output_pretrained(self):
         model = AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy")
         model = model.to(torch_device)
         model.eval()
+        if isinstance(model, ModelMixin):
+            model._mps_warmup(1, sample_posterior=True)
 
         torch.manual_seed(0)
         if torch.cuda.is_available():


Code review comments:
diff_hunk: @@ -585,6 +585,37 @@ def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool
         else:
             return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
 
+    def _mps_warmup_inputs(self, batch_size=None) -> Optional[Tuple]:
+        r"""
+        Temporary procedure to run a one-time forward pass on some models, when using the `mps` device.
+
+        It has been observed that the output of some models (`unet`, `vae`) is different the first time they run than the
+        rest, for the same inputs. We are investigating the root cause of the problem, but meanwhile these methods cand be
path: src/diffusers/modeling_utils.py
text: @user1:
```suggestion
        rest, for the same inputs. We are investigating the root cause of the problem, but meanwhile these methods can be
```

diff_hunk: @@ -191,7 +191,9 @@ def dummy_input(self, sizes=(32, 32)):
         num_channels = 3
 
         noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
-        time_step = torch.tensor(batch_size * [10]).to(torch_device)
+        time_step = torch.tensor(batch_size * [10]).to(device=torch_device)
+        if torch_device == "mps":
+            time_step = time_step.to(dtype=torch.int32)
path: tests/test_models_unet.py
text: @user1:
think we can always cast it to `int32`

@author:
The reason was to not interfere with the common path. I observed some slight differences when using `int32` instead of `int64` during the Flax tests, but they might not be relevant here.

@user2:
Both work fine with me (Flax will be in another file and surely not run on "mps")

@author:
I meant it as a reference point, not to try to integrate it here :)

diff_hunk: @@ -0,0 +1,46 @@
+## How to use Stable Diffusion in Apple Silicon (M1/M2)
+
+🤗 Diffusers is compatible with Apple silicon for Stable Diffusion inference, using the PyTorch `mps` device. These are the steps you need to follow to use your M1 or M2 computer with Stable Diffusion.
+
+### Requirements
+
+- Mac computer with Apple silicon (M1/M2) hardware.
+- macOS 12.3 or later.
+- arm64 version of Python.
+
+### Install PyTorch Nightly 
+
+Install the latest [Preview (Nightly) build of PyTorch](https://pytorch.org/get-started/locally/) on your Apple silicon Mac. Please, make sure it meets the requirements above.
+
+We have tested inference using PyTorch version `1.13.0.dev20220830`. After installing, please verify that your version is at least that one.
+
+### Inference Pipeline
+
+The snippet shown below demonstrates how to use the `mps` backend using the familiar `to()` interface to move the Stable Diffusion pipeline to your M1 or M2 device.
+
+We recommend to "prime" the pipeline using an additional one-time pass through it. This is a temporary workaround for a weird issue we have detected: the first inference pass produces slightly different results than subsequent ones. You only need to do this pass once, and it's ok to use just one inference step and discard the result.
+
+```python
+# make sure you're logged in with `huggingface-cli login`
+from diffusers import StableDiffusionPipeline
+
+pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=True)
+pipe = pipe.to("mps")
+
+prompt = "a photo of an astronaut riding a horse on mars"
+
+# First-time "warmup" pass (see explanation above)
+_ = pipe(prompt, num_inference_steps=1)
+
+# Results match those from the CPU device after the warmup pass.
+image = pipe(prompt)["sample"][0]  
path: README_mps.md
text: @user1:
(nit)
```suggestion
image = pipe(prompt).images[0]
```

@author:
Sure, this became outdated with respect to `main`, I'll merge and adapt. Thanks!

diff_hunk: @@ -135,7 +135,9 @@ def forward(
         if not torch.is_tensor(timesteps):
             timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device)
         elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:
-            timesteps = timesteps[None].to(sample.device)
+            if sample.device.type == "mps":
+                timesteps = timesteps.to(dtype=torch.float32)
path: src/diffusers/models/unet_2d_condition.py
text: @user2:
shouldn't this be handled in the `if` block also ?

also can we always cast it to `float32` cc @user1 ?

@author:
I think I debugged the other branch and it worked fine (at least when using an `int`), I'll double check to make sure.

@user1:
Yes I think we can always cast here indeed