1ST-PLACE-WINNER/MiniMax-H3

1# SPDX-License-Identifier: Apache-2.02# Pixel normalization transforms for the MiniMax H3 visual VAE.3from typing import Tuple4from torchvision.transforms import Normalize5 6 7NORM_CONFIGS = {8    "imagenet": {9        "mean": (0.485, 0.456, 0.406),10        "std": (0.229, 0.224, 0.225),11    },12    "simple": {13        "mean": (0.5, 0.5, 0.5),14        "std": (0.5, 0.5, 0.5),15    },16    "raw": {17        "mean": (0.0, 0.0, 0.0),18        "std": (1.0, 1.0, 1.0),19    },20}21 22 23def get_norm_constants(norm_type: str = "imagenet") -> Tuple[Tuple[float, ...], Tuple[float, ...]]:24    if norm_type not in NORM_CONFIGS:25        raise ValueError(f"Unknown norm_type: {norm_type}. Must be one of {list(NORM_CONFIGS.keys())}")26    config = NORM_CONFIGS[norm_type]27    return config["mean"], config["std"]28 29 30def get_normalize_transform(norm_type: str = "imagenet") -> Normalize:31    mean, std = get_norm_constants(norm_type)32    return Normalize(mean, std)33 34 35def get_denormalize_transform(norm_type: str = "imagenet") -> Normalize:36    mean, std = get_norm_constants(norm_type)37    inv_mean = tuple(-m / s for m, s in zip(mean, std))38    inv_std = tuple(1.0 / s for s in std)39    return Normalize(inv_mean, inv_std)40