Update app.py

app.py

@@ -1,584 +1,620 @@
 """
-WAN-Distributed JAX Inference on Hugging Face Spaces
-Each Space runs this app and can be configured as head or worker.
+Sam-large-2 Distributed Inference - HEAD NODE
+Edit the CONFIG below, then deploy.
 """
 
+# ============================================================================
+# ⚙️ CONFIGURATION - EDIT THIS
+# ============================================================================
+
+CONFIG = {
+    # This node's identity
+    "node_id": "head-main",
+    
+    # Which transformer blocks this node runs (0-indexed)
+    # Sam-large-2 has 12 blocks (0-11)
+    "layer_start": 0,
+    "layer_end": 6,  # exclusive, so this runs blocks 0,1,2,3,4,5
+    
+    # Worker Space URLs (in order of execution)
+    # Leave empty [] for standalone mode (all layers on this node)
+    "worker_urls": [
+        # "https://YOUR-WORKER-SPACE.hf.space",
+    ],
+    
+    # Shared secret for worker communication
+    "secret_token": "sam2-distributed-secret-change-me",
+    
+    # Model settings
+    "model_repo": "Smilyai-labs/Sam-large-2",
+    "cache_dir": "./model_cache",
+}
+
+# ============================================================================
+# CPU Optimization - MUST be before TensorFlow import
+# ============================================================================
+
 import os
+NUM_CORES = os.cpu_count() or 4
+
+os.environ['TF_NUM_INTEROP_THREADS'] = str(NUM_CORES)
+os.environ['TF_NUM_INTRAOP_THREADS'] = str(NUM_CORES)
+os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+os.environ['TF_ENABLE_ONEDNN_OPTS'] = '1'
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+
 import json
 import time
 import threading
-import queue
-from typing import Dict, List, Optional, Any
-from dataclasses import dataclass, field
-import hashlib
+import io
+import base64
+from typing import Dict, List, Optional, Tuple, Any
 
 import gradio as gr
 import numpy as np
 import requests
+import tensorflow as tf
+import keras
+from huggingface_hub import hf_hub_download
 
-# Use CPU JAX
-os.environ["JAX_PLATFORMS"] = "cpu"
-import jax
-import jax.numpy as jnp
+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")
 
 # ============================================================================
-# CONFIGURATION
+# Model Architecture
 # ============================================================================
 
-@dataclass
-class NodeConfig:
-    """Node configuration from environment."""
-    role: str = os.environ.get("NODE_ROLE", "worker")  # "head" or "worker"
-    node_id: str = os.environ.get("NODE_ID", hashlib.md5(os.urandom(8)).hexdigest()[:8])
-    head_url: str = os.environ.get("HEAD_URL", "")  # URL of head Space (for workers)
-    secret_token: str = os.environ.get("SECRET_TOKEN", "default-token")
-    port: int = int(os.environ.get("PORT", "7860"))
+@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
+        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):
+        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)
+
+    def call(self, q, k, offset=0):
+        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)
+        return q_embed, k_embed
+
+    def get_config(self):
+        return {**super().get_config(), "dim": self.dim, "max_len": self.max_len, "theta": self.theta}
+
+
+@keras.saving.register_keras_serializable()
+class RMSNorm(keras.layers.Layer):
+    def __init__(self, epsilon=1e-5, **kwargs):
+        super().__init__(**kwargs)
+        self.epsilon = epsilon
+
+    def build(self, input_shape):
+        self.scale = self.add_weight(name="scale", shape=(input_shape[-1],), initializer="ones")
+        super().build(input_shape)
+
+    def call(self, x):
+        variance = tf.reduce_mean(tf.square(x), axis=-1, keepdims=True)
+        return x * tf.math.rsqrt(variance + self.epsilon) * self.scale
+
+    def get_config(self):
+        return {**super().get_config(), "epsilon": self.epsilon}
+
+
+@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
+        self.n_heads = n_heads
+        self.ff_dim = ff_dim
+        self.dropout_rate = dropout
+        self.max_len = max_len
+        self.rope_theta = rope_theta
+        self.head_dim = d_model // n_heads
+        self.layer_idx = layer_idx
+
+    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")
+        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")
+        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):
+        B, T = tf.shape(x)[0], tf.shape(x)[1]
+        dtype = x.dtype
+
+        res = x
+        y = self.pre_attn_norm(x)
+
+        q = tf.transpose(tf.reshape(self.q_proj(y), [B, T, self.n_heads, self.head_dim]), [0, 2, 1, 3])
+        k = tf.transpose(tf.reshape(self.k_proj(y), [B, T, self.n_heads, self.head_dim]), [0, 2, 1, 3])
+        v = tf.transpose(tf.reshape(self.v_proj(y), [B, T, self.n_heads, self.head_dim]), [0, 2, 1, 3])
+
+        past_len = tf.shape(past_kv[0])[2] if past_kv is not None else 0
+        q, k = self.rope(q, k, offset=past_len)
+
+        if past_kv is not None:
+            k = tf.concat([past_kv[0], k], axis=2)
+            v = tf.concat([past_kv[1], v], axis=2)
+
+        new_kv = (k, v) if use_cache else None
+
+        scores = tf.matmul(q, k, transpose_b=True) / tf.sqrt(tf.cast(self.head_dim, dtype))
+        full_len = tf.shape(k)[2]
+        q_pos = tf.range(past_len, past_len + T)
+        k_pos = tf.range(full_len)
+        mask = tf.where(q_pos[:, None] >= k_pos[None, :], 0.0, -1e9)
+        scores = scores + tf.cast(mask[None, None, :, :], dtype)
+
+        attn = tf.nn.softmax(scores, axis=-1)
+        attn_out = tf.reshape(tf.transpose(tf.matmul(attn, v), [0, 2, 1, 3]), [B, T, self.d_model])
+        x = res + self.dropout(self.out_proj(attn_out), training=training)
+
+        res = x
+        y = self.pre_ffn_norm(x)
+        ffn = self.down_proj(keras.activations.silu(self.gate_proj(y)) * self.up_proj(y))
+        return res + self.dropout(ffn, training=training), new_kv
+
+    def get_config(self):
+        return {**super().get_config(), "d_model": self.d_model, "n_heads": self.n_heads,
+                "ff_dim": self.ff_dim, "dropout": self.dropout_rate, "max_len": self.max_len,
+                "rope_theta": self.rope_theta, "layer_idx": self.layer_idx}
 
 
-CONFIG = NodeConfig()
+# ============================================================================
+# State
+# ============================================================================
+
+class ModelState:
+    def __init__(self):
+        self.config = None
+        self.tokenizer = None
+        self.eos_token_id = 50256
+        
+        # Model components
+        self.embedding = None
+        self.blocks: List = []
+        self.final_norm = None
+        self.lm_head = None
+        
+        self.my_block_start = 0
+        self.my_block_end = 0
 
