class SplitInferenceModule(nn.Module):
    r"""
    A wrapper module class that splits inputs along a specified dimension before performing a forward pass.

    This module is useful when you need to perform inference on large tensors in a memory-efficient way by breaking
    them into smaller chunks, processing each chunk separately, and then reassembling the results.

    Args:
        module (`nn.Module`):
            The underlying PyTorch module that will be applied to each chunk of split inputs.
        split_size (`int`, defaults to `1`):
            The size of each chunk after splitting the input tensor.
        split_dim (`int`, defaults to `0`):
            The dimension along which the input tensors are split.
        input_kwargs_to_split (`List[str]`, defaults to `["hidden_states"]`):
            A list of keyword arguments (strings) that represent the input tensors to be split.

    Workflow:
        1. The keyword arguments specified in `input_kwargs_to_split` are split into smaller chunks using
        `torch.split()` along the dimension `split_dim` and with a chunk size of `split_size`.
        2. The `module` is invoked once for each split with both the split inputs and any unchanged arguments
        that were passed.
        3. The output tensors from each split are concatenated back together along `split_dim` before returning.

    Example:
        ```python
        >>> import torch
        >>> import torch.nn as nn

        >>> model = nn.Linear(1000, 1000)
        >>> split_module = SplitInferenceModule(model, split_size=2, split_dim=0, input_kwargs_to_split=["input"])

        >>> input_tensor = torch.randn(42, 1000)
        >>> # Will split the tensor into 21 slices of shape [2, 1000].
        >>> output = split_module(input=input_tensor)
        ```

    It is also possible to nest `SplitInferenceModule` across different split dimensions for more complex
    multi-dimensional splitting.
    """

    def __init__(
        self,
        module: nn.Module,
        split_size: int = 1,
        split_dim: int = 0,
        input_kwargs_to_split: List[str] = ["hidden_states"],
    ) -> None:
        super().__init__()

        self.module = module
        self.split_size = split_size
        self.split_dim = split_dim
        self.input_kwargs_to_split = set(input_kwargs_to_split)

    def forward(self, *args, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor]]:
        r"""Forward method for the `SplitInferenceModule`.

        This method processes the input by splitting specified keyword arguments along a given dimension, running the
        underlying module on each split, and then concatenating the results. The splitting is controlled by the
        `split_size` and `split_dim` parameters specified during initialization.

        Args:
            *args (`Any`):
                Positional arguments that are passed directly to the `module` without modification.
            **kwargs (`Dict[str, torch.Tensor]`):
                Keyword arguments passed to the underlying `module`. Only keyword arguments whose names match the
                entries in `input_kwargs_to_split` and are of type `torch.Tensor` will be split. The remaining keyword
                arguments are passed unchanged.

        Returns:
            `Union[torch.Tensor, Tuple[torch.Tensor]]`:
                The outputs obtained from `SplitInferenceModule` are the same as if the underlying module was inferred
                without it.
                - If the underlying module returns a single tensor, the result will be a single concatenated tensor
                along the same `split_dim` after processing all splits.
                - If the underlying module returns a tuple of tensors, each element of the tuple will be concatenated
                along the `split_dim` across all splits, and the final result will be a tuple of concatenated tensors.
        """
        split_inputs = {}

        # 1. Split inputs that were specified during initialization and also present in passed kwargs
        for key in list(kwargs.keys()):
            if key not in self.input_kwargs_to_split or not torch.is_tensor(kwargs[key]):
                continue
            split_inputs[key] = torch.split(kwargs[key], self.split_size, self.split_dim)
            kwargs.pop(key)

        # 2. Invoke forward pass across each split
        results = []
        for split_input in zip(*split_inputs.values()):
            inputs = dict(zip(split_inputs.keys(), split_input))
            inputs.update(kwargs)

            intermediate_tensor_or_tensor_tuple = self.module(*args, **inputs)
            results.append(intermediate_tensor_or_tensor_tuple)

        # 3. Concatenate split restuls to obtain final outputs
        if isinstance(results[0], torch.Tensor):
            return torch.cat(results, dim=self.split_dim)
        elif isinstance(results[0], tuple):
            return tuple([torch.cat(x, dim=self.split_dim) for x in zip(*results)])
        else:
            raise ValueError(
                "In order to use the SplitInferenceModule, it is necessary for the underlying `module` to either return a torch.Tensor or a tuple of torch.Tensor's."
            )
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/pipelines/free_noise_utils.py