Unstable flux transformer prediction

### Describe the bug

When feeding the same input (same noise, same prompt embed, same guidance scale and same timestep)to the flux.transformer, i got different output. I tried to fix the seed but nothing changed.

### Reproduction

import torch
import math

from diffusers import FluxPipeline
from torch import Tensor

import random
import numpy as np

def setup_seed(seed=1234):
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

def get_schedule(
num_steps: int,
image_seq_len: int,
base_shift: float = 0.5,
max_shift: float = 1.15,
shift: bool = True,
) -> list:
def get_lin_function(
x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
):
m = (y2 - y1) / (x2 - x1)
b = y1 - m * x1
return lambda x: m * x + b
def time_shift(mu: float, sigma: float, t: Tensor):
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)

# extra step for zero
timesteps = torch.linspace(1, 0, num_steps + 1, dtype=torch.float32)

# shifting the schedule to favor high timesteps for higher signal images
if shift:
# estimate mu based on linear estimation between two points
mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
timesteps = time_shift(mu, 1.0, timesteps)

return timesteps.tolist()

if __name__ == '__main__':
setup_seed(1234) # fix the random seed

DTYPE = torch.bfloat16
device = 'cuda'
pipe = FluxPipeline.from_pretrained('PATH/TO/MODEL', torch_dtype=DTYPE)
pipe.to(device)

with torch.no_grad():
shape = (1, 16, 128, 128)
rand_noise = torch.randn(shape, device=device, dtype=DTYPE)
rand_noise = pipe._pack_latents(
rand_noise,
batch_size=shape[0],
num_channels_latents=shape[1],
height=shape[2],
width=shape[3],
)

t_vec = torch.full((shape[0],), torch.tensor(1.), dtype=DTYPE, device=device)
guidance_vec = torch.full((shape[0],), 2., device=device, dtype=DTYPE)
latent_image_ids = pipe._prepare_latent_image_ids(
shape[0], shape[2], shape[3], device, DTYPE,
)

prompt = 'A cat holding a sign that says hello world'
prompt_embeds, pooled_prompt_embeds, text_ids = pipe.encode_prompt( # null text
prompt=prompt,
prompt_2=prompt
)
while 1:
flux_velocity = pipe.transformer(
hidden_states=rand_noise,
timestep=t_vec,
guidance=guidance_vec,
pooled_projections=pooled_prompt_embeds,
encoder_hidden_states=prompt_embeds,
txt_ids=text_ids,
img_ids=latent_image_ids,
joint_attention_kwargs=None,
return_dict=pipe,
)[0]
print(flux_velocity)

### Logs

