From ec2a161c09540a374792278eca5e028b2c954d03 Mon Sep 17 00:00:00 2001 From: Devam0311 Date: Tue, 1 Sep 2026 17:03:29 +0530 Subject: [PATCH] Run cleanup and honor return_dict for SDXL latent output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For `output_type="latent"` the SDXL inpaint and instruct-pix2pix pipelines returned `StableDiffusionXLPipelineOutput(images=latents)` immediately. That skipped `maybe_free_model_hooks()`, so model offload cleanup never ran, and bypassed the `return_dict` handling, so `return_dict=False` returned the wrong type. Assign `image = latents` instead and guard watermarking, postprocessing and the padding-mask overlay behind `output_type != "latent"`, so the shared cleanup and return handling at the end runs for every output type — matching the text-to-image pipeline. Applied to the three ControlNet and PAG inpaint pipelines carrying the same early return. Add a regression test asserting cleanup runs and `return_dict=False` returns a tuple of latents. Ref #13610 (Issue 3) Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XBYeq5vB4DNZDEaqUsroKR --- .../pipeline_controlnet_inpaint_sd_xl.py | 17 ++++++++++------- ...pipeline_controlnet_union_inpaint_sd_xl.py | 17 ++++++++++------- .../pag/pipeline_pag_sd_xl_inpaint.py | 17 ++++++++++------- .../pipeline_stable_diffusion_xl_inpaint.py | 17 ++++++++++------- ...ne_stable_diffusion_xl_instruct_pix2pix.py | 11 ++++++----- .../test_stable_diffusion_xl_inpaint.py | 19 +++++++++++++++++++ 6 files changed, 65 insertions(+), 33 deletions(-) diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py index f27fcd8aa26f..935842dcd955 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py @@ -1866,16 +1866,19 @@ def denoising_value_valid(dnv): if not output_type == "latent": image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] else: - return StableDiffusionXLPipelineOutput(images=latents) + 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) - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] + if padding_mask_crop is not None: + image = [ + self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image + ] # Offload all models self.maybe_free_model_hooks() diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py index 6d7e1f62beb3..4e5443cf5d41 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py @@ -1878,16 +1878,19 @@ def denoising_value_valid(dnv): if not output_type == "latent": image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] else: - return StableDiffusionXLPipelineOutput(images=latents) + 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) - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] + if padding_mask_crop is not None: + image = [ + self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image + ] # Offload all models self.maybe_free_model_hooks() diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py index 9e70a7779f1e..d6c22c8f8a41 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py @@ -1752,16 +1752,19 @@ def denoising_value_valid(dnv): if needs_upcasting: self.vae.to(dtype=torch.float16) else: - return StableDiffusionXLPipelineOutput(images=latents) + 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) - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] + if padding_mask_crop is not None: + image = [ + self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image + ] # Offload all models self.maybe_free_model_hooks() diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py index 3f18cbe21d0f..9eeb09fd4587 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py @@ -1714,16 +1714,19 @@ def denoising_value_valid(dnv): if needs_upcasting: self.vae.to(dtype=torch.float16) else: - return StableDiffusionXLPipelineOutput(images=latents) + 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) - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] + if padding_mask_crop is not None: + image = [ + self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image + ] # Offload all models self.maybe_free_model_hooks() diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py index bcd337414bac..eb0dfab2c8a4 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py @@ -960,13 +960,14 @@ def __call__( if needs_upcasting: self.vae.to(dtype=torch.float16) else: - return StableDiffusionXLPipelineOutput(images=latents) + 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_inpaint.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_inpaint.py index f973d5afed8f..1052f8a81154 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_inpaint.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_inpaint.py @@ -692,6 +692,25 @@ def test_stable_diffusion_xl_inpaint_2_images(self): image_slice2 = images[1, -1, -3:, -3:] assert (image_slice1 - image_slice2).abs().max() > 1e-2 + def test_latent_output_runs_cleanup_and_honors_return_dict(self): + # Regression: `output_type="latent"` returned `StableDiffusionXLPipelineOutput(images=latents)` + # immediately, which skipped `maybe_free_model_hooks()` and ignored `return_dict=False`. + sd_pipe = self.get_pipeline() + + called = {"cleanup": False} + sd_pipe.maybe_free_model_hooks = lambda: called.__setitem__("cleanup", True) + + inputs = self.get_dummy_inputs() + inputs["output_type"] = "latent" + inputs["return_dict"] = False + output = sd_pipe(**inputs) + + assert called["cleanup"], "`maybe_free_model_hooks()` should still run for latent output." + assert isinstance(output, tuple), "`return_dict=False` should return a tuple." + assert torch.is_tensor(output[0]) + # Raw latents keep the VAE latent channel count; postprocess would have produced 3 channels. + assert output[0].shape[1] == sd_pipe.unet.config.in_channels + def test_pipeline_interrupt(self): sd_pipe = self.get_pipeline().to(torch_device)