Repository: huggingface/diffusers
Pull request: support bf16 for stable diffusion
Currently stable diffusion (or diffusers in general) doesn't work with `bf16` as  nearest upsampling in torch is not supported in `bf16`. 

Minimal code to reproduce:

```python
import torch
import torch.nn.functional as F

image = torch.randn(1, 4, 32, 32).to(device="cuda", dtype=torch.bfloat16)
out = F.interpolate(image, size=(64, 64), mode="nearest")
```

Addinatly, in `pipelines` we need to cast the images to `fp32` as `bf16` is not yet supported in `numpy`.

This is a draft PR to enable `bf16` training/inference for stable diffusion by casting `input` to `fp32` where `bf16` is not supported.

Not sure if this is the right way, curious to you hear your feedback @patrickvonplaten @NouamaneTazi 

fixes #771

Problem statement:
BF16 doesn't work with dreambooth
### Describe the bug

Using bf16 with the dreambooth script throw an exception

### Reproduction

_No response_

### Logs

```shell
RuntimeError: "upsample_nearest2d_out_frame" not implemented for 'BFloat16'
```


### System Info

- `diffusers` version: 0.4.0.dev0
- Platform: Linux-5.15.0-47-generic-x86_64-with-glibc2.35
- Python version: 3.10.0
- PyTorch version (GPU?): 1.12.1 (True)
- Huggingface_hub version: 0.9.1
- Transformers version: 4.22.2
- Using GPU in script?: Yes (3090)
- Using distributed or parallel set-up in script?:  Nope

Hints from issue discussion:
Interesting! @patil-suraj should we maybe add a "if-else" statement for this? 
Hi! Thanks for the issue, I've opened a draft PR to address this #792. Also think this issue should be filed in pytorch as the torch upsample op is not  supported in `bf16`. Minimal code to reproduce.

```python
import torch
import torch.nn.functional as F

image = torch.randn(1, 4, 32, 32).to(device="cuda", dtype=torch.bfloat16)
out = F.interpolate(image, size=(64, 64), mode="nearest")
```

relevant stack-trace 
```
File ~/.pyenv/versions/3.9.14/envs/diffusers-env/lib/python3.9/site-packages/torch/nn/functional.py:3910, in interpolate(input, size, scale_factor, mode, align_corners, recompute_scale_factor, antialias)
   3908     return torch._C._nn.upsample_nearest1d(input, output_size, scale_factors)
   3909 if input.dim() == 4 and mode == "nearest":
-> 3910     return torch._C._nn.upsample_nearest2d(input, output_size, scale_factors)
   3911 if input.dim() == 5 and mode == "nearest":
   3912     return torch._C._nn.upsample_nearest3d(input, output_size, scale_factors)

RuntimeError: "upsample_nearest2d_out_frame" not implemented for 'BFloat16'
```

Commit message: fix typo

Patch to review:
diff --git a/src/diffusers/models/resnet.py b/src/diffusers/models/resnet.py
index b9718e67f279..5b1f88e5416e 100644
--- a/src/diffusers/models/resnet.py
+++ b/src/diffusers/models/resnet.py
@@ -40,6 +40,11 @@ def forward(self, hidden_states, output_size=None):
         if self.use_conv_transpose:
             return self.conv(hidden_states)
 
+        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
+        dtype = hidden_states.dtype
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(torch.float32)
+
         # if `output_size` is passed we force the interpolation output
         # size and do not make use of `scale_factor=2`
         if output_size is None:
@@ -47,6 +52,10 @@ def forward(self, hidden_states, output_size=None):
         else:
             hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
 
+        # If the input is bfloat16, we cast back to bfloat16
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(dtype)
+
         # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed
         if self.use_conv:
             if self.name == "conv":
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
index 00e72de6551a..335d442437f9 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
@@ -327,7 +327,13 @@ def __call__(
         image = self.vae.decode(latents).sample
 
         image = (image / 2 + 0.5).clamp(0, 1)
-        image = image.cpu().permute(0, 2, 3, 1).numpy()
+        image = image.cpu().permute(0, 2, 3, 1)
+
+        # cast to float32 to as numpy doesn't support bfloat16
+        if image.dtype == torch.bfloat16:
+            image = image.to(torch.float32).numpy()
+        else:
+            image = image.numpy()
 
         safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(self.device)
         image, has_nsfw_concept = self.safety_checker(
diff --git a/src/diffusers/pipelines/stable_diffusion/safety_checker.py b/src/diffusers/pipelines/stable_diffusion/safety_checker.py
index 773a7d4b2107..68e4fef6bb6e 100644
--- a/src/diffusers/pipelines/stable_diffusion/safety_checker.py
+++ b/src/diffusers/pipelines/stable_diffusion/safety_checker.py
@@ -36,8 +36,16 @@ def forward(self, clip_input, images):
         pooled_output = self.vision_model(clip_input)[1]  # pooled_output
         image_embeds = self.visual_projection(pooled_output)
 
-        special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds).cpu().numpy()
-        cos_dist = cosine_distance(image_embeds, self.concept_embeds).cpu().numpy()
+        special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds).cpu()
+        cos_dist = cosine_distance(image_embeds, self.concept_embeds).cpu()
+
+        # cast to float32 to as numpy does not support bfloat16
+        if image_embeds.dtype == torch.bfloat16:
+            special_cos_dist = special_cos_dist.float().numpy()
+            cos_dist = cos_dist.float().numpy()
+        else:
+            special_cos_dist = special_cos_dist.numpy()
+            cos_dist = cos_dist.numpy()
 
         result = []
         batch_size = image_embeds.shape[0]


Code review comments:
diff_hunk: @@ -36,8 +36,16 @@ def forward(self, clip_input, images):
         pooled_output = self.vision_model(clip_input)[1]  # pooled_output
         image_embeds = self.visual_projection(pooled_output)
 
-        special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds).cpu().numpy()
-        cos_dist = cosine_distance(image_embeds, self.concept_embeds).cpu().numpy()
+        special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds).cpu()
+        cos_dist = cosine_distance(image_embeds, self.concept_embeds).cpu()
+
+        # cast to float32 to as numpy does not support bfloat16
+        if image_embeds.dtype == torch.bfloat16:
path: src/diffusers/pipelines/stable_diffusion/safety_checker.py
text: @user1:
I don't think we need an `if` statement here as `.float()` should always work and be correct (e.g. we can't do `fp16` on CPU either and doing `numpy()` will move it to `fp32` anyways

@user2:
@user1 why would `numpy()` move the arrays to `fp32`? I thought that numpy arrays support `fp16`?
I believe we should use `.half()` instead which should work on CPU also?

diff_hunk: @@ -327,7 +327,13 @@ def __call__(
         image = self.vae.decode(latents).sample
 
         image = (image / 2 + 0.5).clamp(0, 1)
-        image = image.cpu().permute(0, 2, 3, 1).numpy()
+        image = image.cpu().permute(0, 2, 3, 1)
path: src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
text: @user1:
```suggestion
        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
```

instead of the if-else should be good here IMO

@user2:
Or probably:
```python
image = image.cpu().permute(0, 2, 3, 1).half().numpy()
```

as `self.numpy_to_pil()` only multiplies by 255 then rounds, so I dont think we need a full fp32 precision in here.

@author:
Sticking with `float` here as this is a one time operation and is done on CPU.