+STATE = ModelState()
+stop_generation = False
 
 # ============================================================================
-# SHARED STATE
+# Serialization
 # ============================================================================
 
-class ClusterState:
-    """Shared state for the cluster."""
-    
-    def __init__(self):
-        self.workers: Dict[str, Dict] = {}  # worker_id -> info
-        self.shards: Dict[str, np.ndarray] = {}  # shard_name -> data
-        self.lock = threading.Lock()
-        self.is_initialized = False
-        self.pending_results: Dict[str, Any] = {}
-        self.request_queue: queue.Queue = queue.Queue()
-    
-    def register_worker(self, worker_id: str, url: str, info: Dict) -> bool:
-        with self.lock:
-            self.workers[worker_id] = {
-                "url": url,
-                "info": info,
-                "registered_at": time.time(),
-                "last_seen": time.time(),
-                "status": "active"
-            }
-            return True
-    
-    def get_workers(self) -> List[Dict]:
-        with self.lock:
-            return [
-                {"worker_id": wid, **winfo}
-                for wid, winfo in self.workers.items()
-                if winfo.get("status") == "active"
-            ]
-    
-    def store_shard(self, name: str, data: np.ndarray):
-        with self.lock:
-            self.shards[name] = data
-    
-    def get_shard(self, name: str) -> Optional[np.ndarray]:
-        with self.lock:
-            return self.shards.get(name)
-    
-    def heartbeat(self, worker_id: str):
-        with self.lock:
-            if worker_id in self.workers:
-                self.workers[worker_id]["last_seen"] = time.time()
-
-
-STATE = ClusterState()
+def serialize_tensor(tensor: tf.Tensor) -> str:
+    buffer = io.BytesIO()
+    np.save(buffer, tensor.numpy(), allow_pickle=False)
+    return base64.b64encode(buffer.getvalue()).decode('utf-8')
+
+def deserialize_tensor(data: str) -> tf.Tensor:
+    buffer = io.BytesIO(base64.b64decode(data))
+    return tf.constant(np.load(buffer, allow_pickle=False))
 
+def serialize_kv_cache(past_kv):
+    if past_kv is None:
+        return None
+    return [{"k": serialize_tensor(k), "v": serialize_tensor(v)} if k is not None else None for k, v in past_kv]
+
+def deserialize_kv_cache(data):
+    if data is None:
+        return None
+    return [(deserialize_tensor(item["k"]), deserialize_tensor(item["v"])) if item else None for item in data]
 
 # ============================================================================
-# HTTP COMMUNICATION LAYER
+# HTTP Communication
 # ============================================================================
 
-def make_request(url: str, endpoint: str, data: Dict, timeout: int = 30) -> Optional[Dict]:
-    """Make HTTP request to another Space."""
+def call_worker(url: str, hidden_states: tf.Tensor, past_kv=None, use_cache=False) -> Tuple[tf.Tensor, Any]:
+    """Send hidden states to worker and get result."""
     try:
