Expose group offload pinning options in API

Author: Akilesh-Sahaj

src/diffusers/models/modeling_utils.py

@@ -531,6 +531,8 @@ def enable_group_offload(
         record_stream: bool = False,
         low_cpu_mem_usage=False,
         offload_to_disk_path: Optional[str] = None,
+        pin_groups: Optional[Union[str, Callable]] = None,
+        pin_first_last: bool = False,
     ) -> None:
         r"""
         Activates group offloading for the current model.
@@ -583,6 +585,8 @@ def enable_group_offload(
             low_cpu_mem_usage=low_cpu_mem_usage,
             offload_to_disk_path=offload_to_disk_path,
             block_modules=block_modules,
+            pin_groups=pin_groups,
+            pin_first_last=pin_first_last,
         )
 
     def set_attention_backend(self, backend: str) -> None:

src/diffusers/pipelines/pipeline_utils.py

@@ -1342,6 +1342,8 @@ def enable_group_offload(
         low_cpu_mem_usage=False,
         offload_to_disk_path: Optional[str] = None,
         exclude_modules: Optional[Union[str, List[str]]] = None,
+        pin_groups: Optional[Union[str, Callable]] = None,
+        pin_first_last: bool = False,
     ) -> None:
         r"""
         Applies group offloading to the internal layers of a torch.nn.Module. To understand what group offloading is,
@@ -1402,6 +1404,11 @@ def enable_group_offload(
                 This option only matters when using streamed CPU offloading (i.e. `use_stream=True`). This can be
                 useful when the CPU memory is a bottleneck but may counteract the benefits of using streams.
             exclude_modules (`Union[str, List[str]]`, defaults to `None`): List of modules to exclude from offloading.
+            pin_groups (`\"first_last\"` | `\"all\"` | `Callable`, *optional*):
+                Optionally keep selected groups on the onload device permanently. See `ModelMixin.enable_group_offload`
+                for details.
+            pin_first_last (`bool`, *optional*, defaults to `False`):
+                Deprecated alias for `pin_groups=\"first_last\"`.
 
         Example:
             ```python
@@ -1442,6 +1449,8 @@ def enable_group_offload(
             "record_stream": record_stream,
             "low_cpu_mem_usage": low_cpu_mem_usage,
             "offload_to_disk_path": offload_to_disk_path,
+            "pin_groups": pin_groups,
+            "pin_first_last": pin_first_last,
         }
         for name, component in self.components.items():
             if name not in exclude_modules and isinstance(component, torch.nn.Module):