Flux inference with torch.half outputs NaN values

[latentCall145]

Describe the bug

When running Flux inference with torch.half (i.e. for faster inference on Turing GPUs), there are NaN outputs in the resulting image (will appear as black in matplotlib).

Reproduction

from diffusers import FluxPipeline
import matplotlib.pyplot as plt
import torch
torch.backends.cudnn.benchmark = True

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.bfloat16,
)

pipe.enable_sequential_cpu_offload()
pipe.vae.enable_tiling()
pipe.to(torch.half) # for Turing GPUs, torch_dtype=torch.half in the pipeline constructor uses unnecessary RAM since CPU offload hasn't been enabled yet

images = pipe(
    "an astronaut riding a horse on mars",
    num_inference_steps=1,
    num_images_per_prompt=1,
    guidance_scale=0.0,
    height=1024,
    width=1024,
).images

for img_idx, image in enumerate(images):
    plt.imshow(image)
    plt.show() # will show black images

Logs

/opt/conda/lib/python3.10/site-packages/diffusers/image_processor.py:111: RuntimeWarning: invalid value encountered in cast
  images = (images * 255).round().astype("uint8")

System Info

• 🤗 Diffusers version: 0.30.0
• Platform: Linux-5.15.154+-x86_64-with-glibc2.31
• Running on Google Colab?: No
• Python version: 3.10.13
• PyTorch version (GPU?): 2.1.2 (True)
• Flax version (CPU?/GPU?/TPU?): 0.8.4 (gpu)
• Jax version: 0.4.26
• JaxLib version: 0.4.26.dev20240504
• Huggingface_hub version: 0.23.4
• Transformers version: 4.42.3
• Accelerate version: 0.32.1
• PEFT version: not installed
• Bitsandbytes version: not installed
• Safetensors version: 0.4.3
• xFormers version: not installed
• Accelerator: Tesla T4, 15360 MiB
  Tesla T4, 15360 MiB
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: no

Who can help?

No response