Update app.py

app.py

@@ -1,3 +1,480 @@
+# import os
+# import sys
+# import json
+# import time
+# import importlib.util
+# from pathlib import Path
+# from flask import Flask, request, jsonify, Response, stream_with_context
+# from flask_cors import CORS
+# import torch
+# from transformers import AutoTokenizer
+
+# app = Flask(__name__, static_folder='static', static_url_path='/static')
+# CORS(app)
+
+# # Global state
+# model = None
+# tokenizer = None
+# config = None
+# device = None
+# DiffusionLLM = None
+# chat_function = None
+
+
+# def find_file(filename, search_dirs=None):
+#     """Find a file in current directory or parent directories."""
+#     if search_dirs is None:
+#         search_dirs = [
+#             os.path.dirname(__file__),  # Current directory
+#             os.path.dirname(os.path.dirname(__file__)),  # Parent directory
+#             os.getcwd(),  # Working directory
+#         ]
+
+#     for directory in search_dirs:
+#         filepath = os.path.join(directory, filename)
+#         if os.path.exists(filepath):
+#             print(f"Found {filename} at: {filepath}")
+#             return filepath
+
+#     return None
+
+
+# def try_import_module(filepath, module_name):
+#     """Dynamically import a Python file as a module."""
+#     if not filepath or not os.path.exists(filepath):
+#         return None
+
+#     try:
+#         # Add the directory to sys.path
+#         module_dir = os.path.dirname(filepath)
+#         if module_dir not in sys.path:
+#             sys.path.insert(0, module_dir)
+
+#         spec = importlib.util.spec_from_file_location(module_name, filepath)
+#         if spec is None:
+#             print(f"Could not create spec for {filepath}")
+#             return None
+
+#         module = importlib.util.module_from_spec(spec)
+#         sys.modules[module_name] = module
+#         spec.loader.exec_module(module)
+
+#         print(f"Successfully imported {module_name} from {filepath}")
+#         return module
+#     except Exception as e:
+#         print(f"Error importing {filepath}: {e}")
+#         import traceback
+#         traceback.print_exc()
+#         return None
+
+
+# def load_model_internal():
+#     """Load the model and tokenizer."""
+#     global model, tokenizer, config, device, DiffusionLLM, chat_function
+
+#     if model is not None:
+#         return True
+
+#     try:
+#         print("=" * 60)
+#         print("Starting model loading process...")
+#         print("=" * 60)
+
+#         # Find and import infer-base.py
+#         base_path = find_file("infer-base.py")
+#         if base_path is None:
+#             raise RuntimeError("Could not find infer-base.py. Make sure it's in the same directory as app.py or parent directory.")
+
+#         print(f"\nImporting infer-base.py from: {base_path}")
+#         base_mod = try_import_module(base_path, "infer_base")
+
+#         if base_mod is None:
+#             raise RuntimeError("Failed to import infer-base.py")
+
+#         # Check for DiffusionLLM class
+#         if not hasattr(base_mod, 'DiffusionLLM'):
+#             print("Available attributes in infer_base:", dir(base_mod))
+#             raise RuntimeError("DiffusionLLM class not found in infer-base.py")
+
+#         DiffusionLLM = base_mod.DiffusionLLM
+#         print("✓ Successfully loaded DiffusionLLM class")
+
+#         # Find and import infer-chat.py
+#         chat_path = find_file("infer-chat.py")
+#         if chat_path is None:
+#             raise RuntimeError("Could not find infer-chat.py")
+
+#         print(f"\nImporting infer-chat.py from: {chat_path}")
+#         chat_mod = try_import_module(chat_path, "infer_chat")
+
+#         if chat_mod is None or not hasattr(chat_mod, 'chat'):
+#             raise RuntimeError("Failed to import chat function from infer-chat.py")
+
+#         chat_function = chat_mod.chat
+#         print("✓ Successfully loaded chat function")
+
+#         # Setup pickling workaround for torch.load
+#         try:
+#             if hasattr(base_mod, 'ModelConfig'):
+#                 sys.modules['__main__'].ModelConfig = base_mod.ModelConfig
+#             sys.modules['__main__'].DiffusionLLM = DiffusionLLM
+#             print("✓ Configured pickle support for model loading")
+#         except Exception as e:
+#             print(f"Warning: Could not setup pickle workaround: {e}")
+
+#         # Set device
+#         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#         print(f"\n✓ Using device: {device}")
+
+#         # Load tokenizer
+#         print("\nLoading tokenizer...")
+#         tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
+#         if tokenizer.pad_token is None:
+#             tokenizer.pad_token = tokenizer.eos_token
+#         print("✓ Tokenizer loaded")
+
+#         # Find model checkpoint
+#         checkpoint_dirs = [
+#             "checkpoints",
+#             "../checkpoints",
+#             "./checkpoints",
+#             os.path.join(os.path.dirname(__file__), "checkpoints"),
+#             os.path.join(os.path.dirname(__file__), "../checkpoints"),
+#         ]
+
+#         model_path = None
+#         for checkpoint_dir in checkpoint_dirs:
+#             best_path = os.path.join(checkpoint_dir, "best_model.pt")
+#             fp32_path = os.path.join(checkpoint_dir, "model_fp32.pt")
+
+#             if os.path.exists(best_path):
+#                 model_path = best_path
+#                 break
+#             elif os.path.exists(fp32_path):
+#                 model_path = fp32_path
+#                 break
+
+#         if model_path is None:
+#             raise RuntimeError(
+#                 "Could not find model checkpoint. Looking for:\n"
+#                 "  - checkpoints/best_model.pt\n"
+#                 "  - checkpoints/model_fp32.pt\n"
+#                 f"Searched directories: {checkpoint_dirs}"
+#             )
+
+#         print(f"\n✓ Found model checkpoint: {model_path}")
+#         print("Loading model weights (this may take a minute)...")
+
+#         # Load model
+#         checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+#         config = checkpoint['config']
+
+#         print("Creating model...")
+#         model = DiffusionLLM(config)
+
+#         print("Loading state dict...")
+#         state_dict = checkpoint['model_state']
+#         state_dict = {k: v.float() for k, v in state_dict.items()}
+#         model.load_state_dict(state_dict)
+
+#         model = model.to(device)
+#         model.eval()
+
+#         num_params = sum(p.numel() for p in model.parameters()) / 1e6
+#         print(f"\n{'=' * 60}")
+#         print(f"✓✓✓ MODEL LOADED SUCCESSFULLY ✓✓✓")
+#         print(f"{'=' * 60}")
+#         print(f"Parameters: {num_params:.1f}M")
+#         if 'step' in checkpoint:
+#             print(f"Training steps: {checkpoint['step']}")
+#         if 'best_val_loss' in checkpoint:
+#             print(f"Best validation loss: {checkpoint['best_val_loss']:.4f}")
+#         print(f"{'=' * 60}\n")
+
+#         return True
+
+#     except Exception as e:
+#         print("\n" + "=" * 60)
+#         print("ERROR LOADING MODEL")
+#         print("=" * 60)
+#         print(f"Error: {e}")
+#         import traceback
+#         traceback.print_exc()
+#         print("=" * 60 + "\n")
+#         return False
+
+
+# def create_streaming_visualizer():
+#     """Create a visualizer that yields SSE events instead of printing to terminal."""
+#     def visualizer(tok, context_ids, mask_blocks, is_masked_list, cfg, clear=True):
+#         # Normalize inputs to lists
+#         if not isinstance(mask_blocks, list):
+#             mask_blocks = [mask_blocks]
+#             is_masked_list = [is_masked_list]
+
+#         # Decode context
+#         try:
+#             context_text = tok.decode(context_ids[0], skip_special_tokens=True).replace('\n', ' ')
+#         except Exception:
+#             context_text = str(context_ids[0].tolist())
+
+#         # Build blocks visualization
+#         all_blocks = []
+#         for block_idx, (mask_block, is_masked) in enumerate(zip(mask_blocks, is_masked_list)):
+#             block_tokens = mask_block[0].tolist()
+#             block_data = []
+
+#             for i, token_id in enumerate(block_tokens):
+#                 if is_masked[0, i]:
+#                     block_data.append({
+#                         'type': 'masked',
+#                         'text': '███'
+#                     })
+#                 else:
+#                     try:
+#                         token_text = tok.decode([token_id], skip_special_tokens=False)
+#                     except Exception:
+#                         token_text = str(int(token_id))
+#                     block_data.append({
+#                         'type': 'revealed',
+#                         'text': token_text
+#                     })
+
+#             all_blocks.append({
+#                 'block_index': block_idx,
+#                 'tokens': block_data
+#             })
+
+#         # Return data structure that will be sent as SSE
+#         return {
+#             'context': context_text,
+#             'blocks': all_blocks,
+#             'num_blocks': len(mask_blocks)
+#         }
+
+#     return visualizer
+
+
+# @app.route('/')
+# def index():
+#     """Serve the main HTML page."""
+#     return app.send_static_file('index.html')
+
+
+# @app.route('/api/load', methods=['POST'])
+# def load_model_endpoint():
+#     """Load the model."""
+#     data = request.json or {}
+#     check_only = data.get('check_only', False)
+
+#     global model
+
+#     if check_only:
+#         return jsonify({
+#             'loaded': model is not None,
+#             'message': 'Model is loaded' if model is not None else 'Model not loaded'
+#         })
+
+#     if model is not None:
+#         return jsonify({
+#             'loaded': True,
+#             'message': 'Model already loaded'
+#         })
+
+#     success = load_model_internal()
+
+#     if success:
+#         return jsonify({
+#             'loaded': True,
+#             'message': 'Model loaded successfully'
+#         })
+#     else:
+#         return jsonify({
+#             'loaded': False,
+#             'message': 'Failed to load model. Check server logs for details.'
+#         }), 500
+
+
+# @app.route('/api/generate', methods=['POST'])
+# def generate():
+#     """Generate response without streaming."""
+#     global model, tokenizer, config, device, chat_function
+
+#     if model is None:
+#         return jsonify({'error': 'Model not loaded'}), 400
+
+#     if chat_function is None:
+#         return jsonify({'error': 'Chat function not available'}), 400
+
+#     data = request.json
+#     instruction = data.get('instruction', '')
+#     steps = data.get('steps', 64)
+#     block_size = data.get('block_size', 128)
+#     max_new_tokens = data.get('max_new_tokens', 128)
+#     parallel_blocks = data.get('parallel_blocks', 1)
+
+#     if not instruction:
+#         return jsonify({'error': 'No instruction provided'}), 400
+
+#     try:
+#         # Generate response
+#         raw_output, response = chat_function(
+#             model,
+#             tokenizer,
+#             instruction,
+#             steps=steps,
+#             block_size=block_size,
+#             max_new_tokens=max_new_tokens,
+#             temperature=0.8,
+#             top_k=50,
+#             top_p=0.9,
+#             repetition_penalty=1.2,
+#             no_repeat_ngram_size=3,
+#             verbose=False,
+#             visualize_fn=None,
+#             parallel_blocks=parallel_blocks,
+#         )
+
+#         return jsonify({
+#             'response': response,
+#             'raw_output': raw_output
+#         })
+#     except Exception as e:
+#         import traceback
+#         traceback.print_exc()
+#         return jsonify({'error': str(e)}), 500
+
+
+# @app.route('/api/generate-stream', methods=['POST'])
+# def generate_stream():
+#     """Generate response with streaming visualization."""
+#     global model, tokenizer, config, device, chat_function
+
+#     if model is None:
+#         return jsonify({'error': 'Model not loaded'}), 400
+
+#     if chat_function is None:
+#         return jsonify({'error': 'Chat function not available'}), 400
+
+#     data = request.json
+#     instruction = data.get('instruction', '')
+#     steps = data.get('steps', 64)
+#     block_size = data.get('block_size', 128)
+#     max_new_tokens = data.get('max_new_tokens', 128)
+#     parallel_blocks = data.get('parallel_blocks', 1)
+
+#     if not instruction:
+#         return jsonify({'error': 'No instruction provided'}), 400
+
+#     def generate_events():
+#         try:
+#             # Import threading to allow yielding from callback
+#             import queue
+#             event_queue = queue.Queue()
+#             generation_complete = {'done': False, 'result': None}
+
+#             def streaming_visualizer(tok, context_ids, mask_blocks, is_masked_list, cfg, clear=True):
+#                 """This gets called during generation - we need to send events immediately"""
+#                 visualizer = create_streaming_visualizer()
+#                 data = visualizer(tok, context_ids, mask_blocks, is_masked_list, cfg, clear)
+#                 # Put the update in the queue so it can be yielded immediately
+#                 event_queue.put({'type': 'update', 'data': data})
+
+#             # Start generation in a separate thread so we can yield events as they come
+#             import threading
+
+#             def run_generation():
+#                 try:
+#                     raw_output, response = chat_function(
+#                         model,
+#                         tokenizer,
+#                         instruction,
+#                         steps=steps,
+#                         block_size=block_size,
+#                         max_new_tokens=max_new_tokens,
+#                         temperature=0.8,
+#                         top_k=50,
+#                         top_p=0.9,
+#                         repetition_penalty=1.2,
+#                         no_repeat_ngram_size=3,
+#                         verbose=False,
+#                         visualize_fn=streaming_visualizer,
+#                         parallel_blocks=parallel_blocks,
+#                     )
+#                     generation_complete['result'] = (raw_output, response)
+#                 except Exception as e:
+#                     generation_complete['result'] = ('error', str(e))
+#                 finally:
+#                     generation_complete['done'] = True
+#                     event_queue.put(None)  # Signal completion
+
+#             # Start generation thread
+#             gen_thread = threading.Thread(target=run_generation)
+#             gen_thread.daemon = True
+#             gen_thread.start()
+
+#             # Yield start event
+#             yield f"data: {json.dumps({'type': 'start', 'message': 'Generation started'})}\n\n"
+
+#             # Yield events as they come from the queue
+#             while not generation_complete['done'] or not event_queue.empty():
+#                 try:
+#                     event = event_queue.get(timeout=0.1)
+#                     if event is None:  # Completion signal
+#                         break
+#                     yield f"data: {json.dumps(event)}\n\n"
+#                 except queue.Empty:
+#                     continue
+
+#             # Wait for thread to finish
+#             gen_thread.join(timeout=1.0)
+
+#             # Send final response
+#             if generation_complete['result']:
+#                 raw_output, response = generation_complete['result']
+#                 if raw_output == 'error':
+#                     yield f"data: {json.dumps({'type': 'error', 'error': response})}\n\n"
+#                 else:
+#                     yield f"data: {json.dumps({'type': 'complete', 'response': response, 'raw_output': raw_output})}\n\n"
+
+#         except Exception as e:
+#             import traceback
+#             traceback.print_exc()
+#             yield f"data: {json.dumps({'type': 'error', 'error': str(e)})}\n\n"
+
+#     return Response(
+#         stream_with_context(generate_events()),
+#         mimetype='text/event-stream',
+#         headers={
+#             'Cache-Control': 'no-cache',
+#             'X-Accel-Buffering': 'no'
+#         }
+#     )
+
+
+# @app.route('/api/test-stream', methods=['GET'])
+# def test_stream():
+#     """Test streaming endpoint."""
+#     def generate():
+#         for i in range(10):
+#             yield f"data: {json.dumps({'message': f'Test message {i+1}'})}\n\n"
+#             time.sleep(0.5)
+#         yield f"data: {json.dumps({'message': 'Stream complete'})}\n\n"
+
+#     return Response(
+#         stream_with_context(generate()),
+#         mimetype='text/event-stream',
+#         headers={
+#             'Cache-Control': 'no-cache',
+#             'X-Accel-Buffering': 'no'
+#         }
+#     )
+
+
+# if __name__ == '__main__':
+#     app.run(debug=True, host='0.0.0.0', port=7860, threaded=True)
+
+
 import os
 import sys
 import json
