Support dynamic-shifting schedulers in SD3 ControlNet pipelines

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)):

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)):

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

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."""