diff --git a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py index 4530a424adb4..28bf00d5cc16 100644 --- a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py +++ b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py @@ -80,6 +80,20 @@ """ +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, @@ -852,6 +866,7 @@ def __call__( callback_on_step_end: Callable[[int, int], None] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, + mu: float | None = None, ): r""" Function invoked when calling the pipeline for generation. @@ -963,6 +978,7 @@ def __call__( will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. + mu (`float`, *optional*): `mu` value used for `dynamic_shifting`. Examples: @@ -1101,18 +1117,7 @@ def __call__( else: assert False - # 4. Prepare timesteps - if XLA_AVAILABLE: - timestep_device = "cpu" - else: - timestep_device = device - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas - ) - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # 5. Prepare latent variables + # 4. Prepare latent variables num_channels_latents = self.transformer.config.in_channels latents = self.prepare_latents( batch_size * num_images_per_prompt, @@ -1125,6 +1130,34 @@ def __call__( latents, ) + # 5. Prepare timesteps + scheduler_kwargs = {} + if self.scheduler.config.get("use_dynamic_shifting", None) and mu is None: + _, _, latent_height, latent_width = latents.shape + image_seq_len = (latent_height // self.transformer.config.patch_size) * ( + latent_width // self.transformer.config.patch_size + ) + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.16), + ) + scheduler_kwargs["mu"] = mu + elif mu is not None: + scheduler_kwargs["mu"] = mu + + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas, **scheduler_kwargs + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + # 6. Create tensor stating which controlnets to keep controlnet_keep = [] for i in range(len(timesteps)): diff --git a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py index d2890d55811c..c418ade5c6e0 100644 --- a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py +++ b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py @@ -103,6 +103,20 @@ """ +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, @@ -1020,6 +1034,7 @@ def __call__( callback_on_step_end: Callable[[int, int], None] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, + mu: float | None = None, ): r""" Function invoked when calling the pipeline for generation. @@ -1135,6 +1150,7 @@ def __call__( will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. + mu (`float`, *optional*): `mu` value used for `dynamic_shifting`. Examples: @@ -1272,18 +1288,7 @@ def __call__( else: controlnet_pooled_projections = controlnet_pooled_projections or pooled_prompt_embeds - # 4. Prepare timesteps - if XLA_AVAILABLE: - timestep_device = "cpu" - else: - timestep_device = device - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas - ) - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # 5. Prepare latent variables + # 4. Prepare latent variables num_channels_latents = self.transformer.config.in_channels latents = self.prepare_latents( batch_size * num_images_per_prompt, @@ -1296,6 +1301,34 @@ def __call__( latents, ) + # 5. Prepare timesteps + scheduler_kwargs = {} + if self.scheduler.config.get("use_dynamic_shifting", None) and mu is None: + _, _, latent_height, latent_width = latents.shape + image_seq_len = (latent_height // self.transformer.config.patch_size) * ( + latent_width // self.transformer.config.patch_size + ) + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.16), + ) + scheduler_kwargs["mu"] = mu + elif mu is not None: + scheduler_kwargs["mu"] = mu + + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas, **scheduler_kwargs + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + # 6. Create tensor stating which controlnets to keep controlnet_keep = [] for i in range(len(timesteps)): diff --git a/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py b/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py index 554fbc150e09..968384c15079 100644 --- a/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py +++ b/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py @@ -195,6 +195,29 @@ def test_controlnet_inpaint_sd3(self): assert_tensors_close(image_slice.flatten().cpu(), expected_slice, atol=1e-2) + def test_dynamic_shifting_scheduler(self): + # Regression: this pipeline called `retrieve_timesteps()` without computing `mu`, so any + # scheduler with `use_dynamic_shifting=True` (the SD3.5 style configs) raised + # "`mu` must be passed when `use_dynamic_shifting` is set to be `True`" before inference. + components = self.get_dummy_components() + components["scheduler"] = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True) + pipe = self.get_pipeline(**components).to(torch_device, dtype=torch.float32) + + image = pipe(**self.get_dummy_inputs()).images + + assert image.shape == (1, *self.output_shape) + + def test_dynamic_shifting_scheduler_accepts_explicit_mu(self): + components = self.get_dummy_components() + components["scheduler"] = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True) + pipe = self.get_pipeline(**components).to(torch_device, dtype=torch.float32) + + inputs = self.get_dummy_inputs() + inputs["mu"] = 0.7 + image = pipe(**inputs).images + + assert image.shape == (1, *self.output_shape) + class TestStableDiffusion3ControlNetInpaintingPipelineMemory( StableDiffusion3ControlNetInpaintingPipelineTesterConfig, MemoryTesterMixin diff --git a/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py b/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py index 5fa6770dffdf..cdc75214e5be 100644 --- a/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py +++ b/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py @@ -212,6 +212,29 @@ def test_controlnet_sd35(self): # fmt: on self._run_and_check_slice(components, expected_slice) + def test_dynamic_shifting_scheduler(self): + # Regression: this pipeline called `retrieve_timesteps()` without computing `mu`, so any + # scheduler with `use_dynamic_shifting=True` (the SD3.5 style configs) raised + # "`mu` must be passed when `use_dynamic_shifting` is set to be `True`" before inference. + components = self.get_dummy_components() + components["scheduler"] = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True) + pipe = self.get_pipeline(**components).to(torch_device, dtype=torch.float32) + + image = pipe(**self.get_dummy_inputs()).images + + assert image.shape == (1, *self.output_shape) + + def test_dynamic_shifting_scheduler_accepts_explicit_mu(self): + components = self.get_dummy_components() + components["scheduler"] = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True) + pipe = self.get_pipeline(**components).to(torch_device, dtype=torch.float32) + + inputs = self.get_dummy_inputs() + inputs["mu"] = 0.7 + image = pipe(**inputs).images + + assert image.shape == (1, *self.output_shape) + class TestStableDiffusion3ControlNetPipelineMemory(StableDiffusion3ControlNetPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD3 ControlNet pipeline."""