nvidia/C-RADIO

1# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto.  Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8 9from typing import Union, Tuple10 11import torch12from torch import nn13 14 15norm_t = Union[Tuple[float, float, float], torch.Tensor]16 17class InputConditioner(nn.Module):18    def __init__(self,19                 input_scale: float,20                 norm_mean: norm_t,21                 norm_std: norm_t,22                 dtype: torch.dtype = None,23    ):24        super().__init__()25 26        self.dtype = dtype27 28        self.register_buffer("norm_mean", _to_tensor(norm_mean) / input_scale)29        self.register_buffer("norm_std", _to_tensor(norm_std) / input_scale)30 31    def forward(self, x: torch.Tensor):32        y = (x - self.norm_mean) / self.norm_std33        if self.dtype is not None:34            y = y.to(self.dtype)35        return y36 37 38def get_default_conditioner():39    from timm.data.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD40 41    return InputConditioner(42        input_scale=1.0,43        norm_mean=OPENAI_CLIP_MEAN,44        norm_std=OPENAI_CLIP_STD,45    )46 47 48def _to_tensor(v: norm_t):49    return torch.as_tensor(v, dtype=torch.float32).view(-1, 1, 1)50