@@ -8,6 +485,8 @@ from flask import Flask, request, jsonify, Response, stream_with_context
 from flask_cors import CORS
 import torch
 from transformers import AutoTokenizer
+import threading
+import queue
 
 app = Flask(__name__, static_folder='static', static_url_path='/static')
 CORS(app)
@@ -19,457 +498,581 @@ config = None
 device = None
 DiffusionLLM = None
 chat_function = None
+optimized_pipeline = None  # For ONNX/IPEX
 
+# ==================== CONFIGURATION ====================
+USE_ONNX_RUNTIME = False  # Set True for ONNX (fastest)
+USE_IPEX = False  # Set True for Intel CPUs
+USE_TORCH_COMPILE = True  # Set True for PyTorch 2.0+ (good default)
+QUANTIZE_MODEL = True  # INT8/BF16 quantization
+WARMUP_ITERATIONS = 3  # Warmup for stable performance
+# =======================================================
 
 def find_file(filename, search_dirs=None):
     """Find a file in current directory or parent directories."""
     if search_dirs is None:
         search_dirs = [
-            os.path.dirname(__file__),  # Current directory
-            os.path.dirname(os.path.dirname(__file__)),  # Parent directory
-            os.getcwd(),  # Working directory
+            os.path.dirname(__file__),
+            os.path.dirname(os.path.dirname(__file__)),
+            os.getcwd(),
         ]
