class MotionModules(nn.Module):
    def __init__(
        self,
        in_channels: int,
        layers_per_block: int = 2,
        transformer_layers_per_block: Union[int, Tuple[int]] = 8,
        num_attention_heads: Union[int, Tuple[int]] = 8,
        attention_bias: bool = False,
        cross_attention_dim: Optional[int] = None,
        activation_fn: str = "geglu",
        norm_num_groups: int = 32,
        max_seq_length: int = 32,
    ):
        super().__init__()
        self.motion_modules = nn.ModuleList([])

        if isinstance(transformer_layers_per_block, int):
            transformer_layers_per_block = (transformer_layers_per_block,) * layers_per_block
        elif len(transformer_layers_per_block) != layers_per_block:
            raise ValueError(
                f"The number of transformer layers per block must match the number of layers per block, "
                f"got {layers_per_block} and {len(transformer_layers_per_block)}"
            )

        for i in range(layers_per_block):
            self.motion_modules.append(
                AnimateDiffTransformer3D(
                    in_channels=in_channels,
                    num_layers=transformer_layers_per_block[i],
                    norm_num_groups=norm_num_groups,
                    cross_attention_dim=cross_attention_dim,
                    activation_fn=activation_fn,
                    attention_bias=attention_bias,
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=in_channels // num_attention_heads,
                    positional_embeddings="sinusoidal",
                    num_positional_embeddings=max_seq_length,
                )
            )
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/unets/unet_motion_model.py