```shell
tensor([[[-1.7812, 0.5273, -0.5156, ..., 0.6719, 0.5898, 0.5234],
[-0.5664, -0.9258, 0.0598, ..., 1.1719, -0.4688, -0.3555],
[-1.2188, -0.2441, -0.4238, ..., -0.3906, -0.2373, 1.0391],
...,
[-0.0537, -0.2383, -0.0073, ..., -1.3750, -1.1250, 0.8750],
[-0.2812, 0.9062, -0.6250, ..., -0.5195, 0.6797, 1.0391],
[-0.7109, -0.5664, -0.7461, ..., 0.2129, -0.5430, 0.6367]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5312, -0.5156, ..., 0.6641, 0.5742, 0.5156],
[-0.5703, -0.9297, 0.0569, ..., 1.1719, -0.4766, -0.3594],
[-1.2188, -0.2480, -0.4238, ..., -0.3945, -0.2451, 1.0391],
...,
[-0.0574, -0.2412, -0.0081, ..., -1.3750, -1.1328, 0.8750],
[-0.2793, 0.8984, -0.6289, ..., -0.5234, 0.6719, 1.0391],
[-0.7188, -0.5664, -0.7500, ..., 0.2031, -0.5547, 0.6250]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7891, 0.5273, -0.5156, ..., 0.6641, 0.5820, 0.5156],
[-0.5625, -0.9258, 0.0588, ..., 1.1719, -0.4668, -0.3594],
[-1.2188, -0.2461, -0.4219, ..., -0.3906, -0.2422, 1.0391],
...,
[-0.0564, -0.2373, -0.0037, ..., -1.3750, -1.1250, 0.8750],
[-0.2793, 0.9062, -0.6211, ..., -0.5234, 0.6758, 1.0391],
[-0.7109, -0.5625, -0.7461, ..., 0.2090, -0.5508, 0.6328]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7891, 0.5273, -0.5195, ..., 0.6680, 0.5820, 0.5195],
[-0.5664, -0.9297, 0.0598, ..., 1.1719, -0.4688, -0.3594],
[-1.2188, -0.2451, -0.4219, ..., -0.3906, -0.2422, 1.0391],
...,
[-0.0559, -0.2412, -0.0044, ..., -1.3828, -1.1250, 0.8711],
[-0.2793, 0.8984, -0.6328, ..., -0.5273, 0.6719, 1.0391],
[-0.7109, -0.5664, -0.7500, ..., 0.2041, -0.5508, 0.6289]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5273, -0.5195, ..., 0.6680, 0.5820, 0.5234],
[-0.5625, -0.9219, 0.0603, ..., 1.1797, -0.4688, -0.3594],
[-1.2188, -0.2471, -0.4199, ..., -0.3926, -0.2422, 1.0391],
...,
[-0.0552, -0.2393, -0.0046, ..., -1.3750, -1.1250, 0.8750],
[-0.2734, 0.9062, -0.6250, ..., -0.5234, 0.6758, 1.0391],
[-0.7109, -0.5664, -0.7500, ..., 0.2051, -0.5508, 0.6328]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7891, 0.5273, -0.5195, ..., 0.6680, 0.5781, 0.5195],
[-0.5664, -0.9258, 0.0581, ..., 1.1797, -0.4707, -0.3594],
[-1.2188, -0.2471, -0.4258, ..., -0.3926, -0.2441, 1.0391],
...,
[-0.0542, -0.2393, -0.0028, ..., -1.3750, -1.1250, 0.8711],
[-0.2812, 0.8984, -0.6328, ..., -0.5234, 0.6758, 1.0312],
[-0.7148, -0.5664, -0.7461, ..., 0.2080, -0.5469, 0.6328]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5312, -0.5195, ..., 0.6680, 0.5820, 0.5234],
[-0.5625, -0.9258, 0.0593, ..., 1.1797, -0.4688, -0.3574],
[-1.2188, -0.2461, -0.4219, ..., -0.3926, -0.2412, 1.0469],
...,
[-0.0574, -0.2393, -0.0076, ..., -1.3750, -1.1250, 0.8750],
[-0.2793, 0.9023, -0.6289, ..., -0.5234, 0.6758, 1.0391],
[-0.7148, -0.5625, -0.7461, ..., 0.2080, -0.5469, 0.6328]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5312, -0.5156, ..., 0.6680, 0.5781, 0.5195],
[-0.5664, -0.9258, 0.0588, ..., 1.1797, -0.4707, -0.3594],
[-1.2188, -0.2451, -0.4238, ..., -0.3887, -0.2412, 1.0469],
...,
[-0.0513, -0.2402, -0.0075, ..., -1.3828, -1.1328, 0.8711],
[-0.2773, 0.9023, -0.6250, ..., -0.5234, 0.6758, 1.0312],
[-0.7188, -0.5664, -0.7461, ..., 0.2100, -0.5508, 0.6289]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5312, -0.5195, ..., 0.6680, 0.5859, 0.5273],
[-0.5625, -0.9258, 0.0549, ..., 1.1797, -0.4668, -0.3555],
[-1.2188, -0.2422, -0.4199, ..., -0.3867, -0.2383, 1.0469],
...,
[-0.0544, -0.2402, -0.0097, ..., -1.3750, -1.1250, 0.8789],
[-0.2754, 0.9023, -0.6289, ..., -0.5234, 0.6797, 1.0391],
[-0.7188, -0.5664, -0.7422, ..., 0.2129, -0.5469, 0.6328]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5312, -0.5156, ..., 0.6680, 0.5781, 0.5273],
[-0.5664, -0.9258, 0.0598, ..., 1.1797, -0.4668, -0.3574],
[-1.2109, -0.2451, -0.4238, ..., -0.3887, -0.2412, 1.0391],
...,
[-0.0566, -0.2373, -0.0042, ..., -1.3750, -1.1250, 0.8750],
[-0.2773, 0.8984, -0.6250, ..., -0.5234, 0.6797, 1.0391],
[-0.7188, -0.5664, -0.7461, ..., 0.2090, -0.5508, 0.6250]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5273, -0.5195, ..., 0.6680, 0.5820, 0.5195],
[-0.5664, -0.9297, 0.0605, ..., 1.1719, -0.4688, -0.3594],
[-1.2188, -0.2490, -0.4219, ..., -0.3926, -0.2432, 1.0391],
...,
[-0.0586, -0.2412, -0.0067, ..., -1.3750, -1.1250, 0.8750],
[-0.2773, 0.8984, -0.6289, ..., -0.5234, 0.6719, 1.0391],
[-0.7148, -0.5664, -0.7461, ..., 0.2070, -0.5547, 0.6328]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5312, -0.5195, ..., 0.6680, 0.5781, 0.5195],
[-0.5664, -0.9258, 0.0608, ..., 1.1719, -0.4707, -0.3594],
[-1.2188, -0.2471, -0.4238, ..., -0.3906, -0.2402, 1.0391],
...,
[-0.0583, -0.2393, -0.0055, ..., -1.3750, -1.1250, 0.8711],
[-0.2773, 0.9023, -0.6250, ..., -0.5195, 0.6758, 1.0391],
[-0.7148, -0.5664, -0.7461, ..., 0.2070, -0.5469, 0.6289]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5273, -0.5195, ..., 0.6641, 0.5781, 0.5195],
[-0.5664, -0.9297, 0.0579, ..., 1.1797, -0.4727, -0.3594],
[-1.2188, -0.2461, -0.4238, ..., -0.3906, -0.2412, 1.0391],
...,
[-0.0544, -0.2383, -0.0058, ..., -1.3750, -1.1250, 0.8750],
[-0.2793, 0.9062, -0.6250, ..., -0.5195, 0.6758, 1.0391],
[-0.7109, -0.5742, -0.7500, ..., 0.2051, -0.5508, 0.6289]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5273, -0.5117, ..., 0.6641, 0.5781, 0.5195],
[-0.5664, -0.9297, 0.0579, ..., 1.1719, -0.4727, -0.3633],
[-1.2188, -0.2490, -0.4219, ..., -0.3926, -0.2471, 1.0391],
...,
[-0.0554, -0.2412, -0.0082, ..., -1.3828, -1.1328, 0.8750],
[-0.2793, 0.9023, -0.6250, ..., -0.5234, 0.6719, 1.0312],
[-0.7188, -0.5703, -0.7422, ..., 0.2041, -0.5508, 0.6250]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5312, -0.5156, ..., 0.6680, 0.5781, 0.5234],
[-0.5664, -0.9258, 0.0576, ..., 1.1719, -0.4707, -0.3613],
[-1.2188, -0.2461, -0.4238, ..., -0.3906, -0.2432, 1.0391],
...,
[-0.0562, -0.2422, -0.0058, ..., -1.3750, -1.1250, 0.8750],
[-0.2773, 0.9023, -0.6289, ..., -0.5273, 0.6719, 1.0391],
[-0.7188, -0.5664, -0.7500, ..., 0.2100, -0.5469, 0.6250]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5312, -0.5195, ..., 0.6641, 0.5859, 0.5234],
[-0.5664, -0.9258, 0.0593, ..., 1.1719, -0.4688, -0.3574],
[-1.2188, -0.2471, -0.4219, ..., -0.3867, -0.2412, 1.0469],
...,
[-0.0554, -0.2422, -0.0103, ..., -1.3750, -1.1250, 0.8750],
[-0.2793, 0.9062, -0.6289, ..., -0.5234, 0.6758, 1.0312],
[-0.7188, -0.5664, -0.7461, ..., 0.2080, -0.5547, 0.6250]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812e+00, 5.3125e-01, -5.1562e-01, ..., 6.6797e-01,
5.8203e-01, 5.2344e-01],
[-5.6641e-01, -9.2969e-01, 5.8594e-02, ..., 1.1797e+00,
-4.6680e-01, -3.5938e-01],
[-1.2188e+00, -2.4414e-01, -4.2188e-01, ..., -3.9062e-01,
-2.3926e-01, 1.0469e+00],
...,
[-5.6641e-02, -2.3730e-01, -9.3460e-04, ..., -1.3750e+00,
-1.1250e+00, 8.7500e-01],
[-2.8125e-01, 9.0625e-01, -6.2891e-01, ..., -5.1953e-01,
6.7578e-01, 1.0469e+00],
[-7.0703e-01, -5.6250e-01, -7.4609e-01, ..., 2.1094e-01,
-5.4688e-01, 6.3672e-01]]], device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7891, 0.5273, -0.5156, ..., 0.6641, 0.5820, 0.5156],
[-0.5664, -0.9297, 0.0559, ..., 1.1719, -0.4707, -0.3633],
[-1.2188, -0.2461, -0.4219, ..., -0.3926, -0.2402, 1.0391],
...,
[-0.0535, -0.2393, -0.0095, ..., -1.3828, -1.1250, 0.8750],
[-0.2773, 0.9023, -0.6250, ..., -0.5195, 0.6758, 1.0391],
[-0.7109, -0.5703, -0.7500, ..., 0.2061, -0.5508, 0.6328]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5312, -0.5156, ..., 0.6641, 0.5781, 0.5234],
[-0.5625, -0.9258, 0.0569, ..., 1.1719, -0.4727, -0.3613],
[-1.2188, -0.2500, -0.4219, ..., -0.3926, -0.2451, 1.0391],
...,
[-0.0569, -0.2412, -0.0060, ..., -1.3828, -1.1250, 0.8711],
[-0.2812, 0.8984, -0.6250, ..., -0.5234, 0.6719, 1.0312],
[-0.7188, -0.5664, -0.7500, ..., 0.1992, -0.5586, 0.6289]]],
device='cuda:3', dtype=torch.bfloat16)
tensor([[[-1.7812, 0.5312, -0.5156, ..., 0.6641, 0.5742, 0.5234],
[-0.5664, -0.9258, 0.0596, ..., 1.1797, -0.4746, -0.3633],
[-1.2188, -0.2480, -0.4238, ..., -0.3926, -0.2432, 1.0391],
...,
[-0.0576, -0.2393, -0.0052, ..., -1.3828, -1.1328, 0.8750],
[-0.2773, 0.8984, -0.6328, ..., -0.5312, 0.6680, 1.0312],
[-0.7188, -0.5703, -0.7500, ..., 0.2002, -0.5586, 0.6289]]],
device='cuda:3', dtype=torch.bfloat16)
```

### System Info

- 🤗 Diffusers version: 0.31.0
- Platform: Linux-5.4.0-144-generic-x86_64-with-glibc2.35
- Running on Google Colab?: No
- Python version: 3.10.14
- PyTorch version (GPU?): 2.0.0+cu117 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.25.2
- Transformers version: 4.45.2
- Accelerate version: 1.0.1
- PEFT version: not installed
- Bitsandbytes version: not installed
- Safetensors version: 0.4.5
- xFormers version: not installed
- Accelerator: NVIDIA A800-SXM4-80GB, 81920 MiB
NVIDIA A800-SXM4-80GB, 81920 MiB
NVIDIA A800-SXM4-80GB, 81920 MiB
NVIDIA A800-SXM4-80GB, 81920 MiB
NVIDIA A800-SXM4-80GB, 81920 MiB
NVIDIA A800-SXM4-80GB, 81920 MiB
NVIDIA A800-SXM4-80GB, 81920 MiB
NVIDIA A800-SXM4-80GB, 81920 MiB
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

@DN6 @yiyixuxu @sayakpaul