""" ArcisVLM API — Dependency injection for model, agents, and camera manager. Provides a singleton ModelManager that lazy-loads the VLM and agent framework on first use, and exposes it via get_model_manager(). """ from __future__ import annotations import logging import os import time import threading from typing import Optional logger = logging.getLogger(__name__) # Singleton instance _manager: Optional["ModelManager"] = None _lock = threading.Lock() class ModelManager: """ Manages the lifecycle of the VLM model, agent orchestrator, and camera manager. All heavy objects are lazy-loaded on initialize() to keep import time fast. """ def __init__( self, config_path: str = "configs/scale_1.3b.yaml", checkpoint_path: Optional[str] = None, device: str = "cpu", ) -> None: self.config_path = config_path self.checkpoint_path = checkpoint_path self.device = device self._model = None self._tokenizer = None self._orchestrator = None self._camera_manager = None self._initialized = False self._alert_history: list[dict] = [] self.latency_history: list[float] = [] self.inference_count: int = 0 self.queries_per_sec: float = 0.0 self.start_time: float = 0.0 self.model_params: int = 0 self._gemma_backbone = None # Gemma 4 E2B backbone (when available) @property def is_initialized(self) -> bool: return self._initialized @property def model(self): return self._model @property def tokenizer(self): return self._tokenizer @property def orchestrator(self): return self._orchestrator @property def camera_manager(self): return self._camera_manager @property def gemma_backbone(self): return self._gemma_backbone @property def alert_history(self) -> list[dict]: return self._alert_history def add_alert_event(self, event: dict) -> None: self._alert_history.append(event) # Keep only last 1000 events if len(self._alert_history) > 1000: self._alert_history = self._alert_history[-1000:] def initialize(self) -> None: """Load model, tokenizer, and set up the agent orchestrator.""" if self._initialized: return logger.info("ModelManager: initializing (config=%s, device=%s)", self.config_path, self.device) try: import yaml import torch # Load config with open(self.config_path) as f: config = yaml.safe_load(f) # Build model from model.vlm import VLJEPAModel self._model = VLJEPAModel(config) # Load checkpoint if provided if self.checkpoint_path: ckpt = torch.load(self.checkpoint_path, map_location=self.device, weights_only=False) state = ckpt.get("model_state_dict", ckpt) self._model.load_state_dict(state, strict=False) logger.info("Loaded checkpoint: %s", self.checkpoint_path) self._model = self._model.to(self.device) self._model.eval() # Load tokenizer from model.tokenizer import BPETokenizer tok_path = config.get("tokenizer", {}).get("path", "checkpoints/tokenizer_32k.json") self._tokenizer = BPETokenizer(vocab_size=config.get("tokenizer", {}).get("vocab_size", 32768)) try: self._tokenizer.load(tok_path) logger.info("Loaded tokenizer from %s", tok_path) except FileNotFoundError: logger.warning("Tokenizer file not found at %s — using untrained tokenizer", tok_path) # Set up orchestrator with agents self._setup_agents() self.start_time = time.time() if self._model: self.model_params = sum(p.numel() for p in self._model.parameters()) # Try loading Gemma 4 E2B backbone (if USE_GEMMA=1 env var set) if os.environ.get("USE_GEMMA", "").strip() == "1": try: from model.gemma_backbone import GemmaBackbone self._gemma_backbone = GemmaBackbone(device=self.device) self._gemma_backbone.load() self.model_params = sum( p.numel() for p in self._gemma_backbone._model.parameters() ) logger.info("Gemma 4 E2B backbone loaded successfully") except Exception as e: logger.warning(f"Gemma backbone not available: {e}. Using custom VL-JEPA.") self._gemma_backbone = None self._initialized = True logger.info("ModelManager: initialization complete") except Exception as e: logger.error("ModelManager: initialization failed — %s", e) # Still mark as partially initialized so API can respond with degraded status self._initialized = True logger.info("ModelManager: running in degraded mode (no model)") def _setup_agents(self) -> None: """Create the MotherOrchestrator and register child agents.""" from agents.mother import MotherOrchestrator from agents.vqa import VQAAgent from agents.detect import DetectAgent from agents.alert import AlertAgent from agents.caption import CaptionAgent from agents.track import TrackingAgent from agents.count import CountingAgent from agents.ocr import OCRAgent from agents.reason import ReasoningAgent self._orchestrator = MotherOrchestrator(max_agents_per_type=4) # Create all 8 agent types backed by the shared model agent_classes = [ ("vqa-0001", "vqa", VQAAgent), ("detect-0001", "detect", DetectAgent), ("alert-0001", "alert", AlertAgent), ("caption-0001", "caption", CaptionAgent), ("track-0001", "track", TrackingAgent), ("count-0001", "count", CountingAgent), ("ocr-0001", "ocr", OCRAgent), ("reason-0001", "reason", ReasoningAgent), ] for agent_id, expert_type, cls in agent_classes: try: agent = cls(agent_id=agent_id, model=self._model, tokenizer=self._tokenizer) self._orchestrator.register_agent(agent) except Exception as e: logger.warning("Failed to register %s agent: %s", expert_type, e) def submit_task(self, task): """Submit a task to the orchestrator (synchronous path).""" if self._orchestrator is None: from agents.base import Result return Result(answer="", error="Agent orchestrator not initialized") start = time.time() result = self._orchestrator.submit(task) elapsed_ms = (time.time() - start) * 1000 self.latency_history.append(elapsed_ms) self.inference_count += 1 if len(self.latency_history) > 10000: self.latency_history = self.latency_history[-10000:] return result def shutdown(self) -> None: """Clean up all resources.""" if self._orchestrator: self._orchestrator.shutdown() if self._camera_manager: try: self._camera_manager.stop_all() except Exception: pass self._initialized = False logger.info("ModelManager: shutdown complete") def get_model_manager() -> ModelManager: """Return the global ModelManager singleton, creating it if needed.""" global _manager if _manager is None: with _lock: if _manager is None: _manager = ModelManager() return _manager def configure_model_manager( config_path: str = "configs/scale_1.3b.yaml", checkpoint_path: Optional[str] = None, device: str = "cpu", ) -> ModelManager: """Create and set the global ModelManager with specific options.""" global _manager with _lock: _manager = ModelManager( config_path=config_path, checkpoint_path=checkpoint_path, device=device, ) return _manager