Update app.py

app.py

@@ -10,7 +10,6 @@ os.environ['TF_NUM_INTRAOP_THREADS'] = str(NUM_CORES)
 os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # Force CPU only
 os.environ['TF_ENABLE_ONEDNN_OPTS'] = '1'  # Intel optimization
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'   # Reduce TF logging
-os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '0'  # We'll handle precision manually
 
 import gradio as gr
 import tensorflow as tf
@@ -27,9 +26,11 @@ tf.config.threading.set_inter_op_parallelism_threads(NUM_CORES)
 tf.config.threading.set_intra_op_parallelism_threads(NUM_CORES)
 
 # Enable XLA JIT compilation for CPU
-tf.config.optimizer.set_jit(True)
-
-print(f"✅ CPU optimized: {NUM_CORES} threads, oneDNN enabled, XLA JIT enabled")
+try:
+    tf.config.optimizer.set_jit(True)
+    print(f"✅ CPU optimized: {NUM_CORES} threads, oneDNN enabled, XLA JIT enabled")
+except:
+    print(f"✅ CPU optimized: {NUM_CORES} threads, oneDNN enabled")
 
 # ============================================================================
 # 🎊 FESTIVE MODE TOGGLE 🎊
@@ -46,53 +47,45 @@ MODEL_REPO = "Smilyai-labs/Sam-large-2"
 CACHE_DIR = "./model_cache"
 
 # ============================================================================
-# Optimized Model Architecture with KV-Cache
+# Model Architecture - MUST MATCH CHECKPOINT STRUCTURE
 # ============================================================================
 
 @keras.saving.register_keras_serializable()
 class RotaryEmbedding(keras.layers.Layer):
-    """Optimized RoPE with pre-computed cache."""
+    """RoPE with pre-computed cache (no trainable weights - compatible with checkpoint)."""
     
     def __init__(self, dim, max_len=2048, theta=10000, **kwargs):
         super().__init__(**kwargs)
         self.dim = dim
         self.max_len = max_len
         self.theta = theta
+        self.built_cache = False
         self.cos_cached = None
         self.sin_cached = None
 
     def build(self, input_shape):
-        # Pre-compute RoPE cache during build
-        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)
-        
-        # Store as non-trainable weights for better graph optimization
-        self.cos_cached = self.add_weight(
-            name="cos_cache",
-            shape=emb.shape,
-            initializer=keras.initializers.Constant(np.cos(emb)),
-            trainable=False
-        )
-        self.sin_cached = self.add_weight(
-            name="sin_cache", 
-            shape=emb.shape,
-            initializer=keras.initializers.Constant(np.sin(emb)),
-            trainable=False
-        )
         super().build(input_shape)
 
-    @tf.function(reduce_retracing=True)
+    def _build_cache(self):
+        if not self.built_cache:
+            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)
+            self.cos_cached = tf.constant(np.cos(emb), dtype=tf.float32)
+            self.sin_cached = tf.constant(np.sin(emb), dtype=tf.float32)
+            self.built_cache = True
+
     def call(self, q, k, offset=0):
         """Apply rotary embeddings with position offset for KV-cache."""
+        self._build_cache()
         seq_len = tf.shape(q)[2]
         dtype = q.dtype
 
         cos = tf.cast(self.cos_cached[offset:offset + seq_len, :], dtype)[None, None, :, :]
         sin = tf.cast(self.sin_cached[offset:offset + seq_len, :], dtype)[None, None, :, :]
 
-        # Fused rotate_half operation
+        # Fused rotate_half
         x1_q, x2_q = tf.split(q, 2, axis=-1)
         x1_k, x2_k = tf.split(k, 2, axis=-1)
         
@@ -109,8 +102,6 @@ class RotaryEmbedding(keras.layers.Layer):
 
 @keras.saving.register_keras_serializable()
 class RMSNorm(keras.layers.Layer):
-    """Optimized RMSNorm."""
-    
     def __init__(self, epsilon=1e-5, **kwargs):
         super().__init__(**kwargs)
         self.epsilon = epsilon
