diff --git a/src/diffusers/callbacks.py b/src/diffusers/callbacks.py index 087a6b7fee56..cc43a1832be3 100644 --- a/src/diffusers/callbacks.py +++ b/src/diffusers/callbacks.py @@ -86,12 +86,14 @@ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[s ) if step_index == cutoff_step: - prompt_embeds = callback_kwargs[self.tensor_inputs[0]] - prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens. + if pipeline.do_classifier_free_guidance: + # the CFG batches are `[negative, conditional]`; keep the conditional half of each tensor + batch_size = callback_kwargs[self.tensor_inputs[0]].shape[0] // 2 + for name in self.tensor_inputs: + callback_kwargs[name] = callback_kwargs[name][-batch_size:] pipeline._guidance_scale = 0.0 - callback_kwargs[self.tensor_inputs[0]] = prompt_embeds return callback_kwargs @@ -119,21 +121,14 @@ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[s ) if step_index == cutoff_step: - prompt_embeds = callback_kwargs[self.tensor_inputs[0]] - prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens. - - add_text_embeds = callback_kwargs[self.tensor_inputs[1]] - add_text_embeds = add_text_embeds[-1:] # "-1" denotes the embeddings for conditional pooled text tokens - - add_time_ids = callback_kwargs[self.tensor_inputs[2]] - add_time_ids = add_time_ids[-1:] # "-1" denotes the embeddings for conditional added time vector + if pipeline.do_classifier_free_guidance: + # the CFG batches are `[negative, conditional]`; keep the conditional half of each tensor + batch_size = callback_kwargs[self.tensor_inputs[0]].shape[0] // 2 + for name in self.tensor_inputs: + callback_kwargs[name] = callback_kwargs[name][-batch_size:] pipeline._guidance_scale = 0.0 - callback_kwargs[self.tensor_inputs[0]] = prompt_embeds - callback_kwargs[self.tensor_inputs[1]] = add_text_embeds - callback_kwargs[self.tensor_inputs[2]] = add_time_ids - return callback_kwargs @@ -162,26 +157,15 @@ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[s ) if step_index == cutoff_step: - prompt_embeds = callback_kwargs[self.tensor_inputs[0]] - prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens. - - add_text_embeds = callback_kwargs[self.tensor_inputs[1]] - add_text_embeds = add_text_embeds[-1:] # "-1" denotes the embeddings for conditional pooled text tokens - - add_time_ids = callback_kwargs[self.tensor_inputs[2]] - add_time_ids = add_time_ids[-1:] # "-1" denotes the embeddings for conditional added time vector - - # For Controlnet - image = callback_kwargs[self.tensor_inputs[3]] - image = image[-1:] + if pipeline.do_classifier_free_guidance: + # the CFG batches are `[negative, conditional]`; keep the conditional half of each tensor. In guess + # mode the controlnet `image` is not duplicated, so the slice keeps it whole. + batch_size = callback_kwargs[self.tensor_inputs[0]].shape[0] // 2 + for name in self.tensor_inputs: + callback_kwargs[name] = callback_kwargs[name][-batch_size:] pipeline._guidance_scale = 0.0 - callback_kwargs[self.tensor_inputs[0]] = prompt_embeds - callback_kwargs[self.tensor_inputs[1]] = add_text_embeds - callback_kwargs[self.tensor_inputs[2]] = add_time_ids - callback_kwargs[self.tensor_inputs[3]] = image - return callback_kwargs @@ -229,16 +213,12 @@ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[s ) if step_index == cutoff_step: - prompt_embeds = callback_kwargs[self.tensor_inputs[0]] - prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens. - - pooled_prompt_embeds = callback_kwargs[self.tensor_inputs[1]] - pooled_prompt_embeds = pooled_prompt_embeds[ - -1: - ] # "-1" denotes the embeddings for conditional pooled text tokens. + if pipeline.do_classifier_free_guidance: + # the CFG batches are `[negative, conditional]`; keep the conditional half of each tensor + batch_size = callback_kwargs[self.tensor_inputs[0]].shape[0] // 2 + for name in self.tensor_inputs: + callback_kwargs[name] = callback_kwargs[name][-batch_size:] pipeline._guidance_scale = 0.0 - callback_kwargs[self.tensor_inputs[0]] = prompt_embeds - callback_kwargs[self.tensor_inputs[1]] = pooled_prompt_embeds return callback_kwargs diff --git a/tests/pipelines/controlnet/test_controlnet_sdxl.py b/tests/pipelines/controlnet/test_controlnet_sdxl.py index d7e442df612a..c6ea1b7f6c36 100644 --- a/tests/pipelines/controlnet/test_controlnet_sdxl.py +++ b/tests/pipelines/controlnet/test_controlnet_sdxl.py @@ -31,6 +31,7 @@ StableDiffusionXLImg2ImgPipeline, UNet2DConditionModel, ) +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback, SDXLControlnetCFGCutoffCallback from diffusers.models.unets.unet_2d_blocks import UNetMidBlock2D from diffusers.pipelines.controlnet.pipeline_controlnet import MultiControlNetModel from diffusers.utils.torch_utils import randn_tensor @@ -252,6 +253,42 @@ def test_controlnet_sdxl_lcm(self): image_slice = image[0, -1, -3:, -3:] assert_tensors_close(image_slice.flatten().cpu(), self.expected_lcm_slice, atol=1e-2) + def test_cfg_cutoff_callback(self): + # after the cutoff the callback must keep one conditional embedding per sample, not only the last row + cutoff_callback = SDXLControlnetCFGCutoffCallback(cutoff_step_ratio=None, cutoff_step_index=1) + + class CheckBatchCallback(PipelineCallback): + tensor_inputs = ["latents", "prompt_embeds"] + + def callback_fn(self, pipeline, step_index, timestep, callback_kwargs): + if step_index >= 1: + assert callback_kwargs["prompt_embeds"].shape[0] == callback_kwargs["latents"].shape[0] + return callback_kwargs + + # Run on CPU: guess mode calls `torch.logspace`, which has no MPS kernel. + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs["prompt"] = [inputs["prompt"], "a different prompt"] + inputs["num_inference_steps"] = 3 + inputs["callback_on_step_end"] = MultiPipelineCallbacks([cutoff_callback, CheckBatchCallback()]) + image = pipe(**inputs).images + + assert image.shape == (2, *self.output_shape) + assert pipe.guidance_scale == 0.0 + + # without CFG there is no negative batch to drop + inputs["guidance_scale"] = 1.0 + image = pipe(**inputs).images + + assert image.shape == (2, *self.output_shape) + + # in guess mode the conditioning image is not duplicated for CFG and has to be kept whole + inputs["guidance_scale"] = 6.0 + inputs["guess_mode"] = True + image = pipe(**inputs).images + + assert image.shape == (2, *self.output_shape) + class TestStableDiffusionXLControlNetPipeline( StableDiffusionXLControlNetPipelineTesterConfig, StableDiffusionXLControlNetPipelineTests, PipelineTesterMixin diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_stable_diffusion.py index f647f009ec58..a0dab9d23a9e 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion.py @@ -45,6 +45,7 @@ UNet2DConditionModel, logging, ) +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback, SDCFGCutoffCallback from diffusers.utils.import_utils import is_accelerate_available from ...models.testing_utils.lora import check_if_lora_correctly_set @@ -603,6 +604,34 @@ def callback_on_step_end(pipe, i, t, callback_kwargs): # they should be the same assert torch.allclose(intermediate_latent, output_interrupted, atol=1e-4) + def test_cfg_cutoff_callback(self): + # after the cutoff the callback must keep one conditional embedding per sample, not only the last row + cutoff_callback = SDCFGCutoffCallback(cutoff_step_ratio=None, cutoff_step_index=1) + + class CheckBatchCallback(PipelineCallback): + tensor_inputs = ["latents", "prompt_embeds"] + + def callback_fn(self, pipeline, step_index, timestep, callback_kwargs): + if step_index >= 1: + assert callback_kwargs["prompt_embeds"].shape[0] == callback_kwargs["latents"].shape[0] + return callback_kwargs + + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["prompt"] = [inputs["prompt"], "a different prompt"] + inputs["num_inference_steps"] = 3 + inputs["callback_on_step_end"] = MultiPipelineCallbacks([cutoff_callback, CheckBatchCallback()]) + image = pipe(**inputs).images + + assert image.shape == (2, *self.output_shape) + assert pipe.guidance_scale == 0.0 + + # without CFG there is no negative batch to drop + inputs["guidance_scale"] = 1.0 + image = pipe(**inputs).images + + assert image.shape == (2, *self.output_shape) + def test_pipeline_accept_tuple_type_unet_sample_size(self): # the purpose of this test is to see whether the pipeline would accept a unet with the tuple-typed sample size sd_repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py index 4b8952188efa..0469fa03d2e8 100644 --- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py +++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py @@ -13,6 +13,7 @@ ) from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel, StableDiffusion3Pipeline +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback, SD3CFGCutoffCallback from ...testing_utils import ( assert_tensors_close, @@ -205,6 +206,34 @@ def test_skip_guidance_layers(self): assert not torch.allclose(output_full, output_skip, atol=1e-5), "Outputs should differ when layers are skipped" assert output_full.shape == output_skip.shape, "Outputs should have the same shape" + def test_cfg_cutoff_callback(self): + # after the cutoff the callback must keep one conditional embedding per sample, not only the last row + cutoff_callback = SD3CFGCutoffCallback(cutoff_step_ratio=None, cutoff_step_index=1) + + class CheckBatchCallback(PipelineCallback): + tensor_inputs = ["latents", "prompt_embeds"] + + def callback_fn(self, pipeline, step_index, timestep, callback_kwargs): + if step_index >= 1: + assert callback_kwargs["prompt_embeds"].shape[0] == callback_kwargs["latents"].shape[0] + return callback_kwargs + + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["prompt"] = [inputs["prompt"], "a different prompt"] + inputs["num_inference_steps"] = 3 + inputs["callback_on_step_end"] = MultiPipelineCallbacks([cutoff_callback, CheckBatchCallback()]) + image = pipe(**inputs).images + + assert image.shape == (2, *self.output_shape) + assert pipe.guidance_scale == 0.0 + + # without CFG there is no negative batch to drop + inputs["guidance_scale"] = 1.0 + image = pipe(**inputs).images + + assert image.shape == (2, *self.output_shape) + class TestStableDiffusion3PipelineMemory(StableDiffusion3PipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD3 pipeline.""" diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py index b4dd632c0128..b288a500d99f 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py @@ -40,6 +40,7 @@ UNet2DConditionModel, UniPCMultistepScheduler, ) +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback, SDXLCFGCutoffCallback from diffusers.utils import logging from ...models.testing_utils.lora import check_if_lora_correctly_set @@ -993,6 +994,34 @@ def callback_on_step_end(pipe, i, t, callback_kwargs): # they should be the same assert_tensors_close(intermediate_latent, output_interrupted, atol=1e-4) + def test_cfg_cutoff_callback(self): + # after the cutoff the callback must keep one conditional embedding per sample, not only the last row + cutoff_callback = SDXLCFGCutoffCallback(cutoff_step_ratio=None, cutoff_step_index=1) + + class CheckBatchCallback(PipelineCallback): + tensor_inputs = ["latents", "prompt_embeds"] + + def callback_fn(self, pipeline, step_index, timestep, callback_kwargs): + if step_index >= 1: + assert callback_kwargs["prompt_embeds"].shape[0] == callback_kwargs["latents"].shape[0] + return callback_kwargs + + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() + inputs["prompt"] = [inputs["prompt"], "a different prompt"] + inputs["num_inference_steps"] = 3 + inputs["callback_on_step_end"] = MultiPipelineCallbacks([cutoff_callback, CheckBatchCallback()]) + image = pipe(**inputs).images + + assert image.shape == (2, *self.output_shape) + assert pipe.guidance_scale == 0.0 + + # without CFG there is no negative batch to drop + inputs["guidance_scale"] = 1.0 + image = pipe(**inputs).images + + assert image.shape == (2, *self.output_shape) + class TestStableDiffusionXLPipelineMemory(StableDiffusionXLPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SDXL pipeline."""