Factor Studios commited on
Commit
a31170e
·
verified ·
1 Parent(s): e41bfbc

Upload 3 files

Browse files
Files changed (2) hide show
  1. ai_http.py +6 -0
  2. http_storage.py +12 -0
ai_http.py CHANGED
@@ -71,6 +71,12 @@ class AIAccelerator:
71
  self._vram_allocated += size_bytes
72
  return True
73
 
 
 
 
 
 
 
74
  # Initialize model, tensor, and tokenizer tracking
75
  self.model_registry: Dict[str, Dict[str, Any]] = {} # Track loaded models
76
  self.tensor_registry: Dict[str, Dict[str, Any]] = {} # Track tensor metadata
 
71
  self._vram_allocated += size_bytes
72
  return True
73
 
74
+ def has_model(self, model_id: str) -> bool:
75
+ """Check if a model is loaded"""
76
+ if not self.storage:
77
+ return False
78
+ return self.storage.is_model_loaded(model_id)
79
+
80
  # Initialize model, tensor, and tokenizer tracking
81
  self.model_registry: Dict[str, Dict[str, Any]] = {} # Track loaded models
82
  self.tensor_registry: Dict[str, Dict[str, Any]] = {} # Track tensor metadata
http_storage.py CHANGED
@@ -48,6 +48,18 @@ class LocalStorage:
48
 
49
  self.lock = threading.Lock()
50
  self._closing = False
 
 
 
 
 
 
 
 
 
 
 
 
51
  self.model_registry = {}
52
  self._connected = True
53
 
 
48
 
49
  self.lock = threading.Lock()
50
  self._closing = False
51
+
52
+ # Initialize resource monitoring
53
+ self.resource_monitor = {
54
+ 'vram_used': 0,
55
+ 'active_tensors': 0,
56
+ 'loaded_models': set(),
57
+ 'last_updated': time.time()
58
+ }
59
+
60
+ # Initialize model registry and connection state
61
+ self.model_registry = {}
62
+ self._connected = True
63
  self.model_registry = {}
64
  self._connected = True
65