mld/models/architectures/mld_denoiser.py

from typing import Optional, Union

import torch
import torch.nn as nn

from mld.models.operator.embeddings import TimestepEmbedding, Timesteps
from mld.models.operator.attention import (SkipTransformerEncoder,
                                           SkipTransformerDecoder,
                                           TransformerDecoder,
                                           TransformerDecoderLayer,
                                           TransformerEncoder,
                                           TransformerEncoderLayer)
from mld.models.operator.moe import MoeTransformerEncoderLayer, MoeTransformerDecoderLayer
from mld.models.operator.utils import get_clones, get_activation_fn, zero_module
from mld.models.operator.position_encoding import build_position_encoding


def load_balancing_loss_func(router_logits: tuple, num_experts: int = 4, topk: int = 2):
    router_logits = torch.cat(router_logits, dim=0)
    routing_weights = torch.nn.functional.softmax(router_logits, dim=-1)
    _, selected_experts = torch.topk(routing_weights, topk, dim=-1)
    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
    tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
    router_prob_per_expert = torch.mean(routing_weights, dim=0)
    overall_loss = num_experts * torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
    return overall_loss


class MldDenoiser(nn.Module):

    def __init__(self,
                 latent_dim: list = [1, 256],
                 hidden_dim: Optional[int] = None,
                 text_dim: int = 768,
                 time_dim: int = 768,
                 ff_size: int = 1024,
                 num_layers: int = 9,
                 num_heads: int = 4,
                 dropout: float = 0.1,
                 normalize_before: bool = False,
                 norm_eps: float = 1e-5,
                 activation: str = "gelu",
                 norm_post: bool = True,
                 activation_post: Optional[str] = None,
                 flip_sin_to_cos: bool = True,
                 freq_shift: float = 0,
                 time_act_fn: str = 'silu',
                 time_post_act_fn: Optional[str] = None,
                 position_embedding: str = "learned",
                 arch: str = "trans_enc",
                 add_mem_pos: bool = True,
                 force_pre_post_proj: bool = False,
                 text_act_fn: str = 'relu',
                 time_cond_proj_dim: Optional[int] = None,
                 zero_init_cond: bool = True,
                 is_controlnet: bool = False,
                 controlnet_embed_dim: Optional[int] = None,
                 controlnet_act_fn: str = 'silu',
                 moe: bool = False,
                 moe_num_experts: int = 4,
                 moe_topk: int = 2,
                 moe_loss_weight: float = 1e-2,
                 moe_jitter_noise: Optional[float] = None
                 ) -> None:
        super(MldDenoiser, self).__init__()

        self.latent_dim = latent_dim[-1] if hidden_dim is None else hidden_dim
        add_pre_post_proj = force_pre_post_proj or (hidden_dim is not None and hidden_dim != latent_dim[-1])
        self.latent_pre = nn.Linear(latent_dim[-1], self.latent_dim) if add_pre_post_proj else nn.Identity()
        self.latent_post = nn.Linear(self.latent_dim, latent_dim[-1]) if add_pre_post_proj else nn.Identity()

        self.arch = arch
        self.time_cond_proj_dim = time_cond_proj_dim

        self.moe_num_experts = moe_num_experts
        self.moe_topk = moe_topk
        self.moe_loss_weight = moe_loss_weight

        self.time_proj = Timesteps(time_dim, flip_sin_to_cos, freq_shift)
        self.time_embedding = TimestepEmbedding(time_dim, self.latent_dim, time_act_fn, post_act_fn=time_post_act_fn,
                                                cond_proj_dim=time_cond_proj_dim, zero_init_cond=zero_init_cond)
        self.emb_proj = nn.Sequential(get_activation_fn(text_act_fn), nn.Linear(text_dim, self.latent_dim))

        self.query_pos = build_position_encoding(self.latent_dim, position_embedding=position_embedding)
        if self.arch == "trans_enc":
            if moe:
                encoder_layer = MoeTransformerEncoderLayer(
                    self.latent_dim, num_heads, moe_num_experts, moe_topk, ff_size,
                    dropout, activation, normalize_before, norm_eps, moe_jitter_noise)
            else:
                encoder_layer = TransformerEncoderLayer(
                    self.latent_dim, num_heads, ff_size, dropout,
                    activation, normalize_before, norm_eps)

            encoder_norm = nn.LayerNorm(self.latent_dim, eps=norm_eps) if norm_post and not is_controlnet else None
            self.encoder = SkipTransformerEncoder(encoder_layer, num_layers, encoder_norm, activation_post,
                                                  is_controlnet=is_controlnet, is_moe=moe)

        elif self.arch == 'trans_dec':
            if add_mem_pos:
                self.mem_pos = build_position_encoding(self.latent_dim, position_embedding=position_embedding)
            else:
                self.mem_pos = None
            if moe:
                decoder_layer = MoeTransformerDecoderLayer(
                    self.latent_dim, num_heads, moe_num_experts, moe_topk, ff_size,
                    dropout, activation, normalize_before, norm_eps, moe_jitter_noise)
            else:
                decoder_layer = TransformerDecoderLayer(
                    self.latent_dim, num_heads, ff_size, dropout,
                    activation, normalize_before, norm_eps)

            decoder_norm = nn.LayerNorm(self.latent_dim, eps=norm_eps) if norm_post and not is_controlnet else None
            self.decoder = SkipTransformerDecoder(decoder_layer, num_layers, decoder_norm, activation_post,
                                                  is_controlnet=is_controlnet, is_moe=moe)
        else:
            raise ValueError(f"Not supported architecture: {self.arch}!")

        self.is_controlnet = is_controlnet
        if self.is_controlnet:
            embed_dim = controlnet_embed_dim if controlnet_embed_dim is not None else self.latent_dim
            modules = [
                nn.Linear(latent_dim[-1], embed_dim),
                get_activation_fn(controlnet_act_fn) if controlnet_act_fn else None,
                nn.Linear(embed_dim, embed_dim),
                get_activation_fn(controlnet_act_fn) if controlnet_act_fn else None,
                zero_module(nn.Linear(embed_dim, latent_dim[-1]))
            ]
            self.controlnet_cond_embedding = nn.Sequential(*[m for m in modules if m is not None])

            self.controlnet_down_mid_blocks = nn.ModuleList([
                zero_module(nn.Linear(self.latent_dim, self.latent_dim)) for _ in range(num_layers)])

    def forward(self,
                sample: torch.Tensor,
                timestep: torch.Tensor,
                encoder_hidden_states: torch.Tensor,
                timestep_cond: Optional[torch.Tensor] = None,
                controlnet_cond: Optional[torch.Tensor] = None,
                controlnet_residuals: Optional[list[torch.Tensor]] = None
                ) -> tuple:

        # 0. check if controlnet
        if self.is_controlnet:
            sample = sample + self.controlnet_cond_embedding(controlnet_cond)

        # 1. dimension matching (pre)
        sample = sample.permute(1, 0, 2)
        sample = self.latent_pre(sample)

        # 2. time_embedding
        timesteps = timestep.expand(sample.shape[1]).clone()
        time_emb = self.time_proj(timesteps)
        time_emb = time_emb.to(dtype=sample.dtype)
        # [1, bs, latent_dim] <= [bs, latent_dim]
        time_emb = self.time_embedding(time_emb, timestep_cond).unsqueeze(0)

        # 3. condition + time embedding
        # text_emb [seq_len, batch_size, text_dim] <= [batch_size, seq_len, text_dim]
        encoder_hidden_states = encoder_hidden_states.permute(1, 0, 2)
        # text embedding projection
        text_emb_latent = self.emb_proj(encoder_hidden_states)
        emb_latent = torch.cat((time_emb, text_emb_latent), 0)

        # 4. transformer
        if self.arch == "trans_enc":
            xseq = torch.cat((sample, emb_latent), axis=0)
            xseq = self.query_pos(xseq)
            tokens, intermediates, router_logits = self.encoder(xseq, controlnet_residuals=controlnet_residuals)
        elif self.arch == 'trans_dec':
            sample = self.query_pos(sample)
            if self.mem_pos:
                emb_latent = self.mem_pos(emb_latent)
            tokens, intermediates, router_logits = self.decoder(sample, emb_latent,
                                                                controlnet_residuals=controlnet_residuals)
        else:
            raise TypeError(f"{self.arch} is not supported")

        router_loss = None
        if router_logits is not None:
            router_loss = load_balancing_loss_func(router_logits, self.moe_num_experts, self.moe_topk)
            router_loss = self.moe_loss_weight * router_loss

        if self.is_controlnet:
            control_res_samples = []
            for res, block in zip(intermediates, self.controlnet_down_mid_blocks):
                r = block(res)
                control_res_samples.append(r)
            return control_res_samples, router_loss
        elif self.arch == "trans_enc":
            sample = tokens[:sample.shape[0]]
        elif self.arch == 'trans_dec':
            sample = tokens
        else:
            raise TypeError(f"{self.arch} is not supported")

        # 5. dimension matching (post)
        sample = self.latent_post(sample)
        sample = sample.permute(1, 0, 2)
        return sample, router_loss