From 5406aed0e70647b8754383a2d55f90cc8b0805d3 Mon Sep 17 00:00:00 2001 From: Devam0311 Date: Tue, 1 Sep 2026 16:42:28 +0530 Subject: [PATCH] Do not watermark or postprocess SDXL latent output With `output_type="latent"` the SDXL img2img pipeline sets `image = latents` but still called `watermark.apply_watermark()` and `image_processor.postprocess()` on it, so raw latents were treated as decoded RGB. The text-to-image pipeline already guards both behind `output_type != "latent"`. Apply the same guard in the img2img pipeline, the PAG img2img variant that shares the pattern, and the modular SDXL decoder step. Add a regression test that installs a watermarker which fails if called and asserts the latent output keeps its latent channel count. Ref #13610 (Issue 4) Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XBYeq5vB4DNZDEaqUsroKR --- .../stable_diffusion_xl/decoders.py | 13 +++++++------ .../pag/pipeline_pag_sd_xl_img2img.py | 9 +++++---- .../pipeline_stable_diffusion_xl_img2img.py | 9 +++++---- .../test_stable_diffusion_xl_img2img.py | 19 +++++++++++++++++++ 4 files changed, 36 insertions(+), 14 deletions(-) diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py index b4f15df8b411..86be4ed37931 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py @@ -129,13 +129,14 @@ def __call__(self, components, state: PipelineState) -> PipelineState: else: block_state.images = block_state.latents - # apply watermark if available - if hasattr(components, "watermark") and components.watermark is not None: - block_state.images = components.watermark.apply_watermark(block_state.images) + if not block_state.output_type == "latent": + # apply watermark if available + if hasattr(components, "watermark") and components.watermark is not None: + block_state.images = components.watermark.apply_watermark(block_state.images) - block_state.images = components.image_processor.postprocess( - block_state.images, output_type=block_state.output_type - ) + block_state.images = components.image_processor.postprocess( + block_state.images, output_type=block_state.output_type + ) self.set_block_state(state, block_state) diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py index 201d16a86f8a..0be92a54aeab 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py @@ -1521,11 +1521,12 @@ def denoising_value_valid(dnv): else: image = latents - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) + if not output_type == "latent": + # apply watermark if available + if self.watermark is not None: + image = self.watermark.apply_watermark(image) - image = self.image_processor.postprocess(image, output_type=output_type) + image = self.image_processor.postprocess(image, output_type=output_type) # Offload all models self.maybe_free_model_hooks() diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py index c7a13ca02524..7899025d623a 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py @@ -1477,11 +1477,12 @@ def denoising_value_valid(dnv): else: image = latents - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) + if not output_type == "latent": + # apply watermark if available + if self.watermark is not None: + image = self.watermark.apply_watermark(image) - image = self.image_processor.postprocess(image, output_type=output_type) + image = self.image_processor.postprocess(image, output_type=output_type) # Offload all models self.maybe_free_model_hooks() diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_img2img.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_img2img.py index 15c271950764..06b8a323f35e 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_img2img.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_img2img.py @@ -357,6 +357,25 @@ def test_stable_diffusion_xl_img2img_negative_conditions(self): assert (image_slice_with_no_neg_conditions - image_slice_with_neg_conditions).abs().max() > 1e-4 + def test_latent_output_skips_watermark_and_postprocess(self): + # Regression: `output_type="latent"` sets `image = latents`, but this pipeline still ran the + # watermarker and `image_processor.postprocess` over it, treating raw latents as decoded RGB. + # The text-to-image pipeline already guards both behind `output_type != "latent"`. + class SentinelWatermark: + def apply_watermark(self, images): + raise AssertionError("watermark must not be applied to latent output") + + sd_pipe = self.get_pipeline() + sd_pipe.watermark = SentinelWatermark() + + inputs = self.get_dummy_inputs() + inputs["output_type"] = "latent" + output = sd_pipe(**inputs).images + + assert torch.is_tensor(output) + # Raw latents keep the VAE latent channel count; postprocess would have produced 3 channels. + assert output.shape[1] == sd_pipe.unet.config.in_channels + def test_pipeline_interrupt(self): sd_pipe = self.get_pipeline().to(torch_device)