From 0d1cd52e4ccb44f2062d5ea2f16e4f4beceed17e Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 24 Jun 2025 14:17:37 +0530 Subject: [PATCH] raise as early as possible in group offloading --- src/diffusers/hooks/group_offloading.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py index 1f3a36b5d1..59c32bfb3e 100644 --- a/src/diffusers/hooks/group_offloading.py +++ b/src/diffusers/hooks/group_offloading.py @@ -63,6 +63,9 @@ class ModuleGroup: onload_self: bool = True, offload_to_disk_path: Optional[str] = None, ) -> None: + if stream is None and record_stream: + raise ValueError("`record_stream` cannot be True when `stream` is None.") + self.modules = modules self.offload_device = offload_device self.onload_device = onload_device @@ -96,9 +99,6 @@ class ModuleGroup: else: self.cpu_param_dict = self._init_cpu_param_dict() - if self.stream is None and self.record_stream: - raise ValueError("`record_stream` cannot be True when `stream` is None.") - def _init_cpu_param_dict(self): cpu_param_dict = {} if self.stream is None: @@ -513,6 +513,9 @@ def apply_group_offloading( else: raise ValueError("Using streams for data transfer requires a CUDA device, or an Intel XPU device.") + if not use_stream and record_stream: + raise ValueError("`record_stream` cannot be True when `use_stream=False`.") + _raise_error_if_accelerate_model_or_sequential_hook_present(module) if offload_type == "block_level":