Update app.py

app.py

@@ -10,6 +10,7 @@ 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
@@ -19,12 +20,16 @@ import json
 from tokenizers import Tokenizer
 import numpy as np
 import time
+from functools import lru_cache
 
 # Configure TF threading
 tf.config.threading.set_inter_op_parallelism_threads(NUM_CORES)
 tf.config.threading.set_intra_op_parallelism_threads(NUM_CORES)
 
-print(f"✅ CPU optimized: {NUM_CORES} threads, oneDNN enabled")
+# Enable XLA JIT compilation for CPU
+tf.config.optimizer.set_jit(True)
+
+print(f"✅ CPU optimized: {NUM_CORES} threads, oneDNN enabled, XLA JIT enabled")
 
 # ============================================================================
 # 🎊 FESTIVE MODE TOGGLE 🎊
@@ -41,48 +46,59 @@ MODEL_REPO = "Smilyai-labs/Sam-large-2"
 CACHE_DIR = "./model_cache"
 
 # ============================================================================
-# Model Architecture Definitions (Optimized with KV-Cache)
+# Optimized Model Architecture with KV-Cache
 # ============================================================================
 
 @keras.saving.register_keras_serializable()
 class RotaryEmbedding(keras.layers.Layer):
+    """Optimized RoPE with pre-computed cache."""
+    
     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)
 
-    def _build_cache(self):
-        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(np.cos(emb.numpy()), dtype=tf.float32)
-            self.sin_cached = tf.constant(np.sin(emb.numpy()), dtype=tf.float32)
-            self.built_cache = True
-
-    def rotate_half(self, x):
-        x1, x2 = tf.split(x, 2, axis=-1)
-        return tf.concat([-x2, x1], axis=-1)
-
+    @tf.function(reduce_retracing=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, :, :]
 
-        q_embed = (q * cos) + (self.rotate_half(q) * sin)
-        k_embed = (k * cos) + (self.rotate_half(k) * sin)
+        # Fused rotate_half operation
+        x1_q, x2_q = tf.split(q, 2, axis=-1)
+        x1_k, x2_k = tf.split(k, 2, axis=-1)
+        
+        q_embed = (q * cos) + (tf.concat([-x2_q, x1_q], axis=-1) * sin)
+        k_embed = (k * cos) + (tf.concat([-x2_k, x1_k], axis=-1) * sin)
+        
         return q_embed, k_embed
 
     def get_config(self):
@@ -93,6 +109,8 @@ 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
@@ -102,7 +120,9 @@ 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
 
@@ -114,6 +134,8 @@ class RMSNorm(keras.layers.Layer):
 
 @keras.saving.register_keras_serializable()
 class TransformerBlock(keras.layers.Layer):
+    """Optimized transformer block with efficient attention."""
+    
     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
@@ -124,52 +146,40 @@ class TransformerBlock(keras.layers.Layer):
         self.rope_theta = rope_theta
         self.head_dim = d_model // n_heads
         self.layer_idx = layer_idx
+        self.scale = 1.0 / np.sqrt(self.head_dim)
 
     def build(self, input_shape):
         self.pre_attn_norm = RMSNorm(name="pre_attn_norm")
         self.pre_ffn_norm = RMSNorm(name="pre_ffn_norm")
-        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")
+        
+        # Fused QKV projection for better memory access
+        self.qkv_proj = keras.layers.Dense(self.d_model * 3, use_bias=False, name="qkv_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)
-        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")
+        
+        # Fused gate/up projection
+        self.gate_up_proj = keras.layers.Dense(self.ff_dim * 2, use_bias=False, name="gate_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)
         super().build(input_shape)
 
     def call(self, x, training=None, past_kv=None, use_cache=False):
-        """
-        Args:
-            x: input tensor [B, T, D] (T=1 during cached generation)
-            past_kv: tuple of (past_k, past_v) each [B, n_heads, past_len, head_dim]
-            use_cache: whether to return updated kv cache
-        Returns:
-            output, (new_k, new_v) if use_cache else output, None
-        """
         B = tf.shape(x)[0]
         T = tf.shape(x)[1]
-        dtype = x.dtype
 
         res = x
         y = self.pre_attn_norm(x)
 
-        # Project Q, K, V for current input
-        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])
+        # 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]
 
         # Determine position offset for RoPE
