BiliSakura/BitDance-Tokenizer-diffusers

1from __future__ import annotations2 3import math4from typing import Optional5 6import torch7import torch.nn.functional as F8from torch import nn9 10from diffusers.configuration_utils import ConfigMixin, register_to_config11from diffusers.models.modeling_utils import ModelMixin12 13try:14    from flash_attn import flash_attn_func  # type: ignore15 16    _FLASH_ATTN_AVAILABLE = True17except ImportError:18    flash_attn_func = None  # type: ignore19    _FLASH_ATTN_AVAILABLE = False20 21 22def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000, time_factor: float = 1000.0) -> torch.Tensor:23    half = dim // 224    t = time_factor * t.float()25    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half)26 27    args = t[:, None] * freqs[None]28    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)29    if dim % 2:30        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)31    if torch.is_floating_point(t):32        embedding = embedding.to(t)33    return embedding34 35 36def time_shift_func(t: torch.Tensor, flow_shift: float = 1.0, sigma: float = 1.0) -> torch.Tensor:37    return (1.0 / flow_shift) / ((1.0 / flow_shift) + (1.0 / t - 1.0) ** sigma)38 39 40def get_score_from_velocity(velocity: torch.Tensor, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:41    alpha_t, d_alpha_t = t, 142    sigma_t, d_sigma_t = 1 - t, -143    mean = x44    reverse_alpha_ratio = alpha_t / d_alpha_t45    var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t46    score = (reverse_alpha_ratio * velocity - mean) / var47    return score48 49 50def get_velocity_from_cfg(velocity: torch.Tensor, cfg: float, cfg_mult: int) -> torch.Tensor:51    if cfg_mult == 2:52        cond_v, uncond_v = torch.chunk(velocity, 2, dim=0)53        velocity = uncond_v + cfg * (cond_v - uncond_v)54    return velocity55 56 57def _randn_like(x: torch.Tensor, generator: Optional[torch.Generator]) -> torch.Tensor:58    if generator is None:59        return torch.randn_like(x)60    return torch.randn(x.shape, device=x.device, dtype=x.dtype, generator=generator)61 62 63def euler_step(x: torch.Tensor, v: torch.Tensor, dt: float, cfg: float, cfg_mult: int) -> torch.Tensor:64    with torch.amp.autocast("cuda", enabled=False):65        v = v.to(torch.float32)66        v = get_velocity_from_cfg(v, cfg, cfg_mult)67        x = x + v * dt68    return x69 70 71def euler_maruyama_step(72    x: torch.Tensor,73    v: torch.Tensor,74    t: torch.Tensor,75    dt: float,76    cfg: float,77    cfg_mult: int,78    generator: Optional[torch.Generator],79) -> torch.Tensor:80    with torch.amp.autocast("cuda", enabled=False):81        v = v.to(torch.float32)82        v = get_velocity_from_cfg(v, cfg, cfg_mult)83        score = get_score_from_velocity(v, x, t)84        drift = v + (1 - t) * score85        noise_scale = (2.0 * (1.0 - t) * dt) ** 0.586        x = x + drift * dt + noise_scale * _randn_like(x, generator=generator)87    return x88 89 90def euler_maruyama(91    input_dim: int,92    forward_fn,93    c: torch.Tensor,94    cfg: float = 1.0,95    num_sampling_steps: int = 20,96    last_step_size: float = 0.05,97    time_shift: float = 1.0,98    generator: Optional[torch.Generator] = None,99) -> torch.Tensor:100    cfg_mult = 1101    if cfg > 1.0:102        cfg_mult += 1103 104    x_shape = list(c.shape)105    x_shape[0] = x_shape[0] // cfg_mult106    x_shape[-1] = input_dim107    x = torch.randn(x_shape, device=c.device, dtype=c.dtype, generator=generator)108 109    t_all = torch.linspace(0, 1 - last_step_size, num_sampling_steps + 1, device=c.device, dtype=torch.float32)110    t_all = time_shift_func(t_all, time_shift)111    dt = t_all[1:] - t_all[:-1]112 113    t = torch.tensor(0.0, device=c.device, dtype=torch.float32)114    t_batch = torch.zeros(c.shape[0], device=c.device, dtype=c.dtype)115    for i in range(num_sampling_steps):116        t_batch[:] = t117        combined = torch.cat([x] * cfg_mult, dim=0)118        output = forward_fn(combined, t_batch, c)119        if output.dim() == 2:120            v = (output - combined) / (1 - t_batch.view(-1, 1)).clamp_min(0.05)121        elif output.dim() == 3:122            v = (output - combined) / (1 - t_batch.view(-1, 1, 1)).clamp_min(0.05)123        else:124            raise ValueError(f"Unsupported output rank from diffusion head: {output.dim()}")125 126        x = euler_maruyama_step(x, v, t, float(dt[i]), cfg, cfg_mult, generator=generator)127        t += dt[i]128 129    combined = torch.cat([x] * cfg_mult, dim=0)130    t_batch[:] = 1 - last_step_size131    output = forward_fn(combined, t_batch, c)132    if output.dim() == 2:133        v = (output - combined) / (1 - t_batch.view(-1, 1)).clamp_min(0.05)134    elif output.dim() == 3:135        v = (output - combined) / (1 - t_batch.view(-1, 1, 1)).clamp_min(0.05)136    else:137        raise ValueError(f"Unsupported output rank from diffusion head: {output.dim()}")138 139    x = euler_step(x, v, last_step_size, cfg, cfg_mult)140    return torch.cat([x] * cfg_mult, dim=0)141 142 143class TimestepEmbedder(nn.Module):144    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256) -> None:145        super().__init__()146        self.mlp = nn.Sequential(147            nn.Linear(frequency_embedding_size, hidden_size, bias=True),148            nn.SiLU(),149            nn.Linear(hidden_size, hidden_size, bias=True),150        )151        self.frequency_embedding_size = frequency_embedding_size152 153    def forward(self, t: torch.Tensor) -> torch.Tensor:154        t_freq = timestep_embedding(t, self.frequency_embedding_size)155        t_freq = t_freq.to(self.mlp[0].weight.dtype)156        return self.mlp(t_freq)157 158 159class FinalLayer(nn.Module):160    def __init__(self, channels: int, out_channels: int) -> None:161        super().__init__()162        self.norm_final = nn.LayerNorm(channels, eps=1e-6, elementwise_affine=False)163        self.ada_ln_modulation = nn.Linear(channels, channels * 2, bias=True)164        self.linear = nn.Linear(channels, out_channels, bias=True)165 166    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:167        scale, shift = self.ada_ln_modulation(y).chunk(2, dim=-1)168        x = self.norm_final(x) * (1.0 + scale) + shift169        return self.linear(x)170 171 172class Attention(nn.Module):173    def __init__(self, dim: int, n_head: int) -> None:174        super().__init__()175        if dim % n_head != 0:176            raise ValueError(f"dim ({dim}) must be divisible by n_head ({n_head}).")177 178        self.dim = dim179        self.head_dim = dim // n_head180        self.n_head = n_head181        total_kv_dim = (self.n_head * 3) * self.head_dim182        self.wqkv = nn.Linear(dim, total_kv_dim, bias=True)183        self.wo = nn.Linear(dim, dim, bias=True)184 185    def forward(self, x: torch.Tensor) -> torch.Tensor:186        bsz, seqlen, _ = x.shape187        xq, xk, xv = self.wqkv(x).chunk(3, dim=-1)188 189        xq = xq.view(bsz, seqlen, self.n_head, self.head_dim)190        xk = xk.view(bsz, seqlen, self.n_head, self.head_dim)191        xv = xv.view(bsz, seqlen, self.n_head, self.head_dim)192 193        if _FLASH_ATTN_AVAILABLE and xq.is_cuda:194            output = flash_attn_func(xq, xk, xv, causal=False)195        else:196            xq = xq.transpose(1, 2)197            xk = xk.transpose(1, 2)198            xv = xv.transpose(1, 2)199            output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=0.0, is_causal=False)200            output = output.transpose(1, 2).contiguous()201 202        output = output.view(bsz, seqlen, self.dim)203        return self.wo(output)204 205 206class TransBlock(nn.Module):207    def __init__(self, channels: int, use_swiglu: bool = False) -> None:208        super().__init__()209        self.channels = channels210        self.norm1 = nn.LayerNorm(channels, eps=1e-6, elementwise_affine=True)211        self.attn = Attention(channels, n_head=channels // 128)212 213        self.norm2 = nn.LayerNorm(channels, eps=1e-6, elementwise_affine=True)214        hidden_dim = int(channels * 1.5)215        self.use_swiglu = use_swiglu216        if not use_swiglu:217            self.mlp = nn.Sequential(218                nn.Linear(channels, hidden_dim),219                nn.SiLU(),220                nn.Linear(hidden_dim, channels),221            )222        else:223            self.w1 = nn.Linear(channels, hidden_dim * 2, bias=True)224            self.w2 = nn.Linear(hidden_dim, channels, bias=True)225 226    def forward(227        self,228        x: torch.Tensor,229        scale1: torch.Tensor,230        shift1: torch.Tensor,231        gate1: torch.Tensor,232        scale2: torch.Tensor,233        shift2: torch.Tensor,234        gate2: torch.Tensor,235    ) -> torch.Tensor:236        h = self.norm1(x) * (1 + scale1) + shift1237        h = self.attn(h)238        x = x + h * gate1239 240        h = self.norm2(x) * (1 + scale2) + shift2241        if not self.use_swiglu:242            h = self.mlp(h)243        else:244            h1, h2 = self.w1(h).chunk(2, dim=-1)245            h = self.w2(F.silu(h1) * h2)246        return x + h * gate2247 248 249class TransEncoder(nn.Module):250    def __init__(251        self,252        in_channels: int,253        model_channels: int,254        z_channels: int,255        num_res_blocks: int,256        num_ada_ln_blocks: int = 2,257        grad_checkpointing: bool = False,258        parallel_num: int = 4,259        use_swiglu: bool = False,260    ) -> None:261        super().__init__()262        self.in_channels = in_channels263        self.model_channels = model_channels264        self.out_channels = in_channels265        self.num_res_blocks = num_res_blocks266        self.grad_checkpointing = grad_checkpointing267        self.parallel_num = parallel_num268 269        self.time_embed = TimestepEmbedder(model_channels)270        self.cond_embed = nn.Linear(z_channels, model_channels)271        self.input_proj = nn.Linear(in_channels, model_channels)272 273        self.res_blocks = nn.ModuleList([TransBlock(model_channels, use_swiglu) for _ in range(num_res_blocks)])274        self.ada_ln_blocks = nn.ModuleList(275            [nn.Linear(model_channels, model_channels * 6, bias=True) for _ in range(num_ada_ln_blocks)]276        )277        self.ada_ln_switch_freq = max(1, num_res_blocks // num_ada_ln_blocks)278        if (num_res_blocks % self.ada_ln_switch_freq) != 0:279            raise ValueError("num_res_blocks must be divisible by num_ada_ln_blocks")280 281        self.final_layer = FinalLayer(model_channels, self.out_channels)282        self.initialize_weights()283 284    def initialize_weights(self) -> None:285        def _basic_init(module: nn.Module) -> None:286            if isinstance(module, nn.Linear):287                nn.init.xavier_uniform_(module.weight)288                if module.bias is not None:289                    nn.init.constant_(module.bias, 0)290 291        self.apply(_basic_init)292        nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02)293        nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02)294 295        for block in self.ada_ln_blocks:296            nn.init.constant_(block.weight, 0)297            nn.init.constant_(block.bias, 0)298 299        nn.init.constant_(self.final_layer.ada_ln_modulation.weight, 0)300        nn.init.constant_(self.final_layer.ada_ln_modulation.bias, 0)301        nn.init.constant_(self.final_layer.linear.weight, 0)302        nn.init.constant_(self.final_layer.linear.bias, 0)303 304    def forward(self, x: torch.Tensor, t: torch.Tensor, c: torch.Tensor) -> torch.Tensor:305        dtype = next(self.parameters()).dtype306        x = x.to(dtype)307        t = t.to(dtype)308        c = c.to(dtype)309        x = self.input_proj(x)310        t = self.time_embed(t).unsqueeze(1)311        c = self.cond_embed(c)312        y = F.silu(t + c)313 314        scale1, shift1, gate1, scale2, shift2, gate2 = self.ada_ln_blocks[0](y).chunk(6, dim=-1)315        for i, block in enumerate(self.res_blocks):316            if i > 0 and i % self.ada_ln_switch_freq == 0:317                ada_ln_block = self.ada_ln_blocks[i // self.ada_ln_switch_freq]318                scale1, shift1, gate1, scale2, shift2, gate2 = ada_ln_block(y).chunk(6, dim=-1)319            x = block(x, scale1, shift1, gate1, scale2, shift2, gate2)320 321        output = self.final_layer(x, y)322        return 2 * torch.sigmoid(output) - 1323 324 325class BitDanceDiffusionHead(ModelMixin, ConfigMixin):326    @register_to_config327    def __init__(328        self,329        ch_target: int,330        ch_cond: int,331        ch_latent: int,332        depth_latent: int,333        depth_adanln: int,334        grad_checkpointing: bool = False,335        time_shift: float = 1.0,336        time_schedule: str = "logit_normal",337        P_mean: float = 0.0,338        P_std: float = 1.0,339        parallel_num: int = 4,340        diff_batch_mul: int = 1,341        use_swiglu: bool = False,342    ) -> None:343        super().__init__()344        self.ch_target = ch_target345        self.time_shift = time_shift346        self.time_schedule = time_schedule347        self.P_mean = P_mean348        self.P_std = P_std349        self.diff_batch_mul = diff_batch_mul350 351        self.net = TransEncoder(352            in_channels=ch_target,353            model_channels=ch_latent,354            z_channels=ch_cond,355            num_res_blocks=depth_latent,356            num_ada_ln_blocks=depth_adanln,357            grad_checkpointing=grad_checkpointing,358            parallel_num=parallel_num,359            use_swiglu=use_swiglu,360        )361 362    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:363        with torch.autocast(device_type="cuda", enabled=False):364            with torch.no_grad():365                if self.time_schedule == "logit_normal":366                    t = (torch.randn((x.shape[0]), device=x.device) * self.P_std + self.P_mean).sigmoid()367                    if self.time_shift != 1.0:368                        t = time_shift_func(t, self.time_shift)369                elif self.time_schedule == "uniform":370                    t = torch.rand((x.shape[0]), device=x.device)371                    if self.time_shift != 1.0:372                        t = time_shift_func(t, self.time_shift)373                else:374                    raise NotImplementedError(f"Unknown time_schedule={self.time_schedule}")375 376                e = torch.randn_like(x)377                ti = t.view(-1, 1, 1)378                z = (1.0 - ti) * e + ti * x379                v = (x - z) / (1 - ti).clamp_min(0.05)380 381        if self.diff_batch_mul > 1:382            chunks = self.diff_batch_mul383            x_pred_list = []384            z_chunks = torch.chunk(z, chunks, dim=0)385            t_chunks = torch.chunk(t, chunks, dim=0)386            cond_chunks = torch.chunk(cond, chunks, dim=0)387            for z_i, t_i, cond_i in zip(z_chunks, t_chunks, cond_chunks):388                x_pred_list.append(self.net(z_i, t_i, cond_i))389            x_pred = torch.cat(x_pred_list, dim=0)390        else:391            x_pred = self.net(z, t, cond)392 393        v_pred = (x_pred - z) / (1 - ti).clamp_min(0.05)394        with torch.autocast(device_type="cuda", enabled=False):395            v_pred = v_pred.float()396            loss = torch.mean((v - v_pred) ** 2, dim=2)397        return loss398 399    def sample(400        self,401        z: torch.Tensor,402        cfg: float,403        num_sampling_steps: int,404        generator: Optional[torch.Generator] = None,405    ) -> torch.Tensor:406        return euler_maruyama(407            self.ch_target,408            self.net.forward,409            z,410            cfg,411            num_sampling_steps=num_sampling_steps,412            time_shift=self.time_shift,413            generator=generator,414        )415 416    def initialize_weights(self) -> None:417        self.net.initialize_weights()418