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)