1ST-PLACE-WINNER/MiniMax-H3

1# SPDX-License-Identifier: Apache-2.02# Module helpers for the MiniMax H3 visual VAE (inference-only bundle).3 4 5def apply_spatial_parallel(module, enabled, chunk_dim=-1):6    from .conv import SpatialParallelConv3d7    from .norm import FusedGroupNorm3D, SpatialParallelGroupNorm8 9    if hasattr(module, "set_spatial_parallel"):10        module.set_spatial_parallel(enabled)11    for m in module.modules():12        if enabled and isinstance(m, FusedGroupNorm3D):13            raise NotImplementedError("FusedGroupNorm3D is incompatible with SP")14        if isinstance(m, SpatialParallelGroupNorm):15            m.spatial_parallel = enabled16        elif isinstance(m, SpatialParallelConv3d):17            m.spatial_parallel = enabled18            m.chunk_dim = chunk_dim19