-
     for directory in search_dirs:
         filepath = os.path.join(directory, filename)
         if os.path.exists(filepath):
             print(f"Found {filename} at: {filepath}")
             return filepath
-
     return None
 
-
 def try_import_module(filepath, module_name):
     """Dynamically import a Python file as a module."""
     if not filepath or not os.path.exists(filepath):
         return None
-
     try:
-        # Add the directory to sys.path
         module_dir = os.path.dirname(filepath)
         if module_dir not in sys.path:
             sys.path.insert(0, module_dir)
-
         spec = importlib.util.spec_from_file_location(module_name, filepath)
         if spec is None:
             print(f"Could not create spec for {filepath}")
             return None
-
         module = importlib.util.module_from_spec(spec)
         sys.modules[module_name] = module
         spec.loader.exec_module(module)
-
         print(f"Successfully imported {module_name} from {filepath}")
         return module
     except Exception as e:
         print(f"Error importing {filepath}: {e}")
-        import traceback
-        traceback.print_exc()
+        if __debug__:
+            import traceback
+            traceback.print_exc()
         return None
 
+def configure_cpu_optimization():
+    """Configure CPU for maximum performance."""
+    print("\n" + "=" * 60)
+    print("CPU OPTIMIZATION CONFIGURATION")
+    print("=" * 60)
+    
+    # Get CPU info
+    cpu_count = os.cpu_count()
+    physical_cores = cpu_count // 2 if cpu_count else cpu_count
+    
+    # Optimal thread settings
+    threads = physical_cores or cpu_count or 1
+    torch.set_num_threads(threads)
+    torch.set_num_interop_threads(threads)
+    
+    # Environment variables for MKL/OMP
+    os.environ["OMP_NUM_THREADS"] = str(threads)
+    os.environ["MKL_NUM_THREADS"] = str(threads)
+    os.environ["NUMEXPR_NUM_THREADS"] = str(threads)
+    
+    print(f"✓ CPU Cores: {cpu_count} ({physical_cores} physical)")
+    print(f"✓ Threads: {threads}")
+    print(f"✓ OMP/MKL threads: {threads}")
+    
+    # Intel-specific optimizations
+    if "intel" in torch.__version__.lower() or USE_IPEX:
+        torch.backends.quantized.engine = 'fbgemm'
+        print("✓ Intel FBGEMM backend enabled")
+    
+    print("=" * 60 + "\n")
+
+def quantize_model(model):
+    """Apply quantization for faster inference."""
+    if not QUANTIZE_MODEL:
+        return model
+    
+    print("\nApplying quantization...")
+    try:
+        # Use torch.compile with quantization if available
+        if hasattr(torch, 'ao') and hasattr(torch.ao, 'quantization'):
+            # Dynamic quantization (fastest, no calibration needed)
+            model = torch.quantization.quantize_dynamic(
+                model,
+                {torch.nn.Linear, torch.nn.Conv1d, torch.nn.Conv2d},
+                dtype=torch.qint8
+            )
+            print("✓ Applied INT8 dynamic quantization")
+        elif USE_IPEX:
+            # IPEX BF16 optimization
+            model = model.to(torch.bfloat16)
+            print("✓ Applied BF16 precision (IPEX)")
+        
+        return model
+    except Exception as e:
+        print(f"⚠ Quantization failed: {e}")
+        return model
+
+def compile_model(model):
+    """Compile model for maximum speed."""
+    global USE_TORCH_COMPILE, USE_ONNX_RUNTIME
+    
+    print("\nCompiling model...")
+    
+    # Option 1: ONNX Runtime (BEST performance)
+    if USE_ONNX_RUNTIME:
+        try:
+            import onnxruntime as ort
+            # Export to ONNX (one-time cost, but worth it)
+            onnx_path = "model_optimized.onnx"
+            if not os.path.exists(onnx_path):
+                print("Exporting to ONNX format...")
+                dummy_input = torch.randint(0, 100, (1, 128))  # Adjust shape as needed
+                torch.onnx.export(
+                    model,
+                    dummy_input,
+                    onnx_path,
+                    input_names=['input_ids'],
+                    output_names=['logits'],
+                    dynamic_axes={'input_ids': {0: 'batch', 1: 'sequence'}},
+                    opset_version=16
+                )
+            
+            # Create ONNX Runtime session
+            sess_options = ort.SessionOptions()
+            sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+            sess_options.enable_cpu_mem_arena = True
+            sess_options.enable_mem_pattern = True
+            
+            # Use all cores
+            sess_options.intra_op_num_threads = torch.get_num_threads()
+            
+            provider = 'CPUExecutionProvider'
+            compiled_model = ort.InferenceSession(onnx_path, sess_options, providers=[provider])
+            
+            print("✓ ONNX Runtime compilation complete")
+            return compiled_model
+        except Exception as e:
+            print(f"⚠ ONNX Runtime failed: {e}, using torch.compile")
+            USE_ONNX_RUNTIME = False
+    
+    # Option 2: Intel IPEX
+    if USE_IPEX:
+        try:
+            import intel_extension_for_pytorch as ipex
+            model = ipex.optimize(model, dtype=torch.bfloat16, level="O3")
+            print("✓ Intel IPEX O3 optimization applied")
+            return model
+        except Exception as e:
+            print(f"⚠ IPEX failed: {e}")
+            USE_IPEX = False
+    
+    # Option 3: torch.compile (PyTorch 2.0+)
+    if USE_TORCH_COMPILE and hasattr(torch, 'compile'):
+        try:
+            # Use "max-autotune" for best performance
+            model = torch.compile(
+                model,
+                mode="max-autotune",
+                fullgraph=True,
+                backend="inductor"
+            )
+            print("✓ torch.compile (max-autotune) applied")
+        except Exception as e:
+            print(f"⚠ torch.compile failed: {e}, using eager mode")
+    
+    return model
+
+def warmup_model(model, tokenizer, chat_func):
+    """Warmup the model for consistent performance."""
+    if WARMUP_ITERATIONS == 0:
+        return
+    
+    print("\nWarming up model...")
+    start_time = time.time()
+    
+    try:
+        with torch.inference_mode():
+            for i in range(WARMUP_ITERATIONS):
+                _ = chat_func(
+                    model, tokenizer, "Hello world",
+                    steps=8, block_size=32, max_new_tokens=16,
+                    temperature=0.8, top_k=50, top_p=0.9,
+                    repetition_penalty=1.2, no_repeat_ngram_size=3,
+                    verbose=False, visualize_fn=None, parallel_blocks=2
+                )
+                print(f"  Warmup {i+1}/{WARMUP_ITERATIONS} complete")
+    except Exception as e:
+        print(f"⚠ Warmup failed: {e}")
+    
+    print(f"✓ Warmup finished in {time.time() - start_time:.2f}s")
 
 def load_model_internal():
