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