-        if past_kv is not None:
-            past_len = tf.shape(past_kv[0])[2]
-        else:
-            past_len = 0
+        past_len = tf.shape(past_kv[0])[2] if past_kv is not None else 0
 
         # Apply RoPE with position offset
         q, k = self.rope(q, k, offset=past_len)
@@ -181,15 +191,14 @@ class TransformerBlock(keras.layers.Layer):
 
         new_kv = (k, v) if use_cache else None
 
-        # Attention
+        # Scaled dot-product attention
         full_len = tf.shape(k)[2]
-        scores = tf.matmul(q, k, transpose_b=True) / tf.sqrt(tf.cast(self.head_dim, dtype))
+        scores = tf.matmul(q, k, transpose_b=True) * self.scale
 
-        # Causal mask
+        # 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, :], dtype)
-        mask = tf.where(mask == 0, tf.constant(-1e9, dtype=dtype), tf.constant(0.0, dtype=dtype))
+        mask = tf.cast(q_positions[:, None] < k_positions[None, :], q.dtype) * -1e9
         scores = scores + mask[None, None, :, :]
 
         attn = tf.nn.softmax(scores, axis=-1)
@@ -199,10 +208,12 @@ class TransformerBlock(keras.layers.Layer):
 
         x = res + self.dropout(self.out_proj(attn_out), training=training)
 
-        # FFN
+        # Optimized FFN with fused gate/up
         res = x
         y = self.pre_ffn_norm(x)
-        ffn = self.down_proj(keras.activations.silu(self.gate_proj(y)) * self.up_proj(y))
+        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)
         output = res + self.dropout(ffn, training=training)
 
         return output, new_kv
@@ -223,6 +234,8 @@ 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):
@@ -248,16 +261,11 @@ 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):
-        """
-        Args:
-            input_ids: [B, T]
-            past_kv: list of (k, v) tuples, one per layer
-            use_cache: whether to return updated cache
-        Returns:
-            logits, new_past_kv (or None)
-        """
         x = self.embed(input_ids)
 
         new_past_kv = [] if use_cache else None
@@ -271,12 +279,85 @@ 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
         return base_config
 
 
