BF16 doesn't work with dreambooth

[MalumaDev]

Describe the bug

Using bf16 with the dreambooth script throw an exception

Reproduction

No response

Logs

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

[patrickvonplaten]

Interesting! @patil-suraj should we maybe add a "if-else" statement for this?

[patil-suraj]

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.

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'