From 9e689bdc51e39187b25a82f566b52b3a3a39e84c Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Fri, 1 May 2026 06:14:24 +0200 Subject: [PATCH] update --- src/diffusers/pipelines/flux2/pipeline_flux2.py | 4 +++- src/diffusers/pipelines/flux2/pipeline_flux2_klein.py | 4 +++- src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py | 4 +++- src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py | 4 +++- 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2.py b/src/diffusers/pipelines/flux2/pipeline_flux2.py index 4b60c6042d..b1645b4ae2 100644 --- a/src/diffusers/pipelines/flux2/pipeline_flux2.py +++ b/src/diffusers/pipelines/flux2/pipeline_flux2.py @@ -611,7 +611,9 @@ class Flux2Pipeline(DiffusionPipeline, Flux2LoraLoaderMixin): 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 diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py b/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py index 1f3b5c3c4f..9a3468525c 100644 --- a/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py +++ b/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py @@ -467,7 +467,9 @@ class Flux2KleinPipeline(DiffusionPipeline, Flux2LoraLoaderMixin): 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 diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py b/src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py index f4aecc1876..fd9467003a 100644 --- a/src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py +++ b/src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py @@ -547,7 +547,9 @@ class Flux2KleinInpaintPipeline(DiffusionPipeline, Flux2LoraLoaderMixin): 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 diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py b/src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py index 671953be63..78ed42f20a 100644 --- a/src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py +++ b/src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py @@ -477,7 +477,9 @@ class Flux2KleinKVPipeline(DiffusionPipeline, Flux2LoraLoaderMixin): 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