from typing import *
import logging
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint
import comfy.ops
import comfy.model_management
from .model import LayerNorm32, GroupNorm32, ChannelLayerNorm32, zero_module, pixel_shuffle_3d, str_to_dtype
from .sparse import SparseTensor, VarLenTensor, SparseDownsample, SparseUpsample, SparseSpatial2Channel, SparseChannel2Spatial, sparse_cat, SparseActivation
from .ops_sparse import manual_cast as sparse_ops
ops = comfy.ops.manual_cast


# ============================================================================
# Section 1: Dense 3D VAE (from sparse_structure_vae.py)
# ============================================================================

def norm_layer(norm_type: str, *args, dtype=None, device=None, **kwargs) -> nn.Module:
    """
    Return a normalization layer.
    """
    if norm_type == "group":
        return GroupNorm32(32, *args, dtype=dtype, device=device, **kwargs)
    elif norm_type == "layer":
        return ChannelLayerNorm32(*args, dtype=dtype, device=device, **kwargs)
    else:
        raise ValueError(f"Invalid norm type {norm_type}")


class ResBlock3d(nn.Module):
    def __init__(
        self,
        channels: int,
        out_channels: Optional[int] = None,
        norm_type: Literal["group", "layer"] = "layer",
        dtype=None, device=None, operations=ops,
    ):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels

        self.norm1 = norm_layer(norm_type, channels, dtype=dtype, device=device)
        self.norm2 = norm_layer(norm_type, self.out_channels, dtype=dtype, device=device)
        self.conv1 = operations.Conv3d(channels, self.out_channels, 3, padding=1, dtype=dtype, device=device)
        self.conv2 = zero_module(operations.Conv3d(self.out_channels, self.out_channels, 3, padding=1, dtype=dtype, device=device))
        self.skip_connection = operations.Conv3d(channels, self.out_channels, 1, dtype=dtype, device=device) if channels != self.out_channels else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.norm1(x)
        h = F.silu(h)
        h = self.conv1(h)
        h = self.norm2(h)
        h = F.silu(h)
        h = self.conv2(h)
        h = h + self.skip_connection(x)
        return h


class DownsampleBlock3d(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        mode: Literal["conv", "avgpool"] = "conv",
        dtype=None, device=None, operations=ops,
    ):
        assert mode in ["conv", "avgpool"], f"Invalid mode {mode}"

        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels

        if mode == "conv":
            self.conv = operations.Conv3d(in_channels, out_channels, 2, stride=2, dtype=dtype, device=device)
        elif mode == "avgpool":
            assert in_channels == out_channels, "Pooling mode requires in_channels to be equal to out_channels"

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if hasattr(self, "conv"):
            return self.conv(x)
        else:
            return F.avg_pool3d(x, 2)


class UpsampleBlock3d(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        mode: Literal["conv", "nearest"] = "conv",
        dtype=None, device=None, operations=ops,
    ):
        assert mode in ["conv", "nearest"], f"Invalid mode {mode}"

        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels

        if mode == "conv":
            self.conv = operations.Conv3d(in_channels, out_channels*8, 3, padding=1, dtype=dtype, device=device)
        elif mode == "nearest":
            assert in_channels == out_channels, "Nearest mode requires in_channels to be equal to out_channels"

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if hasattr(self, "conv"):
            x = self.conv(x)
            return pixel_shuffle_3d(x, 2)
        else:
            return F.interpolate(x, scale_factor=2, mode="nearest")


