Upload 3 files

model_index.json

@@ -0,0 +1,97 @@
+{
+  "model_type": "nexus-worldmodel",
+  "version": "2.0",
+  "architecture_type": "cognitive-dynamic",
+  "framework": "pytorch",
+  "library_name": "nexus-worldmodel",
+  
+  "description": "NEXUS-WorldModel v2.0 - Neural World Simulator with Full Cognitive Architecture",
+  
+  "features": [
+    "EARCP (Ensemble Auto-Regulated Coherence Protocol)",
+    "LPOL Memory with 9 world-specific domains",
+    "GQA (Grouped Query Attention) for efficient memory access",
+    "Neurogenesis (dynamic neuron growth)",
+    "Energy System for cognitive resource management",
+    "Dream Phase with prioritized replay",
+    "Multi-World Buffers (physical, social, abstract, temporal)",
+    "Physics Prior (Mixture Density Network)",
+    "RoPE Attention (Rotary Positional Encoding)",
+    "VAE for visual encoding/decoding"
+  ],
+  
+  "files": {
+    "config": "config.json",
+    "weights": "pytorch_model.bin",
+    "cognitive_state": "cognitive_state.json",
+    "training_state": "training_state.json"
+  },
+  
+  "optional_files": {
+    "optimizer": "optimizer.pt",
+    "source_code": "NEXUS_WorldModel_v2.py",
+    "config_module": "nexus_worldmodel_config.py"
+  },
+  
+  "load_instructions": {
+    "python": "from nexus_worldmodel_config import load_nexus_worldmodel\nfrom NEXUS_WorldModel_v2 import NexusWorldModel\n\nmodel, config, cognitive_state, warnings = load_nexus_worldmodel(\n    NexusWorldModel,\n    'amewebstudio/nexus-worldmodel-v2',\n    device='cuda'\n)",
+    "note": "Use strict=False to handle dynamic architecture size mismatches"
+  },
+  
+  "dynamic_components": {
+    "experts": {
+      "description": "Expert count per EARCP layer can grow during training",
+      "initial": 6,
+      "max": 12,
+      "growth_trigger": "coherence < 0.3"
+    },
+    "neurons": {
+      "description": "Neuron count in neurogenesis layer can change",
+      "initial": 64,
+      "min": 32,
+      "max": 256,
+      "birth_trigger": "coherence > 0.8",
+      "death_trigger": "usage < 0.05"
+    }
+  },
+  
+  "cognitive_systems": {
+    "energy": {
+      "description": "Manages cognitive resource consumption",
+      "think_cost": 0.02,
+      "dream_cost": 0.1,
+      "regeneration": 0.05
+    },
+    "dream": {
+      "description": "Memory consolidation through prioritized replay",
+      "cycle_length": 50,
+      "duration": 10
+    },
+    "memory": {
+      "lpol_domains": 9,
+      "episodic_slots": 256,
+      "multi_scale": true
+    }
+  },
+  
+  "training_info": {
+    "recommended_epochs": 15,
+    "recommended_batch_size": 32,
+    "learning_rate": 0.0001,
+    "optimizer": "AdamW",
+    "scheduler": "CosineAnnealingLR",
+    "mixed_precision": true
+  },
+  
+  "author": {
+    "name": "Mike Amega (Logo)",
+    "organization": "Ame Web Studio",
+    "email": "contact@amewebstudio.com"
+  },
+  
+  "license": "Apache-2.0",
+  
+  "citation": {
+    "bibtex": "@software{nexus_worldmodel_2025,\n  author = {Amega, Mike},\n  title = {NEXUS-WorldModel: Neural World Simulator with Cognitive Architecture},\n  year = {2025},\n  url = {https://huggingface.co/amewebstudio/nexus-worldmodel-v2}\n}"
+  }
+}

nexus_worldmodel_config.py

