Ref2VA/video_vae/normalize.py
| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # Pixel normalization transforms for the MiniMax H3 visual VAE. |
| 3 | from typing import Tuple |
| 4 | from torchvision.transforms import Normalize |
| 5 | |
| 6 | |
| 7 | NORM_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 | |
| 23 | def 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 | |
| 30 | def get_normalize_transform(norm_type: str = "imagenet") -> Normalize: |
| 31 | mean, std = get_norm_constants(norm_type) |
| 32 | return Normalize(mean, std) |
| 33 | |
| 34 | |
| 35 | def 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 | |