@@ -120,9 +111,7 @@ class RMSNorm(keras.layers.Layer):
         self.scale = self.add_weight(name="scale", shape=(input_shape[-1],), initializer="ones")
         super().build(input_shape)
 
-    @tf.function(reduce_retracing=True)
     def call(self, x):
-        # Fused computation
         variance = tf.reduce_mean(tf.square(x), axis=-1, keepdims=True)
         return x * tf.math.rsqrt(variance + self.epsilon) * self.scale
 
@@ -134,7 +123,7 @@ class RMSNorm(keras.layers.Layer):
 
 @keras.saving.register_keras_serializable()
 class TransformerBlock(keras.layers.Layer):
-    """Optimized transformer block with efficient attention."""
+    """Transformer block - MATCHES ORIGINAL CHECKPOINT STRUCTURE."""
     
     def __init__(self, d_model, n_heads, ff_dim, dropout, max_len, rope_theta, layer_idx=0, **kwargs):
         super().__init__(**kwargs)
@@ -149,17 +138,21 @@ class TransformerBlock(keras.layers.Layer):
         self.scale = 1.0 / np.sqrt(self.head_dim)
 
     def build(self, input_shape):
+        # MUST use same layer names as checkpoint
         self.pre_attn_norm = RMSNorm(name="pre_attn_norm")
         self.pre_ffn_norm = RMSNorm(name="pre_ffn_norm")
         
-        # Fused QKV projection for better memory access
-        self.qkv_proj = keras.layers.Dense(self.d_model * 3, use_bias=False, name="qkv_proj")
+        # Separate Q, K, V projections (matches checkpoint)
+        self.q_proj = keras.layers.Dense(self.d_model, use_bias=False, name="q_proj")
+        self.k_proj = keras.layers.Dense(self.d_model, use_bias=False, name="k_proj")
+        self.v_proj = keras.layers.Dense(self.d_model, use_bias=False, name="v_proj")
         self.out_proj = keras.layers.Dense(self.d_model, use_bias=False, name="o_proj")
         
         self.rope = RotaryEmbedding(self.head_dim, max_len=self.max_len, theta=self.rope_theta)
         
-        # Fused gate/up projection
-        self.gate_up_proj = keras.layers.Dense(self.ff_dim * 2, use_bias=False, name="gate_up_proj")
+        # Separate gate, up, down projections (matches checkpoint)
+        self.gate_proj = keras.layers.Dense(self.ff_dim, use_bias=False, name="gate_proj")
+        self.up_proj = keras.layers.Dense(self.ff_dim, use_bias=False, name="up_proj")
         self.down_proj = keras.layers.Dense(self.d_model, use_bias=False, name="down_proj")
         
         self.dropout = keras.layers.Dropout(self.dropout_rate)
@@ -172,11 +165,15 @@ class TransformerBlock(keras.layers.Layer):
         res = x
         y = self.pre_attn_norm(x)
 
-        # Fused QKV projection
-        qkv = self.qkv_proj(y)
-        qkv = tf.reshape(qkv, [B, T, 3, self.n_heads, self.head_dim])
-        qkv = tf.transpose(qkv, [2, 0, 3, 1, 4])  # [3, B, n_heads, T, head_dim]
-        q, k, v = qkv[0], qkv[1], qkv[2]
+        # Separate Q, K, V projections
+        q = tf.reshape(self.q_proj(y), [B, T, self.n_heads, self.head_dim])
+        q = tf.transpose(q, [0, 2, 1, 3])  # [B, n_heads, T, head_dim]
+
+        k = tf.reshape(self.k_proj(y), [B, T, self.n_heads, self.head_dim])
+        k = tf.transpose(k, [0, 2, 1, 3])
+
+        v = tf.reshape(self.v_proj(y), [B, T, self.n_heads, self.head_dim])
+        v = tf.transpose(v, [0, 2, 1, 3])
 
         # Determine position offset for RoPE
         past_len = tf.shape(past_kv[0])[2] if past_kv is not None else 0
@@ -198,7 +195,7 @@ class TransformerBlock(keras.layers.Layer):
         # Optimized causal mask
         q_positions = tf.range(past_len, past_len + T)
         k_positions = tf.range(full_len)
