wusize/Harmon-0_5B

1import torch2import torch.nn as nn3from torch.utils.checkpoint import checkpoint4import math5 6from .misc import create_diffusion7 8 9class DiffLoss(nn.Module):10    """Diffusion Loss"""11    def __init__(self, target_channels, z_channels, depth, width, num_sampling_steps, grad_checkpointing=False):12        super(DiffLoss, self).__init__()13        self.in_channels = target_channels14        self.net = SimpleMLPAdaLN(15            in_channels=target_channels,16            model_channels=width,17            out_channels=target_channels * 2,  # for vlb loss18            z_channels=z_channels,19            num_res_blocks=depth,20            grad_checkpointing=grad_checkpointing21        )22 23        self.train_diffusion = create_diffusion(timestep_respacing="", noise_schedule="cosine")24        self.gen_diffusion = create_diffusion(timestep_respacing=num_sampling_steps, noise_schedule="cosine")25 26    def forward(self, target, z, mask=None):27        t = torch.randint(0, self.train_diffusion.num_timesteps, (target.shape[0],), device=target.device)28        model_kwargs = dict(c=z)29        loss_dict = self.train_diffusion.training_losses(self.net, target, t, model_kwargs)30        loss = loss_dict["loss"]31        if mask is not None:32            loss = (loss * mask).sum() / mask.sum()33        return loss.mean()34 35    def sample(self, z, temperature=1.0, cfg=1.0):36        # diffusion loss sampling37        if not cfg == 1.0:38            noise = torch.randn(z.shape[0] // 2, self.in_channels).cuda()39            noise = torch.cat([noise, noise], dim=0)40            model_kwargs = dict(c=z, cfg_scale=cfg)41            sample_fn = self.net.forward_with_cfg42        else:43            noise = torch.randn(z.shape[0], self.in_channels).cuda()44            model_kwargs = dict(c=z)45            sample_fn = self.net.forward46 47        sampled_token_latent = self.gen_diffusion.p_sample_loop(48            sample_fn, noise.shape, noise, clip_denoised=False, model_kwargs=model_kwargs, progress=False,49            temperature=temperature50        )51 52        return sampled_token_latent53 54 55def modulate(x, shift, scale):56    return x * (1 + scale) + shift57 58 59class TimestepEmbedder(nn.Module):60    """61    Embeds scalar timesteps into vector representations.62    """63    def __init__(self, hidden_size, frequency_embedding_size=256):64        super().__init__()65        self.mlp = nn.Sequential(66            nn.Linear(frequency_embedding_size, hidden_size, bias=True),67            nn.SiLU(),68            nn.Linear(hidden_size, hidden_size, bias=True),69        )70        self.frequency_embedding_size = frequency_embedding_size71 72    @staticmethod73    def timestep_embedding(t, dim, max_period=10000):74        """75        Create sinusoidal timestep embeddings.76        :param t: a 1-D Tensor of N indices, one per batch element.77                          These may be fractional.78        :param dim: the dimension of the output.79        :param max_period: controls the minimum frequency of the embeddings.80        :return: an (N, D) Tensor of positional embeddings.81        """82        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py83        half = dim // 284        freqs = torch.exp(85            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half86        ).to(device=t.device)87        args = t[:, None].float() * freqs[None]88        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)89        if dim % 2:90            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)91        return embedding92 93    def forward(self, t):94        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)95        t_emb = self.mlp(t_freq.to(self.mlp[0].weight.data.dtype))96        return t_emb97 98 99class ResBlock(nn.Module):100    """101    A residual block that can optionally change the number of channels.102    :param channels: the number of input channels.103    """104 105    def __init__(106        self,107        channels108    ):109        super().__init__()110        self.channels = channels111 112        self.in_ln = nn.LayerNorm(channels, eps=1e-6)113        self.mlp = nn.Sequential(114            nn.Linear(channels, channels, bias=True),115            nn.SiLU(),116            nn.Linear(channels, channels, bias=True),117        )118 119        self.adaLN_modulation = nn.Sequential(120            nn.SiLU(),121            nn.Linear(channels, 3 * channels, bias=True)122        )123 124    def forward(self, x, y):125        shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1)126        h = modulate(self.in_ln(x), shift_mlp, scale_mlp)127        h = self.mlp(h)128        return x + gate_mlp * h129 130 131class FinalLayer(nn.Module):132    """133    The final layer adopted from DiT.134    """135    def __init__(self, model_channels, out_channels):136        super().__init__()137        self.norm_final = nn.LayerNorm(model_channels, elementwise_affine=False, eps=1e-6)138        self.linear = nn.Linear(model_channels, out_channels, bias=True)139        self.adaLN_modulation = nn.Sequential(140            nn.SiLU(),141            nn.Linear(model_channels, 2 * model_channels, bias=True)142        )143 144    def forward(self, x, c):145        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)146        x = modulate(self.norm_final(x), shift, scale)147        x = self.linear(x)148        return x149 150 151class SimpleMLPAdaLN(nn.Module):152    """153    The MLP for Diffusion Loss.154    :param in_channels: channels in the input Tensor.155    :param model_channels: base channel count for the model.156    :param out_channels: channels in the output Tensor.157    :param z_channels: channels in the condition.158    :param num_res_blocks: number of residual blocks per downsample.159    """160 161    def __init__(162        self,163        in_channels,164        model_channels,165        out_channels,166        z_channels,167        num_res_blocks,168        grad_checkpointing=False169    ):170        super().__init__()171 172        self.in_channels = in_channels173        self.model_channels = model_channels174        self.out_channels = out_channels175        self.num_res_blocks = num_res_blocks176        self.grad_checkpointing = grad_checkpointing177 178        self.time_embed = TimestepEmbedder(model_channels)179        self.cond_embed = nn.Linear(z_channels, model_channels)180 181        self.input_proj = nn.Linear(in_channels, model_channels)182 183        res_blocks = []184        for i in range(num_res_blocks):185            res_blocks.append(ResBlock(186                model_channels,187            ))188 189        self.res_blocks = nn.ModuleList(res_blocks)190        self.final_layer = FinalLayer(model_channels, out_channels)191 192        self.initialize_weights()193 194    def initialize_weights(self):195        def _basic_init(module):196            if isinstance(module, nn.Linear):197                torch.nn.init.xavier_uniform_(module.weight)198                if module.bias is not None:199                    nn.init.constant_(module.bias, 0)200        self.apply(_basic_init)201 202        # Initialize timestep embedding MLP203        nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02)204        nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02)205 206        # Zero-out adaLN modulation layers207        for block in self.res_blocks:208            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)209            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)210 211        # Zero-out output layers212        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)213        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)214        nn.init.constant_(self.final_layer.linear.weight, 0)215        nn.init.constant_(self.final_layer.linear.bias, 0)216 217    def forward(self, x, t, c):218        """219        Apply the model to an input batch.220        :param x: an [N x C] Tensor of inputs.221        :param t: a 1-D batch of timesteps.222        :param c: conditioning from AR transformer.223        :return: an [N x C] Tensor of outputs.224        """225        # import pdb; pdb.set_trace()226        x = self.input_proj(x.to(self.input_proj.weight.data.dtype))227        t = self.time_embed(t)228        c = self.cond_embed(c.to(self.cond_embed.weight.data.dtype))229 230        y = t + c231 232        if self.grad_checkpointing and not torch.jit.is_scripting():233            for block in self.res_blocks:234                x = checkpoint(block, x, y)235        else:236            for block in self.res_blocks:237                x = block(x, y)238 239        return self.final_layer(x, y)240 241    def forward_with_cfg(self, x, t, c, cfg_scale):242        half = x[: len(x) // 2]243        combined = torch.cat([half, half], dim=0)244        model_out = self.forward(combined, t, c)245        eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]246        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)247        half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)248        eps = torch.cat([half_eps, half_eps], dim=0)249        return torch.cat([eps, rest], dim=1)250