Division by zero in rescale_noise_cfg can produce NaNs during inference

[Akash504-ai]

Describe the bug

Bug: Division by zero in rescale_noise_cfg can produce NaNs

Description

The function rescale_noise_cfg performs a division by std_cfg without any numerical stability guard:

noise_pred_rescaled = noise_cfg * (std_text / std_cfg)

If std_cfg becomes zero, this leads to NaN or inf values, which can silently corrupt the diffusion process and produce invalid outputs.

This can happen in edge cases where noise_cfg has zero variance.

Reproduction

import torch

noise_cfg = torch.zeros(1, 4, 64, 64)  # std = 0
noise_pred_text = torch.randn_like(noise_cfg)

std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)

result = noise_cfg * (std_text / std_cfg)
print(result)

Logs

System Info

• Diffusers: main
• PyTorch: any
• OS: any

Who can help?

No response