Update handler.py

handler.py

@@ -1,532 +1,530 @@
-from transformers import PreTrainedTokenizerFast
-
-token2vec = PreTrainedTokenizerFast.from_pretrained("cl100k_tokenizer")
-
-from typing import Dict, Any
-import numpy as np
-import tensorflow as tf
-from tensorflow.keras.layers import Dense, LayerNormalization, Conv2D, UpSampling2D, Embedding, MultiHeadAttention
-from tensorflow.keras.saving import register_keras_serializable
-import tensorflow as tf
-
-token2vec = tiktoken.encoding_for_model("gpt-3.5-turbo")
-
-# @title Config
-def small_config():
-    T = 500
-    beta = np.linspace(1e-4, 0.02, T)
-    alpha = 1 - beta
-    a = np.cumprod(alpha)
-
-    return {
-        "filters": [128, 256],
-        "hidden_dim": 384,
-        "heads": 6,
-        "layers": 8,
-        "patch_size": 4,
-        "batch_size": 64,
-        "T": T,
-        "context_size": 8,
-        "image_size": 128,
-        "latent_shape": (32, 32, 4),
-        "beta": beta,
-        "alpha": alpha,
-        "a": a}
-
-def med_config():
-    T = 1000
-    beta = np.linspace(1e-4, 0.02, T)
-    alpha = 1 - beta
-    a = np.cumprod(alpha)
-
-    return {
-        "filters": [128, 256],
-        "hidden_dim": 768,
-        "heads": 12,
-        "layers": 12,
-        "patch_size": 4,
-        "batch_size": 64,
-        "T": T,
-        "context_size": 8,
-        "image_size": 128,
-        "latent_shape": (32, 32, 4),
-        "beta": beta,
-        "alpha": alpha,
-        "a": a}
-
-def large_config():
-    T = 1000
-    beta = np.linspace(1e-4, 0.02, T)
-    alpha = 1 - beta
-    a = np.cumprod(alpha)
-
-    return {
-        "filters": [128, 256],
-        "hidden_dim": 1024,
-        "heads": 16,
-        "layers": 24,
-        "patch_size": 4,
-        "batch_size": 64,
-        "T": T,
-        "context_size": 8,
-        "image_size": 128,
-        "latent_shape": (32, 32, 4),
-        "beta": beta,
-        "alpha": alpha,
-        "a": a}
-
-config = med_config()
-
-filters = config['filters']
-hidden_dim = config['hidden_dim']
-heads = config['heads']
-layers = config['layers']
-patch_size = config['patch_size']
-batch_size = config['batch_size']
-T = config['T']
-context_size = config['context_size']
-image_size = config['image_size']
-latent_shape = config['latent_shape']
-beta = config['beta']
-alpha = config['alpha']
-a = config['a']
-
-# @title ResBlock, UpBlock, DownBlock
-@register_keras_serializable()
-class ResBlock(tf.keras.layers.Layer):
-  def __init__(self, filters, p, **kwargs):
-    super(ResBlock, self).__init__(**kwargs)
-    self.filters = filters
-    self.p = p
-    self.reshape = Conv2D(filters, kernel_size=1, strides=1, padding="same")
-    #self.norm = BatchNormalization(center=False, scale=False)
-    self.conv1 = Conv2D(filters, kernel_size=p, strides=1, padding="same", activation="swish")
-    self.conv2 = Conv2D(filters, kernel_size=p, strides=1, padding="same")
-
-  def call(self, x):
-    x = self.reshape(x)
-    resid = x
-    #resid = self.norm(resid)
-    resid = self.conv1(resid)
-    resid = self.conv2(resid)
-    x = x + resid
-    return x
-
-  def get_config(self):
-    config = super().get_config()
-    config.update({
-        "filters": self.filters,
-        "p": self.p})
-    return config
-
-@register_keras_serializable()
-class DownBlock(tf.keras.layers.Layer):
-  def __init__(self, filters, **kwargs):
-    super(DownBlock, self).__init__(**kwargs)
-    self.filters = filters
-    self.resBlocks = [ResBlock(f, p=3) for f in filters]
-    self.pool = tf.keras.layers.MaxPool2D(pool_size=(2, 2))
-
-  def call(self, x):
-    for resBlock in self.resBlocks:
-      x = resBlock(x)
-    x = self.pool(x)
-    return x
-
-  def get_config(self):
-    config = super().get_config()
-    config.update({
-        "filters": self.filters})
-    return config
-
-@register_keras_serializable()
-class UpBlock(tf.keras.layers.Layer):
-  def __init__(self, filters, **kwargs):
-    super(UpBlock, self).__init__(**kwargs)
-    self.filters = filters
-    self.resBlocks = [ResBlock(f, p=3) for f in filters]
-    self.upSample = UpSampling2D(size=2, interpolation="bilinear")
-
-  def call(self, x):
-    x = self.upSample(x)
-    for resBlock in self.resBlocks:
-      x = resBlock(x)
-    return x
-
-  def get_config(self):
-    config = super().get_config()
-    config.update({
-        "filters": self.filters})
-    return config
-
-# @title Encoder, Decoder
-@register_keras_serializable()
-class Encoder(tf.keras.Model):
-  def __init__(self, filters, latent_dim, **kwargs):
-    super(Encoder, self).__init__(**kwargs)
-    self.filters = filters
-    self.latent_dim = latent_dim
-    self.downBlocks = [DownBlock([f,f]) for f in filters]
-    self.latent_proj = Conv2D(latent_dim * 2, kernel_size=1, strides=1, padding="same", activation="linear")
-
-  @tf.function
-  def sample(self, mu, logvar):
-    eps = tf.random.normal(shape=tf.shape(mu))
-    return eps * tf.exp(logvar * .5) + mu
-
-  def call(self, x, training=1):
-    for downBlock in self.downBlocks:
-      x = downBlock(x)
-    x = self.latent_proj(x)
-    mu, logvar = tf.split(x, 2, axis=-1)
-    z = self.sample(mu, logvar)
-    return z, mu, logvar
-
-  def get_config(self):
-    config = super().get_config()
-    config.update({
-        "filters": self.filters,
-        "latent_dim": self.latent_dim})
-    return config
-
-  def compute_output_shape(self, input_shape):
-    return (input_shape[0], self.latent_dim), (input_shape[0], self.latent_dim), (input_shape[0], self.latent_dim)
-
-@register_keras_serializable()
-class Decoder(tf.keras.Model):
-  def __init__(self, filters, img_size, **kwargs):
-    super(Decoder, self).__init__(**kwargs)
-    self.filters = filters[::-1]
-    self.img_size = img_size
-    self.undo_latent_proj = Conv2D(filters[0], kernel_size=1, strides=1, padding="same")
-    self.upBlocks = [UpBlock([f,f]) for f in filters]
-    self.conv_proj = Conv2D(3, kernel_size=3, padding="same", activation="linear")
-
-  def call(self, z, training=1):
-    z = self.undo_latent_proj(z)
-    for upBlock in self.upBlocks:
-      z = upBlock(z)
-    x = self.conv_proj(z)
-    return x
-
-  def get_config(self):
-    config = super().get_config()
-    config.update({
-        "filters": self.filters[::-1],
-        "img_size": self.img_size})
-    return config
-
-  def compute_output_shape(self, input_shape):
-    return (input_shape[0], self.img_size, self.img_size, 3)
-
-# @title Helper Functions
-def process_text(text):
-  import tiktoken
-  tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
-  tokens = tokenizer.encode(text)
-  while len(tokens) < context_size:
-      tokens.append(0)
-  return tokens[:context_size]
-
-def normalise_img(img_tensor): # Maps [-1,1] to [0,1]
-  img = img_tensor
-  img *= 0.5
-  img += 0.5
-  return img
-
-def prep_img(img_tensor): # Maps [0,255] to [-1,1]
-  img = img_tensor.copy()
-  img = img / 127.5
-  img -= 1
-  return img
-
-def noisify_img(img_tensor, t, a):  # Returns x_t and the noise used
-  epsilon = np.random.normal(0, 1, img_tensor.shape).astype(np.float32)  # Standard normal
-  sqrt_alpha_bar = np.sqrt(a[t])
-  sqrt_one_minus_alpha_bar = np.sqrt(1 - a[t])
-  x_t = sqrt_alpha_bar * img_tensor + sqrt_one_minus_alpha_bar * epsilon
-  return x_t, epsilon
-
-def denoise_step(x_t, eps_hat, t, a, beta):
-  """
-  Reverse one DDPM step: x_t → x_{t-1}
-  """
-  a_bar_t = tf.convert_to_tensor(a[t], dtype=tf.float32)
-  a_bar_prev = tf.convert_to_tensor(a[t - 1] if t > 0 else 1.0, dtype=tf.float32)
-  a_t = a_bar_t / a_bar_prev
-  beta_t = tf.convert_to_tensor(beta[t], dtype=tf.float32)
-
-  # Avoid NaNs with clamping
-  sqrt_recip_a_t = tf.math.rsqrt(tf.maximum(a_t, 1e-5))
-  sqrt_one_minus_ab = tf.sqrt(tf.maximum(1. - a_bar_t, 1e-5))
-
-  eps_term = (beta_t / sqrt_one_minus_ab) * eps_hat
-  mean = sqrt_recip_a_t * (x_t - eps_term)
-
-  if t > 1:
-      noise = tf.random.normal(shape=x_t.shape)
-      sigma = tf.sqrt(tf.maximum(beta_t, 1e-5))
-      x_prev = mean + sigma * noise
-  else:
-      x_prev = mean
-
-  return x_prev
-
-# @title Transformer Block
-@register_keras_serializable()
-class TransformerBlock(tf.keras.Layer):
-  def __init__(self, context_size, head_no, latent_dim, **kwargs):
-    super().__init__(**kwargs)
-    self.context_size = context_size
-    self.head_no = head_no
-    self.latent_dim = latent_dim
-    self.attn = MultiHeadAttention(num_heads=head_no, key_dim=latent_dim//head_no, output_shape=latent_dim)
-    self.mlp_up = Dense(latent_dim*4, activation="gelu")
-    self.mlp_down = Dense(latent_dim)
-    self.norm1 = LayerNormalization()
-    self.norm2 = LayerNormalization()
-
-  def call(self, x):
-    normed = self.norm1(x)
-    x = x + self.attn(normed, normed, normed)
-    normed = self.norm2(x)
-    dx = self.mlp_up(normed)
-    x = x + self.mlp_down(dx)
-    return x
-
-  def build(self, input_shape):
-    super().build(input_shape)
-
-  def compute_output_shape(self, input_shape):
-    return input_shape
-
-  def get_config(self):
-    config = super().get_config()
-    config.update({
-        "context_size": self.context_size,
-        "head_no": self.head_no,
-        "latent_dim": self.latent_dim})
-    return config
-
-# @title AdaLN-Zero
-@register_keras_serializable()
-class AdaptiveLayerNorm(tf.keras.Layer):
-  def __init__(self, eps=1e-6,**kwargs):
-    self.layernorm = LayerNormalization(epsilon=eps,center=False, scale=False)
-    super(AdaptiveLayerNorm, self).__init__(**kwargs)
-
-  def build(self, input_shape):
-    #B, num_patches, hidden_dim
-    self.M = Dense(input_shape[2], use_bias=True, kernel_initializer='glorot_uniform', activation="linear")
-    self.b = Dense(input_shape[2], use_bias=True, kernel_initializer='glorot_uniform', activation="linear")
-
-  def call(self, x, cond):
-    gamma = self.M(cond)
-    beta = self.b(cond)
-    x = self.layernorm(x)
-    x = x * (1 + tf.expand_dims(gamma, 1)) + tf.expand_dims(beta, 1)
-    return x
-
-  def get_config(self):
-    config = super().get_config()
-    return config
-
-# @title Image Embedder, Unembedder
-@register_keras_serializable()
-class ImageEmbedder(tf.keras.Layer):
-  def __init__(self, latent_size, patch_size, emb_dim,**kwargs):
-    super().__init__(**kwargs)
-    self.emb_dim = emb_dim
-    self.patch_size = patch_size
-    self.latent_size = latent_size
-    self.pos_emb = Embedding(input_dim=(latent_size // patch_size)**2 , output_dim=emb_dim, embeddings_initializer="glorot_uniform")
-    self.reshaper = Dense(emb_dim, kernel_initializer="glorot_uniform")
-    self.conv_expansion = Conv2D(emb_dim, kernel_size=patch_size, strides=patch_size, padding="same")
-
-  def call(self, x):
-    x = self.reshaper(x)
-    x = self.conv_expansion(x)
-    x = tf.reshape(x, shape=[tf.shape(x)[0], tf.shape(x)[1]*tf.shape(x)[2], tf.shape(x)[3]])
-    positions = tf.range(start=0, limit=(self.latent_size // self.patch_size)**2, delta=1)
-    embeddings = self.pos_emb(positions)
-    x = embeddings + x
-    return x
-
-  def get_config(self):
-    config = super().get_config()
-    config.update({
-        "latent_size" : self.latent_size,
-        "patch_size": self.patch_size,
-        "emb_dim": self.emb_dim})
-    return config
-
-@register_keras_serializable()
-class ImageUnembedder(tf.keras.Layer):
-  def __init__(self, latent_size, patch_size, latent_dim, **kwargs):
-    super().__init__(**kwargs)
-    self.latent_dim = latent_dim
-    self.patch_size = patch_size
-    self.latent_size = latent_size
-    self.AdaLN = AdaptiveLayerNorm()
-    self.reshape_to_latent = Dense(patch_size*patch_size*latent_dim, kernel_initializer="glorot_uniform")
-
-  def call(self, x, cond):
-    x = self.AdaLN(x, cond)
-    x = self.reshape_to_latent(x)
-    x = tf.reshape(x, shape=
-     [tf.shape(x)[0],
-      self.latent_size // self.patch_size,
-      self.latent_size // self.patch_size,
-      self.latent_dim*(self.patch_size**2)])
-    x = tf.nn.depth_to_space(x, block_size=self.patch_size)
-    return x
-
-  def get_config(self):
-    config = super().get_config()
-    config.update({
-        "latent_size" : self.latent_size,
-        "patch_size": self.patch_size,
-        "latent_dim": self.latent_dim})
-    return config
-
-# @title LEGACY Prompt and Timestep Embedder
-@register_keras_serializable()
-class ConditioningEmbedder(tf.keras.layers.Layer):
-  def __init__(self, emb_dim, T, context_size, vocab_size=100266, **kwargs):
-    super().__init__(**kwargs)
-    self.emb_dim = emb_dim
-    self.T = T
-    self.context_size = context_size
-    self.vocab_size = vocab_size
-    positions = tf.range(T, dtype=tf.float32)[:, tf.newaxis]
-    frequencies = tf.constant(10000 ** (-tf.range(0, emb_dim, 2, dtype=tf.float32) / emb_dim))
-    angle_rates = positions * frequencies  # (T, emb_dim/2)
-    sin_part = tf.sin(angle_rates)
-    cos_part = tf.cos(angle_rates)
-    emb = tf.stack([sin_part, cos_part], axis=-1)  # (T, emb_dim/2, 2)
-    emb = tf.reshape(emb, [T, emb_dim])  # (T, emb_dim)
-    self.t_embeddings = tf.constant(emb, dtype=tf.float32)
-
-    self.prompt_emb = self.add_weight(shape=(vocab_size, emb_dim), initializer='glorot_uniform', name='prompt_emb', trainable=True)
-    self.CLS = self.add_weight(shape=(emb_dim,), initializer='glorot_uniform', name='CLS', trainable=True)
-    self.prompt_pos_enc = self.add_weight(shape=(1, context_size+1, emb_dim), initializer='glorot_uniform', name='prompt_pos_enc', trainable=True)
-    self.transformer = TransformerBlock(context_size+1, head_no=6, latent_dim=emb_dim)
-
-  def call(self, x):
-    t, prompt_tokens = x
-
-    # ── timestep embedding ───────────────────────────
-    t = tf.cast(tf.squeeze(t, axis=-1), tf.int32)        # (batch,)
-    embedded_t = tf.gather(self.t_embeddings, t)         # (batch, emb_dim)
-    embedded_t = embedded_t[:, tf.newaxis, :]            # (batch, 1, emb_dim)
-
-    # ── prompt embedding path ─────────────────────────
-    embedded_prompt = tf.nn.embedding_lookup(
-        self.prompt_emb, prompt_tokens)                  # (batch, seq_len, emb_dim)
-
-    cls_tok = tf.tile(self.CLS[None, None, :],
-                      [tf.shape(embedded_prompt)[0], 1, 1])
-    embedded_prompt = tf.concat([cls_tok, embedded_prompt], axis=1)
-    embedded_prompt += self.prompt_pos_enc
-    embedded_prompt = self.transformer(embedded_prompt)  # (batch, seq_len+1, emb_dim)
-
-    # add t-embedding to every token (broadcasts along axis-1)
-    embedded_prompt += embedded_t
-
-    # return CLS (keep singleton axis if you need it)
-    return embedded_prompt[:, 0, :]                     # (batch, 1, emb_dim)
-
-
-  def get_config(self):
-    config = super().get_config()
-    config.update({
-        "emb_dim": self.emb_dim,
-        "T": self.T,
-        "context_size": self.context_size,
-        "vocab_size": self.vocab_size})
-    return config
-
-# @title DiT Block
-class Gain(tf.keras.layers.Layer):
-  def __init__(self):
-    super(Gain, self).__init__()
-
-  def build(self, input_shape):
-    self.M = Dense(input_shape[2], use_bias=True,kernel_initializer='glorot_uniform')
-
-  def call(self, x, cond):
-    scale = self.M(cond)
-    x *= tf.expand_dims(scale, 1)
-    return x
-
-@register_keras_serializable()
-class DiTBlock(tf.keras.layers.Layer):
-  def __init__(self, hidden_dim, heads, context_size, **kwargs):
-    super().__init__(**kwargs)
-    self.emb_dim = hidden_dim
-    self.heads = heads
-    self.context_size = context_size
-    self.gain1 = Gain()
-    self.gain2 = Gain()
-    self.adaLN1 = AdaptiveLayerNorm()
-
-    self.attn = MultiHeadAttention(num_heads=self.heads, key_dim=self.emb_dim//self.heads, output_shape=self.emb_dim)
-    self.adaLN2 = AdaptiveLayerNorm()
-    self.mlp_up = Dense(self.emb_dim*4, activation="gelu")
-    self.mlp_down = Dense(self.emb_dim)
-
-  def call(self, x, cond):
-    R = self.adaLN1(x, cond)
-    R = self.gain1(self.attn(R, R, R), cond)
-    x = x + R
-    R = self.adaLN2(x, cond)
-    R = self.mlp_up(R)
-    R = self.gain2(self.mlp_down(R), cond)
-    x = x + R
-    return x
-
-  def get_config(self):
-    config = super().get_config()
-    config.update({"hidden_dim": self.emb_dim,
-                   "heads": self.heads,
-                   "context_size": self.context_size})
-    return config
-
-encoder = tf.keras.models.load_model("encoder.keras")
-decoder = tf.keras.models.load_model("decoder.keras")
-diffuser = tf.keras.models.load_model("diffusion-med-coco.keras")
-
-def inference(prompts):
-    N = len(prompts)
-    x_t = tf.random.normal(shape=(N, 32, 32, 4))
-    texts = tf.convert_to_tensor([process_text(p) for p in prompts])
-    t_shape = (N, 1)
-
-    for t in reversed(range(T)):
-        t_batch = tf.convert_to_tensor([[t]] * N)
-        eps_hat = diffuser([x_t, texts, t_batch])
-        x_t = tf.convert_to_tensor(denoise_step(x_t.numpy(), eps_hat.numpy(), t, a, beta), dtype=tf.float32)
-
-    x_0 = x_t.numpy()
-    imgs = decoder(x_0)
-    return imgs
-
-class EndpointHandler:
-    def __init__(self, path="."):
-        pass  # models already loaded above
-
-    def __call__(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
-        prompts = inputs["inputs"]
-        N = len(prompts)
-        x_t = tf.random.normal(shape=(N, *latent_shape))
-        texts = tf.convert_to_tensor([process_text(p) for p in prompts])
-        
-        for t in reversed(range(T)):
-            t_batch = tf.convert_to_tensor([[t]] * N)
-            eps_hat = diffuser([x_t, texts, t_batch])
-            x_t = tf.convert_to_tensor(
-                denoise_step(x_t.numpy(), eps_hat.numpy(), t, a, beta), dtype=tf.float32
-            )
-
-        imgs = decoder(x_t)
+from transformers import PreTrainedTokenizerFast
+
+token2vec = PreTrainedTokenizerFast.from_pretrained("cl100k_tokenizer")
+
+from typing import Dict, Any
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.layers import Dense, LayerNormalization, Conv2D, UpSampling2D, Embedding, MultiHeadAttention
+from tensorflow.keras.saving import register_keras_serializable
+import tensorflow as tf
+
+# @title Config
+def small_config():
+    T = 500
+    beta = np.linspace(1e-4, 0.02, T)
+    alpha = 1 - beta
+    a = np.cumprod(alpha)
+
+    return {
+        "filters": [128, 256],
+        "hidden_dim": 384,
+        "heads": 6,
+        "layers": 8,
+        "patch_size": 4,
+        "batch_size": 64,
+        "T": T,
+        "context_size": 8,
+        "image_size": 128,
+        "latent_shape": (32, 32, 4),
+        "beta": beta,
+        "alpha": alpha,
+        "a": a}
+
+def med_config():
+    T = 1000
+    beta = np.linspace(1e-4, 0.02, T)
+    alpha = 1 - beta
+    a = np.cumprod(alpha)
+
+    return {
+        "filters": [128, 256],
+        "hidden_dim": 768,
+        "heads": 12,
+        "layers": 12,
+        "patch_size": 4,
+        "batch_size": 64,
+        "T": T,
+        "context_size": 8,
+        "image_size": 128,
+        "latent_shape": (32, 32, 4),
+        "beta": beta,
+        "alpha": alpha,
+        "a": a}
+
+def large_config():
+    T = 1000
+    beta = np.linspace(1e-4, 0.02, T)
+    alpha = 1 - beta
+    a = np.cumprod(alpha)
+
+    return {
+        "filters": [128, 256],
+        "hidden_dim": 1024,
+        "heads": 16,
+        "layers": 24,
+        "patch_size": 4,
+        "batch_size": 64,
+        "T": T,
+        "context_size": 8,
+        "image_size": 128,
+        "latent_shape": (32, 32, 4),
+        "beta": beta,
+        "alpha": alpha,
+        "a": a}
+
+config = med_config()
+
+filters = config['filters']
+hidden_dim = config['hidden_dim']
+heads = config['heads']
+layers = config['layers']
+patch_size = config['patch_size']
+batch_size = config['batch_size']
+T = config['T']
+context_size = config['context_size']
+image_size = config['image_size']
+latent_shape = config['latent_shape']
+beta = config['beta']
+alpha = config['alpha']
+a = config['a']
+
+# @title ResBlock, UpBlock, DownBlock
+@register_keras_serializable()
+class ResBlock(tf.keras.layers.Layer):
+  def __init__(self, filters, p, **kwargs):
+    super(ResBlock, self).__init__(**kwargs)
+    self.filters = filters
+    self.p = p
+    self.reshape = Conv2D(filters, kernel_size=1, strides=1, padding="same")
+    #self.norm = BatchNormalization(center=False, scale=False)
+    self.conv1 = Conv2D(filters, kernel_size=p, strides=1, padding="same", activation="swish")
+    self.conv2 = Conv2D(filters, kernel_size=p, strides=1, padding="same")
+
+  def call(self, x):
+    x = self.reshape(x)
+    resid = x
+    #resid = self.norm(resid)
+    resid = self.conv1(resid)
+    resid = self.conv2(resid)
+    x = x + resid
+    return x
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({
+        "filters": self.filters,
+        "p": self.p})
+    return config
+
+@register_keras_serializable()
+class DownBlock(tf.keras.layers.Layer):
+  def __init__(self, filters, **kwargs):
+    super(DownBlock, self).__init__(**kwargs)
+    self.filters = filters
+    self.resBlocks = [ResBlock(f, p=3) for f in filters]
+    self.pool = tf.keras.layers.MaxPool2D(pool_size=(2, 2))
+
+  def call(self, x):
+    for resBlock in self.resBlocks:
+      x = resBlock(x)
+    x = self.pool(x)
+    return x
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({
+        "filters": self.filters})
+    return config
+
+@register_keras_serializable()
+class UpBlock(tf.keras.layers.Layer):
+  def __init__(self, filters, **kwargs):
+    super(UpBlock, self).__init__(**kwargs)
+    self.filters = filters
+    self.resBlocks = [ResBlock(f, p=3) for f in filters]
+    self.upSample = UpSampling2D(size=2, interpolation="bilinear")
+
+  def call(self, x):
+    x = self.upSample(x)
+    for resBlock in self.resBlocks:
+      x = resBlock(x)
+    return x
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({
+        "filters": self.filters})
+    return config
+
+# @title Encoder, Decoder
+@register_keras_serializable()
+class Encoder(tf.keras.Model):
+  def __init__(self, filters, latent_dim, **kwargs):
+    super(Encoder, self).__init__(**kwargs)
+    self.filters = filters
+    self.latent_dim = latent_dim
+    self.downBlocks = [DownBlock([f,f]) for f in filters]
+    self.latent_proj = Conv2D(latent_dim * 2, kernel_size=1, strides=1, padding="same", activation="linear")
+
+  @tf.function
+  def sample(self, mu, logvar):
+    eps = tf.random.normal(shape=tf.shape(mu))
+    return eps * tf.exp(logvar * .5) + mu
+
+  def call(self, x, training=1):
+    for downBlock in self.downBlocks:
+      x = downBlock(x)
+    x = self.latent_proj(x)
+    mu, logvar = tf.split(x, 2, axis=-1)
+    z = self.sample(mu, logvar)
+    return z, mu, logvar
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({
+        "filters": self.filters,
+        "latent_dim": self.latent_dim})
+    return config
+
+  def compute_output_shape(self, input_shape):
+    return (input_shape[0], self.latent_dim), (input_shape[0], self.latent_dim), (input_shape[0], self.latent_dim)
+
+@register_keras_serializable()
+class Decoder(tf.keras.Model):
+  def __init__(self, filters, img_size, **kwargs):
+    super(Decoder, self).__init__(**kwargs)
+    self.filters = filters[::-1]
+    self.img_size = img_size
+    self.undo_latent_proj = Conv2D(filters[0], kernel_size=1, strides=1, padding="same")
+    self.upBlocks = [UpBlock([f,f]) for f in filters]
+    self.conv_proj = Conv2D(3, kernel_size=3, padding="same", activation="linear")
+
+  def call(self, z, training=1):
+    z = self.undo_latent_proj(z)
+    for upBlock in self.upBlocks:
+      z = upBlock(z)
+    x = self.conv_proj(z)
+    return x
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({
+        "filters": self.filters[::-1],
+        "img_size": self.img_size})
+    return config
+
+  def compute_output_shape(self, input_shape):
+    return (input_shape[0], self.img_size, self.img_size, 3)
+
+# @title Helper Functions
+def process_text(text):
+  import tiktoken
+  tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
+  tokens = tokenizer.encode(text)
+  while len(tokens) < context_size:
+      tokens.append(0)
+  return tokens[:context_size]
+
+def normalise_img(img_tensor): # Maps [-1,1] to [0,1]
+  img = img_tensor
+  img *= 0.5
+  img += 0.5
+  return img
+
+def prep_img(img_tensor): # Maps [0,255] to [-1,1]
+  img = img_tensor.copy()
+  img = img / 127.5
+  img -= 1
+  return img
+
+def noisify_img(img_tensor, t, a):  # Returns x_t and the noise used
+  epsilon = np.random.normal(0, 1, img_tensor.shape).astype(np.float32)  # Standard normal
+  sqrt_alpha_bar = np.sqrt(a[t])
+  sqrt_one_minus_alpha_bar = np.sqrt(1 - a[t])
+  x_t = sqrt_alpha_bar * img_tensor + sqrt_one_minus_alpha_bar * epsilon
+  return x_t, epsilon
+
+def denoise_step(x_t, eps_hat, t, a, beta):
+  """
+  Reverse one DDPM step: x_t → x_{t-1}
+  """
+  a_bar_t = tf.convert_to_tensor(a[t], dtype=tf.float32)
+  a_bar_prev = tf.convert_to_tensor(a[t - 1] if t > 0 else 1.0, dtype=tf.float32)
+  a_t = a_bar_t / a_bar_prev
+  beta_t = tf.convert_to_tensor(beta[t], dtype=tf.float32)
+
+  # Avoid NaNs with clamping
+  sqrt_recip_a_t = tf.math.rsqrt(tf.maximum(a_t, 1e-5))
+  sqrt_one_minus_ab = tf.sqrt(tf.maximum(1. - a_bar_t, 1e-5))
+
+  eps_term = (beta_t / sqrt_one_minus_ab) * eps_hat
+  mean = sqrt_recip_a_t * (x_t - eps_term)
+
+  if t > 1:
+      noise = tf.random.normal(shape=x_t.shape)
+      sigma = tf.sqrt(tf.maximum(beta_t, 1e-5))
+      x_prev = mean + sigma * noise
+  else:
+      x_prev = mean
+
+  return x_prev
+
+# @title Transformer Block
+@register_keras_serializable()
+class TransformerBlock(tf.keras.Layer):
+  def __init__(self, context_size, head_no, latent_dim, **kwargs):
+    super().__init__(**kwargs)
+    self.context_size = context_size
+    self.head_no = head_no
+    self.latent_dim = latent_dim
+    self.attn = MultiHeadAttention(num_heads=head_no, key_dim=latent_dim//head_no, output_shape=latent_dim)
+    self.mlp_up = Dense(latent_dim*4, activation="gelu")
+    self.mlp_down = Dense(latent_dim)
+    self.norm1 = LayerNormalization()
+    self.norm2 = LayerNormalization()
+
+  def call(self, x):
+    normed = self.norm1(x)
+    x = x + self.attn(normed, normed, normed)
+    normed = self.norm2(x)
+    dx = self.mlp_up(normed)
+    x = x + self.mlp_down(dx)
+    return x
+
+  def build(self, input_shape):
+    super().build(input_shape)
+
+  def compute_output_shape(self, input_shape):
+    return input_shape
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({
+        "context_size": self.context_size,
+        "head_no": self.head_no,
+        "latent_dim": self.latent_dim})
+    return config
+
+# @title AdaLN-Zero
+@register_keras_serializable()
+class AdaptiveLayerNorm(tf.keras.Layer):
+  def __init__(self, eps=1e-6,**kwargs):
+    self.layernorm = LayerNormalization(epsilon=eps,center=False, scale=False)
+    super(AdaptiveLayerNorm, self).__init__(**kwargs)
+
+  def build(self, input_shape):
+    #B, num_patches, hidden_dim
+    self.M = Dense(input_shape[2], use_bias=True, kernel_initializer='glorot_uniform', activation="linear")
+    self.b = Dense(input_shape[2], use_bias=True, kernel_initializer='glorot_uniform', activation="linear")
+
+  def call(self, x, cond):
+    gamma = self.M(cond)
+    beta = self.b(cond)
+    x = self.layernorm(x)
+    x = x * (1 + tf.expand_dims(gamma, 1)) + tf.expand_dims(beta, 1)
+    return x
+
+  def get_config(self):
+    config = super().get_config()
+    return config
+
+# @title Image Embedder, Unembedder
+@register_keras_serializable()
+class ImageEmbedder(tf.keras.Layer):
+  def __init__(self, latent_size, patch_size, emb_dim,**kwargs):
+    super().__init__(**kwargs)
+    self.emb_dim = emb_dim
+    self.patch_size = patch_size
+    self.latent_size = latent_size
+    self.pos_emb = Embedding(input_dim=(latent_size // patch_size)**2 , output_dim=emb_dim, embeddings_initializer="glorot_uniform")
+    self.reshaper = Dense(emb_dim, kernel_initializer="glorot_uniform")
+    self.conv_expansion = Conv2D(emb_dim, kernel_size=patch_size, strides=patch_size, padding="same")
+
+  def call(self, x):
+    x = self.reshaper(x)
+    x = self.conv_expansion(x)
+    x = tf.reshape(x, shape=[tf.shape(x)[0], tf.shape(x)[1]*tf.shape(x)[2], tf.shape(x)[3]])
+    positions = tf.range(start=0, limit=(self.latent_size // self.patch_size)**2, delta=1)
+    embeddings = self.pos_emb(positions)
+    x = embeddings + x
+    return x
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({
+        "latent_size" : self.latent_size,
+        "patch_size": self.patch_size,
+        "emb_dim": self.emb_dim})
+    return config
+
+@register_keras_serializable()
+class ImageUnembedder(tf.keras.Layer):
+  def __init__(self, latent_size, patch_size, latent_dim, **kwargs):
+    super().__init__(**kwargs)
+    self.latent_dim = latent_dim
+    self.patch_size = patch_size
+    self.latent_size = latent_size
+    self.AdaLN = AdaptiveLayerNorm()
+    self.reshape_to_latent = Dense(patch_size*patch_size*latent_dim, kernel_initializer="glorot_uniform")
+
+  def call(self, x, cond):
+    x = self.AdaLN(x, cond)
+    x = self.reshape_to_latent(x)
+    x = tf.reshape(x, shape=
+     [tf.shape(x)[0],
+      self.latent_size // self.patch_size,
+      self.latent_size // self.patch_size,
+      self.latent_dim*(self.patch_size**2)])
+    x = tf.nn.depth_to_space(x, block_size=self.patch_size)
+    return x
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({
+        "latent_size" : self.latent_size,
+        "patch_size": self.patch_size,
+        "latent_dim": self.latent_dim})
+    return config
+
+# @title LEGACY Prompt and Timestep Embedder
+@register_keras_serializable()
+class ConditioningEmbedder(tf.keras.layers.Layer):
+  def __init__(self, emb_dim, T, context_size, vocab_size=100266, **kwargs):
+    super().__init__(**kwargs)
+    self.emb_dim = emb_dim
+    self.T = T
+    self.context_size = context_size
+    self.vocab_size = vocab_size
+    positions = tf.range(T, dtype=tf.float32)[:, tf.newaxis]
+    frequencies = tf.constant(10000 ** (-tf.range(0, emb_dim, 2, dtype=tf.float32) / emb_dim))
+    angle_rates = positions * frequencies  # (T, emb_dim/2)
+    sin_part = tf.sin(angle_rates)
+    cos_part = tf.cos(angle_rates)
+    emb = tf.stack([sin_part, cos_part], axis=-1)  # (T, emb_dim/2, 2)
+    emb = tf.reshape(emb, [T, emb_dim])  # (T, emb_dim)
+    self.t_embeddings = tf.constant(emb, dtype=tf.float32)
+
+    self.prompt_emb = self.add_weight(shape=(vocab_size, emb_dim), initializer='glorot_uniform', name='prompt_emb', trainable=True)
+    self.CLS = self.add_weight(shape=(emb_dim,), initializer='glorot_uniform', name='CLS', trainable=True)
+    self.prompt_pos_enc = self.add_weight(shape=(1, context_size+1, emb_dim), initializer='glorot_uniform', name='prompt_pos_enc', trainable=True)
+    self.transformer = TransformerBlock(context_size+1, head_no=6, latent_dim=emb_dim)
+
+  def call(self, x):
+    t, prompt_tokens = x
+
+    # ── timestep embedding ───────────────────────────
+    t = tf.cast(tf.squeeze(t, axis=-1), tf.int32)        # (batch,)
+    embedded_t = tf.gather(self.t_embeddings, t)         # (batch, emb_dim)
+    embedded_t = embedded_t[:, tf.newaxis, :]            # (batch, 1, emb_dim)
+
+    # ── prompt embedding path ─────────────────────────
+    embedded_prompt = tf.nn.embedding_lookup(
+        self.prompt_emb, prompt_tokens)                  # (batch, seq_len, emb_dim)
+
+    cls_tok = tf.tile(self.CLS[None, None, :],
+                      [tf.shape(embedded_prompt)[0], 1, 1])
+    embedded_prompt = tf.concat([cls_tok, embedded_prompt], axis=1)
+    embedded_prompt += self.prompt_pos_enc
+    embedded_prompt = self.transformer(embedded_prompt)  # (batch, seq_len+1, emb_dim)
+
+    # add t-embedding to every token (broadcasts along axis-1)
+    embedded_prompt += embedded_t
+
+    # return CLS (keep singleton axis if you need it)
+    return embedded_prompt[:, 0, :]                     # (batch, 1, emb_dim)
+
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({
+        "emb_dim": self.emb_dim,
+        "T": self.T,
+        "context_size": self.context_size,
+        "vocab_size": self.vocab_size})
+    return config
+
+# @title DiT Block
+class Gain(tf.keras.layers.Layer):
+  def __init__(self):
+    super(Gain, self).__init__()
+
+  def build(self, input_shape):
+    self.M = Dense(input_shape[2], use_bias=True,kernel_initializer='glorot_uniform')
+
+  def call(self, x, cond):
+    scale = self.M(cond)
+    x *= tf.expand_dims(scale, 1)
+    return x
+
+@register_keras_serializable()
+class DiTBlock(tf.keras.layers.Layer):
+  def __init__(self, hidden_dim, heads, context_size, **kwargs):
+    super().__init__(**kwargs)
+    self.emb_dim = hidden_dim
+    self.heads = heads
+    self.context_size = context_size
+    self.gain1 = Gain()
+    self.gain2 = Gain()
+    self.adaLN1 = AdaptiveLayerNorm()
+
+    self.attn = MultiHeadAttention(num_heads=self.heads, key_dim=self.emb_dim//self.heads, output_shape=self.emb_dim)
+    self.adaLN2 = AdaptiveLayerNorm()
+    self.mlp_up = Dense(self.emb_dim*4, activation="gelu")
+    self.mlp_down = Dense(self.emb_dim)
+
+  def call(self, x, cond):
+    R = self.adaLN1(x, cond)
+    R = self.gain1(self.attn(R, R, R), cond)
+    x = x + R
+    R = self.adaLN2(x, cond)
+    R = self.mlp_up(R)
+    R = self.gain2(self.mlp_down(R), cond)
+    x = x + R
+    return x
+
+  def get_config(self):
+    config = super().get_config()
+    config.update({"hidden_dim": self.emb_dim,
+                   "heads": self.heads,
+                   "context_size": self.context_size})
+    return config
+
+encoder = tf.keras.models.load_model("encoder.keras")
+decoder = tf.keras.models.load_model("decoder.keras")
+diffuser = tf.keras.models.load_model("diffusion-med-coco.keras")
+
+def inference(prompts):
+    N = len(prompts)
+    x_t = tf.random.normal(shape=(N, 32, 32, 4))
+    texts = tf.convert_to_tensor([process_text(p) for p in prompts])
+    t_shape = (N, 1)
+
+    for t in reversed(range(T)):
+        t_batch = tf.convert_to_tensor([[t]] * N)
+        eps_hat = diffuser([x_t, texts, t_batch])
+        x_t = tf.convert_to_tensor(denoise_step(x_t.numpy(), eps_hat.numpy(), t, a, beta), dtype=tf.float32)
+
+    x_0 = x_t.numpy()
+    imgs = decoder(x_0)
+    return imgs
+
+class EndpointHandler:
+    def __init__(self, path="."):
+        pass  # models already loaded above
+
+    def __call__(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        prompts = inputs["inputs"]
+        N = len(prompts)
+        x_t = tf.random.normal(shape=(N, *latent_shape))
+        texts = tf.convert_to_tensor([process_text(p) for p in prompts])
+        
+        for t in reversed(range(T)):
+            t_batch = tf.convert_to_tensor([[t]] * N)
+            eps_hat = diffuser([x_t, texts, t_batch])
+            x_t = tf.convert_to_tensor(
+                denoise_step(x_t.numpy(), eps_hat.numpy(), t, a, beta), dtype=tf.float32
+            )
+
+        imgs = decoder(x_t)
         return {"outputs": imgs.numpy().tolist()}