Spaces:
Runtime error
Runtime error
Factor Studios
commited on
Upload 3 files
Browse files- ai_http.py +6 -0
- 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 |
|