-    """Load the model and tokenizer."""
-    global model, tokenizer, config, device, DiffusionLLM, chat_function
-
+    """Load the model with ultra-fast optimizations."""
+    global model, tokenizer, config, device, DiffusionLLM, chat_function, optimized_pipeline
+    
     if model is not None:
         return True
-
+    
     try:
         print("=" * 60)
-        print("Starting model loading process...")
+        print("ULTRA-FAST CPU MODEL LOADING")
         print("=" * 60)
-
-        # Find and import infer-base.py
+        
+        # Configure CPU
+        configure_cpu_optimization()
+        
+        # Import modules
         base_path = find_file("infer-base.py")
         if base_path is None:
-            raise RuntimeError("Could not find infer-base.py. Make sure it's in the same directory as app.py or parent directory.")
-
-        print(f"\nImporting infer-base.py from: {base_path}")
+            raise RuntimeError("Could not find infer-base.py")
+        
         base_mod = try_import_module(base_path, "infer_base")
-
-        if base_mod is None:
-            raise RuntimeError("Failed to import infer-base.py")
-
-        # Check for DiffusionLLM class
-        if not hasattr(base_mod, 'DiffusionLLM'):
-            print("Available attributes in infer_base:", dir(base_mod))
-            raise RuntimeError("DiffusionLLM class not found in infer-base.py")
-
+        if base_mod is None or not hasattr(base_mod, 'DiffusionLLM'):
+            raise RuntimeError("DiffusionLLM class not found")
+        
         DiffusionLLM = base_mod.DiffusionLLM
