declare-lab/tango2

1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14import math15from typing import Optional16 17import numpy as np18import torch19from torch import nn20 21 22def get_timestep_embedding(23    timesteps: torch.Tensor,24    embedding_dim: int,25    flip_sin_to_cos: bool = False,26    downscale_freq_shift: float = 1,27    scale: float = 1,28    max_period: int = 10000,29):30    """31    This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.32 33    :param timesteps: a 1-D Tensor of N indices, one per batch element.34                      These may be fractional.35    :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the36    embeddings. :return: an [N x dim] Tensor of positional embeddings.37    """38    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"39 40    half_dim = embedding_dim // 241    exponent = -math.log(max_period) * torch.arange(42        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device43    )44    exponent = exponent / (half_dim - downscale_freq_shift)45 46    emb = torch.exp(exponent)47    emb = timesteps[:, None].float() * emb[None, :]48 49    # scale embeddings50    emb = scale * emb51 52    # concat sine and cosine embeddings53    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)54 55    # flip sine and cosine embeddings56    if flip_sin_to_cos:57        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)58 59    # zero pad60    if embedding_dim % 2 == 1:61        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))62    return emb63 64 65def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):66    """67    grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or68    [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)69    """70    grid_h = np.arange(grid_size, dtype=np.float32)71    grid_w = np.arange(grid_size, dtype=np.float32)72    grid = np.meshgrid(grid_w, grid_h)  # here w goes first73    grid = np.stack(grid, axis=0)74 75    grid = grid.reshape([2, 1, grid_size, grid_size])76    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)77    if cls_token and extra_tokens > 0:78        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)79    return pos_embed80 81 82def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):83    if embed_dim % 2 != 0:84        raise ValueError("embed_dim must be divisible by 2")85 86    # use half of dimensions to encode grid_h87    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)88    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)89 90    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)91    return emb92 93 94def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):95    """96    embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)97    """98    if embed_dim % 2 != 0:99        raise ValueError("embed_dim must be divisible by 2")100 101    omega = np.arange(embed_dim // 2, dtype=np.float64)102    omega /= embed_dim / 2.0103    omega = 1.0 / 10000**omega  # (D/2,)104 105    pos = pos.reshape(-1)  # (M,)106    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product107 108    emb_sin = np.sin(out)  # (M, D/2)109    emb_cos = np.cos(out)  # (M, D/2)110 111    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)112    return emb113 114 115class PatchEmbed(nn.Module):116    """2D Image to Patch Embedding"""117 118    def __init__(119        self,120        height=224,121        width=224,122        patch_size=16,123        in_channels=3,124        embed_dim=768,125        layer_norm=False,126        flatten=True,127        bias=True,128    ):129        super().__init__()130 131        num_patches = (height // patch_size) * (width // patch_size)132        self.flatten = flatten133        self.layer_norm = layer_norm134 135        self.proj = nn.Conv2d(136            in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias137        )138        if layer_norm:139            self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)140        else:141            self.norm = None142 143        pos_embed = get_2d_sincos_pos_embed(embed_dim, int(num_patches**0.5))144        self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False)145 146    def forward(self, latent):147        latent = self.proj(latent)148        if self.flatten:149            latent = latent.flatten(2).transpose(1, 2)  # BCHW -> BNC150        if self.layer_norm:151            latent = self.norm(latent)152        return latent + self.pos_embed153 154 155class TimestepEmbedding(nn.Module):156    def __init__(157        self,158        in_channels: int,159        time_embed_dim: int,160        act_fn: str = "silu",161        out_dim: int = None,162        post_act_fn: Optional[str] = None,163        cond_proj_dim=None,164    ):165        super().__init__()166 167        self.linear_1 = nn.Linear(in_channels, time_embed_dim)168 169        if cond_proj_dim is not None:170            self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)171        else:172            self.cond_proj = None173 174        if act_fn == "silu":175            self.act = nn.SiLU()176        elif act_fn == "mish":177            self.act = nn.Mish()178        elif act_fn == "gelu":179            self.act = nn.GELU()180        else:181            raise ValueError(f"{act_fn} does not exist. Make sure to define one of 'silu', 'mish', or 'gelu'")182 183        if out_dim is not None:184            time_embed_dim_out = out_dim185        else:186            time_embed_dim_out = time_embed_dim187        self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out)188 189        if post_act_fn is None:190            self.post_act = None191        elif post_act_fn == "silu":192            self.post_act = nn.SiLU()193        elif post_act_fn == "mish":194            self.post_act = nn.Mish()195        elif post_act_fn == "gelu":196            self.post_act = nn.GELU()197        else:198            raise ValueError(f"{post_act_fn} does not exist. Make sure to define one of 'silu', 'mish', or 'gelu'")199 200    def forward(self, sample, condition=None):201        if condition is not None:202            sample = sample + self.cond_proj(condition)203        sample = self.linear_1(sample)204 205        if self.act is not None:206            sample = self.act(sample)207 208        sample = self.linear_2(sample)209 210        if self.post_act is not None:211            sample = self.post_act(sample)212        return sample213 214 215class Timesteps(nn.Module):216    def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float):217        super().__init__()218        self.num_channels = num_channels219        self.flip_sin_to_cos = flip_sin_to_cos220        self.downscale_freq_shift = downscale_freq_shift221 222    def forward(self, timesteps):223        t_emb = get_timestep_embedding(224            timesteps,225            self.num_channels,226            flip_sin_to_cos=self.flip_sin_to_cos,227            downscale_freq_shift=self.downscale_freq_shift,228        )229        return t_emb230 231 232class GaussianFourierProjection(nn.Module):233    """Gaussian Fourier embeddings for noise levels."""234 235    def __init__(236        self, embedding_size: int = 256, scale: float = 1.0, set_W_to_weight=True, log=True, flip_sin_to_cos=False237    ):238        super().__init__()239        self.weight = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False)240        self.log = log241        self.flip_sin_to_cos = flip_sin_to_cos242 243        if set_W_to_weight:244            # to delete later245            self.W = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False)246 247            self.weight = self.W248 249    def forward(self, x):250        if self.log:251            x = torch.log(x)252 253        x_proj = x[:, None] * self.weight[None, :] * 2 * np.pi254 255        if self.flip_sin_to_cos:256            out = torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1)257        else:258            out = torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)259        return out260 261 262class ImagePositionalEmbeddings(nn.Module):263    """264    Converts latent image classes into vector embeddings. Sums the vector embeddings with positional embeddings for the265    height and width of the latent space.266 267    For more details, see figure 10 of the dall-e paper: https://arxiv.org/abs/2102.12092268 269    For VQ-diffusion:270 271    Output vector embeddings are used as input for the transformer.272 273    Note that the vector embeddings for the transformer are different than the vector embeddings from the VQVAE.274 275    Args:276        num_embed (`int`):277            Number of embeddings for the latent pixels embeddings.278        height (`int`):279            Height of the latent image i.e. the number of height embeddings.280        width (`int`):281            Width of the latent image i.e. the number of width embeddings.282        embed_dim (`int`):283            Dimension of the produced vector embeddings. Used for the latent pixel, height, and width embeddings.284    """285 286    def __init__(287        self,288        num_embed: int,289        height: int,290        width: int,291        embed_dim: int,292    ):293        super().__init__()294 295        self.height = height296        self.width = width297        self.num_embed = num_embed298        self.embed_dim = embed_dim299 300        self.emb = nn.Embedding(self.num_embed, embed_dim)301        self.height_emb = nn.Embedding(self.height, embed_dim)302        self.width_emb = nn.Embedding(self.width, embed_dim)303 304    def forward(self, index):305        emb = self.emb(index)306 307        height_emb = self.height_emb(torch.arange(self.height, device=index.device).view(1, self.height))308 309        # 1 x H x D -> 1 x H x 1 x D310        height_emb = height_emb.unsqueeze(2)311 312        width_emb = self.width_emb(torch.arange(self.width, device=index.device).view(1, self.width))313 314        # 1 x W x D -> 1 x 1 x W x D315        width_emb = width_emb.unsqueeze(1)316 317        pos_emb = height_emb + width_emb318 319        # 1 x H x W x D -> 1 x L xD320        pos_emb = pos_emb.view(1, self.height * self.width, -1)321 322        emb = emb + pos_emb[:, : emb.shape[1], :]323 324        return emb325 326 327class LabelEmbedding(nn.Module):328    """329    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.330 331    Args:332        num_classes (`int`): The number of classes.333        hidden_size (`int`): The size of the vector embeddings.334        dropout_prob (`float`): The probability of dropping a label.335    """336 337    def __init__(self, num_classes, hidden_size, dropout_prob):338        super().__init__()339        use_cfg_embedding = dropout_prob > 0340        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)341        self.num_classes = num_classes342        self.dropout_prob = dropout_prob343 344    def token_drop(self, labels, force_drop_ids=None):345        """346        Drops labels to enable classifier-free guidance.347        """348        if force_drop_ids is None:349            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob350        else:351            drop_ids = torch.tensor(force_drop_ids == 1)352        labels = torch.where(drop_ids, self.num_classes, labels)353        return labels354 355    def forward(self, labels: torch.LongTensor, force_drop_ids=None):356        use_dropout = self.dropout_prob > 0357        if (self.training and use_dropout) or (force_drop_ids is not None):358            labels = self.token_drop(labels, force_drop_ids)359        embeddings = self.embedding_table(labels)360        return embeddings361 362 363class CombinedTimestepLabelEmbeddings(nn.Module):364    def __init__(self, num_classes, embedding_dim, class_dropout_prob=0.1):365        super().__init__()366 367        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1)368        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)369        self.class_embedder = LabelEmbedding(num_classes, embedding_dim, class_dropout_prob)370 371    def forward(self, timestep, class_labels, hidden_dtype=None):372        timesteps_proj = self.time_proj(timestep)373        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype))  # (N, D)374 375        class_labels = self.class_embedder(class_labels)  # (N, D)376 377        conditioning = timesteps_emb + class_labels  # (N, D)378 379        return conditioning380