Skip to content

Fix missing latents_bn_std dtype cast in VAE normalization #13299

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sayakpaul merged 5 commits into from
May 1, 2026
Merged

Fix missing latents_bn_std dtype cast in VAE normalization #13299

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
14b1403
Corrected casting of latents_bn_std
adi776borate Mar 21, 2026
9e7b26e
Merge branch 'main' into fix/flux2-vae-device-casting
yiyixuxu Apr 13, 2026
8af5ee2
Merge branch 'main' into fix/flux2-vae-device-casting
adi776borate Apr 18, 2026
e0d2e1b
Propagated the fix to the klein inpaint pipeline
adi776borate Apr 18, 2026
f0b0b13
Merge branch 'main' into fix/flux2-vae-device-casting
DN6 May 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion src/diffusers/pipelines/flux2/pipeline_flux2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -611,7 +611,9 @@ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
image_latents = self._patchify_latents(image_latents)

latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype)
latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps)
latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
image_latents.device, image_latents.dtype
)
image_latents = (image_latents - latents_bn_mean) / latents_bn_std

return image_latents
Expand Down
4 changes: 3 additions & 1 deletion src/diffusers/pipelines/flux2/pipeline_flux2_klein.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -467,7 +467,9 @@ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
image_latents = self._patchify_latents(image_latents)

latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype)
latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps)
latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
image_latents.device, image_latents.dtype
)
image_latents = (image_latents - latents_bn_mean) / latents_bn_std

return image_latents
Expand Down
4 changes: 3 additions & 1 deletion src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -477,7 +477,9 @@ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
image_latents = self._patchify_latents(image_latents)

latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype)
latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps)
latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
image_latents.device, image_latents.dtype
)
image_latents = (image_latents - latents_bn_mean) / latents_bn_std

return image_latents
Expand Down
Loading