File size: 7,248 Bytes
31dc8dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | import jax
import jax.numpy as jnp
import flax.linen as nn
from modules.layers import (
Attention, BottleneckTextProj, FinalLayer, RMSNorm, SwiGLUFFN,
TextRotaryEmbeddingFast, TimestepEmbedder,
DEFAULT_KERNEL_INIT, DEFAULT_BIAS_INIT, NORMAL_INIT_002,
)
class ELFBlock(nn.Module):
"""ELF Transformer block."""
hidden_size: int
num_heads: int
mlp_ratio: float = 4.0
attn_drop: float = 0.0
proj_drop: float = 0.0
@nn.compact
def __call__(self, x, rope_fn=None, attention_mask=None, deterministic=True):
mlp_hidden_dim = int(self.hidden_size * self.mlp_ratio)
x_normed = RMSNorm(self.hidden_size, eps=1e-6, name='norm1')(x)
attn_out = Attention(
self.hidden_size, self.num_heads, qkv_bias=True, qk_norm=True,
attn_drop=self.attn_drop, proj_drop=self.proj_drop, name='attn',
)(x_normed, rope_fn, attention_mask=attention_mask, deterministic=deterministic)
x = x + attn_out
x_normed = RMSNorm(self.hidden_size, eps=1e-6, name='norm2')(x)
mlp_out = SwiGLUFFN(self.hidden_size, mlp_hidden_dim, drop=self.proj_drop, name='mlp')(
x_normed, deterministic=deterministic,
)
x = x + mlp_out
return x
class ELF(nn.Module):
"""Text ELF Transformer."""
text_encoder_dim: int
max_length: int
hidden_size: int = 1024
depth: int = 24
num_heads: int = 16
mlp_ratio: float = 4.0
attn_drop: float = 0.0
proj_drop: float = 0.0
bottleneck_dim: int = 128
num_time_tokens: int = 4 # Number of in-context time conditioning tokens
num_self_cond_cfg_tokens: int = 4 # Number of in-context self-cond CFG tokens
num_model_mode_tokens: int = 0 # If > 0, prepend learnable model-mode tokens that signal decoding mode
vocab_size: int = 0 # Vocabulary size for decoder unembedding
def build_context(self, t, self_cond_cfg_scale=None):
prefix_tokens = []
B = t.shape[0]
def _make_prefix(emb, n_tokens, param_name):
tokens = self.param(param_name, NORMAL_INIT_002, (1, n_tokens, self.hidden_size))
return jnp.tile(tokens, (B, 1, 1)) + jnp.expand_dims(emb, 1)
if self.num_time_tokens <= 0:
raise ValueError("num_time_tokens must be positive for prefix time conditioning")
time_emb = TimestepEmbedder(self.hidden_size, name='t_embedder')(t)
prefix_tokens.append(_make_prefix(time_emb, self.num_time_tokens, 't_emb_tokens'))
if self_cond_cfg_scale is not None:
sc_emb = TimestepEmbedder(self.hidden_size, name='self_cond_cfg_embedder')(self_cond_cfg_scale)
if self.num_self_cond_cfg_tokens > 0:
prefix_tokens.append(_make_prefix(sc_emb, self.num_self_cond_cfg_tokens, 'self_cond_cfg_tokens'))
return prefix_tokens
@nn.compact
def __call__(
self, x, t, attention_mask=None, deterministic=True,
self_cond_cfg_scale=None, decoder_step_active=None,
):
"""x: (N, S, C) or (N, S, 2C) with self-cond. t: (N,). attention_mask: (N, S), 1=valid."""
patch_size = 1
head_dim = self.hidden_size // self.num_heads
B = x.shape[0]
# Self-conditioning: input is [z, x_pred] when 2x encoder dim
if x.shape[-1] == 2 * self.text_encoder_dim:
x = nn.Dense(
self.text_encoder_dim, use_bias=True,
kernel_init=DEFAULT_KERNEL_INIT, bias_init=DEFAULT_BIAS_INIT, name='self_cond_proj',
)(x)
# Text projection (with bottleneck)
x = BottleneckTextProj(
self.text_encoder_dim, self.hidden_size, self.bottleneck_dim, name='text_proj',
)(x)
# Prepend learnable model-mode tokens (gated: zero unless decoder_step_active=True)
model_mode_offset = 0
if self.num_model_mode_tokens > 0:
mode_tokens = jnp.tile(
self.param('mode_tokens', NORMAL_INIT_002,
(1, self.num_model_mode_tokens, self.hidden_size)),
(B, 1, 1),
)
active_gate = jnp.array(False) if decoder_step_active is None else decoder_step_active
mode_tokens = mode_tokens * active_gate.astype(mode_tokens.dtype)
x = jnp.concatenate([mode_tokens, x], axis=1)
model_mode_offset = self.num_model_mode_tokens
if attention_mask is not None:
mode_mask = jnp.ones((B, self.num_model_mode_tokens), dtype=attention_mask.dtype)
attention_mask = jnp.concatenate([mode_mask, attention_mask], axis=1)
prefix_len = 0
context_prefix_tokens = self.build_context(t, self_cond_cfg_scale)
if context_prefix_tokens:
prefix_tokens = jnp.concatenate(context_prefix_tokens, axis=1)
prefix_len = prefix_tokens.shape[1]
x = jnp.concatenate([prefix_tokens, x], axis=1)
if attention_mask is not None:
prefix_mask = jnp.ones((B, prefix_len), dtype=attention_mask.dtype)
attention_mask = jnp.concatenate([prefix_mask, attention_mask], axis=1)
feat_rope = TextRotaryEmbeddingFast(
dim=head_dim, pt_seq_len=self.max_length,
num_empty_token=prefix_len + model_mode_offset, name='feat_rope',
)
q1, q3 = self.depth // 4, self.depth // 4 * 3
for i in range(self.depth):
in_drop_range = q3 > i >= q1
block = ELFBlock(
self.hidden_size, self.num_heads, mlp_ratio=self.mlp_ratio,
attn_drop=self.attn_drop if in_drop_range else 0.0,
proj_drop=self.proj_drop if in_drop_range else 0.0,
name=f'blocks_{i}',
)
x = block(x, rope_fn=feat_rope, attention_mask=attention_mask, deterministic=deterministic)
x = x[:, prefix_len + model_mode_offset:]
# Factored decoder unembedding: hidden -> text_encoder_dim -> vocab
decoder_logits = None
bn = self.text_encoder_dim
proj_kernel = self.param('proj_kernel', DEFAULT_KERNEL_INIT, (self.hidden_size, bn))
proj_bias = self.param('proj_bias', DEFAULT_BIAS_INIT, (bn,))
unembed_kernel = self.param('unembed_kernel', DEFAULT_KERNEL_INIT, (bn, self.vocab_size))
unembed_bias = self.param('unembed_bias', DEFAULT_BIAS_INIT, (self.vocab_size,))
if decoder_step_active is not None:
decoder_logits = jax.lax.cond(
decoder_step_active,
lambda xi: jax.nn.gelu(xi @ proj_kernel + proj_bias) @ unembed_kernel + unembed_bias,
lambda xi: jnp.zeros((*xi.shape[:2], self.vocab_size), dtype=xi.dtype),
x,
)
output = FinalLayer(self.hidden_size, patch_size, self.text_encoder_dim, name='final_layer')(x)
return output, decoder_logits
# Model factory functions
def ELF_B(**kwargs): return ELF(depth=12, hidden_size=768, num_heads=12, **kwargs)
def ELF_M(**kwargs): return ELF(depth=24, hidden_size=1056, num_heads=16, **kwargs)
def ELF_L(**kwargs): return ELF(depth=32, hidden_size=1280, num_heads=16, **kwargs)
ELF_models = {
'ELF-B': ELF_B, 'ELF-M': ELF_M, 'ELF-L': ELF_L,
}
|