Merge remote-tracking branch 'upstream/main' into sefi-image-diffusers

Author: jmliu206

src/diffusers/modular_pipelines/components_manager.py

@@ -23,6 +23,7 @@
 import torch
 
 from ..hooks import ModelHook
+from ..hooks.group_offloading import _is_group_offload_enabled
 from ..utils import (
     is_accelerate_available,
     logging,
@@ -76,12 +77,20 @@ def add_other_hook(self, hook: "UserCustomOffloadHook"):
         self.other_hooks.append(hook)
 
     def init_hook(self, module):
+        # A group offloaded module holds one group at a time and refuses `.to()`. Moving it here would be a
+        # silent no-op that leaves this hook recording an offload that never happened.
+        if _is_group_offload_enabled(module):
+            return module
         return module.to("cpu")
 
     def pre_forward(self, module, *args, **kwargs):
         if module.device != self.execution_device:
             if self.other_hooks is not None:
-                hooks_to_offload = [hook for hook in self.other_hooks if hook.model.device == self.execution_device]
+                hooks_to_offload = [
+                    hook
+                    for hook in self.other_hooks
+                    if hook.model.device == self.execution_device and not _is_group_offload_enabled(hook.model)
+                ]
                 # offload all other hooks
                 start_time = time.perf_counter()
                 if self.offload_strategy is not None:
@@ -104,7 +113,10 @@ def pre_forward(self, module, *args, **kwargs):
 
                 if hooks_to_offload:
                     clear_device_cache()
-            module.to(self.execution_device)
+            # The strategy still runs above, so a group offloaded model can make room for itself by moving other
+            # models — it just places itself.
+            if not _is_group_offload_enabled(module):
+                module.to(self.execution_device)
         return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device)
 
 
@@ -336,6 +348,7 @@ def __init__(self):
         self.collections = OrderedDict()  # collection_name -> set of component_names
         self.model_hooks = None
         self._auto_offload_enabled = False
