From c5fd848e6d5300d2ec63d9279186fb53e6a7ad54 Mon Sep 17 00:00:00 2001 From: Devam0311 Date: Tue, 1 Sep 2026 15:35:15 +0530 Subject: [PATCH] Apply the VAE shift factor when decoding in SD3 inpaint `StableDiffusion3InpaintPipeline` encodes VAE latents with `(latents - shift_factor) * scaling_factor` but decoded with only `latents / scaling_factor`, dropping `+ shift_factor`. Encode and decode therefore disagreed within the same pipeline, and the output differed from `StableDiffusion3Pipeline`, `StableDiffusion3Img2ImgPipeline` and both SD3 ControlNet pipelines, which all apply the shift. SD3 VAEs have a nonzero `shift_factor`, so inpaint output was decoded from the wrong latent distribution. Add a regression test asserting the decoded image matches a manual decode with the shift applied, and update the `test_inference` expected slice, which pinned the previous incorrect output. Ref #13611 (Issue 1) Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XBYeq5vB4DNZDEaqUsroKR --- .../pipeline_stable_diffusion_3_inpaint.py | 5 ++-- ...est_pipeline_stable_diffusion_3_inpaint.py | 30 ++++++++++++++++++- 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py index 321e9f8dd80e..5d137a5e47ce 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py @@ -1365,9 +1365,8 @@ def __call__( xm.mark_step() if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 0 - ] + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + image = self.vae.decode(latents, return_dict=False, generator=generator)[0] else: image = latents diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py index 862f08e3c720..75121995bb3d 100644 --- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py +++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py @@ -146,11 +146,39 @@ def test_inference(self): generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) # fmt: off - expected_slice = torch.tensor([0.5076, 0.4238, 0.7243, 0.4664, 0.3933, 0.5421, 0.4952, 0.5001, 0.5716, 0.5092, 0.5091, 0.7205, 0.5442, 0.7069, 0.6149, 0.6009]) + expected_slice = torch.tensor([0.5066, 0.4189, 0.7212, 0.4641, 0.3917, 0.5386, 0.4909, 0.5002, 0.5655, 0.4970, 0.4891, 0.7210, 0.5366, 0.6935, 0.6176, 0.5953]) # fmt: on assert_tensors_close(generated_slice, expected_slice, atol=1e-3, msg="Output does not match expected slice.") + def test_vae_shift_factor_applied_on_decode(self): + # Regression: this pipeline decoded with `latents / scaling_factor` only, dropping + # `+ shift_factor`, so it decoded from a different latent distribution than the one its + # own `_encode_vae_image` encodes into (and than the other SD3 pipelines decode from). + pipe = self.get_pipeline() + + inputs = self.get_dummy_inputs() + inputs["output_type"] = "latent" + latents = pipe(**inputs).images + + image = pipe(**self.get_dummy_inputs()).images + + def decode(scaled_latents): + decoded = pipe.vae.decode(scaled_latents, return_dict=False)[0] + return pipe.image_processor.postprocess( + decoded, output_type="pt", do_denormalize=[True] * decoded.shape[0] + ) + + with_shift = decode((latents / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor) + without_shift = decode(latents / pipe.vae.config.scaling_factor) + + assert not torch.allclose(with_shift, without_shift), ( + "The dummy VAE needs a nonzero `shift_factor` for this regression test to be meaningful." + ) + assert_tensors_close( + image, with_shift, atol=1e-4, msg="Decoded output should have the VAE shift factor applied." + ) + class TestStableDiffusion3InpaintPipelineMemory(StableDiffusion3InpaintPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD3 inpaint pipeline."""