+# ============================================================================
+# Optimized Sampling Functions
+# ============================================================================
+
+@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."""
+    
+    def __init__(self, vocab_size):
+        self.vocab_size = vocab_size
+        self.rng = np.random.default_rng()
+    
+    def sample(self, logits, temperature, top_k, top_p, token_freq, repetition_penalty):
+        """Optimized sampling with vectorized operations."""
+        # Temperature scaling
+        if temperature != 1.0:
+            logits = logits / temperature
+
+        # Vectorized repetition penalty
+        if repetition_penalty != 1.0 and token_freq:
+            freq_tokens = np.array(list(token_freq.keys()), dtype=np.int32)
+            freq_values = np.array(list(token_freq.values()), dtype=np.float32)
+            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)
+
+        # Top-K filtering with partial sort
+        if 0 < top_k < len(logits):
+            top_k_indices = np.argpartition(logits, -top_k)[-top_k:]
+            top_k_logits = logits[top_k_indices]
+        else:
+            top_k_indices = np.arange(len(logits))
+            top_k_logits = logits
+
+        # Stable softmax
+        top_k_logits = top_k_logits - np.max(top_k_logits)
+        exp_logits = np.exp(top_k_logits)
+        top_k_probs = exp_logits / exp_logits.sum()
+
+        # Top-P (nucleus) filtering
+        if top_p < 1.0:
+            sorted_idx = np.argsort(top_k_probs)[::-1]
+            cumsum = np.cumsum(top_k_probs[sorted_idx])
+            cutoff = np.searchsorted(cumsum, top_p) + 1
+            nucleus_idx = sorted_idx[:cutoff]
+            nucleus_probs = top_k_probs[nucleus_idx]
+            nucleus_probs /= nucleus_probs.sum()
+            sampled = self.rng.choice(len(nucleus_probs), p=nucleus_probs)
+            return int(top_k_indices[nucleus_idx[sampled]])
+        else:
+            sampled = self.rng.choice(len(top_k_probs), p=top_k_probs)
+            return int(top_k_indices[sampled])
+
+
 # --- Model and Tokenizer Loading ---
 
 config_path = hf_hub_download(MODEL_REPO, "config.json", cache_dir=CACHE_DIR)
@@ -322,7 +403,7 @@ if use_checkpoint:
         'n_heads': config['num_attention_heads'],
         'ff_mult': config['intermediate_size'] / config['hidden_size'],
         'max_len': config['max_position_embeddings'],
-        'dropout': 0.1,
+        'dropout': 0.0,  # Disable dropout for inference
         'rope_theta': config['rope_theta']
     }
     model = SAM1Model(config=model_config)
@@ -356,11 +437,23 @@ else:
 if model:
     print(f"✅ Model loaded: {config['num_hidden_layers']} layers, {config['vocab_size']} vocab")
 
-# Warm up the model
-print("🔥 Warming up model...")
+# Initialize fast sampler
+sampler = FastSampler(config['vocab_size'])
+
+# Warm up with trace compilation
+print("🔥 Warming up model and compiling traces...")
 warmup_input = tf.constant([[1, 2, 3, 4, 5]], dtype=tf.int32)
-_, _ = model(warmup_input, training=False, use_cache=True)
-print("✅ Model warmed up")
+
+# Warm up prefill
+for _ in range(3):
+    logits, past_kv = model(warmup_input, training=False, use_cache=True)
+
+# Warm up decode step
+single_token = tf.constant([[1]], dtype=tf.int32)
+for _ in range(3):
+    logits, past_kv = model(single_token, training=False, past_kv=past_kv, use_cache=True)
+
+print("✅ Model warmed up and traces compiled")
 
 # ============================================================================
 # Optimized Inference Logic with KV-Cache
@@ -369,45 +462,6 @@ print("✅ Model warmed up")
 stop_generation = False
 
 
-def sample_token(logits, temperature, top_k, top_p, token_freq, repetition_penalty):
-    """Pure NumPy sampling for speed."""
-    # Temperature scaling
-    scaled_logits = logits / temperature
-
-    # Repetition penalty
-    if repetition_penalty != 1.0:
-        for token_id, freq in token_freq.items():
-            if token_id < len(scaled_logits):
-                scaled_logits[token_id] /= (repetition_penalty ** freq)
-
-    # Top-K filtering
-    if top_k > 0 and top_k < len(scaled_logits):
-        top_k_indices = np.argpartition(scaled_logits, -top_k)[-top_k:]
-        top_k_logits = scaled_logits[top_k_indices]
-    else:
-        top_k_indices = np.arange(len(scaled_logits))
-        top_k_logits = scaled_logits
-
-    # Softmax (numerically stable)
-    top_k_logits = top_k_logits - np.max(top_k_logits)
-    top_k_probs = np.exp(top_k_logits)
-    top_k_probs /= top_k_probs.sum()
-
-    # Top-P (nucleus) filtering
-    if top_p < 1.0:
-        sorted_idx = np.argsort(top_k_probs)[::-1]
-        cumsum = np.cumsum(top_k_probs[sorted_idx])
-        cutoff = np.searchsorted(cumsum, top_p) + 1
-        nucleus_idx = sorted_idx[:cutoff]
-        nucleus_probs = top_k_probs[nucleus_idx]
-        nucleus_probs /= nucleus_probs.sum()
-        sampled = np.random.choice(len(nucleus_probs), p=nucleus_probs)
-        return int(top_k_indices[nucleus_idx[sampled]])
-    else:
-        sampled = np.random.choice(len(top_k_probs), p=top_k_probs)
-        return int(top_k_indices[sampled])
-
-
 def generate_stream(
     prompt: str,
     max_tokens: int = 512,
@@ -440,10 +494,9 @@ def generate_stream(
 
     max_context = config['max_position_embeddings']
 
-    start_time = time.time()
+    start_time = time.perf_counter()  # More precise timing
 
     # === PREFILL PHASE ===
-    # Truncate if prompt is too long
     if len(input_ids) > max_context - max_tokens:
         input_ids = input_ids[-(max_context - max_tokens):]
 
@@ -455,22 +508,25 @@ def generate_stream(
         yield f"Error during prefill: {e}"
         return
 
-    # Get logits for last position
+    # Get logits for last position (avoid copy with indexing)
     next_token_logits = logits[0, -1, :].numpy()
 
-    prefill_time = time.time() - start_time
-    print(f"⚡ Prefill: {len(input_ids)} tokens in {prefill_time:.2f}s")
+    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)")
 
     # === GENERATION LOOP ===
-    decode_start = time.time()
+    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]*"
             return
 
-        # Sample next token
-        next_token_id = sample_token(
+        # Sample next token using optimized sampler
+        next_token_id = sampler.sample(
             next_token_logits, temperature, top_k, top_p, token_freq, repetition_penalty
         )
 
@@ -485,7 +541,9 @@ def generate_stream(
         token_text = tokenizer.decode([next_token_id])
         generated_text += token_text
         token_count += 1
-        yield generated_text
+        
+        if step % yield_interval == 0:
+            yield generated_text
 
         # === DECODE PHASE (single token, reuse cache) ===
         next_input = tf.constant([[next_token_id]], dtype=tf.int32)
@@ -498,25 +556,25 @@ def generate_stream(
             
         next_token_logits = logits[0, -1, :].numpy()
 
-        # Truncate cache if too long
-        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  # Keep some buffer
-            past_kv = [
-                (k[:, :, trim_amount:, :], v[:, :, trim_amount:, :])
-                for k, v in past_kv
-            ]
-
-    decode_time = time.time() - decode_start
-    total_time = time.time() - start_time
+        # Truncate cache if too long (less frequent check)
+        if step % 100 == 0:
+            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
+                past_kv = [
+                    (k[:, :, trim_amount:, :], v[:, :, trim_amount:, :])
+                    for k, v in past_kv
+                ]
+
+    decode_time = time.perf_counter() - decode_start
+    total_time = time.perf_counter() - start_time
     
     if token_count > 0:
         decode_tps = token_count / decode_time if decode_time > 0 else 0
-        total_tps = token_count / total_time if total_time > 0 else 0
         
         stats = (
             f"\n\n*[Generated {token_count} tokens in {total_time:.1f}s "
-            f"(prefill: {prefill_time:.1f}s, decode: {decode_tps:.1f} tok/s)]*"
+            f"(prefill: {prefill_time:.2f}s, decode: {decode_tps:.1f} tok/s)]*"
         )
         
         if not stop_generation:
@@ -531,20 +589,21 @@ def generate_stream(
 
 def format_chat_prompt(message: str, history: list, reasoning_enabled: bool) -> str:
     """Format message history and seed <think> if enabled."""
-    prompt = ""
+    prompt_parts = []
+    
     for user_msg, assistant_msg in history:
-        prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
+        prompt_parts.append(f"<|im_start|>user\n{user_msg}<|im_end|>")
         if assistant_msg:
-            # Clean up any stats from previous messages
             clean_msg = assistant_msg.split("\n\n*[")[0]
-            prompt += f"<|im_start|>assistant\n{clean_msg}<|im_end|>\n"
-
-    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
+            prompt_parts.append(f"<|im_start|>assistant\n{clean_msg}<|im_end|>")
 
+    prompt_parts.append(f"<|im_start|>user\n{message}<|im_end|>")
+    prompt_parts.append("<|im_start|>assistant")
+    
     if reasoning_enabled:
-        prompt += "<think>"
+        prompt_parts.append("<think>")
 
-    return prompt
+    return "\n".join(prompt_parts)
 
 
 def chat_stream(
@@ -583,7 +642,6 @@ def chat_stream(
 
         display_response = partial_response
         if should_stop:
-            # Keep the stats portion if present
             stats_start = partial_response.find("\n\n*[")
             if stats_start > earliest_stop:
                 display_response = partial_response[:earliest_stop] + partial_response[stats_start:]
@@ -769,7 +827,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 enabled ⚡
+**Optimization:** KV-Cache + XLA JIT ⚡
 """)
 
     gr.Examples(
@@ -845,7 +903,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
 
 if __name__ == "__main__":
     print("\n" + "=" * 60)
-    print("🚀 Starting Sam-large-2 Chat with KV-Cache Optimization")
+    print("🚀 Starting Sam-large-2 Chat with Optimized Inference")
     print("=" * 60 + "\n")
     demo.queue(max_size=20)
     demo.launch(server_name="0.0.0.0", server_port=7860, share=False, show_error=True)