Update app.py


app.py

@@ -33,18 +33,30 @@ class RotaryEmbedding(keras.layers.Layer):
         self.dim = dim
         self.max_len = max_len
         self.theta = theta
-        self.built_cache = False
     
     def build(self, input_shape):
-        if not self.built_cache:
-            inv_freq = 1.0 / (self.theta ** (tf.range(0, self.dim, 2, dtype=tf.float32) / self.dim))
-            t = tf.range(self.max_len, dtype=tf.float32)
-            freqs = tf.einsum("i,j->ij", t, inv_freq)
-            emb = tf.concat([freqs, freqs], axis=-1)
-            
-            self.cos_cached = tf.constant(tf.cos(emb), dtype=tf.float32)
-            self.sin_cached = tf.constant(tf.sin(emb), dtype=tf.float32)
-            self.built_cache = True
+        # Compute embeddings using numpy, then convert to TF tensors
+        inv_freq = 1.0 / (self.theta ** (np.arange(0, self.dim, 2, dtype=np.float32) / self.dim))
+        t = np.arange(self.max_len, dtype=np.float32)
+        freqs = np.outer(t, inv_freq)
+        emb = np.concatenate([freqs, freqs], axis=-1)
+        
+        # Create non-trainable weights for cos and sin embeddings
+        self.cos_cached = self.add_weight(
+            name="cos_cached",
+            shape=(self.max_len, self.dim),
+            initializer=keras.initializers.Constant(np.cos(emb)),
+            trainable=False,
+            dtype=tf.float32
+        )
+        
+        self.sin_cached = self.add_weight(
+            name="sin_cached",
+            shape=(self.max_len, self.dim),
+            initializer=keras.initializers.Constant(np.sin(emb)),
+            trainable=False,
+            dtype=tf.float32
+        )
         
         super().build(input_shape)
     
@@ -579,6 +591,9 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
         gr.HTML("""
             <div class="header">
                 <div class="celebration">🎉 🎊 ✨ 🎈 🎆</div>
+                <img src="https://cdn-uploads.huggingface.co/production/uploads/64e3486b82fb6ae7a06c749c/yBUDdaTze1L84NaDSpZGf.jpeg" 
+                     alt="SAM-Z-1" 
+                     style="max-width: 400px; border-radius: 12px; margin: 1rem auto; display: block; box-shadow: 0 8px 24px rgba(0,0,0,0.2);">
                 <h1>🤖 SAM-Z-1 Chat 🤖</h1>
                 <p><strong>LATEST RELEASE!</strong> Our fastest non-reasoning model</p>
                 <div class="twin-badge">Twin of SAM-X-1 (Reasoning Model)</div>
@@ -591,6 +606,9 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
     else:
         gr.HTML("""
             <div class="header">
+                <img src="https://cdn-uploads.huggingface.co/production/uploads/64e3486b82fb6ae7a06c749c/yBUDdaTze1L84NaDSpZGf.jpeg" 
+                     alt="SAM-Z-1" 
+                     style="max-width: 300px; border-radius: 12px; margin: 1rem auto; display: block; box-shadow: 0 4px 16px rgba(0,0,0,0.15);">
                 <h1>🤖 SAM-Z-1 Chat</h1>
                 <p>Fast, direct responses without reasoning overhead</p>
                 <p style="font-size: 0.9rem; margin-top: 0.5rem;">