def forward(self, x):
        """Apply local attention by splitting into grids and reconstructing.

        Args:
            x (torch.Tensor): Input tensor of shape (batch, channels, height, width).

        Returns:
            torch.Tensor: Processed tensor of original size.
        """
        b, c, h, w = x.shape
        qkv = self.grids(x)
        out = self.grids_inverse(qkv)
        return out
Apply local attention by splitting into grids and reconstructing.

        Args:
            x (torch.Tensor): Input tensor of shape (batch, channels, height, width).

        Returns:
            torch.Tensor: Processed tensor of original size.
        
examples/community/pipeline_faithdiff_stable_diffusion_xl.py