From 7eacfb3ac4e1929ce086829a136b1a24123f81f2 Mon Sep 17 00:00:00 2001 From: Devam0311 Date: Tue, 1 Sep 2026 16:13:39 +0530 Subject: [PATCH] Support dynamic-shifting schedulers in SD3 ControlNet pipelines MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both SD3 ControlNet pipelines called `retrieve_timesteps()` without any `mu` handling and exposed no `mu` argument, so any scheduler configured 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. This made ControlNet inconsistent with the rest of the SD3 family. Port the `calculate_shift()` helper, the `mu` argument and the `scheduler_kwargs["mu"]` handling from the base SD3 pipelines. `mu` is derived from the latents, so latent preparation now runs before timestep preparation, matching the base SD3 ordering. Neither step depends on the other, and the existing expected slices are unchanged, confirming the reorder preserves behaviour. Add regression tests for the derived and explicitly passed `mu` paths in both pipelines. Ref #13611 (Issue 2) Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XBYeq5vB4DNZDEaqUsroKR --- .../pipeline_stable_diffusion_3_controlnet.py | 57 +++++++++++++++---- ...table_diffusion_3_controlnet_inpainting.py | 57 +++++++++++++++---- .../test_controlnet_inpaint_sd3.py | 23 ++++++++ .../controlnet_sd3/test_controlnet_sd3.py | 23 ++++++++ 4 files changed, 136 insertions(+), 24 deletions(-) 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."""