-        mask = tf.cast(q_positions[:, None] < k_positions[None, :], q.dtype) * -1e9
+        mask = tf.cast(q_positions[:, None] < k_positions[None, :], scores.dtype) * -1e9
         scores = scores + mask[None, None, :, :]
 
         attn = tf.nn.softmax(scores, axis=-1)
@@ -208,12 +205,10 @@ class TransformerBlock(keras.layers.Layer):
 
         x = res + self.dropout(self.out_proj(attn_out), training=training)
 
-        # Optimized FFN with fused gate/up
+        # FFN with SwiGLU
         res = x
         y = self.pre_ffn_norm(x)
-        gate_up = self.gate_up_proj(y)
-        gate, up = tf.split(gate_up, 2, axis=-1)
-        ffn = self.down_proj(tf.nn.silu(gate) * up)
+        ffn = self.down_proj(tf.nn.silu(self.gate_proj(y)) * self.up_proj(y))
         output = res + self.dropout(ffn, training=training)
 
         return output, new_kv
@@ -234,8 +229,6 @@ class TransformerBlock(keras.layers.Layer):
 
 @keras.saving.register_keras_serializable()
 class SAM1Model(keras.Model):
-    """Optimized SAM model with compiled inference."""
-    
     def __init__(self, **kwargs):
         super().__init__()
         if 'config' in kwargs and isinstance(kwargs['config'], dict):
@@ -261,9 +254,6 @@ class SAM1Model(keras.Model):
         ]
         self.norm = RMSNorm(name="final_norm")
         self.lm_head = keras.layers.Dense(self.cfg['vocab_size'], use_bias=False, name="lm_head")
-        
-        self._compiled_prefill = None
-        self._compiled_decode = None
 
     def call(self, input_ids, training=None, past_kv=None, use_cache=False):
         x = self.embed(input_ids)
@@ -279,19 +269,6 @@ class SAM1Model(keras.Model):
         logits = self.lm_head(self.norm(x))
         return logits, new_past_kv
 
-    @tf.function(reduce_retracing=True)
-    def prefill(self, input_ids):
-        """Compiled prefill for initial prompt processing."""
-        return self.call(input_ids, training=False, past_kv=None, use_cache=True)
-
-    @tf.function(reduce_retracing=True, input_signature=[
-        tf.TensorSpec(shape=[1, 1], dtype=tf.int32),
-        tf.TensorSpec(shape=[None], dtype=tf.variant)  # For the list of KV tuples
-    ])
-    def decode_step(self, input_ids, past_kv):
-        """Compiled single-token decode step."""
-        return self.call(input_ids, training=False, past_kv=past_kv, use_cache=True)
-
     def get_config(self):
         base_config = super().get_config()
         base_config['config'] = self.cfg
@@ -299,15 +276,9 @@ class SAM1Model(keras.Model):
 
 
 # ============================================================================
-# Optimized Sampling Functions
+# Optimized Sampling
 # ============================================================================
 
-@lru_cache(maxsize=128)
-def get_top_k_mask(vocab_size, top_k):
-    """Cache top-k masks for common vocab sizes."""
-    return top_k
-
-
 class FastSampler:
     """Vectorized sampler for faster token selection."""
     
@@ -317,6 +288,9 @@ class FastSampler:
     
     def sample(self, logits, temperature, top_k, top_p, token_freq, repetition_penalty):
         """Optimized sampling with vectorized operations."""
+        # Make a copy to avoid modifying original
+        logits = logits.copy()
+        
         # Temperature scaling
         if temperature != 1.0:
             logits = logits / temperature
@@ -328,7 +302,8 @@ class FastSampler:
             valid_mask = freq_tokens < len(logits)
             freq_tokens = freq_tokens[valid_mask]
             freq_values = freq_values[valid_mask]
-            logits[freq_tokens] /= np.power(repetition_penalty, freq_values)
+            if len(freq_tokens) > 0:
+                logits[freq_tokens] /= np.power(repetition_penalty, freq_values)
 
         # Top-K filtering with partial sort
         if 0 < top_k < len(logits):
