Update app.py

app.py

@@ -245,7 +245,7 @@ class TokenManager:
                 self.user_tokens[username] = {
                     "balance": float('inf'),
                     "start_time": time.time(),
-                    "purchases": {"batch_multiplier": 1, "token_limit": 2048},
+                    "purchases": {"batch_size": 512, "max_tokens": 2048},
                     "total_spent": 0.0,
                     "is_owner": True,
                     "username": username
@@ -255,7 +255,7 @@ class TokenManager:
                 self.user_tokens[username] = {
                     "balance": MONTHLY_TOKEN_CREDITS,
                     "start_time": time.time(),
-                    "purchases": {"batch_multiplier": 1, "token_limit": 2048},
+                    "purchases": {"batch_size": 512, "max_tokens": 2048},
                     "total_spent": 0.0,
                     "is_owner": False,
                     "username": username,
@@ -328,62 +328,93 @@ class TokenManager:
             last_reset = self.user_tokens[username].get("last_reset", time.time())
             time_until_reset = 30 * 24 * 60 * 60 - (time.time() - last_reset)
             days_left = int(time_until_reset / (24 * 60 * 60))
-            return False, f"❌ Out of tokens! Resets in {days_left} days. Current balance: 0.00"
+            return False, f" Out of tokens! Resets in {days_left} days. Current balance: 0.00"
         
-        return True, f"✅ Access granted. Balance: {balance:.2f} tokens"
+        return True, f" Access granted. Balance: {balance:.2f} tokens"
     
-    def purchase_batch_upgrade(self, username: str) -> tuple:
-        """Purchase batch size upgrade (exponential cost). Free for owner."""
+    def purchase_batch_upgrade(self, username: str, batch_size: int = 512) -> tuple:
+        """Purchase specific batch size upgrade. Free for owner. Auto-rounds to nearest power of 2."""
         if not username:
-            return False, "❌ Please login first"
+            return False, " Please login first"
         
         self.initialize_user(username)
         
+        # SMART ROUNDING: Round to nearest power of 2 for optimal performance
+        if batch_size <= 0:
+            batch_size = 512
+        
+        # Find nearest power of 2
+        def round_to_power_of_2(n):
+            if n <= 128:
+                return 128
+            # Find the next power of 2
+            power = 1
+            while power < n:
+                power *= 2
+            # Check if previous power of 2 is closer
+            prev_power = power // 2
+            if abs(n - prev_power) <= abs(n - power):
+                return prev_power
+            return power
+        
+        rounded_batch = round_to_power_of_2(batch_size)
+        
         # Owner gets free upgrades
         if self.user_tokens[username].get("is_owner", False):
-            current_mult = self.user_tokens[username]["purchases"]["batch_multiplier"]
-            self.user_tokens[username]["purchases"]["batch_multiplier"] = current_mult * 2
-            new_mult = current_mult * 2
-            logger.info(f"[TOKEN] 👑 OWNER free batch upgrade: {current_mult}x → {new_mult}x")
-            return True, f"👑 FREE UPGRADE! Batch now {new_mult}x!"
+            self.user_tokens[username]["purchases"]["batch_size"] = rounded_batch
+            logger.info(f"[TOKEN] OWNER set batch size to: {rounded_batch} (rounded from {batch_size})")
+            if rounded_batch != batch_size:
+                return True, f" Batch size set to {rounded_batch} (rounded from {batch_size})!"
+            return True, f" Batch size set to {rounded_batch}!"
         
-        current_mult = self.user_tokens[username]["purchases"]["batch_multiplier"]
-        upgrade_level = int(math.log2(current_mult)) if current_mult > 1 else 0
-        cost = BATCH_UPGRADE_BASE_COST * (2 ** upgrade_level)
+        # Cost based on rounded batch size (larger batches cost more)
+        cost = (rounded_batch / 1000) * 0.01  # 0.01 tokens per 1000 batch size
         
         if self.user_tokens[username]["balance"] >= cost:
             self.user_tokens[username]["balance"] -= cost
-            self.user_tokens[username]["purchases"]["batch_multiplier"] = current_mult * 2
-            new_mult = current_mult * 2
-            logger.info(f"[TOKEN] Batch upgrade: {current_mult}x → {new_mult}x | Cost: {cost:.5f}")
-            return True, f"✅ Batch upgraded to {new_mult}x! (-{cost:.5f} tokens)"
+            self.user_tokens[username]["purchases"]["batch_size"] = rounded_batch
+            logger.info(f"[TOKEN] Batch size set to {rounded_batch} (rounded from {batch_size}) | Cost: {cost:.5f}")
+            if rounded_batch != batch_size:
+                return True, f" Batch size set to {rounded_batch} (rounded from {batch_size})! (-{cost:.5f} tokens)"
+            return True, f" Batch size set to {rounded_batch}! (-{cost:.5f} tokens)"
         else:
-            return False, f"❌ Insufficient tokens! Need {cost:.5f}, have {self.user_tokens[username]['balance']:.2f}"
+            return False, f" Insufficient tokens! Need {cost:.5f}, have {self.user_tokens[username]['balance']:.2f}"
     
-    def purchase_token_upgrade(self, username: str, extra_tokens: int = 1000) -> tuple:
-        """Purchase extra response token length. Free for owner."""
+    def purchase_token_upgrade(self, username: str, max_tokens: int = 2048) -> tuple:
+        """Purchase specific max tokens setting. Free for owner. Auto-rounds to nearest 256."""
         if not username:
-            return False, "❌ Please login first"
+            return False, " Please login first"
         
         self.initialize_user(username)
         
+        # SMART ROUNDING: Round to nearest 256 for optimal memory alignment
+        if max_tokens <= 0:
+            max_tokens = 2048
+        
+        # Find nearest multiple of 256
+        rounded_tokens = ((max_tokens + 128) // 256) * 256
+        rounded_tokens = max(256, min(8192, rounded_tokens))  # Clamp between 256-8192
+        
         # Owner gets free upgrades
         if self.user_tokens[username].get("is_owner", False):
-            self.user_tokens[username]["purchases"]["token_limit"] += extra_tokens
-            new_limit = self.user_tokens[username]["purchases"]["token_limit"]
-            logger.info(f"[TOKEN] 👑 OWNER free token upgrade: +{extra_tokens} tokens")
-            return True, f"👑 FREE UPGRADE! Token limit now {new_limit}!"
+            self.user_tokens[username]["purchases"]["max_tokens"] = rounded_tokens
+            logger.info(f"[TOKEN] OWNER set max tokens to: {rounded_tokens} (rounded from {max_tokens})")
+            if rounded_tokens != max_tokens:
+                return True, f" Max tokens set to {rounded_tokens} (rounded from {max_tokens})!"
+            return True, f" Max tokens set to {rounded_tokens}!"
         
-        cost = (extra_tokens / 1000) * TOKEN_UPGRADE_COST_PER_1K
+        # Cost based on rounded max tokens (larger context costs more)
+        cost = (rounded_tokens / 1000) * TOKEN_UPGRADE_COST_PER_1K
         
         if self.user_tokens[username]["balance"] >= cost:
             self.user_tokens[username]["balance"] -= cost
-            self.user_tokens[username]["purchases"]["token_limit"] += extra_tokens
-            new_limit = self.user_tokens[username]["purchases"]["token_limit"]
-            logger.info(f"[TOKEN] Token limit upgrade: +{extra_tokens} tokens | Cost: {cost:.5f}")
-            return True, f"✅ Token limit now {new_limit}! (-{cost:.5f} tokens)"
+            self.user_tokens[username]["purchases"]["max_tokens"] = rounded_tokens
+            logger.info(f"[TOKEN] Max tokens set to {rounded_tokens} (rounded from {max_tokens}) | Cost: {cost:.5f}")
+            if rounded_tokens != max_tokens:
+                return True, f" Max tokens set to {rounded_tokens} (rounded from {max_tokens})! (-{cost:.5f} tokens)"
+            return True, f" Max tokens set to {rounded_tokens}! (-{cost:.5f} tokens)"
         else:
-            return False, f"❌ Insufficient tokens! Need {cost:.5f}, have {self.user_tokens[username]['balance']:.2f}"
+            return False, f" Insufficient tokens! Need {cost:.5f}, have {self.user_tokens[username]['balance']:.2f}"
     
     def get_balance(self, username: str) -> float:
         """Get user's current token balance"""
@@ -418,7 +449,7 @@ class TokenManager:
             stats = self.user_tokens[username]
             
             if stats.get("is_owner", False):
-                return f"👑 Owner session ended. Welcome back anytime, {stats['username']}!"
+                return f" OWNER session ended. Welcome back anytime, {stats['username']}!"
             
             logger.info(f"[TOKEN] Session ended: Spent {stats['total_spent']:.2f}, Remaining {stats['balance']:.2f}")
             return f"Session ended. You spent {stats['total_spent']:.2f} tokens this session. Balance: {stats['balance']:.2f}"
@@ -560,7 +591,7 @@ class ZeroEngine:
                 logger.info(f"[PREPROCESS] Tokenizing {len(text)} chars in background...")
                 tokens = self.llm.tokenize(text.encode("utf-8"))
                 self.preprocessed_tokens = tokens
-                logger.info(f"[PREPROCESS] ✅ Ready: {len(tokens)} tokens cached")
+                logger.info(f"[PREPROCESS] Ready: {len(tokens)} tokens cached")
             except Exception as e:
                 logger.error(f"[PREPROCESS] Failed: {e}")
                 self.preprocessed_tokens = None
@@ -637,7 +668,7 @@ class ZeroEngine:
         """HYPER-OPTIMIZED Boot kernel with format auto-detection and Gemma fixes"""
         try:
             if not repo or not filename:
-                return "🔴 ERROR: Repository or filename missing"
+                return " ERROR: Repository or filename missing"
             
             logger.info(f"[BOOT] Starting download: {filename} from {repo}")
             
@@ -658,17 +689,17 @@ class ZeroEngine:
                 logger.info(f"[BOOT] Download complete: {path}")
             except Exception as e:
                 logger.error(f"[BOOT] Download failed: {e}")
-                return f"🔴 DOWNLOAD FAILED: {str(e)}"
+                return f" DOWNLOAD FAILED: {str(e)}"
             
             # Check if model is cached
             is_cached = model_cache.is_cached(path)
-            cache_status = "🎯 CACHED" if is_cached else "🆕 NEW"
+            cache_status = " CACHED" if is_cached else " NEW"
             
             # Validate before loading
             valid, msg = ResourceMonitor.validate_deployment(path)
             if not valid:
                 logger.warning(f"[BOOT] Validation failed: {msg}")
-                return f"🔴 VALIDATION FAILED: {msg}"
+                return f" VALIDATION FAILED: {msg}"
             
             logger.info(f"[BOOT] Validation passed ({cache_status}), applying {quant_config['type']} optimizations for {model_format.upper()}...")
             
@@ -676,13 +707,13 @@ class ZeroEngine:
             with self.kernel_lock:
                 # WRECK OLD MODEL
                 if self.llm:
-                    logger.info("[BOOT] 💣 WRECKING old model...")
+                    logger.info("[BOOT]  WRECKING old model...")
                     try:
                         model_cache.wreck_old_model_cache()
                         del self.llm
                         self.llm = None
                         nuclear_ram_clear()
-                        logger.info("[BOOT] ✅ Old model DESTROYED")
+                        logger.info("[BOOT]  Old model DESTROYED")
                     except Exception as e:
                         logger.warning(f"[BOOT] Cleanup warning: {e}")
                 
@@ -693,13 +724,13 @@ class ZeroEngine:
                 # CPU-OPTIMIZED BATCH CALCULATION - Very aggressive for 16GB RAM
                 # Base calculation: use more RAM for batching on CPU
                 base_batch = int(512 * available_ram_gb / 8)  # More aggressive base
-                optimal_batch = int(base_batch * quant_config["batch_multiplier"])
+                optimal_batch = base_batch
                 
-                # Apply user's batch multiplier from token purchases
+                # Apply user's batch size from token purchases
                 if session_id:
-                    user_batch_mult = token_manager.get_purchases(session_id)["batch_multiplier"]
-                    optimal_batch = int(optimal_batch * user_batch_mult)
-                    logger.info(f"[TOKEN] User batch multiplier: {user_batch_mult}x")
+                    user_batch_size = token_manager.get_purchases(session_id)["batch_size"]
+                    optimal_batch = user_batch_size
+                    logger.info(f"[TOKEN] User batch size: {user_batch_size}")
                 
                 # CPU can handle larger batches with quantized models
                 optimal_batch = max(1024, min(8192, optimal_batch))  # 1024-8192 range for CPU
@@ -772,19 +803,19 @@ class ZeroEngine:
                     except:
                         pass
                     
-                    logger.info("[BOOT] 🚀 CPU-OPTIMIZED MODEL READY!")
-                    return f"🟢 {model_format.upper()} {quant_config['type']} {cache_status} | CPU:{optimal_threads}T | B:{optimal_batch} | Ctx:{optimal_ctx}"
+                    logger.info("[BOOT]  CPU-OPTIMIZED MODEL READY!")
+                    return f" {model_format.upper()} {quant_config['type']} {cache_status} | CPU:{optimal_threads}T | B:{optimal_batch} | Ctx:{optimal_ctx}"
                     
                 except Exception as e:
                     logger.error(f"[BOOT] Model loading failed: {e}")
                     self.llm = None
                     nuclear_ram_clear()
-                    return f"🔴 LOAD FAILED: {str(e)}"
+                    return f" LOAD FAILED: {str(e)}"
                     
         except Exception as e:
             logger.error(f"[BOOT] Unexpected error: {e}")
             nuclear_ram_clear()
-            return f"🔴 BOOT FAILURE: {str(e)}"
+            return f" BOOT FAILURE: {str(e)}"
 
     def stitch_cache(self, ghost_text: str) -> str:
         """Prime KV cache with ghost context"""
@@ -794,6 +825,7 @@ class ZeroEngine:
         def _bg_eval():
             self.is_prefilling = True
             try:
+                logger.info(f"[PREPROCESS] Tokenizing {len(ghost_text)} chars in background...")
                 tokens = self.llm.tokenize(ghost_text.encode("utf-8"))
                 self.llm.eval(tokens)
                 logger.info(f"Ghost cache primed: {len(tokens)} tokens")
@@ -804,7 +836,7 @@ class ZeroEngine:
                 self.is_prefilling = False
         
         threading.Thread(target=_bg_eval, daemon=True).start()
-        return "⚡ Primed"
+        return " Primed"
 
     def inference_generator(self, prompt: str, history: List[Dict], ghost_context: str, repo: str, quant: str, profile: gr.OAuthProfile | None) -> Generator:
         username = profile.username if profile else "anonymous"
@@ -817,7 +849,7 @@ class ZeroEngine:
         # AUTO-BOOT: If model not loaded, auto-boot default model
         if not self.llm:
             logger.info("[AUTO-BOOT] No model loaded, initiating auto-boot...")
-            history.append({"role": "assistant", "content": "🔄 Auto-booting model, please wait..."})
+            history.append({"role": "assistant", "content": " Auto-booting model, please wait..."})
             yield history
             
             # Use provided repo/quant or fallback to defaults
@@ -826,12 +858,12 @@ class ZeroEngine:
             
             boot_result = self.boot_kernel(boot_repo, boot_quant)
             
-            if "🔴" in boot_result or "FAILED" in boot_result:
-                history[-1]["content"] = f"❌ Auto-boot failed: {boot_result}\n\nPlease manually SCAN and BOOT a model."
+            if " " in boot_result or "FAILED" in boot_result:
+                history[-1]["content"] = f" Auto-boot failed: {boot_result}\n\nPlease manually SCAN and BOOT a model."
                 yield history
                 return
             
-            history[-1]["content"] = f"✅ {boot_result}\n\nProcessing your request..."
+            history[-1]["content"] = f" {boot_result}\n\nProcessing your request..."
             yield history
             time.sleep(0.5)  # Brief pause for user to see the message
 
@@ -839,7 +871,7 @@ class ZeroEngine:
         cache_key = f"{ghost_context}:{prompt}"
         if cache_key in self.prompt_cache:
             self.perf_stats["cache_hits"] += 1
-            logger.info("⚡ CACHE HIT - Instant response!")
+            logger.info(" CACHE HIT - Instant response!")
             history.append({"role": "user", "content": prompt})
             history.append({"role": "assistant", "content": self.prompt_cache[cache_key]})
             yield history
@@ -863,7 +895,7 @@ class ZeroEngine:
             # Get max tokens from user purchases
             max_tokens = 2048
             if username:
-                max_tokens = token_manager.get_purchases(username)["token_limit"]
+                max_tokens = token_manager.get_purchases(username)["max_tokens"]
             
             # HYPER-OPTIMIZED CPU INFERENCE SETTINGS
             stream = self.llm(
@@ -1125,12 +1157,14 @@ with gr.Blocks(title="ZeroEngine V0.2", css=CUSTOM_CSS) as demo:
             
             gr.Markdown("---")
             
-            # Token Purchases
-            gr.Markdown("### 💎 Token Upgrades")
+            # Performance Settings
+            gr.Markdown("### 💎 Performance Settings")
+            batch_size_input = gr.Number(label="Batch Size", value=512, minimum=128, maximum=8192, step=128)
+            max_tokens_input = gr.Number(label="Max Tokens", value=2048, minimum=512, maximum=8192, step=256)
             with gr.Row():
-                batch_upgrade_btn = gr.Button("🚀 Batch x2", size="sm", variant="secondary")
-                token_upgrade_btn = gr.Button("📈 +1K Tokens", size="sm", variant="secondary")
-            purchase_status = gr.Markdown("Ready to upgrade!")
+                batch_upgrade_btn = gr.Button("🚀 Set Batch Size", size="sm", variant="secondary")
+                token_upgrade_btn = gr.Button("📈 Set Max Tokens", size="sm", variant="secondary")
+            purchase_status = gr.Markdown("Ready to configure!")
             
             gr.Markdown("---")
             
@@ -1185,13 +1219,13 @@ with gr.Blocks(title="ZeroEngine V0.2", css=CUSTOM_CSS) as demo:
         result = kernel.boot_kernel(repo, file, username)
         return result
     
-    def on_batch_upgrade():
-        success, msg = token_manager.purchase_batch_upgrade(session_id)
+    def on_batch_upgrade(batch_size):
+        success, msg = token_manager.purchase_batch_upgrade(session_id, int(batch_size))
         balance = token_manager.get_balance(session_id)
         return msg, f"{balance}"
     
-    def on_token_upgrade():
-        success, msg = token_manager.purchase_token_upgrade(session_id, 1000)
+    def on_token_upgrade(max_tokens):
+        success, msg = token_manager.purchase_token_upgrade(session_id, int(max_tokens))
         balance = token_manager.get_balance(session_id)
         return msg, f"{balance}"
     
@@ -1215,8 +1249,8 @@ with gr.Blocks(title="ZeroEngine V0.2", css=CUSTOM_CSS) as demo:
     boot_btn.click(on_boot, [repo_input, quant_dropdown], [boot_status])
     
     # Token purchases
-    batch_upgrade_btn.click(on_batch_upgrade, None, [purchase_status, token_balance])
-    token_upgrade_btn.click(on_token_upgrade, None, [purchase_status, token_balance])
+    batch_upgrade_btn.click(on_batch_upgrade, [batch_size_input], [purchase_status, token_balance])
+    token_upgrade_btn.click(on_token_upgrade, [max_tokens_input], [purchase_status, token_balance])
     end_session_btn.click(on_end_session, None, [session_status])
     
     # Custom parameter updates