clip model is running
MODIFY
examples/stable_diffusion.py
examples/stable_diffusion.py
@@ -364,10 +364,114 @@ class UNetModel:
x = run(x, bb)
return x.sequential(self.out)
+class CLIPMLP:
+ def __init__(self):
+ self.fc1 = Linear(768, 3072)
+ self.fc2 = Linear(3072, 768)
+
+ def __call__(self, hidden_states):
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = hidden_states.quick_gelu()
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+class CLIPAttention:
+ def __init__(self):
+ self.embed_dim = 768
+ self.num_heads = 12
+ self.head_dim = self.embed_dim // self.num_heads
+ self.scale = self.head_dim**-0.5
+ self.k_proj = Linear(self.embed_dim, self.embed_dim)
+ self.v_proj = Linear(self.embed_dim, self.embed_dim)
+ self.q_proj = Linear(self.embed_dim, self.embed_dim)
+ self.out_proj = Linear(self.embed_dim, self.embed_dim)
+
+ def _shape(self, tensor, seq_len: int, bsz: int):
+ return tensor.reshape(bsz, seq_len, self.num_heads, self.head_dim).permute(0,2,1,3)
+
+ def __call__(self, hidden_states):
+ bsz, tgt_len, embed_dim = hidden_states.shape
+
+ query_states = self.q_proj(hidden_states) * self.scale
+ key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
+ value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+
+ proj_shape = (bsz * self.num_heads, -1, self.head_dim)
+ query_states = self._shape(query_states, tgt_len, bsz).reshape(*proj_shape)
+ key_states = key_states.reshape(*proj_shape)
+ value_states = value_states.reshape(*proj_shape)
+
+ attn_weights = query_states @ key_states.permute(0,2,1)
+ attn_weights = attn_weights.softmax()
+
+ attn_output = attn_weights @ value_states
+
+ attn_output = attn_output.reshape(bsz, self.num_heads, tgt_len, self.head_dim)
+ attn_output = attn_output.permute(0,2,1,3)
+ attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)
+
+ attn_output = self.out_proj(attn_output)
+ return attn_output
+
+class CLIPEncoderLayer:
+ def __init__(self):
+ self.self_attn = CLIPAttention()
+ self.layer_norm1 = Normalize(768, num_groups=1)
+ self.mlp = CLIPMLP()
+ self.layer_norm2 = Normalize(768, num_groups=1)
+
+ def __call__(self, hidden_states):
+ residual = hidden_states
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states = self.self_attn(hidden_states)
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+class CLIPEncoder:
+ def __init__(self):
+ self.layers = [CLIPEncoderLayer() for i in range(12)]
+
+ def __call__(self, hidden_states):
+ return hidden_states.sequential(self.layers)
+
+class CLIPTextEmbeddings:
+ def __init__(self):
+ self.position_ids = Tensor.empty(1, 77)
+ self.token_embedding = {"weight": Tensor.empty(49408, 768)}
+ self.position_embedding = {"weight": Tensor.empty(77, 768)}
+
+ def __call__(self, input_ids, position_ids):
+ # TODO: actually support batches
+ inputs = np.zeros((1, len(input_ids), 49408))
+ positions = np.zeros((1, len(position_ids), 77))
+ for i,x in enumerate(input_ids): inputs[0][i][x] = 1
+ for i,x in enumerate(position_ids): positions[0][i][x] = 1
+ inputs_embeds = Tensor(inputs, device=self.token_embedding['weight'].device) @ self.token_embedding['weight']
+ position_embeddings = Tensor(positions, device=self.position_embedding['weight'].device) @ self.position_embedding['weight']
+ return inputs_embeds + position_embeddings
+
+class CLIPTextTransformer:
+ def __init__(self):
+ self.embeddings = CLIPTextEmbeddings()
+ self.encoder = CLIPEncoder()
+ self.final_layer_norm = Normalize(768, num_groups=1)
+
+ def __call__(self, input_ids):
+ x = self.embeddings(input_ids, list(range(len(input_ids))))
+ x = self.encoder(x)
+ return self.final_layer_norm(x)
+
class StableDiffusion:
def __init__(self):
- self.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel())
- self.first_stage_model = AutoencoderKL()
+ #self.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel())
+ #self.first_stage_model = AutoencoderKL()
+ self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = CLIPTextTransformer()))
def __call__(self, x):
context = Tensor.uniform(1, 77, 768)
@@ -414,6 +518,11 @@ if __name__ == "__main__":
assert w.shape == v.shape
w.assign(v.astype(np.float32))
+ outs = model.cond_stage_model.transformer.text_model([1,2,3])
+ print(outs.numpy())
+
+ exit(0)
+
# load apple latent space
nz = Tensor(np.load("datasets/stable_diffusion_apple.npy"))

MODIFY
tinygrad/nn/optim.py
tinygrad/nn/optim.py
@@ -6,8 +6,7 @@ class Optimizer:
self.params = [x for x in params if x.requires_grad]
def zero_grad(self):
- for param in self.params:
- param.grad = None
+ for param in self.params: param.grad = None
def realize(self, extra=None):
# TODO: corealize
@@ -19,8 +18,7 @@ class SGD(Optimizer):
self.lr = lr
def step(self):
- for t in self.params:
- t.assign(t.detach() - t.grad * self.lr)
+ for t in self.params: t.assign(t.detach() - t.grad * self.lr)
self.realize()
class RMSprop(Optimizer):

MODIFY
tinygrad/shapetracker.py
tinygrad/shapetracker.py
@@ -112,7 +112,7 @@ class ShapeTracker:
def permute(self, *axis):
assert all([isinstance(x, int) and x >= 0 and x < len(self.shape) for x in axis])
- assert len(set(axis)) == len(axis) and len(axis) == len(self.shape)
+ assert len(set(axis)) == len(axis) and len(axis) == len(self.shape), f"can't permute {self.shape} with {axis}"
self.views[-1] = View([self.shape[a] for a in axis], [self.strides[a] for a in axis], self.offset)
# TODO: this is a special case of slice with strides, remove it

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -180,6 +180,7 @@ class Tensor:
# (padding_left, padding_right, padding_top, padding_bottom)
def pad2d(self, padding:Tuple[int, ...]): return self[:, :, -padding[2]:self.shape[2]+padding[3], -padding[0]:self.shape[3]+padding[1]]
+ # TODO: this is totally not transpose
def transpose(self, order=(1,0)): return self.permute(order=order)
def flatten(self, start_dim=0): return self.reshape(shape=tuple(list(self.shape[0:start_dim]) + [-1]))
@@ -247,6 +248,7 @@ class Tensor:
def hardswish(self): return self * (self+3).relu6() * (1/6)
def tanh(self): return 2.0 * ((2.0 * self).sigmoid()) - 1.0
def gelu(self): return 0.5 * self * (1 + (self * 0.7978845608 * (1 + 0.044715 * self * self)).tanh())
+ def quick_gelu(self): return self * (self * 1.702).sigmoid()
def leakyrelu(self, neg_slope=0.01): return self.relu() - (-neg_slope*self).relu()
def mish(self): return self * self.softplus().tanh()
def softplus(self, limit=20, beta=1): return (1/beta) * (1 + (self*beta).exp()).log()

clip model is running