-        full_url = f"{url.rstrip('/')}/api/{endpoint}"
-        headers = {"Authorization": f"Bearer {CONFIG.secret_token}"}
-        
         response = requests.post(
-            full_url,
-            json=data,
-            headers=headers,
-            timeout=timeout
+            f"{url.rstrip('/')}/api/forward",
+            json={
+                "hidden_states": serialize_tensor(hidden_states),
+                "past_kv": serialize_kv_cache(past_kv),
+                "use_cache": use_cache,
+            },
+            headers={"Authorization": f"Bearer {CONFIG['secret_token']}"},
+            timeout=120
         )
         
         if response.status_code == 200:
-            return response.json()
+            result = response.json()
+            output = deserialize_tensor(result["hidden_states"])
+            new_kv = deserialize_kv_cache(result.get("past_kv"))
+            return output, new_kv
         else:
-            print(f"Request failed: {response.status_code} - {response.text}")
-            return None
+            raise RuntimeError(f"Worker returned {response.status_code}")
     except Exception as e:
-        print(f"Request error: {e}")
-        return None
-
+        raise RuntimeError(f"Worker call failed: {e}")
 
 # ============================================================================
-# WORKER LOGIC
+# Model Loading
 # ============================================================================
 
-def worker_register_with_head():
-    """Register this worker with the head node."""
-    if not CONFIG.head_url:
-        print("No HEAD_URL configured, cannot register")
-        return False
-    
-    # Get this Space's URL from environment or construct it
-    space_url = os.environ.get("SPACE_URL", f"http://localhost:{CONFIG.port}")
-    
-    result = make_request(
-        CONFIG.head_url,
-        "register_worker",
-        {
-            "worker_id": CONFIG.node_id,
-            "worker_url": space_url,
-            "info": {
-                "jax_devices": len(jax.devices()),
-                "platform": jax.default_backend(),
-            }
-        }
-    )
-    
-    if result and result.get("success"):
-        print(f"Registered with head at {CONFIG.head_url}")
-        return True
-    return False
-
-
-def worker_heartbeat_loop():
-    """Send periodic heartbeats to head."""
-    while True:
-        time.sleep(30)
-        if CONFIG.head_url:
-            make_request(
-                CONFIG.head_url,
-                "heartbeat",
-                {"worker_id": CONFIG.node_id}
-            )
-
-
-def worker_forward_pass(input_data: np.ndarray) -> np.ndarray:
-    """Run forward pass on local shards."""
-    x = jnp.array(input_data)
-    
-    # Apply each stored shard (simple linear layers for demo)
-    for name, weight in sorted(STATE.shards.items()):
-        if weight.ndim == 2:
-            # Matrix multiply for weight matrices
-            if x.shape[-1] == weight.shape[0]:
-                x = x @ weight
-        elif weight.ndim == 1:
-            # Add for biases
-            if x.shape[-1] == weight.shape[0]:
-                x = x + weight
-    
-    # Apply simple activation
-    x = jax.nn.relu(x)
-    
-    return np.array(x)
-
+def load_model():
+    """Load model and extract components for this node."""
+    print("🚀 Loading model...")
+    
+    # Load config
+    config_path = hf_hub_download(CONFIG["model_repo"], "config.json", cache_dir=CONFIG["cache_dir"])
+    with open(config_path, 'r') as f:
+        model_config = json.load(f)
+    STATE.config = model_config
+    
+    # Load tokenizer
+    from transformers import AutoTokenizer
+    from tokenizers import Tokenizer
+    
+    hf_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    hf_tokenizer.add_special_tokens({"additional_special_tokens": 
+        ["<|im_start|>", "<|im_end|>", "<think>", "</think>", "<CONTINUE>", "<im end for model tun>"]})
+    os.makedirs("./temp_tokenizer", exist_ok=True)
+    hf_tokenizer.save_pretrained("./temp_tokenizer")
+    STATE.tokenizer = Tokenizer.from_file("./temp_tokenizer/tokenizer.json")
+    STATE.eos_token_id = model_config.get('eos_token_id', 50256)
+    
+    # Load weights
+    weights_path = hf_hub_download(CONFIG["model_repo"], "ckpt.weights.h5", cache_dir=CONFIG["cache_dir"])
+    
+    # Build full model to load weights
+    n_layers = model_config['num_hidden_layers']
+    d_model = model_config['hidden_size']
+    n_heads = model_config['num_attention_heads']
+    ff_dim = model_config['intermediate_size']
+    max_len = model_config['max_position_embeddings']
+    rope_theta = model_config['rope_theta']
+    vocab_size = model_config['vocab_size']
+    
+    # Temporary full model
+    embedding = keras.layers.Embedding(vocab_size, d_model, name="embed_tokens")
+    blocks = [TransformerBlock(d_model, n_heads, ff_dim, 0.0, max_len, rope_theta, i, name=f"block_{i}") 
+              for i in range(n_layers)]
+    final_norm = RMSNorm(name="final_norm")
+    lm_head = keras.layers.Dense(vocab_size, use_bias=False, name="lm_head")
+    
+    # Build
+    dummy = tf.zeros((1, 16), dtype=tf.int32)
+    x = embedding(dummy)
+    for block in blocks:
+        x, _ = block(x)
+    x = final_norm(x)
+    _ = lm_head(x)
+    
+    # Load weights into a temp model structure
+    class TempModel(keras.Model):
+        def __init__(self):
+            super().__init__()
+            self.embed = embedding
+            self.blocks = blocks
+            self.norm = final_norm
+            self.lm_head = lm_head
+        def call(self, x):
+            x = self.embed(x)
+            for b in self.blocks:
+                x, _ = b(x)
+            return self.lm_head(self.norm(x))
+    
+    temp_model = TempModel()
+    temp_model(dummy)
+    temp_model.load_weights(weights_path)
+    print("✅ Weights loaded")
+    
+    # Extract components for this node
+    STATE.my_block_start = CONFIG["layer_start"]
+    STATE.my_block_end = CONFIG["layer_end"] if CONFIG["layer_end"] > 0 else n_layers
+    
+    # HEAD always has embedding
+    STATE.embedding = embedding
+    
+    # Extract our blocks
+    STATE.blocks = blocks[STATE.my_block_start:STATE.my_block_end]
+    print(f"✅ Loaded blocks {STATE.my_block_start} to {STATE.my_block_end - 1}")
+    
+    # HEAD has final norm and lm_head only if no workers OR we handle last block
+    has_workers = len(CONFIG["worker_urls"]) > 0
+    if not has_workers:
+        STATE.final_norm = final_norm
+        STATE.lm_head = lm_head
+        print("✅ Loaded final norm and LM head (standalone mode)")
+    
+    # Warmup
+    print("🔥 Warming up...")
+    dummy = tf.constant([[1, 2, 3]], dtype=tf.int32)
+    x = STATE.embedding(dummy)
+    for block in STATE.blocks:
+        x, _ = block(x, use_cache=False)
+    if STATE.lm_head:
+        _ = STATE.lm_head(STATE.final_norm(x))
+    
+    print("✅ Model ready!")
+    return True
 
 # ============================================================================
