16 lines
448 B
Python
16 lines
448 B
Python
import mlx.core as mx
|
|
from mlx import nn
|
|
|
|
|
|
class TextEmbedder(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear_1 = nn.Linear(768, 3072)
|
|
self.linear_2 = nn.Linear(3072, 3072)
|
|
|
|
def __call__(self, caption: mx.array) -> mx.array:
|
|
hidden_states = self.linear_1(caption)
|
|
hidden_states = nn.silu(hidden_states)
|
|
hidden_states = self.linear_2(hidden_states)
|
|
return hidden_states
|