Repository: huggingface/diffusers
File: src/diffusers/hooks/group_offloading.py (Python)

Code before:
        self.offload_to_disk_path = offload_to_disk_path
        self._is_offloaded_to_disk = False

        if self.offload_to_disk_path:
            self._group_id = _group_id
            short_hash = _compute_group_hash(self._group_id)
            self.safetensors_file_path = os.path.join(self.offload_to_disk_path, f"group_{short_hash}.safetensors")

            all_tensors = []
            for module in self.modules:
                all_tensors.extend(list(module.parameters()))
                all_tensors.extend(list(module.buffers()))
            all_tensors.extend(self.parameters)
            all_tensors.extend(self.buffers)
            all_tensors = list(dict.fromkeys(all_tensors))  # Remove duplicates

            self.tensor_to_key = {tensor: f"tensor_{i}" for i, tensor in enumerate(all_tensors)}
            self.key_to_tensor = {v: k for k, v in self.tensor_to_key.items()}
            self.cpu_param_dict = {}
        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:
            return cpu_param_dict

        for module in self.modules:
            for param in module.parameters():
                cpu_param_dict[param] = param.data.cpu() if self.low_cpu_mem_usage else param.data.cpu().pin_memory()
            for buffer in module.buffers():
                cpu_param_dict[buffer] = (
                    buffer.data.cpu() if self.low_cpu_mem_usage else buffer.data.cpu().pin_memory()
                )

        for param in self.parameters:
            cpu_param_dict[param] = param.data.cpu() if self.low_cpu_mem_usage else param.data.cpu().pin_memory()

        for buffer in self.buffers:
            cpu_param_dict[buffer] = buffer.data.cpu() if self.low_cpu_mem_usage else buffer.data.cpu().pin_memory()

        return cpu_param_dict

    @contextmanager
    def _pinned_memory_tensors(self):
        pinned_dict = {}
        try:
            for param, tensor in self.cpu_param_dict.items():

Review comment (suggestion):
The passed in `group_id` argument should be unique no? I don't think we need to compute a hash. 

Code after:
        low_cpu_mem_usage: bool = False,
        onload_self: bool = True,
        offload_to_disk_path: Optional[str] = None,
        group_id: Optional[int] = None,
    ) -> None:
        self.modules = modules
        self.offload_device = offload_device
        self.onload_device = onload_device
        self.offload_leader = offload_leader
        self.onload_leader = onload_leader
        self.parameters = parameters or []
        self.buffers = buffers or []
        self.non_blocking = non_blocking or stream is not None
        self.stream = stream
        self.record_stream = record_stream
        self.onload_self = onload_self
        self.low_cpu_mem_usage = low_cpu_mem_usage

        self.offload_to_disk_path = offload_to_disk_path
        self._is_offloaded_to_disk = False

        if self.offload_to_disk_path:
            # Instead of `group_id or str(id(self))` we do this because `group_id` can be "" as well.
            self.group_id = group_id if group_id is not None else str(id(self))
            short_hash = _compute_group_hash(self.group_id)
            self.safetensors_file_path = os.path.join(self.offload_to_disk_path, f"group_{short_hash}.safetensors")

            all_tensors = []
            for module in self.modules:
                all_tensors.extend(list(module.parameters()))
                all_tensors.extend(list(module.buffers()))
            all_tensors.extend(self.parameters)
            all_tensors.extend(self.buffers)
            all_tensors = list(dict.fromkeys(all_tensors))  # Remove duplicates

            self.tensor_to_key = {tensor: f"tensor_{i}" for i, tensor in enumerate(all_tensors)}
            self.key_to_tensor = {v: k for k, v in self.tensor_to_key.items()}
            self.cpu_param_dict = {}
        else:
            self.cpu_param_dict = self._init_cpu_param_dict()

    def _init_cpu_param_dict(self):
        cpu_param_dict = {}
        if self.stream is None:
            return cpu_param_dict

        for module in self.modules:
            for param in module.parameters():
                cpu_param_dict[param] = param.data.cpu() if self.low_cpu_mem_usage else param.data.cpu().pin_memory()
            for buffer in module.buffers():
                cpu_param_dict[buffer] = (