class SparseStructureEncoder(nn.Module):
    """
    Encoder for Sparse Structure (E_S in the paper Sec. 3.3).

    Args:
        in_channels (int): Channels of the input.
        latent_channels (int): Channels of the latent representation.
        num_res_blocks (int): Number of residual blocks at each resolution.
        channels (List[int]): Channels of the encoder blocks.
        num_res_blocks_middle (int): Number of residual blocks in the middle.
        norm_type (Literal["group", "layer"]): Type of normalization layer.
        use_fp16 (bool): Whether to use FP16.
    """
    def __init__(
        self,
        in_channels: int,
        latent_channels: int,
        num_res_blocks: int,
        channels: List[int],
        num_res_blocks_middle: int = 2,
        norm_type: Literal["group", "layer"] = "layer",
        use_fp16: bool = False,
        dtype=None, device=None, operations=ops,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.latent_channels = latent_channels
        self.num_res_blocks = num_res_blocks
        self.channels = channels
        self.num_res_blocks_middle = num_res_blocks_middle
        self.norm_type = norm_type
        self.use_fp16 = use_fp16

        self.input_layer = operations.Conv3d(in_channels, channels[0], 3, padding=1, dtype=dtype, device=device)

        self.blocks = nn.ModuleList([])
        for i, ch in enumerate(channels):
            self.blocks.extend([
                ResBlock3d(ch, ch, dtype=dtype, device=device, operations=operations)
                for _ in range(num_res_blocks)
            ])
            if i < len(channels) - 1:
                self.blocks.append(
                    DownsampleBlock3d(ch, channels[i+1], dtype=dtype, device=device, operations=operations)
                )

        self.middle_block = nn.Sequential(*[
            ResBlock3d(channels[-1], channels[-1], dtype=dtype, device=device, operations=operations)
            for _ in range(num_res_blocks_middle)
        ])

        self.out_layer = nn.Sequential(
            norm_layer(norm_type, channels[-1], dtype=dtype, device=device),
            nn.SiLU(),
            operations.Conv3d(channels[-1], latent_channels*2, 3, padding=1, dtype=dtype, device=device)
        )

    @property
    def device(self) -> torch.device:
        return next(self.parameters()).device

    @device.setter
    def device(self, value):
        pass

    def forward(self, x: torch.Tensor, sample_posterior: bool = False, return_raw: bool = False) -> torch.Tensor:
        h = self.input_layer(x)

        for block in self.blocks:
            h = block(h)
        h = self.middle_block(h)

        h = self.out_layer(h)

        mean, logvar = h.chunk(2, dim=1)

        if sample_posterior:
            std = torch.exp(0.5 * logvar)
            z = mean + std * torch.randn_like(std)
        else:
            z = mean

        if return_raw:
            return z, mean, logvar
        return z


class SparseStructureDecoder(nn.Module):
    """
    Decoder for Sparse Structure (D_S in the paper Sec. 3.3).

    Args:
        out_channels (int): Channels of the output.
        latent_channels (int): Channels of the latent representation.
        num_res_blocks (int): Number of residual blocks at each resolution.
        channels (List[int]): Channels of the decoder blocks.
        num_res_blocks_middle (int): Number of residual blocks in the middle.
        norm_type (Literal["group", "layer"]): Type of normalization layer.
        use_fp16 (bool): Whether to use FP16.
    """
    def __init__(
        self,
        out_channels: int,
        latent_channels: int,
        num_res_blocks: int,
        channels: List[int],
        num_res_blocks_middle: int = 2,
        norm_type: Literal["group", "layer"] = "layer",
        use_fp16: bool = False,
        dtype=None, device=None, operations=ops,
    ):
        super().__init__()
        self.out_channels = out_channels
        self.latent_channels = latent_channels
        self.num_res_blocks = num_res_blocks
        self.channels = channels
        self.num_res_blocks_middle = num_res_blocks_middle
        self.norm_type = norm_type
        self.use_fp16 = use_fp16

        self.input_layer = operations.Conv3d(latent_channels, channels[0], 3, padding=1, dtype=dtype, device=device)

        self.middle_block = nn.Sequential(*[
            ResBlock3d(channels[0], channels[0], dtype=dtype, device=device, operations=operations)
            for _ in range(num_res_blocks_middle)
        ])

        self.blocks = nn.ModuleList([])
        for i, ch in enumerate(channels):
            self.blocks.extend([
                ResBlock3d(ch, ch, dtype=dtype, device=device, operations=operations)
                for _ in range(num_res_blocks)
            ])
            if i < len(channels) - 1:
                self.blocks.append(
                    UpsampleBlock3d(ch, channels[i+1], dtype=dtype, device=device, operations=operations)
                )

        self.out_layer = nn.Sequential(
            norm_layer(norm_type, channels[-1], dtype=dtype, device=device),
            nn.SiLU(),
            operations.Conv3d(channels[-1], out_channels, 3, padding=1, dtype=dtype, device=device)
        )

    @property
    def device(self) -> torch.device:
        return next(self.parameters()).device

    @device.setter
    def device(self, value):
        pass

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.input_layer(x)
        h = self.middle_block(h)
        for block in self.blocks:
            h = block(h)
        h = self.out_layer(h)
        return h


# ============================================================================
# Section 2: Sparse UNet VAE (from sparse_unet_vae.py)
# ============================================================================

# Use the sparse module as 'sp' namespace for SparseTensor operations
import importlib as _importlib
sp = _importlib.import_module('.sparse', __name__.rsplit('.', 1)[0] if '.' in __name__ else __name__)


def _apply_in_chunks(module: nn.Module, x: torch.Tensor, chunk_size: int) -> torch.Tensor:
    if chunk_size <= 0 or x.shape[0] <= chunk_size:
        return module(x)

    outputs = []
    for start in range(0, x.shape[0], chunk_size):
        outputs.append(module(x[start : start + chunk_size]))
    return torch.cat(outputs, dim=0)


class SparseResBlock3d(nn.Module):
    def __init__(
        self,
        channels: int,
        out_channels: Optional[int] = None,
        downsample: bool = False,
        upsample: bool = False,
        resample_mode: Literal['nearest', 'spatial2channel'] = 'nearest',
        use_checkpoint: bool = False,
        dtype=None, device=None, operations=ops, sparse_operations=sparse_ops,
    ):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.downsample = downsample
        self.upsample = upsample
        self.resample_mode = resample_mode
        self.use_checkpoint = use_checkpoint

        assert not (downsample and upsample), "Cannot downsample and upsample at the same time"

        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6, dtype=dtype, device=device)
        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
        if resample_mode == 'nearest':
            self.conv1 = sparse_operations.SparseConv3d(channels, self.out_channels, 3, dtype=dtype, device=device)
        elif resample_mode =='spatial2channel' and not self.downsample:
            self.conv1 = sparse_operations.SparseConv3d(channels, self.out_channels * 8, 3, dtype=dtype, device=device)
        elif resample_mode =='spatial2channel' and self.downsample:
            self.conv1 = sparse_operations.SparseConv3d(channels, self.out_channels // 8, 3, dtype=dtype, device=device)
        self.conv2 = zero_module(sparse_operations.SparseConv3d(self.out_channels, self.out_channels, 3, dtype=dtype, device=device))
        if resample_mode == 'nearest':
            self.skip_connection = sparse_operations.SparseLinear(channels, self.out_channels, dtype=dtype, device=device) if channels != self.out_channels else nn.Identity()
        elif resample_mode =='spatial2channel' and self.downsample:
            self.skip_connection = lambda x: x.replace(x.feats.reshape(x.feats.shape[0], out_channels, channels * 8 // out_channels).mean(dim=-1))
        elif resample_mode =='spatial2channel' and not self.downsample:
            self.skip_connection = lambda x: x.replace(x.feats.repeat_interleave(out_channels // (channels // 8), dim=1))
        self.updown = None
        if self.downsample:
            if resample_mode == 'nearest':
                self.updown = sp.SparseDownsample(2)
            elif resample_mode =='spatial2channel':
                self.updown = sp.SparseSpatial2Channel(2)
        elif self.upsample:
            self.to_subdiv = sparse_operations.SparseLinear(channels, 8, dtype=dtype, device=device)
            if resample_mode == 'nearest':
                self.updown = sp.SparseUpsample(2)
            elif resample_mode =='spatial2channel':
                self.updown = sp.SparseChannel2Spatial(2)

    def _updown(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
        if self.downsample:
            x = self.updown(x)
        elif self.upsample:
            x = self.updown(x, subdiv.replace(subdiv.feats > 0))
        return x

    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
        subdiv = None
        if self.upsample:
            subdiv = self.to_subdiv(x)
        h = x.replace(self.norm1(x.feats))
        h = h.replace(F.silu(h.feats))
        if self.resample_mode == 'spatial2channel':
            h = self.conv1(h)
        h = self._updown(h, subdiv)
        x = self._updown(x, subdiv)
        if self.resample_mode == 'nearest':
            h = self.conv1(h)
        h = h.replace(self.norm2(h.feats))
        h = h.replace(F.silu(h.feats))
        h = self.conv2(h)
        h = h + self.skip_connection(x)
        if self.upsample:
            return h, subdiv
        return h

    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
        if self.use_checkpoint:
            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
        else:
            return self._forward(x)


class SparseResBlockDownsample3d(nn.Module):
    def __init__(
        self,
        channels: int,
        out_channels: Optional[int] = None,
        use_checkpoint: bool = False,
        dtype=None, device=None, operations=ops, sparse_operations=sparse_ops,
    ):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.use_checkpoint = use_checkpoint

        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6, dtype=dtype, device=device)
        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
        self.conv1 = sparse_operations.SparseConv3d(channels, self.out_channels, 3, dtype=dtype, device=device)
        self.conv2 = zero_module(sparse_operations.SparseConv3d(self.out_channels, self.out_channels, 3, dtype=dtype, device=device))
        self.skip_connection = sparse_operations.SparseLinear(channels, self.out_channels, dtype=dtype, device=device) if channels != self.out_channels else nn.Identity()
        self.updown = sp.SparseDownsample(2)

    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
        h = x.replace(self.norm1(x.feats))
        h = h.replace(F.silu(h.feats))
        h = self.updown(h)
        x = self.updown(x)
        h = self.conv1(h)
        h = h.replace(self.norm2(h.feats))
        h = h.replace(F.silu(h.feats))
        h = self.conv2(h)
        h = h + self.skip_connection(x)
        return h

    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
        if self.use_checkpoint:
            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
        else:
            return self._forward(x)


class SparseResBlockUpsample3d(nn.Module):
    def __init__(
        self,
        channels: int,
        out_channels: Optional[int] = None,
        use_checkpoint: bool = False,
        pred_subdiv: bool = True,
        dtype=None, device=None, operations=ops, sparse_operations=sparse_ops,
    ):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.use_checkpoint = use_checkpoint
        self.pred_subdiv = pred_subdiv

        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6, dtype=dtype, device=device)
        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
        self.conv1 = sparse_operations.SparseConv3d(channels, self.out_channels, 3, dtype=dtype, device=device)
        self.conv2 = zero_module(sparse_operations.SparseConv3d(self.out_channels, self.out_channels, 3, dtype=dtype, device=device))
        self.skip_connection = sparse_operations.SparseLinear(channels, self.out_channels, dtype=dtype, device=device) if channels != self.out_channels else nn.Identity()
        if self.pred_subdiv:
            self.to_subdiv = sparse_operations.SparseLinear(channels, 8, dtype=dtype, device=device)
        self.updown = sp.SparseUpsample(2)

    def _forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
        if self.pred_subdiv:
            subdiv = self.to_subdiv(x)
        h = x.replace(self.norm1(x.feats))
        h = h.replace(F.silu(h.feats))
        subdiv_binarized = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None
        h = self.updown(h, subdiv_binarized)
        x = self.updown(x, subdiv_binarized)
        h = self.conv1(h)
        h = h.replace(self.norm2(h.feats))
        h = h.replace(F.silu(h.feats))
        h = self.conv2(h)
        h = h + self.skip_connection(x)
        if self.pred_subdiv:
            return h, subdiv
        else:
            return h

    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
        if self.use_checkpoint:
            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
        else:
            return self._forward(x)


class SparseResBlockS2C3d(nn.Module):
    def __init__(
        self,
        channels: int,
        out_channels: Optional[int] = None,
        use_checkpoint: bool = False,
        dtype=None, device=None, operations=ops, sparse_operations=sparse_ops,
    ):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.use_checkpoint = use_checkpoint

        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6, dtype=dtype, device=device)
        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
        self.conv1 = sparse_operations.SparseConv3d(channels, self.out_channels // 8, 3, dtype=dtype, device=device)
        self.conv2 = zero_module(sparse_operations.SparseConv3d(self.out_channels, self.out_channels, 3, dtype=dtype, device=device))
        self.skip_connection = lambda x: x.replace(x.feats.reshape(x.feats.shape[0], out_channels, channels * 8 // out_channels).mean(dim=-1))
        self.updown = sp.SparseSpatial2Channel(2)

    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
        h = x.replace(self.norm1(x.feats))
        h = h.replace(F.silu(h.feats))
        h = self.conv1(h)
        h = self.updown(h)
        x = self.updown(x)
        h = h.replace(self.norm2(h.feats))
        h = h.replace(F.silu(h.feats))
        h = self.conv2(h)
        h = h + self.skip_connection(x)
        return h

    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
        if self.use_checkpoint:
            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
        else:
            return self._forward(x)


class SparseResBlockC2S3d(nn.Module):
    def __init__(
        self,
        channels: int,
        out_channels: Optional[int] = None,
        use_checkpoint: bool = False,
        pred_subdiv: bool = True,
        dtype=None, device=None, operations=ops, sparse_operations=sparse_ops,
    ):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.use_checkpoint = use_checkpoint
        self.pred_subdiv = pred_subdiv
        self._low_vram = False

        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6, dtype=dtype, device=device)
        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
        self.conv1 = sparse_operations.SparseConv3d(channels, self.out_channels * 8, 3, dtype=dtype, device=device)
        self.conv2 = zero_module(sparse_operations.SparseConv3d(self.out_channels, self.out_channels, 3, dtype=dtype, device=device))
        self.skip_connection = lambda x: x.replace(x.feats.repeat_interleave(out_channels // (channels // 8), dim=1))
        if pred_subdiv:
            self.to_subdiv = sparse_operations.SparseLinear(channels, 8, dtype=dtype, device=device)
        self.updown = sp.SparseChannel2Spatial(2)

    @property
    def low_vram(self) -> bool:
        return self._low_vram

    @low_vram.setter
    def low_vram(self, value: bool) -> None:
        self._low_vram = bool(value)

    def _tiled_conv1_c2s(self, h, x, subdiv_bin):
        """Fused tiled conv1->C2S. Never materializes full conv1 output.

        Returns (h_fine: SparseTensor, x_fine_feats_cpu: Tensor).
        """
        from flex_gemm_ap.ops.spconv.submanifold_conv3d import (
            SubMConv3dFunction, _partition_octree
        )

        _ma = torch.cuda.memory_allocated
        _v = lambda: _ma() // 1048576
        device = h.feats.device
        N = h.feats.shape[0]
        coords = h.coords
        spatial = coords[:, 1:]
        shape = torch.Size([*h.shape, *h.spatial_shape])

        weight = self.conv1.weight
        bias = self.conv1.bias
        Co, Kd, Kh, Kw, Ci = weight.shape
        kernel_size = (Kw, Kh, Kd)
        dilation = self.conv1.dilation
        halo = max((k // 2) * d for k, d in zip(kernel_size, dilation))

        factor = self.updown.factor
        DIM = 3
        factor_cubed = factor ** DIM
        subdiv_feats = subdiv_bin.feats

        x_scale = x._scale

        max_tile = 500_000
        tiles = _partition_octree(spatial, max_tile)
        print(f"[C2S3d] TILED: N={N:,} -> {len(tiles)} tiles, alloc={_v()}MB", flush=True)

        h_feats_cpu = h.feats.cpu()
        x_feats_cpu = x.feats.cpu()
        h.data['feats'] = torch.empty(0, device='cpu')
        x.data['feats'] = torch.empty(0, device='cpu')
        del h, x
        torch.cuda.empty_cache()
        print(f"[C2S3d] TILED feats offloaded, alloc={_v()}MB", flush=True)

        fine_h_list = []
        fine_x_list = []
        fine_coords_list = []

        for ti, (tile_min, tile_max) in enumerate(tiles):
            interior_mask = ((spatial >= tile_min) & (spatial < tile_max)).all(dim=1)
            halo_mask = ((spatial >= tile_min - halo) & (spatial < tile_max + halo)).all(dim=1)

            halo_indices = halo_mask.nonzero(as_tuple=True)[0]
            halo_indices_cpu = halo_indices.cpu()
            tile_feats = h_feats_cpu[halo_indices_cpu].to(device)
            tile_coords = coords[halo_indices].contiguous()

            print(f"[C2S3d] TILED tile {ti+1}/{len(tiles)}: conv1 N_halo={tile_feats.shape[0]:,} "
                  f"alloc={_v()}MB", flush=True)
            tile_cache = SubMConv3dFunction._compute_neighbor_cache(
                tile_coords, shape, kernel_size, dilation)
            tile_conv_out = SubMConv3dFunction._sparse_submanifold_conv_forward(
                tile_feats, tile_cache, weight, bias)
            del tile_feats, tile_cache

            interior_in_halo = interior_mask[halo_indices]
            interior_conv = tile_conv_out[interior_in_halo]
            interior_coords = tile_coords[interior_in_halo]
            del tile_conv_out, tile_coords, halo_indices, halo_indices_cpu

            interior_global = interior_mask.nonzero(as_tuple=True)[0]
            tile_subdiv = subdiv_feats[interior_global]
            tile_x = x_feats_cpu[interior_global.cpu()].to(device)

            N_int = interior_conv.shape[0]
            N_leaf = tile_subdiv.sum(dim=-1)
            subidx = tile_subdiv.nonzero()[:, -1]
            output_size = subidx.shape[0]

            new_coords = interior_coords.clone()
            new_coords[:, 1:] *= factor
            new_coords = torch.repeat_interleave(
                new_coords, N_leaf, dim=0, output_size=output_size)
            for i in range(DIM):
                new_coords[:, i + 1] += subidx // factor ** i % factor

            idx = torch.repeat_interleave(
                torch.arange(N_int, device=device), N_leaf, dim=0,
                output_size=output_size)
            gather = idx * factor_cubed + subidx

            fine_h = interior_conv.reshape(N_int * factor_cubed, -1)[gather]
            fine_x = tile_x.reshape(N_int * factor_cubed, -1)[gather]

            fine_h_list.append(fine_h.cpu())
            fine_x_list.append(fine_x.cpu())
            fine_coords_list.append(new_coords.cpu())

            del interior_conv, tile_x, interior_coords, tile_subdiv
            del fine_h, fine_x, new_coords, interior_mask, halo_mask
            del interior_in_halo, interior_global, idx, gather
            torch.cuda.empty_cache()
            print(f"[C2S3d] TILED tile {ti+1}/{len(tiles)}: N_int={N_int:,} "
                  f"fine={output_size:,} alloc={_v()}MB", flush=True)

        del h_feats_cpu

        fine_coords_gpu = torch.cat(fine_coords_list).to(device)
        del fine_coords_list
        h_fine = sp.SparseTensor(
            torch.cat(fine_h_list).to(device),
            fine_coords_gpu,
        )
        h_fine._scale = tuple(s / factor for s in x_scale)
        del fine_h_list

        x_fine_feats_cpu = torch.cat(fine_x_list)
        del fine_x_list, x_feats_cpu

        print(f"[C2S3d] TILED done: h_fine={h_fine.feats.shape} "
              f"x_fine_cpu={x_fine_feats_cpu.shape} alloc={_v()}MB", flush=True)
        return h_fine, x_fine_feats_cpu

    def _forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
        _ma = torch.cuda.memory_allocated
        _v = lambda: _ma() // 1048576
        if self.pred_subdiv:
            subdiv = self.to_subdiv(x)
        print(f"[C2S3d] after subdiv: alloc={_v()}MB", flush=True)
        h = x.replace(self.norm1(x.feats))
        print(f"[C2S3d] after norm1: alloc={_v()}MB", flush=True)
        h = h.replace(F.silu(h.feats))
        N = h.feats.shape[0]
        print(f"[C2S3d] pre-conv1: N={N:,} Ci={h.feats.shape[1]} Co={self.out_channels*8} alloc={_v()}MB", flush=True)

        subdiv_binarized = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None

        # --- TILED PATH: fused conv1->C2S, never materializes full conv1 output ---
        if self._low_vram and N > 1_000_000:
            h, x_feats_cpu = self._tiled_conv1_c2s(h, x, subdiv_binarized)
            del subdiv_binarized
        else:
            # --- NORMAL PATH ---
            if self._low_vram:
                x_feats_cpu = x.feats.to('cpu', non_blocking=True)
                x = x.replace(torch.empty(0, dtype=x.feats.dtype, device=x.feats.device))
                torch.cuda.current_stream().synchronize()
                torch.cuda.empty_cache()
                print(f"[C2S3d] x.feats offloaded before conv1: alloc={_v()}MB", flush=True)
            h = self.conv1(h)
            print(f"[C2S3d] post-conv1: h={h.feats.shape} alloc={_v()}MB", flush=True)
            if self._low_vram:
                x = x.replace(x_feats_cpu.to(h.feats.device))
                del x_feats_cpu
                print(f"[C2S3d] x.feats restored from CPU: alloc={_v()}MB", flush=True)
            print(f"[C2S3d] pre-C2S(h): alloc={_v()}MB", flush=True)
            h = self.updown(h, subdiv_binarized)
            print(f"[C2S3d] post-C2S(h): h={h.feats.shape} alloc={_v()}MB", flush=True)
            x = self.updown(x, subdiv_binarized)
            print(f"[C2S3d] post-C2S(x): x={x.feats.shape} alloc={_v()}MB", flush=True)
            del subdiv_binarized
            h.clear_spatial_cache()
            x.clear_spatial_cache()
            print(f"[C2S3d] post-C2S: N={h.feats.shape[0]:,} h={h.feats.shape[1]}ch x={x.feats.shape[1]}ch alloc={_v()}MB", flush=True)
            x_feats_cpu = x.feats.to('cpu', non_blocking=True)
            del x
            torch.cuda.current_stream().synchronize()
            torch.cuda.empty_cache()
            print(f"[C2S3d] x offloaded to CPU, pre-conv2: alloc={_v()}MB", flush=True)

        # --- conv2 + skip (shared by both paths) ---
        h.clear_spatial_cache()
        N = h.feats.shape[0]
        _CHUNK = 1_000_000

        if self._low_vram and N > _CHUNK:
            # Offload h.feats to CPU, process norm2+silu in chunks on GPU,
            # then reassemble.  Without offloading first, the full h.feats
            # (2+ GB) stays on GPU and even a small chunk can't fit.
            dev = h.feats.device
            h_cpu = h.feats.to('cpu', non_blocking=True)
            torch.cuda.current_stream().synchronize()
            h = h.replace(torch.empty(0, dtype=h_cpu.dtype, device=dev))
            torch.cuda.empty_cache()
            print(f"[C2S3d] h.feats offloaded for chunked norm2+silu: alloc={_v()}MB", flush=True)
            out_chunks = []
            for i in range(0, N, _CHUNK):
                c = h_cpu[i:i+_CHUNK].to(dev)
                out_chunks.append(F.silu(self.norm2(c)).cpu())
                del c
            h = h.replace(torch.cat(out_chunks).to(dev))
            del out_chunks, h_cpu
            torch.cuda.empty_cache()
            print(f"[C2S3d] after chunked norm2+silu: alloc={_v()}MB", flush=True)
        else:
            h = h.replace(self.norm2(h.feats))
            print(f"[C2S3d] after norm2: alloc={_v()}MB", flush=True)
            h = h.replace(F.silu(h.feats))
            print(f"[C2S3d] after silu: alloc={_v()}MB", flush=True)

        h = self.conv2(h)
        print(f"[C2S3d] post-conv2: alloc={_v()}MB", flush=True)
        h.clear_spatial_cache()
        print(f"[C2S3d] post-conv2 (cache cleared): alloc={_v()}MB", flush=True)

        R = self.out_channels // (self.channels // 8)
        if self._low_vram and N > _CHUNK:
            # Chunked skip to avoid materializing full repeat_interleave
            dev = h.feats.device
            for i in range(0, N, _CHUNK):
                end = min(i + _CHUNK, N)
                skip_chunk = x_feats_cpu[i:end].repeat_interleave(R, dim=1).to(dev)
                h.feats[i:end] += skip_chunk
                del skip_chunk
            del x_feats_cpu
            print(f"[C2S3d] after chunked skip: alloc={_v()}MB", flush=True)
        else:
            skip_feats = x_feats_cpu.repeat_interleave(R, dim=1).to(h.feats.device)
            del x_feats_cpu
            print(f"[C2S3d] after skip expand: alloc={_v()}MB", flush=True)
            h = h.replace(h.feats + skip_feats)
            del skip_feats
            print(f"[C2S3d] post-skip: alloc={_v()}MB", flush=True)
        if self.pred_subdiv:
            return h, subdiv
        else:
            return h

    def forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
        if self.use_checkpoint:
            return torch.utils.checkpoint.checkpoint(self._forward, x, subdiv, use_reentrant=False)
        else:
            return self._forward(x, subdiv)


class SparseConvNeXtBlock3d(nn.Module):
    def __init__(
        self,
        channels: int,
        mlp_ratio: float = 4.0,
        use_checkpoint: bool = False,
        dtype=None, device=None, operations=ops, sparse_operations=sparse_ops,
    ):
        super().__init__()
        self.channels = channels
        self.use_checkpoint = use_checkpoint
        self.low_vram = False
        self.mlp_chunk_size = 8192

        self.norm = LayerNorm32(channels, elementwise_affine=True, eps=1e-6, dtype=dtype, device=device)
        self.conv = sparse_operations.SparseConv3d(channels, channels, 3, dtype=dtype, device=device)
        self.mlp = nn.Sequential(
            operations.Linear(channels, int(channels * mlp_ratio), dtype=dtype, device=device),
            nn.SiLU(),
            zero_module(operations.Linear(int(channels * mlp_ratio), channels, dtype=dtype, device=device)),
        )

    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
        h = self.conv(x)
        _CHUNK = 1_000_000
        N = h.feats.shape[0]
        if self.low_vram and N > _CHUNK:
            # Full CPU-offload path: norm + MLP + residual all chunked.
            # At 18M+ voxels, even one full tensor exceeds available VRAM.
            dev = h.feats.device
            h_cpu = h.feats.to('cpu', non_blocking=True)
            torch.cuda.current_stream().synchronize()
            del h
            torch.cuda.empty_cache()
            if self.low_vram:
                x.clear_spatial_cache()

            # norm -> MLP -> residual, all in 1M chunks on GPU
            x_feats_cpu = x.feats.cpu() if x.feats.device.type == 'cuda' else x.feats
            out_chunks = []
            for i in range(0, N, _CHUNK):
                end = min(i + _CHUNK, N)
                c = h_cpu[i:end].to(dev)
                normed = self.norm(c)
                del c
                mlp_out = self.mlp(normed)
                del normed
                mlp_out += x_feats_cpu[i:end].to(dev)  # residual
                out_chunks.append(mlp_out.cpu())
                del mlp_out
            del h_cpu, x_feats_cpu
            feats = torch.cat(out_chunks).to(dev)
            del out_chunks
            return x.replace(feats)
        else:
            norm_feats = self.norm(h.feats)
            del h  # free conv output feats before MLP expansion
            if self.low_vram:
                x.clear_spatial_cache()
            steps = max(1, (norm_feats.shape[0] + self.mlp_chunk_size - 1) // self.mlp_chunk_size) if self.low_vram else 1
            while True:
                try:
                    chunk_size = (norm_feats.shape[0] + steps - 1) // steps if steps > 1 else 0
                    feats = _apply_in_chunks(self.mlp, norm_feats, chunk_size)
                    break
                except comfy.model_management.OOM_EXCEPTION as e:
                    comfy.model_management.soft_empty_cache(True)
                    steps *= 2
                    if steps > 64:
                        raise e
            del norm_feats
            x_feats = x.feats
            if x_feats.device != feats.device:
                x_feats = x_feats.to(feats.device)
            feats.add_(x_feats)
            del x_feats
            return x.replace(feats)

    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
        if self.use_checkpoint:
            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
        else:
            return self._forward(x)


class SparseUnetVaeEncoder(nn.Module):
    """
    Sparse Swin Transformer Unet VAE model.
    """
    def __init__(
        self,
        in_channels: int,
        model_channels: List[int],
        latent_channels: int,
        num_blocks: List[int],
        block_type: List[str],
        down_block_type: List[str],
        block_args: List[Dict[str, Any]],
        use_fp16: bool = False,
        dtype=None, device=None, operations=ops, sparse_operations=sparse_ops,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.model_channels = model_channels
        self.num_blocks = num_blocks

        self.input_layer = sparse_operations.SparseLinear(in_channels, model_channels[0], dtype=dtype, device=device)
        self.to_latent = sparse_operations.SparseLinear(model_channels[-1], 2 * latent_channels, dtype=dtype, device=device)

        self.blocks = nn.ModuleList([])
        for i in range(len(num_blocks)):
            self.blocks.append(nn.ModuleList([]))
            for j in range(num_blocks[i]):
                self.blocks[-1].append(
                    globals()[block_type[i]](
                        model_channels[i],
                        dtype=dtype,
                        device=device,
                        operations=operations,
                        sparse_operations=sparse_operations,
                        **block_args[i],
                    )
                )
            if i < len(num_blocks) - 1:
                self.blocks[-1].append(
                    globals()[down_block_type[i]](
                        model_channels[i],
                        model_channels[i+1],
                        dtype=dtype,
                        device=device,
                        operations=operations,
                        sparse_operations=sparse_operations,
                        **block_args[i],
                    )
                )

    @property
    def device(self) -> torch.device:
        return next(self.parameters()).device

    @device.setter
    def device(self, value):
        pass

    def convert_to_fp16(self) -> None:
        pass

    def convert_to_fp32(self) -> None:
        pass

    def initialize_weights(self) -> None:
        # Initialize transformer layers:
        def _basic_init(module):
            if isinstance(module, nn.Linear):
                torch.nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)
        self.apply(_basic_init)

    def forward(self, x: sp.SparseTensor, sample_posterior=False, return_raw=False):
        h = self.input_layer(x)
        for i, res in enumerate(self.blocks):
            for j, block in enumerate(res):
                h = block(h)
        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
        h = self.to_latent(h)

        # Sample from the posterior distribution
        mean, logvar = h.feats.chunk(2, dim=-1)
        if sample_posterior:
            std = torch.exp(0.5 * logvar)
            z = mean + std * torch.randn_like(std)
        else:
            z = mean
        z = h.replace(z)

        if return_raw:
            return z, mean, logvar
        else:
            return z


class SparseUnetVaeDecoder(nn.Module):
    """
    Sparse Swin Transformer Unet VAE model.
    """
    def __init__(
        self,
        out_channels: int,
        model_channels: List[int],
        latent_channels: int,
        num_blocks: List[int],
        block_type: List[str],
        up_block_type: List[str],
        block_args: List[Dict[str, Any]],
        use_fp16: bool = False,
        pred_subdiv: bool = True,
        dtype=None, device=None, operations=ops, sparse_operations=sparse_ops,
    ):
        super().__init__()
        self.out_channels = out_channels
        self.model_channels = model_channels
        self.num_blocks = num_blocks
        self.use_fp16 = use_fp16
        self.pred_subdiv = pred_subdiv
        self._low_vram = False

        self.output_layer = sparse_operations.SparseLinear(model_channels[-1], out_channels, dtype=dtype, device=device)
        self.from_latent = sparse_operations.SparseLinear(latent_channels, model_channels[0], dtype=dtype, device=device)

        self.blocks = nn.ModuleList([])
        for i in range(len(num_blocks)):
            self.blocks.append(nn.ModuleList([]))
            for j in range(num_blocks[i]):
                self.blocks[-1].append(
                    globals()[block_type[i]](
                        model_channels[i],
                        dtype=dtype,
                        device=device,
                        operations=operations,
                        sparse_operations=sparse_operations,
                        **block_args[i],
                    )
                )
            if i < len(num_blocks) - 1:
                self.blocks[-1].append(
                    globals()[up_block_type[i]](
                        model_channels[i],
                        model_channels[i+1],
                        pred_subdiv=pred_subdiv,
                        dtype=dtype,
                        device=device,
                        operations=operations,
                        sparse_operations=sparse_operations,
                        **block_args[i],
                    )
                )

    @property
    def device(self) -> torch.device:
        return next(self.parameters()).device

    @device.setter
    def device(self, value):
        pass

    @property
    def low_vram(self) -> bool:
        return self._low_vram

    @low_vram.setter
    def low_vram(self, value: bool) -> None:
        self._low_vram = bool(value)
        for module in self.modules():
            if module is self:
                continue
            if hasattr(module, "low_vram"):
                try:
                    setattr(module, "low_vram", self._low_vram)
                except Exception:
                    pass
        # Tiled sparse conv for large voxel counts
        tile_size = 1_000_000 if self._low_vram else 0
        count = 0
        for module in self.modules():
            if module.__class__.__name__ == 'SparseConv3d':
                module.max_voxels_per_tile = tile_size
                count += 1
        if count:
            print(f"[vae] set max_voxels_per_tile={tile_size} on {count} SparseConv3d modules", flush=True)

    def convert_to_fp16(self) -> None:
        pass

    def convert_to_fp32(self) -> None:
        pass

    def initialize_weights(self) -> None:
        # Initialize transformer layers:
        def _basic_init(module):
            if isinstance(module, nn.Linear):
                torch.nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)
        self.apply(_basic_init)

    def forward(self, x: sp.SparseTensor, guide_subs: Optional[List[sp.SparseTensor]] = None, return_subs: bool = False) -> sp.SparseTensor:
        assert guide_subs is None or self.pred_subdiv == False, "Only decoders with pred_subdiv=False can be used with guide_subs"
        assert return_subs == False or self.pred_subdiv == True, "Only decoders with pred_subdiv=True can be used with return_subs"

        import comfy.utils
        total_blocks = sum(len(res) for res in self.blocks)
        pbar = comfy.utils.ProgressBar(total_blocks + 1)

        h = self.from_latent(x)
        _ma = torch.cuda.memory_allocated

        subs = []
        for i, res in enumerate(self.blocks):
            for j, block in enumerate(res):
                _bname = block.__class__.__name__
                print(f"[UNET] level={i} block={j} {_bname} N={h.feats.shape[0]:,} ch={h.feats.shape[1]} alloc={_ma()//1048576}MB", flush=True)
                if i < len(self.blocks) - 1 and j == len(res) - 1:
                    if self.pred_subdiv:
                        h, sub = block(h)
                        subs.append(sub)
                    else:
                        h = block(h, subdiv=guide_subs[i] if guide_subs is not None else None)
                    h.clear_spatial_cache()
                else:
                    h = block(h)
                pbar.update(1)

        _CHUNK = 1_000_000
        if self.low_vram and h.feats.shape[0] > _CHUNK:
            dev = h.feats.device
            norm_shape = h.feats.shape[-1:]
            h_cpu = h.feats.to('cpu', non_blocking=True)
            torch.cuda.current_stream().synchronize()
            h = h.replace(torch.empty(0, dtype=h_cpu.dtype, device=dev))
            torch.cuda.empty_cache()
            chunks = [F.layer_norm(h_cpu[i:i+_CHUNK].to(dev), norm_shape).cpu()
                      for i in range(0, h_cpu.shape[0], _CHUNK)]
            h = h.replace(torch.cat(chunks).to(dev))
            del chunks, h_cpu
        else:
            h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
        h = self.output_layer(h)

        pbar.update(1)
        if return_subs:
            return h, subs
        return h

    def upsample(self, x: sp.SparseTensor, upsample_times: int) -> torch.Tensor:
        assert self.pred_subdiv == True, "Only decoders with pred_subdiv=True can be used with upsampling"

        h = self.from_latent(x)
        for i, res in enumerate(self.blocks):
            if i == upsample_times:
                return h.coords
            for j, block in enumerate(res):
                if i < len(self.blocks) - 1 and j == len(res) - 1:
                    h, sub = block(h)
                else:
                    h = block(h)


# ============================================================================
# Section 3: FlexiDualGrid VAE (from fdg_vae.py)
# ============================================================================

from o_voxel_vb_ap.convert import tiled_flexible_dual_grid_to_mesh as _tiled_vb
import comfy.model_management as _cmm


class Mesh:
    """Mesh representation (vertices + faces + optional vertex attributes)."""
    def __init__(self, vertices, faces, vertex_attrs=None):
        self.vertices = vertices.float()
        self.faces = faces.int()
        self.vertex_attrs = vertex_attrs

    @property
    def device(self):
        return self.vertices.device

    def to(self, device, non_blocking=False):
        return Mesh(
            self.vertices.to(device, non_blocking=non_blocking),
            self.faces.to(device, non_blocking=non_blocking),
            self.vertex_attrs.to(device, non_blocking=non_blocking) if self.vertex_attrs is not None else None,
        )

    def cuda(self, non_blocking=False):
        return self.to(_cmm.get_torch_device(), non_blocking=non_blocking)

    def cpu(self):
        return self.to('cpu')

    def fill_holes(self, max_hole_perimeter=3e-2):
        import cumesh_vb as cumesh
        device = _cmm.get_torch_device()
        vertices = self.vertices.to(device)
        faces = self.faces.to(device)
        mesh = cumesh.CuMesh()
        mesh.init(vertices, faces)
        mesh.get_edges()
        mesh.get_boundary_info()
        if mesh.num_boundaries == 0:
            return
        mesh.get_vertex_edge_adjacency()
        mesh.get_vertex_boundary_adjacency()
        mesh.get_manifold_boundary_adjacency()
        mesh.read_manifold_boundary_adjacency()
        mesh.get_boundary_connected_components()
        mesh.get_boundary_loops()
        if mesh.num_boundary_loops == 0:
            return
        mesh.fill_holes(max_hole_perimeter=max_hole_perimeter)
        new_vertices, new_faces = mesh.read()
        self.vertices = new_vertices.to(self.device)
        self.faces = new_faces.to(self.device)

    def remove_faces(self, face_mask: torch.Tensor):
        import cumesh_vb as cumesh
        device = _cmm.get_torch_device()
        vertices = self.vertices.to(device)
        faces = self.faces.to(device)
        mesh = cumesh.CuMesh()
        mesh.init(vertices, faces)
        mesh.remove_faces(face_mask)
        new_vertices, new_faces = mesh.read()
        self.vertices = new_vertices.to(self.device)
        self.faces = new_faces.to(self.device)

    def simplify(self, target=1000000, verbose: bool = False, options: dict = {}):
        import cumesh_vb as cumesh
        device = _cmm.get_torch_device()
        vertices = self.vertices.to(device)
        faces = self.faces.to(device)
        mesh = cumesh.CuMesh()
        mesh.init(vertices, faces)
        mesh.simplify(target, verbose=verbose, options=options)
        new_vertices, new_faces = mesh.read()
        self.vertices = new_vertices.to(self.device)
        self.faces = new_faces.to(self.device)


class Voxel:
    """Voxel representation for 3D data."""
    def __init__(self, origin: list, voxel_size: float, coords: torch.Tensor = None,
                 attrs: torch.Tensor = None, layout: Dict = {}, device: torch.device = None):
        if device is None:
            device = _cmm.get_torch_device()
        self.origin = torch.tensor(origin, dtype=torch.float32, device=device)
        self.voxel_size = voxel_size
        self.coords = coords
        self.attrs = attrs
        self.layout = layout
        self.device = device

    @property
    def position(self):
        return (self.coords + 0.5) * self.voxel_size + self.origin[None, :]

    def split_attrs(self):
        return {k: self.attrs[:, self.layout[k]] for k in self.layout}


class MeshWithVoxel(Mesh, Voxel):
    """Combined mesh + voxel representation with PBR attribute querying."""
    def __init__(self, vertices: torch.Tensor, faces: torch.Tensor, origin: list,
                 voxel_size: float, coords: torch.Tensor, attrs: torch.Tensor,
                 voxel_shape: torch.Size, layout: Dict = {}):
        self.vertices = vertices.float()
        self.faces = faces.int()
        self.origin = torch.tensor(origin, dtype=torch.float32, device=self.device)
        self.voxel_size = voxel_size
        self.coords = coords
        self.attrs = attrs
        self.voxel_shape = voxel_shape
        self.layout = layout

    def to(self, device, non_blocking=False):
        return MeshWithVoxel(
            self.vertices.to(device, non_blocking=non_blocking),
            self.faces.to(device, non_blocking=non_blocking),
            self.origin.tolist(), self.voxel_size,
            self.coords.to(device, non_blocking=non_blocking),
            self.attrs.to(device, non_blocking=non_blocking),
            self.voxel_shape, self.layout,
        )

    def query_attrs(self, xyz):
        from flex_gemm_ap.ops.grid_sample import grid_sample_3d
        grid = ((xyz - self.origin) / self.voxel_size).reshape(1, -1, 3)
        vertex_attrs = grid_sample_3d(
            self.attrs,
            torch.cat([torch.zeros_like(self.coords[..., :1]), self.coords], dim=-1),
            self.voxel_shape, grid, mode='trilinear'
        )[0]
        return vertex_attrs

    def query_vertex_attrs(self):
        return self.query_attrs(self.vertices)


class FlexiDualGridVaeEncoder(SparseUnetVaeEncoder):
    def __init__(
        self,
        model_channels: List[int],
        latent_channels: int,
        num_blocks: List[int],
        block_type: List[str],
        down_block_type: List[str],
        block_args: List[Dict[str, Any]],
        use_fp16: bool = False,
        dtype=None,
        device=None,
        operations=ops,
        sparse_operations=sparse_ops,
    ):
        super().__init__(
            6,
            model_channels,
            latent_channels,
            num_blocks,
            block_type,
            down_block_type,
            block_args,
            use_fp16,
            dtype=dtype,
            device=device,
            operations=operations,
            sparse_operations=sparse_operations,
        )

    def forward(self, vertices: sp.SparseTensor, intersected: sp.SparseTensor, sample_posterior=False, return_raw=False):
        x = vertices.replace(torch.cat([
            vertices.feats - 0.5,
            intersected.feats.float() - 0.5,
        ], dim=1))
        return super().forward(x, sample_posterior, return_raw)


class FlexiDualGridVaeDecoder(SparseUnetVaeDecoder):
    def __init__(
        self,
        resolution: int,
        model_channels: List[int],
        latent_channels: int,
        num_blocks: List[int],
        block_type: List[str],
        up_block_type: List[str],
        block_args: List[Dict[str, Any]],
        voxel_margin: float = 0.5,
        use_fp16: bool = False,
        dtype=None,
        device=None,
        operations=ops,
        sparse_operations=sparse_ops,
    ):
        self.resolution = resolution
        self.voxel_margin = voxel_margin

        super().__init__(
            7,
            model_channels,
            latent_channels,
            num_blocks,
            block_type,
            up_block_type,
            block_args,
            use_fp16,
            dtype=dtype,
            device=device,
            operations=operations,
            sparse_operations=sparse_operations,
        )

    def set_resolution(self, resolution: int) -> None:
        self.resolution = resolution

    def forward(self, x: sp.SparseTensor, **kwargs):
        decoded = super().forward(x, **kwargs)
        out_list = list(decoded) if isinstance(decoded, tuple) else [decoded]
        h = out_list[0]

        # Free sparse conv caches and input tensor
        h.clear_spatial_cache()
        x.clear_spatial_cache()
        for item in out_list[1:]:
            if isinstance(item, (list, tuple)):
                for sub in item:
                    if hasattr(sub, 'clear_spatial_cache'):
                        sub.clear_spatial_cache()
            elif hasattr(item, 'clear_spatial_cache'):
                item.clear_spatial_cache()
        del x
        torch.cuda.empty_cache()

        # Extract the 7 output channels with activations, free h
        coords = h.coords[:, 1:]
        vertices = h.replace((1 + 2 * self.voxel_margin) * F.sigmoid(h.feats[..., 0:3]) - self.voxel_margin)
        intersected = h.replace(h.feats[..., 3:6] > 0)
        quad_lerp = h.replace(F.softplus(h.feats[..., 6:7]))
        extra = out_list[1:]
        del h, out_list

        # Free everything possible before mesh extraction
        _ma = torch.cuda.memory_allocated
        _v = lambda: _ma() // 1048576
        print(f"[FlexiDecoder] pre-cleanup: alloc={_v()}MB", flush=True)
        self.to('cpu')
        print(f"[FlexiDecoder] model to CPU: alloc={_v()}MB", flush=True)
        # Move subs to CPU
        for ei, item in enumerate(extra):
            if isinstance(item, (list, tuple)):
                for si, sub in enumerate(item):
                    if hasattr(sub, 'feats') and sub.feats.is_cuda:
                        print(f"[FlexiDecoder] extra[{ei}][{si}] feats={sub.feats.shape} on {sub.feats.device}, {sub.feats.numel()*sub.feats.element_size()//1048576}MB", flush=True)
                        extra[ei][si] = sub.replace(sub.feats.cpu())
            elif hasattr(item, 'feats') and item.feats.is_cuda:
                print(f"[FlexiDecoder] extra[{ei}] feats={item.feats.shape} on {item.feats.device}, {item.feats.numel()*item.feats.element_size()//1048576}MB", flush=True)
                extra[ei] = item.replace(item.feats.cpu())
        print(f"[FlexiDecoder] subs to CPU: alloc={_v()}MB", flush=True)
        # Log what's still on GPU
        print(f"[FlexiDecoder] coords: {coords.shape} {coords.device} {coords.numel()*coords.element_size()//1048576}MB", flush=True)
        print(f"[FlexiDecoder] vertices.feats: {vertices.feats.shape} {vertices.feats.device} {vertices.feats.numel()*vertices.feats.element_size()//1048576}MB", flush=True)
        print(f"[FlexiDecoder] intersected.feats: {intersected.feats.shape} {intersected.feats.device} {intersected.feats.numel()*intersected.feats.element_size()//1048576}MB", flush=True)
        print(f"[FlexiDecoder] quad_lerp.feats: {quad_lerp.feats.shape} {quad_lerp.feats.device} {quad_lerp.feats.numel()*quad_lerp.feats.element_size()//1048576}MB", flush=True)
        torch.cuda.empty_cache()
        print(f"[FlexiDecoder] post-cleanup: alloc={_v()}MB, starting mesh extraction", flush=True)

        mesh = [Mesh(*_tiled_vb(
            coords=coords,
            dual_vertices=v.feats,
            intersected_flag=i.feats,
            split_weight=q.feats,
            aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
            grid_size=self.resolution,
            tile_size=128,
            train=False,
        )) for v, i, q in zip(vertices, intersected, quad_lerp)]

        del coords, vertices, intersected, quad_lerp

        if extra:
            return (mesh, *extra)
        return mesh