@@ -455,6 +430,30 @@ for _ in range(3):
 
 print("✅ Model warmed up and traces compiled")
 
+# ============================================================================
+# Compiled Inference Functions
+# ============================================================================
+
+# Create tf.function wrapped inference for speed
+@tf.function(reduce_retracing=True)
+def model_prefill(input_ids):
+    """Compiled prefill function."""
+    return model(input_ids, training=False, use_cache=True)
+
+
+@tf.function(reduce_retracing=True)
+def model_decode(input_ids, past_kv):
+    """Compiled single-token decode function."""
+    return model(input_ids, training=False, past_kv=past_kv, use_cache=True)
+
+
+# Additional warmup for compiled functions
+print("🔥 Compiling tf.function traces...")
+_ = model_prefill(warmup_input)
+_ = model_decode(single_token, past_kv)
+print("✅ Compiled functions ready")
+
+
 # ============================================================================
 # Optimized Inference Logic with KV-Cache
 # ============================================================================
@@ -494,7 +493,7 @@ def generate_stream(
 
     max_context = config['max_position_embeddings']
 
-    start_time = time.perf_counter()  # More precise timing
+    start_time = time.perf_counter()
 
     # === PREFILL PHASE ===
     if len(input_ids) > max_context - max_tokens:
@@ -503,23 +502,21 @@ def generate_stream(
     input_tensor = tf.constant([input_ids], dtype=tf.int32)
     
     try:
-        logits, past_kv = model(input_tensor, training=False, use_cache=True)
+        logits, past_kv = model_prefill(input_tensor)
     except Exception as e:
         yield f"Error during prefill: {e}"
         return
 
-    # Get logits for last position (avoid copy with indexing)
+    # Get logits for last position
     next_token_logits = logits[0, -1, :].numpy()
 
     prefill_time = time.perf_counter() - start_time
-    print(f"⚡ Prefill: {len(input_ids)} tokens in {prefill_time:.3f}s ({len(input_ids)/prefill_time:.1f} tok/s)")
+    prefill_tps = len(input_ids) / prefill_time if prefill_time > 0 else 0
+    print(f"⚡ Prefill: {len(input_ids)} tokens in {prefill_time:.3f}s ({prefill_tps:.1f} tok/s)")
 
     # === GENERATION LOOP ===
     decode_start = time.perf_counter()
     
-    # Pre-compute constants
-    yield_interval = 1  # Yield every token for streaming
-    
     for step in range(max_tokens):
         if stop_generation:
             yield generated_text + "\n\n*[Generation stopped]*"
@@ -541,23 +538,21 @@ def generate_stream(
         token_text = tokenizer.decode([next_token_id])
         generated_text += token_text
         token_count += 1
-        
-        if step % yield_interval == 0:
-            yield generated_text
+        yield generated_text
 
         # === DECODE PHASE (single token, reuse cache) ===
         next_input = tf.constant([[next_token_id]], dtype=tf.int32)
         
         try:
-            logits, past_kv = model(next_input, training=False, past_kv=past_kv, use_cache=True)
+            logits, past_kv = model_decode(next_input, past_kv)
         except Exception as e:
             yield generated_text + f"\n\n*[Error during generation: {e}]*"
             return
             
         next_token_logits = logits[0, -1, :].numpy()
 
-        # Truncate cache if too long (less frequent check)
-        if step % 100 == 0:
+        # Truncate cache if too long (check less frequently)
+        if step % 100 == 99:
             current_len = past_kv[0][0].shape[2] if past_kv and past_kv[0] is not None else 0
             if current_len > max_context:
                 trim_amount = current_len - max_context + 100
@@ -827,7 +822,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
 **Vocab:** {config['vocab_size']:,}
 **Layers:** {config['num_hidden_layers']}
 **Context:** {config['max_position_embeddings']:,} tokens
-**Optimization:** KV-Cache + XLA JIT ⚡
+**Optimization:** KV-Cache + XLA ⚡
 """)
 
     gr.Examples(