monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor: Transformer blocks

Browse Source
This commit is contained in:
filipstrand 2025-02-02 05:10:04 +01:00
parent 0d9a5d626f
commit 0485cf07c2
4 changed files with 88 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/mflux/models/transformer/ada_layer_norm_zero.py
View File

@ -8,7 +8,7 @@ class AdaLayerNormZero(nn.Module):
self.linear = nn.Linear(3072, 18432)
self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
def __call__(self, x: mx.array, text_embeddings: mx.array):
def __call__(self, hidden_states: mx.array, text_embeddings: mx.array):
text_embeddings = self.linear(nn.silu(text_embeddings))
chunk_size = 18432 // 6
shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
@ -17,5 +17,5 @@ class AdaLayerNormZero(nn.Module):
shift_mlp = text_embeddings[:, 3 * chunk_size : 4 * chunk_size]
scale_mlp = text_embeddings[:, 4 * chunk_size : 5 * chunk_size]
gate_mlp = text_embeddings[:, 5 * chunk_size : 6 * chunk_size]
x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp

6
src/mflux/models/transformer/ada_layer_norm_zero_single.py
View File

@ -8,11 +8,11 @@ class AdaLayerNormZeroSingle(nn.Module):
self.linear = nn.Linear(3072, 3 * 3072)
self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
def __call__(self, x: mx.array, text_embeddings: mx.array):
def __call__(self, hidden_states: mx.array, text_embeddings: mx.array) -> mx.array:
text_embeddings = self.linear(nn.silu(text_embeddings))
chunk_size = 9216 // 3
shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
scale_msa = text_embeddings[:, 1 * chunk_size : 2 * chunk_size]
gate_msa = text_embeddings[:, 2 * chunk_size : 3 * chunk_size]
x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
return x, gate_msa
hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
return hidden_states, gate_msa

71
src/mflux/models/transformer/joint_transformer_block.py
View File

@ -11,12 +11,12 @@ class JointTransformerBlock(nn.Module):
super().__init__()
self.layer = layer
self.norm1 = AdaLayerNormZero()
self.norm2 = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
self.ff = FeedForward(activation_function=nn.gelu)
self.attn = JointAttention()
self.norm1_context = AdaLayerNormZero()
self.ff_context = FeedForward(activation_function=nn.gelu_approx)
self.attn = JointAttention()
self.norm2 = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
self.norm2_context = nn.LayerNorm(dims=1536, eps=1e-6, affine=False)
self.ff = FeedForward(activation_function=nn.gelu)
self.ff_context = FeedForward(activation_function=nn.gelu_approx)
def __call__(
self,
@ -25,30 +25,67 @@ class JointTransformerBlock(nn.Module):
text_embeddings: mx.array,
rotary_embeddings: mx.array,
) -> (mx.array, mx.array):
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, text_embeddings)
# 1a. Compute norm for hidden_states
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
hidden_states=hidden_states,
text_embeddings=text_embeddings
) # fmt: off
# 1b. Compute norm for encoder_hidden_states
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
x=encoder_hidden_states, text_embeddings=text_embeddings
)
hidden_states=encoder_hidden_states,
text_embeddings=text_embeddings
) # fmt: off
# 2. Compute attention
attn_output, context_attn_output = self.attn(
hidden_states=norm_hidden_states,
encoder_hidden_states=norm_encoder_hidden_states,
image_rotary_emb=rotary_embeddings,
)
# 3a. Apply norm and feed forward for hidden states
hidden_states = JointTransformerBlock._apply_norm_and_feed_forward(
hidden_states=hidden_states,
attn_output=attn_output,
gate_mlp=gate_mlp,
gate_msa=gate_msa,
scale_mlp=scale_mlp,
shift_mlp=shift_mlp,
norm_layer=self.norm2,
ff_layer=self.ff,
)
# 3b. Apply norm and feed forward for encoder hidden states
encoder_hidden_states = JointTransformerBlock._apply_norm_and_feed_forward(
hidden_states=encoder_hidden_states,
attn_output=context_attn_output,
gate_mlp=c_gate_mlp,
gate_msa=c_gate_msa,
scale_mlp=c_scale_mlp,
shift_mlp=c_shift_mlp,
norm_layer=self.norm2_context,
ff_layer=self.ff_context,
)
return encoder_hidden_states, hidden_states
@staticmethod
def _apply_norm_and_feed_forward(
hidden_states: mx.array,
attn_output: mx.array,
gate_mlp: mx.array,
gate_msa: mx.array,
scale_mlp: mx.array,
shift_mlp: mx.array,
norm_layer: nn.Module,
ff_layer: nn.Module,
) -> mx.array:
attn_output = mx.expand_dims(gate_msa, axis=1) * attn_output
hidden_states = hidden_states + attn_output
norm_hidden_states = self.norm2(hidden_states)
norm_hidden_states = norm_layer(hidden_states)
norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
ff_output = self.ff(norm_hidden_states)
ff_output = ff_layer(norm_hidden_states)
ff_output = mx.expand_dims(gate_mlp, axis=1) * ff_output
hidden_states = hidden_states + ff_output
context_attn_output = mx.expand_dims(c_gate_msa, axis=1) * context_attn_output
encoder_hidden_states = encoder_hidden_states + context_attn_output
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
context_ff_output = self.ff_context(norm_encoder_hidden_states)
encoder_hidden_states = encoder_hidden_states + mx.expand_dims(c_gate_mlp, axis=1) * context_ff_output
return encoder_hidden_states, hidden_states
return hidden_states

32
src/mflux/models/transformer/single_transformer_block.py
View File

@ -12,8 +12,8 @@ class SingleTransformerBlock(nn.Module):
super().__init__()
self.layer = layer
self.norm = AdaLayerNormZeroSingle()
self.proj_mlp = nn.Linear(3072, 4 * 3072)
self.attn = SingleBlockAttention()
self.proj_mlp = nn.Linear(3072, 4 * 3072)
self.proj_out = nn.Linear(3072 + 4 * 3072, 3072)
def __call__(
@ -22,15 +22,39 @@ class SingleTransformerBlock(nn.Module):
text_embeddings: mx.array,
rotary_embeddings: mx.array,
) -> (mx.array, mx.array):
# 0. Establish residual connection
residual = hidden_states
norm_hidden_states, gate = self.norm(x=hidden_states, text_embeddings=text_embeddings)
mlp_hidden_states = nn.gelu_approx(self.proj_mlp(norm_hidden_states))
# 1. Compute norm for hidden_states
norm_hidden_states, gate = self.norm(
hidden_states=hidden_states,
text_embeddings=text_embeddings
) # fmt: off
# 2. Compute attention
attn_output = self.attn(
hidden_states=norm_hidden_states,
image_rotary_emb=rotary_embeddings,
)
# 3. Apply norm and feed forward for hidden states
hidden_states = self._apply_feed_forward_and_projection(
norm_hidden_states=norm_hidden_states,
attn_output=attn_output,
gate=gate,
)
return residual + hidden_states
def _apply_feed_forward_and_projection(
self,
norm_hidden_states: mx.array,
attn_output: mx.array,
gate: mx.array,
) -> mx.array:
feed_forward = self.proj_mlp(norm_hidden_states)
mlp_hidden_states = nn.gelu_approx(feed_forward)
hidden_states = mx.concatenate([attn_output, mlp_hidden_states], axis=2)
gate = mx.expand_dims(gate, axis=1)
hidden_states = gate * self.proj_out(hidden_states)
hidden_states = residual + hidden_states
return hidden_states