+        self._offload_strategy = None
 
     def _lookup_ids(
         self,
@@ -692,7 +705,12 @@ def matches_pattern(component_id, pattern, exact_match=False):
 
         return get_return_dict(matches, return_dict_with_names)
 
-    def enable_auto_cpu_offload(self, device: str | int | torch.device = None, memory_reserve_margin="3GB"):
+    def enable_auto_cpu_offload(
+        self,
+        device: str | int | torch.device = None,
+        memory_reserve_margin="3GB",
+        offload_strategy=None,
+    ):
         """
         Enable automatic CPU offloading for all components.
 
@@ -703,11 +721,19 @@ def enable_auto_cpu_offload(self, device: str | int | torch.device = None, memor
         4. The system tries to offload the smallest combination of models that frees enough memory
         5. Models stay on the execution device until another model needs memory and forces them off
 
+        A group offloaded model takes part in this but places itself: it can still make room by moving other models
+        aside, and is never moved to make room for them. Either order works — group offload before or after enabling
+        this. `AutoOffloadStrategy` sizes its decisions from model memory footprints, which do not describe a model
+        holding one group at a time, so pass an `offload_strategy` that decides from the workflow instead.
+
         Args:
             device (str | int | torch.device): The execution device where models are moved for forward passes
             memory_reserve_margin (str): The memory reserve margin to use, default is 3GB. This is the amount of
                                         memory to keep free on the device to avoid running out of memory during model
                                         execution (e.g., for intermediate activations, gradients, etc.)
+            offload_strategy: Any callable with the signature `(hooks, model_id, model, execution_device) -> hooks`,
+                              returning which resident models to offload before the incoming one loads. Defaults to
+                              `AutoOffloadStrategy`, which frees the smallest sufficient combination.
         """
         if not is_accelerate_available():
             raise ImportError("Make sure to install accelerate to use auto_cpu_offload")
@@ -732,7 +758,17 @@ def enable_auto_cpu_offload(self, device: str | int | torch.device = None, memor
                 remove_hook_from_module(component, recurse=True)
 
         self.disable_auto_cpu_offload()
-        offload_strategy = AutoOffloadStrategy(memory_reserve_margin=memory_reserve_margin)
+        if offload_strategy is None:
+            offload_strategy = AutoOffloadStrategy(memory_reserve_margin=memory_reserve_margin)
+            if any(
+                isinstance(component, torch.nn.Module) and _is_group_offload_enabled(component)
+                for component in self.components.values()
+            ):
+                logger.warning(
+                    "`AutoOffloadStrategy` decides what to move from model memory footprints, which do not "
+                    "describe a group offloaded model: it holds one group at a time, not its whole weight. Pass "
+                    "an `offload_strategy` that decides from the workflow instead."
+                )
 
         all_hooks = []
         for name, component in self.components.items():
@@ -749,6 +785,23 @@ def enable_auto_cpu_offload(self, device: str | int | torch.device = None, memor
         self.model_hooks = all_hooks
         self._auto_offload_enabled = True
         self._auto_offload_device = device
+        self._offload_strategy = offload_strategy
+
+    def set_offload_strategy(self, offload_strategy):
+        """
+        Replace the offload strategy on all managed models. Only valid while auto CPU offloading is enabled.
+
+        Args:
+            offload_strategy:
+                Any callable with the signature `(hooks, model_id, model, execution_device) -> hooks`: it receives the
+                hooks of the models currently on the device and returns the ones to offload before the incoming model
+                loads. The default is `AutoOffloadStrategy`, which frees the smallest sufficient combination.
+        """
+        if not self._auto_offload_enabled:
+            raise ValueError("Auto CPU offloading is not enabled. Call `enable_auto_cpu_offload` first.")
+        for user_hook in self.model_hooks:
+            user_hook.hook.offload_strategy = offload_strategy
+        self._offload_strategy = offload_strategy
 
     def disable_auto_cpu_offload(self):
         """
@@ -765,6 +818,7 @@ def disable_auto_cpu_offload(self):
             clear_device_cache()
         self.model_hooks = None
         self._auto_offload_enabled = False
+        self._offload_strategy = None
 
     def get_model_info(
         self,

src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py

@@ -459,6 +459,9 @@ def dpm_solver_first_order_update(
                 The direct output from the learned diffusion model.
             sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:
@@ -497,6 +500,9 @@ def multistep_dpm_solver_second_order_update(
                 The direct outputs from learned diffusion model at current and latter timesteps.
             sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:

src/diffusers/schedulers/scheduling_ddim_cogvideox.py

@@ -285,6 +285,8 @@ def set_timesteps(
         Args:
             num_inference_steps (`int`):
                 The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
         """
 
         if num_inference_steps > self.config.num_train_timesteps:

src/diffusers/schedulers/scheduling_ddim_inverse.py

@@ -275,6 +275,8 @@ def set_timesteps(
         Args:
             num_inference_steps (`int`):
                 The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
         """
 
         if num_inference_steps > self.config.num_train_timesteps:

src/diffusers/schedulers/scheduling_dpm_cogvideox.py

@@ -106,8 +106,8 @@ def rescale_zero_terminal_snr(alphas_cumprod):
     Rescales betas to have zero terminal SNR Based on https://huggingface.co/papers/2305.08891 (Algorithm 1)
 
     Args:
-        betas (`torch.Tensor`):
-            the betas that the scheduler is being initialized with.
+        alphas_cumprod (`torch.Tensor`):
+            the cumulative product of alphas that the scheduler is being initialized with.
 
     Returns:
         `torch.Tensor`: rescaled betas with zero terminal SNR

src/diffusers/schedulers/scheduling_dpmsolver_multistep.py

@@ -379,6 +379,9 @@ def set_timesteps(
                 The number of diffusion steps used when generating samples with a pre-trained model.
             device (`str` or `torch.device`, *optional*):
                 The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            mu (`float`, *optional*):
+                Exponent for the dynamic time shift. Requires `use_dynamic_shifting` and a `time_shift_type` of
+                `"exponential"`; when passed, `flow_shift` is set to `exp(mu)`.
             timesteps (`list[int]`, *optional*):
                 Custom timesteps used to support arbitrary timesteps schedule. If `None`, timesteps will be generated
                 based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps` and `sigmas`
@@ -931,6 +934,9 @@ def multistep_dpm_solver_second_order_update(
                 The direct outputs from learned diffusion model at current and latter timesteps.
             sample (`torch.Tensor`, *optional*):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:

src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py

@@ -787,6 +787,9 @@ def multistep_dpm_solver_second_order_update(
                 The direct outputs from learned diffusion model at current and latter timesteps.
             sample (`torch.Tensor`, *optional*):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:

src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py

@@ -342,6 +342,9 @@ def set_timesteps(
                 The number of diffusion steps used when generating samples with a pre-trained model.
             device (`str` or `torch.device`, *optional*):
                 The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            mu (`float`, *optional*):
+                Exponent for the dynamic time shift. Requires `use_dynamic_shifting` and a `time_shift_type` of
+                `"exponential"`; when passed, `flow_shift` is set to `exp(mu)`.
             timesteps (`list[int]`, *optional*):
                 Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
                 timestep spacing strategy of equal spacing between timesteps schedule is used. If `timesteps` is
@@ -776,6 +779,9 @@ def dpm_solver_first_order_update(
                 The previous discrete timestep in the diffusion chain.
             sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:
@@ -841,6 +847,9 @@ def singlestep_dpm_solver_second_order_update(
                 The previous discrete timestep in the diffusion chain.
             sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:
@@ -952,6 +961,9 @@ def singlestep_dpm_solver_third_order_update(
                 The previous discrete timestep in the diffusion chain.
             sample (`torch.Tensor`):
                 A current instance of a sample created by diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:
@@ -1076,6 +1088,9 @@ def singlestep_dpm_solver_update(
                 A current instance of a sample created by diffusion process.
             order (`int`):
                 The solver order at this step.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:

src/diffusers/schedulers/scheduling_helios.py

@@ -386,6 +386,8 @@ def convert_model_output(
                 The current discrete timestep in the diffusion chain.
             sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
+            sigma (`torch.Tensor`, *optional*):
+                The sigma of the current step in the noise schedule.
 
         Returns:
             `torch.Tensor`:
@@ -470,6 +472,10 @@ def multistep_uni_p_bh_update(
                 A current instance of a sample created by the diffusion process.
             order (`int`):
                 The order of UniP at this timestep (corresponds to the *p* in UniPC-p).
+            sigma (`torch.Tensor`, *optional*):
+                The sigma of the current step in the noise schedule.
+            sigma_next (`torch.Tensor`, *optional*):
+                The sigma of the next step in the noise schedule.
 
         Returns:
             `torch.Tensor`:
@@ -607,6 +613,10 @@ def multistep_uni_c_bh_update(
                 The generated sample after the last predictor `x_{t}`.
             order (`int`):
                 The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.
+            sigma_before (`torch.Tensor`, *optional*):
+                The sigma of the previous step in the noise schedule.
+            sigma (`torch.Tensor`, *optional*):
+                The sigma of the current step in the noise schedule.
 
         Returns:
             `torch.Tensor`: