| """Decoder template from the H3 VAE config (no H3 VAE weights needed). |
| |
| The checkpoint holds every decoder parameter: the template is built on meta, then the |
| checkpoint is strictly assigned. QK normalization and rotary embedding follow the Diffusers |
| H3 source (MiniMaxH3VideoAttnProcessor). |
| """ |
| from types import SimpleNamespace |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| def _apply_qk_norm(norm, x): |
| return norm(x.float()).to(x.dtype) |
|
|
|
|
| def _apply_rotary_pos_emb_impl(x, emb): |
| cos, sin = (t.to(x.dtype) for t in emb) |
| dim = cos.shape[-1] |
| a, tail = x[..., :dim], x[..., dim:] |
| first, second = a.chunk(2, dim=-1) |
| rotated = torch.cat((-second, first), dim=-1) |
| return torch.cat((a * cos + rotated * sin, tail), dim=-1) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, dim, hidden): |
| super().__init__() |
| self.w1 = nn.Linear(dim, 2 * hidden) |
| self.w2 = nn.Linear(hidden, dim) |
|
|
| def forward(self, x): |
| gate, up = self.w1(x).chunk(2, dim=-1) |
| return self.w2(F.silu(gate) * up) |
|
|
|
|
| class RGBProcessor: |
| def revert_tensor(self, x): |
| |
| mean = x.new_tensor((0.485, 0.456, 0.406)).view(1, 3, 1, 1, 1) |
| std = x.new_tensor((0.229, 0.224, 0.225)).view(1, 3, 1, 1, 1) |
| return x * std + mean |
|
|
|
|
| def make_template(config): |
| from diffusers.models.autoencoders.autoencoder_kl_minimax_h3 import MiniMaxH3VideoRotaryPosEmbed |
|
|
| heads = config['decoder_num_attention_heads'] |
| dh = config['decoder_attention_head_dim'] |
| dim = heads * dh |
| eps = config['decoder_norm_eps'] |
| |
| rope = MiniMaxH3VideoRotaryPosEmbed( |
| int(dh * config['decoder_rope_dim_ratio']), theta=config['decoder_rope_theta']) |
| mean = torch.tensor(config['latents_mean'], dtype=torch.float32) |
| std = torch.tensor(config['latents_std'], dtype=torch.float32) |
| with torch.device('meta'): |
| dec = nn.Module() |
| dec.x_embedder = nn.Linear(24, dim) |
| dec.register_tokens = nn.Parameter(torch.empty(1, config['decoder_num_register_tokens'], dim)) |
| dec.transformer_blocks = nn.ModuleList() |
| for _ in range(config['decoder_num_layers']): |
| block = nn.Module() |
| block.norm1 = nn.RMSNorm(dim, eps=eps) |
| block.norm2 = nn.RMSNorm(dim, eps=eps) |
| block.scale1 = nn.Parameter(torch.empty(dim)) |
| block.scale2 = nn.Parameter(torch.empty(dim)) |
| block.attn = nn.Module() |
| block.attn.heads, block.attn.dim_head = heads, dh |
| block.attn.to_qkv = nn.Linear(dim, 3 * dim) |
| block.attn.to_out = nn.Linear(dim, dim) |
| block.attn.norm_q = nn.RMSNorm(dh, eps=eps, elementwise_affine=False) |
| block.attn.norm_k = nn.RMSNorm(dh, eps=eps, elementwise_affine=False) |
| block.ff = SwiGLU(dim, dim * config['decoder_ffn_mult']) |
| dec.transformer_blocks.append(block) |
| dec.norm_out = nn.LayerNorm(dim, eps=eps) |
| dec.proj_out = nn.Linear(dim, 3 * 4 * 16 * 16) |
| post_quant_conv = nn.Conv3d(24, 24, 1) |
| dec.pos_embed = rope |
| return SimpleNamespace(decoder=dec, post_quant_conv=post_quant_conv, |
| latents_mean=mean, latents_std=std, processor=RGBProcessor()) |
|
|
