1ST-PLACE-WINNER/MiniMax-H3

1# SPDX-License-Identifier: Apache-2.02# VAE distribution and aggregation helpers for the MiniMax H3 visual VAE.3import torch4 5 6class DiagonalGaussianDistribution(object):7    def __init__(self, parameters, upcast_fp32=True):8        if upcast_fp32:9            parameters = parameters.to(dtype=torch.float32)10 11        self.parameters = parameters12        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)13        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)14        self.std = torch.exp(0.5 * self.logvar)15        self.var = torch.exp(self.logvar)16 17    @torch.compiler.disable18    def sample(self, generator=None):19        noise = torch.randn(self.mean.shape, generator=generator)20        x = self.mean + self.std * noise.to(device=self.parameters.device)21        return x22 23 24class ClsTokenAggregator:25    def __init__(self, vae_model):26        self.vae = vae_model27        self.cls_tokens = []28 29    def __enter__(self):30        return self31 32    def __exit__(self, exc_type, exc_val, exc_tb):33        if self.cls_tokens and hasattr(self.vae.encoder, "loss_info"):34            self.vae.encoder.loss_info["cls_token"] = torch.stack(35                self.cls_tokens, dim=036            ).mean(dim=0)37        return False38 39    def collect(self):40        if (41            hasattr(self.vae.encoder, "loss_info")42            and "cls_token" in self.vae.encoder.loss_info43        ):44            self.cls_tokens.append(self.vae.encoder.loss_info["cls_token"].clone())45 46    def collect_stacked(self, num_tiles, sample_batch_size):47        if (48            hasattr(self.vae.encoder, "loss_info")49            and "cls_token" in self.vae.encoder.loss_info50        ):51            cls_token = self.vae.encoder.loss_info["cls_token"]52            cls_token = cls_token.unflatten(0, (num_tiles, sample_batch_size))53            self.cls_tokens.extend(token.clone() for token in cls_token)54