@@ -0,0 +1,979 @@
+#!/usr/bin/env python3
+"""
+================================================================================
+NEXUS-WorldModel v2.0 - HuggingFace Configuration System
+================================================================================
+
+This module provides a robust configuration system for saving/loading 
+NEXUS-WorldModel with its dynamic cognitive architecture.
+
+Features:
+- Compatible with HuggingFace Hub
+- Handles dynamic components (neurogenesis, expert growth)
+- Preserves cognitive state (energy, dream, memory)
+- Supports incremental retraining
+
+Author: Mike Amega (Logo) - Ame Web Studio
+License: Apache 2.0
+================================================================================
+"""
+
+import os
+import json
+import copy
+from typing import Dict, List, Optional, Any, Union
+from dataclasses import dataclass, field, asdict
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+
+# ==============================================================================
+# CONFIGURATION CLASSES
+# ==============================================================================
+
+@dataclass
+class WorldConfig:
+    """Configuration for the 2D World Simulation"""
+    width: int = 64
+    height: int = 64
+    channels: int = 3
+    gravity: float = 0.1
+    friction: float = 0.98
+    bounce: float = 0.8
+    max_velocity: float = 5.0
+    max_agents: int = 5
+    max_obstacles: int = 10
+    max_zones: int = 3
+    agent_radius: float = 2.0
+    dt: float = 1.0
+    
+    def to_dict(self) -> Dict:
+        return asdict(self)
+    
+    @classmethod
+    def from_dict(cls, d: Dict) -> "WorldConfig":
+        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})
+
+
+@dataclass
+class NexusWorldModelConfig:
+    """
+    NEXUS-WorldModel v2.0 Configuration
+    
+    This configuration class is designed to be:
+    1. Serializable to JSON for HuggingFace Hub
+    2. Robust to architecture changes (neurogenesis, expert growth)
+    3. Compatible with partial loading for transfer learning
+    
+    Architecture Type: cognitive-dynamic
+    - Experts can grow dynamically
+    - Neurons can be born/die
+    - Memory states are persistent
+    """
+    
+    # === Model Identity ===
+    model_type: str = "nexus-worldmodel"
+    version: str = "2.0"
+    codename: str = "WorldSim-Cognitive"
+    architecture_type: str = "cognitive-dynamic"
+    
+    # === Core Dimensions ===
+    d_model: int = 512
+    d_ff: int = 2048
+    n_layers: int = 8
+    n_heads: int = 8
+    dropout: float = 0.1
+    
+    # === Latent Space ===
+    latent_dim: int = 256
+    latent_state_dim: int = 256
+    
+    # === Action Space ===
+    action_dim: int = 5  # [no_op, up, down, left, right]
+    
+    # === Sequence ===
+    max_seq_len: int = 512
+    context_length: int = 16
+    prediction_horizon: int = 8
+    
+    # === VAE ===
+    encoder_channels: List[int] = field(default_factory=lambda: [32, 64, 128, 256])
+    decoder_channels: List[int] = field(default_factory=lambda: [256, 128, 64, 32])
+    kl_weight: float = 0.0001
+    
+    # === LPOL Memory (9 domains) ===
+    use_lpol: bool = True
+    memory_size: int = 512
+    memory_k: int = 8
+    domain_types: List[str] = field(default_factory=lambda: [
+        'physics', 'spatial', 'temporal', 'causal',
+        'agent', 'object', 'zone', 'action', 'reward'
+    ])
+    
+    # === GQA (Grouped Query Attention) ===
+    use_gqa: bool = True
+    gqa_num_heads: int = 8
+    gqa_num_kv_groups: int = 2
+    
+    # === Multi-Scale Memory ===
+    multi_scale_enabled: bool = True
+    st_decay: float = 0.9
+    lt_decay: float = 0.99
+    
+    # === Episodic Memory ===
+    episodic_size: int = 256
+    episodic_dim: int = 128
+    
+    # === Structural State ===
+    structural_dim: int = 64
+    structural_decay: float = 0.95
+    
+    # === EARCP (Ensemble Auto-Regulated Coherence Protocol) ===
+    expert_types: List[str] = field(default_factory=lambda: [
+        'Physics', 'Spatial', 'Temporal', 'Causal', 'Prediction', 'Planning'
+    ])
+    max_experts: int = 12
+    growth_threshold_coherence: float = 0.3
+    growth_patience: int = 10
+    
+    # === Neurogenesis ===
+    neurogenesis_enabled: bool = True
+    min_neurons_per_layer: int = 32
+    max_neurons_per_layer: int = 256
+    neuron_birth_threshold: float = 0.8
+    neuron_death_threshold: float = 0.05
+    
+    # === Energy System ===
+    energy_cost_think: float = 0.02
+    energy_cost_dream: float = 0.1
+    energy_regeneration: float = 0.05
+    
+    # === Dream Phase ===
+    dream_enabled: bool = True
+    dream_cycle_length: int = 50
+    dream_duration: int = 10
+    prioritized_replay: bool = True
+    
+    # === Internal State ===
+    internal_state_dim: int = 128
+    tension_integration_rate: float = 0.1
+    
+    # === Multi-World Buffers ===
+    world_state_dim: int = 128
+    world_update_rate: float = 0.1
+    world_domains: List[str] = field(default_factory=lambda: ['physical', 'social', 'abstract', 'temporal'])
+    
+    # === Physics Prior (MDN) ===
+    physics_components: int = 8
+    physics_hidden: int = 256
+    
+    # === Training ===
+    batch_size: int = 32
+    learning_rate: float = 1e-4
+    epochs: int = 15
+    
+    # === HuggingFace Hub ===
+    push_to_hub: bool = True
+    hub_model_id: str = "amewebstudio/nexus-worldmodel-v2"
+    
+    # === World Simulation ===
+    world: WorldConfig = field(default_factory=WorldConfig)
+    
+    # === Dynamic State Tracking (for save/load) ===
+    # These are updated during training and saved with the model
+    _dynamic_state: Dict = field(default_factory=lambda: {
+        'current_experts_per_layer': [],  # List of expert counts per layer
+        'current_neurons': 64,            # Current neuron count in neurogenesis
+        'total_births': 0,
+        'total_deaths': 0,
+        'total_dreams': 0,
+        'training_epochs_completed': 0,
+        'best_loss': float('inf')
+    })
+    
+    def __post_init__(self):
+        # Initialize dynamic state for layers if empty
+        if not self._dynamic_state.get('current_experts_per_layer'):
+            self._dynamic_state['current_experts_per_layer'] = [
+                len(self.expert_types) for _ in range(self.n_layers)
+            ]
+    
+    def to_dict(self) -> Dict:
+        """Convert config to dictionary for JSON serialization"""
+        d = {}
+        for key, value in self.__dict__.items():
+            if key == 'world':
+                d[key] = value.to_dict()
+            elif isinstance(value, (list, dict, str, int, float, bool, type(None))):
+                d[key] = value
+            else:
+                d[key] = str(value)
+        return d
+    
+    def to_json_string(self) -> str:
+        """Serialize config to JSON string"""
+        return json.dumps(self.to_dict(), indent=2)
+    
+    def save_pretrained(self, save_directory: Union[str, Path]):
+        """Save configuration to directory (HuggingFace standard)"""
+        save_directory = Path(save_directory)
+        save_directory.mkdir(parents=True, exist_ok=True)
+        
+        # Save main config
+        config_path = save_directory / "config.json"
+        with open(config_path, 'w') as f:
+            json.dump(self.to_dict(), f, indent=2)
+        
+        # Save model index for HuggingFace
+        model_index = {
+            "model_type": self.model_type,
+            "version": self.version,
+            "architecture_type": self.architecture_type,
+            "framework": "pytorch",
+            "files": {
+                "config": "config.json",
+                "weights": "pytorch_model.bin",
+                "cognitive_state": "cognitive_state.json"
+            }
+        }
+        with open(save_directory / "model_index.json", 'w') as f:
+            json.dump(model_index, f, indent=2)
+        
+        print(f"✅ Configuration saved to {save_directory}")
+    
+    @classmethod
+    def from_dict(cls, config_dict: Dict) -> "NexusWorldModelConfig":
+        """Create config from dictionary"""
+        # Handle nested WorldConfig
+        if 'world' in config_dict and isinstance(config_dict['world'], dict):
+            config_dict = config_dict.copy()
+            config_dict['world'] = WorldConfig.from_dict(config_dict['world'])
+        
+        # Filter to known fields
+        known_fields = set(cls.__dataclass_fields__.keys())
+        filtered = {k: v for k, v in config_dict.items() if k in known_fields}
+        
+        return cls(**filtered)
+    
+    @classmethod
+    def from_json_file(cls, json_file: Union[str, Path]) -> "NexusWorldModelConfig":
+        """Load config from JSON file"""
+        with open(json_file, 'r') as f:
+            config_dict = json.load(f)
+        return cls.from_dict(config_dict)
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_path: Union[str, Path]) -> "NexusWorldModelConfig":
+        """Load configuration from pretrained directory or HuggingFace Hub"""
+        pretrained_path = Path(pretrained_path)
+        
+        # Try local path first
+        config_file = pretrained_path / "config.json"
+        if config_file.exists():
+            return cls.from_json_file(config_file)
+        
+        # Try HuggingFace Hub
+        try:
+            from huggingface_hub import hf_hub_download
+            config_file = hf_hub_download(
+                repo_id=str(pretrained_path),
+                filename="config.json"
+            )
+            return cls.from_json_file(config_file)
+        except Exception as e:
+            raise ValueError(f"Could not load config from {pretrained_path}: {e}")
+    
+    def update_dynamic_state(self, model: nn.Module):
+        """Update dynamic state from model after training"""
+        # Get expert counts per layer
+        if hasattr(model, 'layers'):
+            self._dynamic_state['current_experts_per_layer'] = [
+                layer.get_expert_count() if hasattr(layer, 'get_expert_count') else len(self.expert_types)
+                for layer in model.layers
+            ]
+        
+        # Get neurogenesis stats
+        if hasattr(model, 'neurogenesis'):
+            stats = model.neurogenesis.get_stats()
+            self._dynamic_state['current_neurons'] = stats.get('total_neurons', 64)
+            self._dynamic_state['total_births'] = stats.get('total_births', 0)
+            self._dynamic_state['total_deaths'] = stats.get('total_deaths', 0)
+        
+        # Get dream stats
+        if hasattr(model, 'dream'):
+            dream_stats = model.dream.get_stats()
+            self._dynamic_state['total_dreams'] = dream_stats.get('total_dreams', 0)
+    
+    def get_architecture_summary(self) -> str:
+        """Get a human-readable architecture summary"""
+        return f"""
+╔══════════════════════════════════════════════════════════════════════╗
+║                    NEXUS-WorldModel v{self.version} Configuration                ║
+╠══════════════════════════════════════════════════════════════════════╣
+║ Architecture Type: {self.architecture_type:47} ║
+║                                                                      ║
+║ Core:                                                                ║
+║   d_model: {self.d_model:5}  |  d_ff: {self.d_ff:5}  |  n_layers: {self.n_layers:2}  |  n_heads: {self.n_heads:2}  ║
+║   latent_dim: {self.latent_dim:3}  |  action_dim: {self.action_dim}                              ║
+║                                                                      ║
+║ Memory:                                                              ║
+║   LPOL: {len(self.domain_types)} domains  |  GQA: {self.gqa_num_heads} heads, {self.gqa_num_kv_groups} KV groups         ║
+║   Episodic: {self.episodic_size} slots  |  Memory size: {self.memory_size}                   ║
+║                                                                      ║
+║ Cognitive:                                                           ║
+║   Experts: {len(self.expert_types)} types (max {self.max_experts})  |  Neurogenesis: {'ON' if self.neurogenesis_enabled else 'OFF':3}        ║
+║   Dream Phase: {'ON' if self.dream_enabled else 'OFF':3}  |  Energy System: ON                        ║
+║   World Buffers: {len(self.world_domains)} domains                                       ║
+║                                                                      ║
+║ Dynamic State:                                                       ║
+║   Current neurons: {self._dynamic_state.get('current_neurons', 64):3}  |  Births: {self._dynamic_state.get('total_births', 0):4}  |  Deaths: {self._dynamic_state.get('total_deaths', 0):4}  ║
+║   Dreams completed: {self._dynamic_state.get('total_dreams', 0):4}                                      ║
+╚══════════════════════════════════════════════════════════════════════╝
+"""
+
+
+# ==============================================================================
+# COGNITIVE STATE MANAGER
+# ==============================================================================
+
+class CognitiveStateManager:
+    """
+    Manages the cognitive state of NEXUS-WorldModel for save/load operations.
+    
+    This is crucial for our dynamic architecture because:
+    1. Expert counts can change during training (growth)
+    2. Neuron counts can change (neurogenesis)
+    3. Memory buffers have persistent state
+    4. Energy and dream systems have state
+    
+    This manager ensures proper serialization and restoration of these
+    dynamic components.
+    """
+    
+    @staticmethod
+    def extract_cognitive_state(model: nn.Module) -> Dict[str, Any]:
+        """Extract all cognitive state from model"""
+        state = {
+            'version': '2.0',
+            'state_type': 'cognitive',
+            'components': {}
+        }
+        
+        # === EARCP Layer States ===
+        if hasattr(model, 'layers'):
+            layer_states = []
+            for i, layer in enumerate(model.layers):
+                layer_state = {
+                    'layer_idx': i,
+                    'expert_count': layer.get_expert_count() if hasattr(layer, 'get_expert_count') else 0,
+                    'low_coh_count': layer.low_coh_count.item() if hasattr(layer, 'low_coh_count') else 0
+                }
+                layer_states.append(layer_state)
+            state['components']['earcp_layers'] = layer_states
+        
+        # === Neurogenesis State ===
+        if hasattr(model, 'neurogenesis'):
+            ng = model.neurogenesis
+            state['components']['neurogenesis'] = {
+                'n_neurons': ng.n_neurons.item(),
+                'usage': ng.usage.tolist(),
+                'lifetime': ng.lifetime.tolist(),
+                'births': ng.births.item(),
+                'deaths': ng.deaths.item()
+            }
+        
+        # === Energy System State ===
+        if hasattr(model, 'energy'):
+            state['components']['energy'] = {
+                'energy': model.energy.energy.item(),
+                'consumed': model.energy.consumed.item()
+            }
+        
+        # === Dream Phase State ===
+        if hasattr(model, 'dream'):
+            dream = model.dream
+            state['components']['dream'] = {
+                'is_dreaming': dream.is_dreaming.item(),
+                'dream_step': dream.dream_step.item(),
+                'cycles_since': dream.cycles_since.item(),
+                'total_dreams': dream.total_dreams.item(),
+                'buffer_size': len(dream.buffer)
+            }
+        
+        # === Self Trace State ===
+        if hasattr(model, 'self_trace'):
+            state['components']['self_trace'] = {
+                'identity_norm': model.self_trace.identity.norm().item(),
+                'n_traces': model.self_trace.n_traces.item()
+            }
+        
+        # === Memory States ===
+        if hasattr(model, 'memory'):
+            mem = model.memory
+            state['components']['memory'] = {
+                'st_norm': mem.st.norm().item(),
+                'lt_norm': mem.lt.norm().item()
+            }
+        
+        # === Internal State ===
+        if hasattr(model, 'internal'):
+            state['components']['internal'] = {
+                'tension': model.internal.tension.integrate(),
+                'discomfort_norm': model.internal.discomfort.norm().item()
+            }
+        
+        # === World Buffers ===
+        if hasattr(model, 'world_buffers'):
+            wb_states = {}
+            for domain, buffer in model.world_buffers.world_buffers.items():
+                wb_states[domain] = {
+                    'state_norm': buffer.state.norm().item(),
+                    'prediction_norm': buffer.prediction.norm().item(),
+                    'surprise': buffer.surprise.item()
+                }
+            state['components']['world_buffers'] = wb_states
+        
+        return state
+    
+    @staticmethod
+    def save_cognitive_state(model: nn.Module, save_path: Union[str, Path]):
+        """Save cognitive state to JSON file"""
+        state = CognitiveStateManager.extract_cognitive_state(model)
+        
+        with open(save_path, 'w') as f:
+            json.dump(state, f, indent=2)
+        
+        print(f"✅ Cognitive state saved to {save_path}")
+    
+    @staticmethod
+    def load_cognitive_state(load_path: Union[str, Path]) -> Dict[str, Any]:
+        """Load cognitive state from JSON file"""
+        with open(load_path, 'r') as f:
+            state = json.load(f)
+        return state
+    
+    @staticmethod
+    def restore_cognitive_state(model: nn.Module, state: Dict[str, Any], 
+                                strict: bool = False) -> List[str]:
+        """
+        Restore cognitive state to model.
+        
+        Args:
+            model: The model to restore state to
+            state: The cognitive state dictionary
+            strict: If True, raise error on mismatch. If False, skip mismatches.
+        
+        Returns:
+            List of warnings/info about restoration
+        """
+        warnings = []
+        components = state.get('components', {})
+        
+        # === Restore Energy ===
+        if 'energy' in components and hasattr(model, 'energy'):
+            model.energy.energy.fill_(components['energy']['energy'])
+            model.energy.consumed.fill_(components['energy']['consumed'])
+            warnings.append("✓ Energy state restored")
+        
+        # === Restore Dream ===
+        if 'dream' in components and hasattr(model, 'dream'):
+            dream_state = components['dream']
+            model.dream.total_dreams.fill_(dream_state['total_dreams'])
+            model.dream.cycles_since.fill_(dream_state['cycles_since'])
+            warnings.append(f"✓ Dream state restored (total dreams: {dream_state['total_dreams']})")
+        
+        # === Restore Self Trace ===
+        if 'self_trace' in components and hasattr(model, 'self_trace'):
+            model.self_trace.n_traces.fill_(components['self_trace']['n_traces'])
+            warnings.append("✓ Self trace state restored")
+        
+        # === Restore Neurogenesis (careful - sizes may differ) ===
+        if 'neurogenesis' in components and hasattr(model, 'neurogenesis'):
+            ng_state = components['neurogenesis']
+            saved_neurons = ng_state['n_neurons']
+            current_neurons = model.neurogenesis.n_neurons.item()
+            
+            if saved_neurons != current_neurons:
+                if strict:
+                    raise ValueError(f"Neurogenesis mismatch: saved {saved_neurons}, current {current_neurons}")
+                warnings.append(f"⚠ Neurogenesis size mismatch: saved {saved_neurons} vs current {current_neurons}")
+            else:
+                model.neurogenesis.births.fill_(ng_state['births'])
+                model.neurogenesis.deaths.fill_(ng_state['deaths'])
+                warnings.append(f"✓ Neurogenesis state restored (neurons: {saved_neurons})")
+        
+        # === Restore EARCP Layers (careful - expert counts may differ) ===
+        if 'earcp_layers' in components and hasattr(model, 'layers'):
+            for layer_state in components['earcp_layers']:
+                idx = layer_state['layer_idx']
+                if idx < len(model.layers):
+                    saved_experts = layer_state['expert_count']
+                    current_experts = model.layers[idx].get_expert_count()
+                    
+                    if saved_experts != current_experts:
+                        warnings.append(f"⚠ Layer {idx} expert mismatch: saved {saved_experts} vs current {current_experts}")
+                    else:
+                        model.layers[idx].low_coh_count.fill_(layer_state['low_coh_count'])
+        
+        return warnings
+
+
+# ==============================================================================
+# MODEL SAVE/LOAD UTILITIES
+# ==============================================================================
+
+def save_nexus_worldmodel(
+    model: nn.Module,
+    config: NexusWorldModelConfig,
+    save_directory: Union[str, Path],
+    save_optimizer: bool = False,
+    optimizer: Optional[torch.optim.Optimizer] = None,
+    epoch: int = 0,
+    loss: float = 0.0
+):
+    """
+    Save NEXUS-WorldModel with full state for HuggingFace Hub.
+    
+    Saves:
+    - config.json: Model configuration
+    - pytorch_model.bin: Model weights
+    - cognitive_state.json: Dynamic cognitive state
+    - model_index.json: HuggingFace model index
+    - training_state.json: Training progress (optional)
+    """
+    save_directory = Path(save_directory)
+    save_directory.mkdir(parents=True, exist_ok=True)
+    
+    # Update config with current dynamic state
+    config.update_dynamic_state(model)
+    config._dynamic_state['training_epochs_completed'] = epoch
+    config._dynamic_state['best_loss'] = loss
+    
+    # 1. Save configuration
+    config.save_pretrained(save_directory)
+    
+    # 2. Save model weights
+    weights_path = save_directory / "pytorch_model.bin"
+    torch.save(model.state_dict(), weights_path)
+    print(f"✅ Model weights saved to {weights_path}")
+    
+    # 3. Save cognitive state
+    cognitive_path = save_directory / "cognitive_state.json"
+    CognitiveStateManager.save_cognitive_state(model, cognitive_path)
+    
+    # 4. Save training state (for resuming)
+    training_state = {
+        'epoch': epoch,
+        'loss': loss,
+        'config_hash': hash(config.to_json_string())
+    }
+    
+    if save_optimizer and optimizer is not None:
+        optimizer_path = save_directory / "optimizer.pt"
+        torch.save(optimizer.state_dict(), optimizer_path)
+        training_state['optimizer_saved'] = True
+        print(f"✅ Optimizer state saved to {optimizer_path}")
+    
+    with open(save_directory / "training_state.json", 'w') as f:
+        json.dump(training_state, f, indent=2)
+    
+    print(f"✅ Full model saved to {save_directory}")
+
+
+def load_nexus_worldmodel(
+    model_class,
+    pretrained_path: Union[str, Path],
+    device: Union[str, torch.device] = 'cpu',
+    strict: bool = False
+) -> tuple:
+    """
+    Load NEXUS-WorldModel from pretrained directory or HuggingFace Hub.
+    
+    IMPORTANT: This function handles dynamic architecture automatically!
+    - Neurogenesis: neurons may have been added/removed during training
+    - Expert Growth: experts may have been added during training
+    
+    The model is automatically resized before loading weights.
+    
+    Args:
+        model_class: The NexusWorldModel class
+        pretrained_path: Local path or HuggingFace repo ID
+        device: Device to load model to
+        strict: If True, require exact architecture match. 
+                Default False (recommended for dynamic arch).
+    
+    Returns:
+        (model, config, cognitive_state, warnings)
+    
+    Usage:
+        from NEXUS_WorldModel_v2 import NexusWorldModel
+        
+        model, config, cog_state, warnings = load_nexus_worldmodel(
+            NexusWorldModel,
+            "amewebstudio/nexus-worldmodel-v2",
+            device="cuda"
+        )
+    """
+    pretrained_path = Path(pretrained_path)
+    warnings = []
+    
+    # Try local path first
+    if pretrained_path.exists():
+        config_file = pretrained_path / "config.json"
+        weights_file = pretrained_path / "pytorch_model.bin"
+        checkpoint_file = pretrained_path / "nexus_worldmodel_v2.pt"
+        cognitive_file = pretrained_path / "cognitive_state.json"
+    else:
+        # Try HuggingFace Hub
+        try:
+            from huggingface_hub import hf_hub_download
+            
+            config_file = hf_hub_download(str(pretrained_path), "config.json")
+            
+            # Try checkpoint first (contains everything)
+            try:
+                checkpoint_file = hf_hub_download(str(pretrained_path), "nexus_worldmodel_v2.pt")
+                weights_file = None
+            except:
+                checkpoint_file = None
+                weights_file = hf_hub_download(str(pretrained_path), "pytorch_model.bin")
+            
+            try:
+                cognitive_file = hf_hub_download(str(pretrained_path), "cognitive_state.json")
+            except:
+                cognitive_file = None
+                warnings.append("⚠ No cognitive_state.json found")
+        except Exception as e:
+            raise ValueError(f"Could not load from {pretrained_path}: {e}")
+    
+    # 1. Load configuration
+    config = NexusWorldModelConfig.from_json_file(config_file)
+    print(f"✅ Configuration loaded")
+    
+    # 2. Create model with default sizes
+    model = model_class(config)
+    model = model.to(device)
+    
+    # 3. Load weights
+    if checkpoint_file and Path(checkpoint_file).exists():
+        checkpoint = torch.load(checkpoint_file, map_location=device)
+        state_dict = checkpoint.get('model', checkpoint)
+        
+        # Extract training info from checkpoint
+        if 'epochs' in checkpoint:
+            warnings.append(f"ℹ Trained for {checkpoint['epochs']} epochs")
+        if 'loss' in checkpoint:
+            warnings.append(f"ℹ Final loss: {checkpoint['loss']:.4f}")
+    elif weights_file and Path(weights_file).exists():
+        state_dict = torch.load(weights_file, map_location=device)
+    else:
+        raise FileNotFoundError(f"No weights found in {pretrained_path}")
+    
+    # 4. CRITICAL: Resize model to match saved dimensions
+    # This handles neurogenesis and expert growth automatically
+    if hasattr(model, 'resize_for_loading'):
+        resize_warnings = model.resize_for_loading(state_dict)
+        warnings.extend(resize_warnings)
+    else:
+        # Manual resize for older model versions
+        warnings.extend(_manual_resize_model(model, state_dict, config))
+    
+    # 5. Load weights (now sizes match!)
+    try:
+        incompatible = model.load_state_dict(state_dict, strict=strict)
+        
+        if hasattr(incompatible, 'missing_keys') and incompatible.missing_keys:
+            warnings.append(f"⚠ Missing {len(incompatible.missing_keys)} keys")
+        if hasattr(incompatible, 'unexpected_keys') and incompatible.unexpected_keys:
+            warnings.append(f"⚠ Unexpected {len(incompatible.unexpected_keys)} keys")
+        
+        print(f"✅ Model weights loaded")
+    except RuntimeError as e:
+        if "size mismatch" in str(e):
+            warnings.append(f"❌ Size mismatch even after resize: {e}")
+            # Try loading what we can
+            model_state = model.state_dict()
+            loadable = {k: v for k, v in state_dict.items() 
+                       if k in model_state and model_state[k].shape == v.shape}
+            model.load_state_dict(loadable, strict=False)
+            warnings.append(f"⚠ Partially loaded {len(loadable)}/{len(state_dict)} tensors")
+        else:
+            raise
+    
+    # 6. Load and restore cognitive state
+    cognitive_state = None
+    if cognitive_file and Path(cognitive_file).exists():
+        cognitive_state = CognitiveStateManager.load_cognitive_state(cognitive_file)
+        restore_warnings = CognitiveStateManager.restore_cognitive_state(
+            model, cognitive_state, strict=strict
+        )
+        warnings.extend(restore_warnings)
+        print(f"✅ Cognitive state restored")
+    
+    return model, config, cognitive_state, warnings
+
+
+def _manual_resize_model(model, state_dict: Dict[str, torch.Tensor], 
+                         config) -> List[str]:
+    """
+    Manually resize model components for older model versions
+    that don't have resize_for_loading method.
+    """
+    warnings = []
+    
+    # Resize neurogenesis
+    for key, tensor in state_dict.items():
+        if 'neurogenesis.weights' in key:
+            saved_neurons = tensor.size(0)
+            if hasattr(model, 'neurogenesis'):
+                current_neurons = model.neurogenesis.n_neurons.item()
+                if current_neurons != saved_neurons:
+                    if hasattr(model.neurogenesis, 'resize'):
+                        model.neurogenesis.resize(saved_neurons)
+                        warnings.append(f"✓ Neurogenesis: {current_neurons} → {saved_neurons}")
+                    else:
+                        warnings.append(f"⚠ Cannot resize neurogenesis: no resize method")
+            break
+    
+    # Resize experts per layer
+    n_layers = config.n_layers if hasattr(config, 'n_layers') else 8
+    for layer_idx in range(n_layers):
+        expert_count = 0
+        for key in state_dict.keys():
+            if f'layers.{layer_idx}.experts.' in key and '.fc1.weight' in key:
+                expert_count += 1
+        
+        if expert_count > 0 and hasattr(model, 'layers') and layer_idx < len(model.layers):
+            current_count = model.layers[layer_idx].get_expert_count()
+            if current_count != expert_count:
+                if hasattr(model.layers[layer_idx], 'resize_experts'):
+                    model.layers[layer_idx].resize_experts(expert_count)
+                    warnings.append(f"✓ Layer {layer_idx}: {current_count} → {expert_count} experts")
+                else:
+                    warnings.append(f"⚠ Cannot resize layer {layer_idx}: no resize_experts method")
+    
+    return warnings
+
+
+# ==============================================================================
+# HUGGINGFACE HUB INTEGRATION
+# ==============================================================================
+
+def push_to_hub(
+    model: nn.Module,
+    config: NexusWorldModelConfig,
+    repo_id: str,
+    token: Optional[str] = None,
+    commit_message: str = "Update NEXUS-WorldModel",
+    private: bool = False
+):
+    """
+    Push NEXUS-WorldModel to HuggingFace Hub.
+    
+    Args:
+        model: The trained model
+        config: Model configuration
+        repo_id: HuggingFace repository ID (e.g., "username/model-name")
+        token: HuggingFace API token
+        commit_message: Commit message
+        private: Whether the repo should be private
+    """
+    from huggingface_hub import HfApi, create_repo, upload_folder
+    import tempfile
+    
+    # Get token
+    if token is None:
+        token = os.environ.get('HF_TOKEN')
+    
+    if token is None:
+        try:
+            from kaggle_secrets import UserSecretsClient
+            token = UserSecretsClient().get_secret("HF_TOKEN")
+        except:
+            pass
+    
+    if token is None:
+        raise ValueError("No HuggingFace token found. Set HF_TOKEN environment variable.")
+    
+    api = HfApi()
+    
+    # Create repo if needed
+    try:
+        create_repo(repo_id, token=token, private=private, exist_ok=True)
+        print(f"✅ Repository ready: {repo_id}")
+    except Exception as e:
+        print(f"⚠ Repository creation: {e}")
+    
+    # Save to temp directory
+    with tempfile.TemporaryDirectory() as tmpdir:
+        save_nexus_worldmodel(
+            model=model,
+            config=config,
+            save_directory=tmpdir,
+            epoch=config._dynamic_state.get('training_epochs_completed', 0),
+            loss=config._dynamic_state.get('best_loss', 0)
+        )
+        
+        # Create README
+        readme_content = create_model_card(model, config)
+        with open(os.path.join(tmpdir, "README.md"), 'w') as f:
+            f.write(readme_content)
+        
+        # Upload
+        api.upload_folder(
+            folder_path=tmpdir,
+            repo_id=repo_id,
+            token=token,
+            commit_message=commit_message
+        )
+    
+    print(f"✅ Model pushed to: https://huggingface.co/{repo_id}")
+
+
+def create_model_card(model: nn.Module, config: NexusWorldModelConfig) -> str:
+    """Create a model card for HuggingFace Hub"""
+    
+    # Get dynamic stats
+    stats = config._dynamic_state
+    
+    return f"""---
+license: apache-2.0
+language:
+  - en
+tags:
+  - nexus-worldmodel
+  - world-model
+  - cognitive-architecture
+  - earcp
+  - lpol
+  - neurogenesis
+  - gqa
+  - pytorch
+pipeline_tag: reinforcement-learning
+library_name: pytorch
+---
+
+# NEXUS-WorldModel v{config.version}
+
+**"Learning to Simulate Reality with Full Cognitive Architecture"**
+
+## Model Description
+
+NEXUS-WorldModel is a neural world simulator that uses a complete cognitive architecture to learn and predict 2D physics environments. Unlike traditional world models, it features:
+
+- **Dynamic Architecture**: Experts and neurons grow during training
+- **Cognitive Systems**: Energy management, dream phases, memory consolidation
+- **Multi-Domain Memory**: 9 specialized LPOL memory domains
+- **GQA**: Grouped Query Attention for efficient memory access
+
+## Architecture Components
+
+| Component | Description |
+|-----------|-------------|
+| **EARCP Module** | Sparse Compression + Gated Integration |
+| **LPOL Memory** | {len(config.domain_types)} domains with GQA |
+| **GQA** | {config.gqa_num_heads} heads, {config.gqa_num_kv_groups} KV groups |
+| **EARCP Layers** | {config.n_layers} layers with dynamic experts |
+| **Neurogenesis** | {'Enabled' if config.neurogenesis_enabled else 'Disabled'} |
+| **Dream Phase** | {'Enabled' if config.dream_enabled else 'Disabled'} |
+| **Physics Prior** | MDN with {config.physics_components} components |
+
+## Training Information
+
+- **Epochs Completed**: {stats.get('training_epochs_completed', 0)}
+- **Best Loss**: {stats.get('best_loss', 'N/A')}
+- **Parameters**: {model.count_params():,}
+- **Current Neurons**: {stats.get('current_neurons', 64)}
+- **Total Births**: {stats.get('total_births', 0)}
+- **Total Deaths**: {stats.get('total_deaths', 0)}
+- **Total Dreams**: {stats.get('total_dreams', 0)}
+
+## Usage
+
+```python
+from nexus_worldmodel_config import load_nexus_worldmodel, NexusWorldModelConfig
+from NEXUS_WorldModel_v2 import NexusWorldModel
+
+# Load from HuggingFace Hub
+model, config, cognitive_state, warnings = load_nexus_worldmodel(
+    NexusWorldModel,
+    "{config.hub_model_id}",
+    device="cuda"
+)
+
+# Or create fresh model
+config = NexusWorldModelConfig()
+model = NexusWorldModel(config)
+
+# Encode observation
+z, mu, logvar, z_d = model.encode(obs_tensor)
+
+# Predict next state
+prediction = model.predict_next(z, action)
+
+# Dream/imagine trajectory
+dream = model.dream_trajectory(z_start, action_sequence)
+```
+
+## Configuration
+
+Key parameters:
+```python
+d_model: {config.d_model}
+n_layers: {config.n_layers}
+latent_dim: {config.latent_dim}
+domain_types: {config.domain_types}
+expert_types: {config.expert_types}
+```
+
+## Files
+
+| File | Description |
+|------|-------------|
+| `config.json` | Model configuration |
+| `pytorch_model.bin` | Model weights |
+| `cognitive_state.json` | Dynamic cognitive state |
+| `model_index.json` | HuggingFace model index |
+
+## Important Notes
+
+⚠️ **Dynamic Architecture**: This model has components that can grow during training:
+- Expert count per layer may vary
+- Neuron count in neurogenesis layer may vary
+- Use `strict=False` when loading to handle size mismatches
+
+## Author
+
+**Mike Amega (Logo)** - Ame Web Studio
+
+## License
+
+Apache 2.0
+"""
+
+
+# ==============================================================================
+# MAIN (Testing)
+# ==============================================================================
+
+if __name__ == "__main__":
+    # Test configuration
+    config = NexusWorldModelConfig()
+    
+    print(config.get_architecture_summary())
+    
+    # Test serialization
+    json_str = config.to_json_string()
+    print("\n📄 Config JSON (truncated):")
+    print(json_str[:500] + "...")
+    
+    # Test save/load cycle
+    import tempfile
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config.save_pretrained(tmpdir)
+        loaded = NexusWorldModelConfig.from_pretrained(tmpdir)
+        print(f"\n✅ Save/Load cycle successful")
+        print(f"   Original d_model: {config.d_model}")
+        print(f"   Loaded d_model: {loaded.d_model}")

nexus_worldmodel_v2.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6f3595ff8e808d62df7f7077c15e8e72542adcac55eb74024945df4479a613df
-size 18509763
+oid sha256:61b77e3c47e91451f07ee84f6e1031561a20088ddb50cdfdaa0563934cb89d11
+size 914665457