-# HEAD NODE LOGIC
+# Distributed Forward
 # ============================================================================
 
-def head_distribute_model(params: Dict[str, np.ndarray]) -> bool:
-    """Distribute model parameters to workers."""
-    workers = STATE.get_workers()
-    if not workers:
-        print("No workers available")
-        return False
-    
-    # Simple round-robin distribution
-    param_list = list(params.items())
-    shards_per_worker = max(1, len(param_list) // len(workers))
-    
-    for i, worker in enumerate(workers):
-        start_idx = i * shards_per_worker
-        end_idx = start_idx + shards_per_worker if i < len(workers) - 1 else len(param_list)
-        
-        worker_shards = dict(param_list[start_idx:end_idx])
-        
-        for shard_name, shard_data in worker_shards.items():
-            result = make_request(
-                worker["url"],
-                "store_shard",
-                {
-                    "name": shard_name,
-                    "data": shard_data.tolist(),
-                    "shape": list(shard_data.shape),
-                    "dtype": str(shard_data.dtype)
-                },
-                timeout=60
-            )
-            
-            if not result or not result.get("success"):
-                print(f"Failed to send shard {shard_name} to worker {worker['worker_id']}")
-                return False
-    
-    print(f"Distributed {len(params)} shards to {len(workers)} workers")
-    return True
-
-
-def head_run_inference(input_data: np.ndarray) -> np.ndarray:
-    """Run distributed inference across workers."""
-    workers = STATE.get_workers()
-    
-    if not workers:
-        # No workers, run locally
-        return worker_forward_pass(input_data)
-    
-    # Pipeline through workers
-    current_data = input_data
-    
-    for worker in workers:
-        result = make_request(
-            worker["url"],
-            "forward",
-            {
-                "data": current_data.tolist(),
-                "shape": list(current_data.shape),
-            },
-            timeout=60
-        )
-        
-        if result and "output" in result:
-            current_data = np.array(result["output"])
-        else:
-            print(f"Worker {worker['worker_id']} failed, using local fallback")
-            current_data = worker_forward_pass(current_data)
+def forward_pass(input_ids: tf.Tensor, past_kv_local=None, past_kv_workers=None, use_cache=False):
+    """
+    Full forward pass through HEAD + all workers.
+    Returns logits and updated KV caches.
+    """
+    # Embedding
+    x = STATE.embedding(input_ids)
+    
+    # Local blocks
+    new_local_kv = [] if use_cache else None
+    for i, block in enumerate(STATE.blocks):
+        block_past = past_kv_local[i] if past_kv_local else None
+        x, kv = block(x, past_kv=block_past, use_cache=use_cache)
+        if use_cache:
+            new_local_kv.append(kv)
+    
+    # Workers
+    new_worker_kv = {} if use_cache else None
+    for worker_url in CONFIG["worker_urls"]:
+        worker_past = past_kv_workers.get(worker_url) if past_kv_workers else None
+        x, worker_kv = call_worker(worker_url, x, worker_past, use_cache)
+        if use_cache:
+            new_worker_kv[worker_url] = worker_kv
+    
+    # Final (only if standalone or last worker returned to us)
+    # In distributed mode, the last worker applies final_norm + lm_head
+    if STATE.lm_head:
+        logits = STATE.lm_head(STATE.final_norm(x))
+    else:
+        # x should already be logits from last worker
+        logits = x
     
-    return current_data
-
+    return logits, new_local_kv, new_worker_kv
 
 # ============================================================================
-# API ENDPOINTS (Gradio doesn't have native API, so we use a simple approach)
+# Generation
 # ============================================================================
 
-def api_handler(endpoint: str, data: Dict) -> Dict:
-    """Handle API requests based on endpoint."""
-    
-    # Verify token
-    # (In production, check Authorization header)
+def sample_token(logits, temperature, top_k, top_p, token_freq, rep_penalty):
+    logits = np.array(logits) / temperature
     
-    if endpoint == "register_worker":
-        success = STATE.register_worker(
-            data["worker_id"],
-            data["worker_url"],
-            data.get("info", {})
-        )
-        return {"success": success, "message": "Worker registered" if success else "Failed"}
-    
-    elif endpoint == "heartbeat":
-        STATE.heartbeat(data.get("worker_id", ""))
-        return {"success": True}
-    
-    elif endpoint == "store_shard":
-        shard_data = np.array(data["data"], dtype=data.get("dtype", "float32"))
-        shard_data = shard_data.reshape(data["shape"])
-        STATE.store_shard(data["name"], shard_data)
-        return {"success": True, "shard": data["name"]}
-    
-    elif endpoint == "forward":
-        input_data = np.array(data["data"]).reshape(data["shape"])
-        output = worker_forward_pass(input_data)
-        return {"output": output.tolist(), "shape": list(output.shape)}
-    
-    elif endpoint == "status":
-        return {
-            "node_id": CONFIG.node_id,
-            "role": CONFIG.role,
-            "workers": len(STATE.get_workers()),
-            "shards": list(STATE.shards.keys()),
-            "jax_devices": len(jax.devices()),
-        }
-    
-    elif endpoint == "get_workers":
-        return {"workers": STATE.get_workers()}
+    for tid, freq in token_freq.items():
+        if tid < len(logits):
+            logits[tid] /= (rep_penalty ** freq)
     
+    if 0 < top_k < len(logits):
+        top_k_idx = np.argpartition(logits, -top_k)[-top_k:]
+        top_k_logits = logits[top_k_idx]
     else:
-        return {"error": f"Unknown endpoint: {endpoint}"}
-
-
-# ============================================================================
-# GRADIO INTERFACE
-# ============================================================================
-
-def create_test_model(num_layers: int = 4, hidden_size: int = 128) -> Dict[str, np.ndarray]:
-    """Create a simple test model."""
-    params = {}
+        top_k_idx = np.arange(len(logits))
+        top_k_logits = logits
+    
+    top_k_logits = top_k_logits - np.max(top_k_logits)
+    probs = np.exp(top_k_logits)
+    probs /= probs.sum()
     
-    for i in range(num_layers):
-        params[f"layer_{i}_weight"] = np.random.randn(hidden_size, hidden_size).astype(np.float32) * 0.02
-        params[f"layer_{i}_bias"] = np.zeros(hidden_size, dtype=np.float32)
+    if top_p < 1.0:
+        sorted_idx = np.argsort(probs)[::-1]
+        cumsum = np.cumsum(probs[sorted_idx])
+        cutoff = np.searchsorted(cumsum, top_p) + 1
+        nucleus_idx = sorted_idx[:cutoff]
+        nucleus_probs = probs[nucleus_idx]
+        nucleus_probs /= nucleus_probs.sum()
+        sampled = np.random.choice(len(nucleus_probs), p=nucleus_probs)
+        return int(top_k_idx[nucleus_idx[sampled]])
     
-    return params
+    return int(top_k_idx[np.random.choice(len(probs), p=probs)])
 
 
-def gradio_run_inference(input_text: str) -> str:
-    """Run inference from Gradio UI."""
-    # Simple tokenization (ASCII values normalized)
-    tokens = np.array([ord(c) / 128.0 for c in input_text[:128]], dtype=np.float32)
+def generate_stream(prompt: str, max_tokens=512, temperature=0.8, top_k=40, top_p=0.9, rep_penalty=1.1):
+    global stop_generation
+    stop_generation = False
     
-    # Pad to fixed size
-    if len(tokens) < 128:
-        tokens = np.pad(tokens, (0, 128 - len(tokens)))
+    input_ids = [i for i in STATE.tokenizer.encode(prompt).ids if i != STATE.eos_token_id]
+    if not input_ids:
+        yield "Error: Empty prompt"
+        return
     
-    # Run inference
-    start_time = time.time()
+    generated = ""
+    token_freq = {}
     
-    if CONFIG.role == "head":
-        output = head_run_inference(tokens)
-    else:
-        output = worker_forward_pass(tokens)
-    
-    latency = (time.time() - start_time) * 1000
-    
-    # Format output
-    result = f"Output shape: {output.shape}\n"
-    result += f"Output mean: {output.mean():.4f}\n"
-    result += f"Output std: {output.std():.4f}\n"
-    result += f"Latency: {latency:.1f}ms\n"
-    result += f"Workers used: {len(STATE.get_workers())}"
-    
-    return result
-
-
-def gradio_get_status() -> str:
-    """Get cluster status for Gradio UI."""
-    status = {
-        "Node ID": CONFIG.node_id,
-        "Role": CONFIG.role,
-        "JAX Devices": len(jax.devices()),
-        "JAX Backend": jax.default_backend(),
-        "Stored Shards": len(STATE.shards),
-        "Shard Names": list(STATE.shards.keys())[:10],  # First 10
-    }
-    
-    if CONFIG.role == "head":
-        workers = STATE.get_workers()
-        status["Connected Workers"] = len(workers)
-        status["Worker List"] = [
-            f"{w['worker_id']} @ {w['url']}"
-            for w in workers
-        ]
-    else:
-        status["Head URL"] = CONFIG.head_url
-        status["Registered"] = STATE.is_initialized
+    stop_ids = {STATE.eos_token_id, STATE.tokenizer.token_to_id("<|im_end|>"), 
+                STATE.tokenizer.token_to_id("<im end for model tun>")}
+    stop_ids.discard(None)
     
-    return json.dumps(status, indent=2)
-
-
-def gradio_init_model(num_layers: int, hidden_size: int) -> str:
-    """Initialize and distribute model."""
-    params = create_test_model(int(num_layers), int(hidden_size))
+    max_ctx = STATE.config['max_position_embeddings']
+    if len(input_ids) > max_ctx - max_tokens:
+        input_ids = input_ids[-(max_ctx - max_tokens):]
     
-    if CONFIG.role == "head":
-        workers = STATE.get_workers()
-        if workers:
-            success = head_distribute_model(params)
-            if success:
-                return f"Distributed {len(params)} shards to {len(workers)} workers"
-            else:
-                return "Failed to distribute model"
-        else:
-            # Store locally
-            for name, data in params.items():
-                STATE.store_shard(name, data)
-            return f"No workers - stored {len(params)} shards locally"
-    else:
-        # Worker stores locally
-        for name, data in params.items():
-            STATE.store_shard(name, data)
-        return f"Stored {len(params)} shards locally"
-
-
-def gradio_register_worker(worker_url: str) -> str:
-    """Manually register a worker (for head node)."""
-    if CONFIG.role != "head":
-        return "Only head node can register workers"
-    
-    # Ping the worker
-    result = make_request(worker_url, "status", {})
+    start = time.time()
     
-    if result:
-        worker_id = result.get("node_id", f"worker_{len(STATE.workers)}")
-        STATE.register_worker(worker_id, worker_url, result)
-        return f"Registered worker {worker_id}"
-    else:
-        return f"Failed to reach worker at {worker_url}"
-
-
-def gradio_api_call(endpoint: str, json_data: str) -> str:
-    """Make API call (for testing)."""
+    # Prefill
+    input_tensor = tf.constant([input_ids], dtype=tf.int32)
     try:
-        data = json.loads(json_data) if json_data else {}
-        result = api_handler(endpoint, data)
-        return json.dumps(result, indent=2)
+        logits, local_kv, worker_kv = forward_pass(input_tensor, None, None, use_cache=True)
     except Exception as e:
-        return f"Error: {e}"
-
-
-# ============================================================================
-# MAIN APP
-# ============================================================================
-
-def create_app():
-    """Create Gradio app based on node role."""
+        yield f"Error: {e}"
+        return
     
-    # Start background tasks
-    if CONFIG.role == "worker" and CONFIG.head_url:
-        # Register with head
-        threading.Thread(target=lambda: time.sleep(5) or worker_register_with_head(), daemon=True).start()
-        # Heartbeat loop
-        threading.Thread(target=worker_heartbeat_loop, daemon=True).start()
+    next_logits = logits[0, -1, :].numpy()
+    prefill_time = time.time() - start
+    print(f"⚡ Prefill: {len(input_ids)} tokens in {prefill_time:.2f}s")
     
-    # Create Gradio interface
-    with gr.Blocks(title=f"WAN-JAX {CONFIG.role.upper()} - {CONFIG.node_id}") as app:
-        gr.Markdown(f"""
-        # 🌐 WAN-Distributed JAX Inference
+    # Generate
+    decode_start = time.time()
+    tokens_generated = 0
+    
+    for _ in range(max_tokens):
+        if stop_generation:
+            yield generated + "\n\n*[Stopped]*"
+            return
         
-        **Node ID:** `{CONFIG.node_id}` | **Role:** `{CONFIG.role.upper()}`
+        next_id = sample_token(next_logits, temperature, top_k, top_p, token_freq, rep_penalty)
         
-        {"This is the **HEAD** node - it coordinates workers and runs inference." if CONFIG.role == "head" else "This is a **WORKER** node - it stores model shards and computes."}
-        """)
+        if next_id in stop_ids:
+            break
         
-        with gr.Tab("Status"):
-            status_output = gr.Textbox(label="Cluster Status", lines=15)
-            refresh_btn = gr.Button("Refresh Status")
-            refresh_btn.click(gradio_get_status, outputs=status_output)
-            
-            # Auto-refresh on load
-            app.load(gradio_get_status, outputs=status_output)
+        token_freq[next_id] = token_freq.get(next_id, 0) + 1
+        generated += STATE.tokenizer.decode([next_id])
+        tokens_generated += 1
+        yield generated
         
-        with gr.Tab("Inference"):
-            with gr.Row():
-                with gr.Column():
-                    input_text = gr.Textbox(
-                        label="Input Text",
-                        placeholder="Enter text to process...",
-                        lines=3
-                    )
-                    infer_btn = gr.Button("Run Inference", variant="primary")
-                
-                with gr.Column():
-                    output_text = gr.Textbox(label="Output", lines=8)
-            
-            infer_btn.click(gradio_run_inference, inputs=input_text, outputs=output_text)
+        # Next step
+        next_input = tf.constant([[next_id]], dtype=tf.int32)
+        try:
+            logits, local_kv, worker_kv = forward_pass(next_input, local_kv, worker_kv, use_cache=True)
+        except Exception as e:
+            yield generated + f"\n\n*[Error: {e}]*"
+            return
         
-        with gr.Tab("Model"):
-            with gr.Row():
-                num_layers = gr.Slider(1, 12, value=4, step=1, label="Number of Layers")
-                hidden_size = gr.Slider(32, 512, value=128, step=32, label="Hidden Size")
-            
-            init_btn = gr.Button("Initialize Model")
-            init_output = gr.Textbox(label="Result")
-            
-            init_btn.click(
-                gradio_init_model,
-                inputs=[num_layers, hidden_size],
-                outputs=init_output
-            )
+        next_logits = logits[0, -1, :].numpy()
+    
+    # Stats
+    if tokens_generated > 0:
+        total = time.time() - start
+        tps = tokens_generated / (time.time() - decode_start)
+        workers = len(CONFIG["worker_urls"])
+        mode = f", {workers} workers" if workers else " standalone"
+        generated += f"\n\n*[{tokens_generated} tokens in {total:.1f}s ({tps:.1f} tok/s){mode}]*"
+    
+    yield generated
+
+
+def format_prompt(message: str, history: list, reasoning: bool) -> str:
+    prompt = ""
+    for user, assistant in history:
+        prompt += f"<|im_start|>user\n{user}<|im_end|>\n"
+        if assistant:
+            prompt += f"<|im_start|>assistant\n{assistant.split('*[')[0].strip()}<|im_end|>\n"
+    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
+    if reasoning:
+        prompt += "<think>"
+    return prompt
+
+
+def chat_respond(message, history, max_tokens, temp, top_k, top_p, rep_pen, reasoning):
+    if not message.strip():
+        yield history
+        return
+    
+    prompt = format_prompt(message, history, reasoning)
+    
+    for text in generate_stream(prompt, max_tokens, temp, top_k, top_p, rep_pen):
+        display = text
+        for tag in ["<|im_end|>", "<im end for model tun>"]:
+            if tag in display:
+                idx = display.find(tag)
+                stats = display.find("\n\n*[")
+                display = display[:idx] + (display[stats:] if stats > idx else "")
         
-        if CONFIG.role == "head":
-            with gr.Tab("Workers"):
-                worker_url_input = gr.Textbox(
-                    label="Worker Space URL",
-                    placeholder="https://username-spacename.hf.space"
-                )
-                register_btn = gr.Button("Register Worker")
-                register_output = gr.Textbox(label="Result")
-                
-                register_btn.click(
-                    gradio_register_worker,
-                    inputs=worker_url_input,
-                    outputs=register_output
-                )
+        if reasoning and '<think>' in display and '</think>' in display:
+            s, e = display.find('<think>'), display.find('</think>')
+            if s < e:
+                thought = display[s+7:e].strip()
+                display = display[:s] + f'<details><summary>🧠 Reasoning</summary><p>{thought}</p></details>' + display[e+8:]
         
-        with gr.Tab("API"):
-            gr.Markdown("""
-            ### Direct API Access
-            Use this tab to test API endpoints directly.
-            
-            **Endpoints:**
-            - `status` - Get node status
-            - `register_worker` - Register a worker (head only)
-            - `store_shard` - Store a model shard
-            - `forward` - Run forward pass
-            - `get_workers` - List workers (head only)
-            """)
-            
-            endpoint_input = gr.Textbox(label="Endpoint", value="status")
-            json_input = gr.Textbox(label="JSON Data", value="{}", lines=5)
-            api_btn = gr.Button("Call API")
-            api_output = gr.Textbox(label="Response", lines=10)
-            
-            api_btn.click(
-                gradio_api_call,
-                inputs=[endpoint_input, json_input],
-                outputs=api_output
-            )
-    
-    return app
+        yield history + [[message, display.strip()]]
+
 
+def stop():
+    global stop_generation
+    stop_generation = True
 
 # ============================================================================
-# FASTAPI MOUNTING FOR TRUE API ACCESS
+# Gradio UI
 # ============================================================================
 
-# Optional: Mount FastAPI for proper API endpoints
-try:
-    from fastapi import FastAPI, Request, HTTPException
-    from fastapi.responses import JSONResponse
+def create_ui():
+    workers = CONFIG["worker_urls"]
+    mode = f"Distributed ({len(workers)} workers)" if workers else "Standalone"
     
-    api_app = FastAPI()
-    
-    @api_app.post("/api/{endpoint}")
-    async def api_endpoint(endpoint: str, request: Request):
-        # Check authorization
-        auth_header = request.headers.get("Authorization", "")
-        if not auth_header.startswith("Bearer "):
-            # Allow without auth for testing, but log it
-            pass
+    with gr.Blocks(title="Sam-large-2 HEAD") as app:
+        gr.Markdown(f"""
+        # 👑 Sam-large-2 - HEAD NODE
+        **Mode:** {mode} | **Blocks:** {CONFIG['layer_start']}-{CONFIG['layer_end']-1} | **ID:** {CONFIG['node_id']}
+        """)
         
-        try:
-            data = await request.json()
-        except:
-            data = {}
+        if workers:
+            gr.Markdown("**Workers:** " + ", ".join(f"`{w}`" for w in workers))
         
-        result = api_handler(endpoint, data)
-        return JSONResponse(result)
-    
-    @api_app.get("/api/status")
-    async def get_status():
-        return JSONResponse(api_handler("status", {}))
-    
-    # Mount Gradio app
-    app = create_app()
-    api_app = gr.mount_gradio_app(api_app, app, path="/")
-    
-    print("Running with FastAPI + Gradio")
+        reasoning = gr.State(False)
+        chatbot = gr.Chatbot(height=500)
+        
+        with gr.Row():
+            reason_btn = gr.Button("💡", size="sm", scale=0)
+            msg = gr.Textbox(placeholder="Type message...", show_label=False, scale=8)
+            send = gr.Button("Send", variant="primary", scale=1)
+            stop_btn = gr.Button("⏹️", scale=0)
+        
+        with gr.Accordion("⚙️ Settings", open=False):
+            max_tok = gr.Slider(50, 1024, 512, label="Max Tokens")
+            temp = gr.Slider(0.1, 2.0, 0.8, label="Temperature")
+            topk = gr.Slider(1, 100, 40, label="Top-K")
+            topp = gr.Slider(0.1, 1.0, 0.9, label="Top-P")
+            rep = gr.Slider(1.0, 2.0, 1.1, label="Repetition Penalty")
+        
+        def toggle(r):
+            return not r, gr.update(variant="primary" if not r else "secondary")
+        
+        reason_btn.click(toggle, [reasoning], [reasoning, reason_btn])
+        
+        inputs = [msg, chatbot, max_tok, temp, topk, topp, rep, reasoning]
+        submit = msg.submit(chat_respond, inputs, chatbot).then(lambda: "", outputs=msg)
+        click = send.click(chat_respond, inputs, chatbot).then(lambda: "", outputs=msg)
+        stop_btn.click(stop, cancels=[submit, click])
+        
+        gr.Button("🗑️ Clear").click(lambda: ([], ""), outputs=[chatbot, msg])
     
-except ImportError:
-    # FastAPI not available, use pure Gradio
-    app = create_app()
-    print("Running with pure Gradio")
-
+    return app
 
 # ============================================================================
-# LAUNCH
+# Main
 # ============================================================================
 
-if __name__ == "__main__":
-    print(f"Starting WAN-JAX Node")
-    print(f"  Node ID: {CONFIG.node_id}")
-    print(f"  Role: {CONFIG.role}")
-    print(f"  Head URL: {CONFIG.head_url}")
-    print(f"  JAX devices: {jax.devices()}")
-    
-    app = create_app()
-    app.launch(server_name="0.0.0.0", server_port=CONFIG.port)
+print("=" * 60)
+print("🚀 Sam-large-2 HEAD Node Starting")
+print(f"   Blocks: {CONFIG['layer_start']} to {CONFIG['layer_end']}")
+print(f"   Workers: {CONFIG['worker_urls'] or 'None (standalone)'}")
+print("=" * 60)
+
+load_model()
+app = create_ui()
+app.queue()
+app.launch(server_name="0.0.0.0", server_port=7860)