-        print("✓ Successfully loaded DiffusionLLM class")
-
-        # Find and import infer-chat.py
+        
         chat_path = find_file("infer-chat.py")
         if chat_path is None:
             raise RuntimeError("Could not find infer-chat.py")
-
-        print(f"\nImporting infer-chat.py from: {chat_path}")
+        
         chat_mod = try_import_module(chat_path, "infer_chat")
-
         if chat_mod is None or not hasattr(chat_mod, 'chat'):
-            raise RuntimeError("Failed to import chat function from infer-chat.py")
-
+            raise RuntimeError("Chat function not found")
+        
         chat_function = chat_mod.chat
-        print("✓ Successfully loaded chat function")
-
-        # Setup pickling workaround for torch.load
-        try:
-            if hasattr(base_mod, 'ModelConfig'):
-                sys.modules['__main__'].ModelConfig = base_mod.ModelConfig
-            sys.modules['__main__'].DiffusionLLM = DiffusionLLM
-            print("✓ Configured pickle support for model loading")
-        except Exception as e:
-            print(f"Warning: Could not setup pickle workaround: {e}")
-
-        # Set device
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"\n✓ Using device: {device}")
-
+        
+        # Device
+        device = torch.device("cpu")
+        
         # Load tokenizer
         print("\nLoading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
+        tokenizer = AutoTokenizer.from_pretrained(
+            "Qwen/Qwen2.5-0.5B",
+            use_fast=True,
+            trust_remote_code=True
+        )
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
-        print("✓ Tokenizer loaded")
-
-        # Find model checkpoint
-        checkpoint_dirs = [
-            "checkpoints",
-            "../checkpoints",
-            "./checkpoints",
-            os.path.join(os.path.dirname(__file__), "checkpoints"),
-            os.path.join(os.path.dirname(__file__), "../checkpoints"),
-        ]
-
+        print("✓ Fast tokenizer loaded")
+        
+        # Find model
+        checkpoint_dirs = ["checkpoints", "../checkpoints", "./checkpoints"]
         model_path = None
         for checkpoint_dir in checkpoint_dirs:
             best_path = os.path.join(checkpoint_dir, "best_model.pt")
             fp32_path = os.path.join(checkpoint_dir, "model_fp32.pt")
-
             if os.path.exists(best_path):
                 model_path = best_path
                 break
             elif os.path.exists(fp32_path):
                 model_path = fp32_path
-                break
-
+        
         if model_path is None:
-            raise RuntimeError(
-                "Could not find model checkpoint. Looking for:\n"
-                "  - checkpoints/best_model.pt\n"
-                "  - checkpoints/model_fp32.pt\n"
-                f"Searched directories: {checkpoint_dirs}"
-            )
-
-        print(f"\n✓ Found model checkpoint: {model_path}")
-        print("Loading model weights (this may take a minute)...")
-
-        # Load model
-        checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+            raise RuntimeError("Model checkpoint not found")
+        
+        print(f"\nLoading model from: {model_path}")
+        
+        # Load checkpoint
+        checkpoint = torch.load(model_path, map_location='cpu', weights_only=False)
         config = checkpoint['config']
-
-        print("Creating model...")
+        
+        # Create and load model
         model = DiffusionLLM(config)
-
-        print("Loading state dict...")
         state_dict = checkpoint['model_state']
-        state_dict = {k: v.float() for k, v in state_dict.items()}
+        
+        # Convert to float32 for CPU
+        if not USE_IPEX:  # Keep FP32 for ONNX, use BF16 for IPEX
+            state_dict = {k: v.float() for k, v in state_dict.items()}
+        
         model.load_state_dict(state_dict)
-
-        model = model.to(device)
         model.eval()
-
+        
+        # Apply optimizations
+        model = quantize_model(model)
+        model = model.to(device)
+        model = compile_model(model)
+        
+        # Warmup
+        warmup_model(model, tokenizer, chat_function)
+        
+        # Print summary
         num_params = sum(p.numel() for p in model.parameters()) / 1e6
         print(f"\n{'=' * 60}")
-        print(f"✓✓✓ MODEL LOADED SUCCESSFULLY ✓✓✓")
+        print(f"✓✓✓ MODEL LOADED & ULTRA-OPTIMIZED FOR CPU ✓✓✓")
         print(f"{'=' * 60}")
+        print(f"Framework: {'ONNX Runtime' if USE_ONNX_RUNTIME else 'IPEX' if USE_IPEX else 'PyTorch'}")
         print(f"Parameters: {num_params:.1f}M")
+        print(f"CPU Threads: {torch.get_num_threads()}")
+        print(f"Quantization: {'INT8' if QUANTIZE_MODEL else 'BF16' if USE_IPEX else 'FP32'}")
         if 'step' in checkpoint:
             print(f"Training steps: {checkpoint['step']}")
-        if 'best_val_loss' in checkpoint:
-            print(f"Best validation loss: {checkpoint['best_val_loss']:.4f}")
         print(f"{'=' * 60}\n")
-
+        
         return True
-
+        
     except Exception as e:
-        print("\n" + "=" * 60)
-        print("ERROR LOADING MODEL")
-        print("=" * 60)
-        print(f"Error: {e}")
-        import traceback
-        traceback.print_exc()
-        print("=" * 60 + "\n")
+        print(f"\nERROR LOADING MODEL: {e}")
+        if __debug__:
+            import traceback
+            traceback.print_exc()
         return False
 
-
 def create_streaming_visualizer():
-    """Create a visualizer that yields SSE events instead of printing to terminal."""
+    """Create optimized visualizer."""
     def visualizer(tok, context_ids, mask_blocks, is_masked_list, cfg, clear=True):
-        # Normalize inputs to lists
         if not isinstance(mask_blocks, list):
             mask_blocks = [mask_blocks]
             is_masked_list = [is_masked_list]
-
-        # Decode context
+        
         try:
+            # Decode only once for efficiency
             context_text = tok.decode(context_ids[0], skip_special_tokens=True).replace('\n', ' ')
         except Exception:
             context_text = str(context_ids[0].tolist())
-
-        # Build blocks visualization
+        
         all_blocks = []
         for block_idx, (mask_block, is_masked) in enumerate(zip(mask_blocks, is_masked_list)):
             block_tokens = mask_block[0].tolist()
             block_data = []
-
+            
+            # Batch decode for speed
+            token_ids_to_decode = []
+            positions = []
+            
+            for i, token_id in enumerate(block_tokens):
+                if not is_masked[0, i]:
+                    token_ids_to_decode.append(token_id)
+                    positions.append(i)
+            
+            try:
+                decoded_tokens = tok.batch_decode(
+                    [[tid] for tid in token_ids_to_decode],
+                    skip_special_tokens=False
+                )
+            except Exception:
+                decoded_tokens = [str(int(tid)) for tid in token_ids_to_decode]
+            
+            # Reconstruct block
+            decoded_idx = 0
             for i, token_id in enumerate(block_tokens):
                 if is_masked[0, i]:
-                    block_data.append({
-                        'type': 'masked',
-                        'text': '███'
-                    })
+                    block_data.append({'type': 'masked', 'text': '███'})
                 else:
-                    try:
-                        token_text = tok.decode([token_id], skip_special_tokens=False)
-                    except Exception:
-                        token_text = str(int(token_id))
                     block_data.append({
                         'type': 'revealed',
-                        'text': token_text
+                        'text': decoded_tokens[decoded_idx]
                     })
-
-            all_blocks.append({
-                'block_index': block_idx,
-                'tokens': block_data
-            })
-
-        # Return data structure that will be sent as SSE
+                    decoded_idx += 1
+            
+            all_blocks.append({'block_index': block_idx, 'tokens': block_data})
+        
         return {
             'context': context_text,
             'blocks': all_blocks,
             'num_blocks': len(mask_blocks)
         }
-
+    
     return visualizer
 
-
 @app.route('/')
 def index():
     """Serve the main HTML page."""
     return app.send_static_file('index.html')
 
-
 @app.route('/api/load', methods=['POST'])
 def load_model_endpoint():
     """Load the model."""
     data = request.json or {}
     check_only = data.get('check_only', False)
-
+    
     global model
-
+    
     if check_only:
         return jsonify({
             'loaded': model is not None,
             'message': 'Model is loaded' if model is not None else 'Model not loaded'
         })
-
+    
     if model is not None:
         return jsonify({
             'loaded': True,
             'message': 'Model already loaded'
         })
-
+    
     success = load_model_internal()
-
     if success:
         return jsonify({
             'loaded': True,
-            'message': 'Model loaded successfully'
+            'message': 'Model loaded with ultra-fast CPU optimizations'
         })
     else:
         return jsonify({
             'loaded': False,
-            'message': 'Failed to load model. Check server logs for details.'
+            'message': 'Failed to load model. Check server logs.'
         }), 500
 
-
 @app.route('/api/generate', methods=['POST'])
 def generate():
-    """Generate response without streaming."""
-    global model, tokenizer, config, device, chat_function
-
+    """Generate response - optimized for minimal latency."""
+    global model, tokenizer, device, chat_function
+    
     if model is None:
         return jsonify({'error': 'Model not loaded'}), 400
-
+    
     if chat_function is None:
         return jsonify({'error': 'Chat function not available'}), 400
-
+    
     data = request.json
     instruction = data.get('instruction', '')
-    steps = data.get('steps', 64)
-    block_size = data.get('block_size', 128)
-    max_new_tokens = data.get('max_new_tokens', 128)
-    parallel_blocks = data.get('parallel_blocks', 1)
-
+    steps = data.get('steps', 16)  # Further reduced for speed
+    block_size = data.get('block_size', 32)
+    max_new_tokens = data.get('max_new_tokens', 64)
+    parallel_blocks = data.get('parallel_blocks', torch.get_num_threads())
+    
     if not instruction:
         return jsonify({'error': 'No instruction provided'}), 400
-
+    
     try:
-        # Generate response
-        raw_output, response = chat_function(
-            model,
-            tokenizer,
-            instruction,
-            steps=steps,
-            block_size=block_size,
-            max_new_tokens=max_new_tokens,
-            temperature=0.8,
-            top_k=50,
-            top_p=0.9,
-            repetition_penalty=1.2,
-            no_repeat_ngram_size=3,
-            verbose=False,
-            visualize_fn=None,
-            parallel_blocks=parallel_blocks,
-        )
-
+        # Fast path: no overhead
+        with torch.inference_mode():
+            raw_output, response = chat_function(
+                model,
+                tokenizer,
+                instruction,
+                steps=steps,
+                block_size=block_size,
+                max_new_tokens=max_new_tokens,
+                temperature=0.7,  # Slightly lower for faster sampling
+                top_k=50,
+                top_p=0.9,
+                repetition_penalty=1.2,
+                no_repeat_ngram_size=3,
+                verbose=False,
+                visualize_fn=None,
+                parallel_blocks=parallel_blocks,
+            )
+        
         return jsonify({
             'response': response,
             'raw_output': raw_output
         })
+        
     except Exception as e:
-        import traceback
-        traceback.print_exc()
+        if __debug__:
+            import traceback
+            traceback.print_exc()
         return jsonify({'error': str(e)}), 500
 
-
 @app.route('/api/generate-stream', methods=['POST'])
 def generate_stream():
-    """Generate response with streaming visualization."""
-    global model, tokenizer, config, device, chat_function
-
+    """Generate with streaming - optimized with queue."""
+    global model, tokenizer, chat_function
+    
     if model is None:
         return jsonify({'error': 'Model not loaded'}), 400
-
+    
     if chat_function is None:
         return jsonify({'error': 'Chat function not available'}), 400
-
+    
     data = request.json
     instruction = data.get('instruction', '')
-    steps = data.get('steps', 64)
-    block_size = data.get('block_size', 128)
-    max_new_tokens = data.get('max_new_tokens', 128)
-    parallel_blocks = data.get('parallel_blocks', 1)
-
+    steps = data.get('steps', 16)
+    block_size = data.get('block_size', 32)
+    max_new_tokens = data.get('max_new_tokens', 64)
+    parallel_blocks = data.get('parallel_blocks', torch.get_num_threads())
+    
     if not instruction:
         return jsonify({'error': 'No instruction provided'}), 400
-
+    
     def generate_events():
-        try:
-            # Import threading to allow yielding from callback
-            import queue
-            event_queue = queue.Queue()
-            generation_complete = {'done': False, 'result': None}
-
-            def streaming_visualizer(tok, context_ids, mask_blocks, is_masked_list, cfg, clear=True):
-                """This gets called during generation - we need to send events immediately"""
+        event_queue = queue.Queue(maxsize=100)  # Limit queue size
+        generation_complete = {'done': False, 'result': None}
+        
+        def streaming_visualizer(tok, context_ids, mask_blocks, is_masked_list, cfg, clear=True):
+            try:
                 visualizer = create_streaming_visualizer()
                 data = visualizer(tok, context_ids, mask_blocks, is_masked_list, cfg, clear)
-                # Put the update in the queue so it can be yielded immediately
-                event_queue.put({'type': 'update', 'data': data})
-
-            # Start generation in a separate thread so we can yield events as they come
-            import threading
-
-            def run_generation():
-                try:
+                event_queue.put({'type': 'update', 'data': data}, block=False)
+            except queue.Full:
+                pass  # Drop updates if queue full (prevents memory issues)
+        
+        def run_generation():
+            try:
+                with torch.inference_mode():
                     raw_output, response = chat_function(
-                        model,
-                        tokenizer,
-                        instruction,
-                        steps=steps,
-                        block_size=block_size,
-                        max_new_tokens=max_new_tokens,
-                        temperature=0.8,
-                        top_k=50,
-                        top_p=0.9,
-                        repetition_penalty=1.2,
-                        no_repeat_ngram_size=3,
-                        verbose=False,
-                        visualize_fn=streaming_visualizer,
+                        model, tokenizer, instruction,
+                        steps=steps, block_size=block_size, max_new_tokens=max_new_tokens,
+                        temperature=0.7, top_k=50, top_p=0.9,
+                        repetition_penalty=1.2, no_repeat_ngram_size=3,
+                        verbose=False, visualize_fn=streaming_visualizer,
                         parallel_blocks=parallel_blocks,
                     )
-                    generation_complete['result'] = (raw_output, response)
-                except Exception as e:
-                    generation_complete['result'] = ('error', str(e))
-                finally:
-                    generation_complete['done'] = True
-                    event_queue.put(None)  # Signal completion
-
-            # Start generation thread
-            gen_thread = threading.Thread(target=run_generation)
-            gen_thread.daemon = True
-            gen_thread.start()
-
-            # Yield start event
-            yield f"data: {json.dumps({'type': 'start', 'message': 'Generation started'})}\n\n"
-
-            # Yield events as they come from the queue
-            while not generation_complete['done'] or not event_queue.empty():
-                try:
-                    event = event_queue.get(timeout=0.1)
-                    if event is None:  # Completion signal
-                        break
-                    yield f"data: {json.dumps(event)}\n\n"
-                except queue.Empty:
-                    continue
-
-            # Wait for thread to finish
-            gen_thread.join(timeout=1.0)
-
-            # Send final response
-            if generation_complete['result']:
-                raw_output, response = generation_complete['result']
-                if raw_output == 'error':
-                    yield f"data: {json.dumps({'type': 'error', 'error': response})}\n\n"
-                else:
-                    yield f"data: {json.dumps({'type': 'complete', 'response': response, 'raw_output': raw_output})}\n\n"
-
-        except Exception as e:
-            import traceback
-            traceback.print_exc()
-            yield f"data: {json.dumps({'type': 'error', 'error': str(e)})}\n\n"
-
+                generation_complete['result'] = (raw_output, response)
+            except Exception as e:
+                generation_complete['result'] = ('error', str(e))
+            finally:
+                generation_complete['done'] = True
+                event_queue.put(None)
+        
+        import threading
+        gen_thread = threading.Thread(target=run_generation)
+        gen_thread.daemon = True
+        gen_thread.start()
+        
+        yield f"data: {json.dumps({'type': 'start'})}\n\n"
+        
+        while not generation_complete['done'] or not event_queue.empty():
+            try:
+                event = event_queue.get(timeout=0.1)
+                if event is None:
+                    break
+                yield f"data: {json.dumps(event)}\n\n"
+            except queue.Empty:
+                continue
+        
+        gen_thread.join(timeout=2.0)
+        
+        if generation_complete['result']:
+            raw_output, response = generation_complete['result']
+            if raw_output == 'error':
+                yield f"data: {json.dumps({'type': 'error', 'error': response})}\n\n"
+            else:
+                yield f"data: {json.dumps({'type': 'complete', 'response': response})}\n\n"
+    
     return Response(
         stream_with_context(generate_events()),
         mimetype='text/event-stream',
         headers={
             'Cache-Control': 'no-cache',
-            'X-Accel-Buffering': 'no'
+            'X-Accel-Buffering': 'no',
+            'Connection': 'keep-alive'
         }
     )
 
-
-@app.route('/api/test-stream', methods=['GET'])
-def test_stream():
-    """Test streaming endpoint."""
-    def generate():
-        for i in range(10):
-            yield f"data: {json.dumps({'message': f'Test message {i+1}'})}\n\n"
-            time.sleep(0.5)
-        yield f"data: {json.dumps({'message': 'Stream complete'})}\n\n"
-
-    return Response(
-        stream_with_context(generate()),
-        mimetype='text/event-stream',
-        headers={
-            'Cache-Control': 'no-cache',
-            'X-Accel-Buffering': 'no'
+@app.route('/api/status', methods=['GET'])
+def status():
+    """Get system status."""
+    return jsonify({
+        'model_loaded': model is not None,
+        'torch_threads': torch.get_num_threads(),
+        'interop_threads': torch.get_num_interop_threads(),
+        'cpu_count': os.cpu_count(),
+        'optimizations': {
+            'onnx_runtime': USE_ONNX_RUNTIME,
+            'ipex': USE_IPEX,
+            'torch_compile': USE_TORCH_COMPILE,
+            'quantization': QUANTIZE_MODEL
         }
-    )
-
+    })
 
 if __name__ == '__main__':
-    app.run(debug=True, host='0.0.0.0', port=7860, threaded=True)
+    print("\n" + "=" * 70)
+    print("ULTRA-FAST CPU INFERENCE SERVER")
+    print("=" * 70)
+    print("Available optimizations:")
+    print("  ✓ ONNX Runtime (best):", USE_ONNX_RUNTIME)
+    print("  ✓ Intel IPEX:", USE_IPEX)
+    print("  ✓ torch.compile:", USE_TORCH_COMPILE)
+    print("  ✓ Quantization:", QUANTIZE_MODEL)
+    print("  ✓ Multi-threading")
+    print("  ✓ Inference mode")
+    print("  ✓ Fast tokenizer")
+    print("  ✓ Memory layout optimization")
+    print("\nTo install ONNX Runtime: pip install onnxruntime")
+    print("To install Intel IPEX: pip install intel-extension-for-pytorch")
+    print("=" * 70 + "\n")
+    
+    app.run(debug=False, host='0.0.0.0', port=7860, threaded=True)