Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import tensorflow as tf
+import keras
 from huggingface_hub import hf_hub_download
 import json
 import os
@@ -25,8 +26,8 @@ CACHE_DIR = "./model_cache"
 # Model Architecture Definitions (Required for Loading)
 # ============================================================================
 
-@tf.keras.saving.register_keras_serializable()
-class RotaryEmbedding(tf.keras.layers.Layer):
+@keras.saving.register_keras_serializable()
+class RotaryEmbedding(keras.layers.Layer):
     def __init__(self, dim, max_len=2048, theta=10000, **kwargs):
         super().__init__(**kwargs)
         self.dim = dim
@@ -68,8 +69,8 @@ class RotaryEmbedding(tf.keras.layers.Layer):
         return config
 
 
-@tf.keras.saving.register_keras_serializable()
-class RMSNorm(tf.keras.layers.Layer):
+@keras.saving.register_keras_serializable()
+class RMSNorm(keras.layers.Layer):
     def __init__(self, epsilon=1e-5, **kwargs):
         super().__init__(**kwargs)
         self.epsilon = epsilon
@@ -87,8 +88,8 @@ class RMSNorm(tf.keras.layers.Layer):
         return config
 
 
-@tf.keras.saving.register_keras_serializable()
-class TransformerBlock(tf.keras.layers.Layer):
+@keras.saving.register_keras_serializable()
+class TransformerBlock(keras.layers.Layer):
     def __init__(self, d_model, n_heads, ff_dim, dropout, max_len, rope_theta, layer_idx=0, **kwargs):
         super().__init__(**kwargs)
         self.d_model = d_model
@@ -103,18 +104,18 @@ class TransformerBlock(tf.keras.layers.Layer):
         self.pre_attn_norm = RMSNorm()
         self.pre_ffn_norm = RMSNorm()
         
-        self.q_proj = tf.keras.layers.Dense(d_model, use_bias=False, name="q_proj")
-        self.k_proj = tf.keras.layers.Dense(d_model, use_bias=False, name="k_proj")
-        self.v_proj = tf.keras.layers.Dense(d_model, use_bias=False, name="v_proj")
-        self.out_proj = tf.keras.layers.Dense(d_model, use_bias=False, name="o_proj")
+        self.q_proj = keras.layers.Dense(d_model, use_bias=False, name="q_proj")
+        self.k_proj = keras.layers.Dense(d_model, use_bias=False, name="k_proj")
+        self.v_proj = keras.layers.Dense(d_model, use_bias=False, name="v_proj")
+        self.out_proj = keras.layers.Dense(d_model, use_bias=False, name="o_proj")
         
         self.rope = RotaryEmbedding(self.head_dim, max_len=max_len, theta=rope_theta)
         
-        self.gate_proj = tf.keras.layers.Dense(ff_dim, use_bias=False, name="gate_proj")
-        self.up_proj = tf.keras.layers.Dense(ff_dim, use_bias=False, name="up_proj")
-        self.down_proj = tf.keras.layers.Dense(d_model, use_bias=False, name="down_proj")
+        self.gate_proj = keras.layers.Dense(ff_dim, use_bias=False, name="gate_proj")
+        self.up_proj = keras.layers.Dense(ff_dim, use_bias=False, name="up_proj")
+        self.down_proj = keras.layers.Dense(d_model, use_bias=False, name="down_proj")
         
-        self.dropout = tf.keras.layers.Dropout(dropout)
+        self.dropout = keras.layers.Dropout(dropout)
     
     def call(self, x, training=None):
         B, T, D = tf.shape(x)[0], tf.shape(x)[1], self.d_model
@@ -146,7 +147,7 @@ class TransformerBlock(tf.keras.layers.Layer):
         # FFN (SwiGLU)
         res = x
         y = self.pre_ffn_norm(x)
-        ffn = self.down_proj(tf.keras.activations.silu(self.gate_proj(y)) * self.up_proj(y))
+        ffn = self.down_proj(keras.activations.silu(self.gate_proj(y)) * self.up_proj(y))
         
         return res + self.dropout(ffn, training=training)
     
@@ -164,8 +165,8 @@ class TransformerBlock(tf.keras.layers.Layer):
         return config
 
 
-@tf.keras.saving.register_keras_serializable()
-class SAM1Model(tf.keras.Model):
+@keras.saving.register_keras_serializable()
+class SAM1Model(keras.Model):
     def __init__(self, **kwargs):
         super().__init__()
         if 'config' in kwargs and isinstance(kwargs['config'], dict):
@@ -175,7 +176,7 @@ class SAM1Model(tf.keras.Model):
         else:
             self.cfg = kwargs.get('cfg', kwargs)
         
-        self.embed = tf.keras.layers.Embedding(self.cfg['vocab_size'], self.cfg['d_model'], name="embed_tokens")
+        self.embed = keras.layers.Embedding(self.cfg['vocab_size'], self.cfg['d_model'], name="embed_tokens")
         
         ff_dim = int(self.cfg['d_model'] * self.cfg['ff_mult'])
         block_args = {
@@ -193,7 +194,7 @@ class SAM1Model(tf.keras.Model):
             self.blocks.append(block)
         
         self.norm = RMSNorm(name="final_norm")
-        self.lm_head = tf.keras.layers.Dense(self.cfg['vocab_size'], use_bias=False, name="lm_head")
+        self.lm_head = keras.layers.Dense(self.cfg['vocab_size'], use_bias=False, name="lm_head")
     
     def call(self, input_ids, training=None):
         x = self.embed(input_ids)
@@ -239,7 +240,7 @@ print(f"   Custom tokens: {custom_tokens}")
 eos_token_id = config.get('eos_token_id', 50256)
 
 # Load model with TF function optimization
-model = tf.keras.models.load_model(model_path, compile=False)
+model = keras.models.load_model(model_path, compile=False)
 
 # Create optimized inference function
 @tf.function(reduce_retracing=True)