Spaces:
Paused
Paused
File size: 2,509 Bytes
e6ed91e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | import os
import tensorflow as tf
import threading
class ModelLoader:
def __init__(self):
# Model is stored directly in the Space
self.local_model_path = "models/saved_model"
self.model = None
self._model_ready = False
self._lock = threading.Lock() # Add thread safety for web/mobile concurrent requests
print("ModelLoader initialized - ready for lazy loading")
def get_model(self):
"""Get model, loading it if necessary (lazy loading)"""
with self._lock: # Prevent race conditions between web/mobile requests
if self.model is None or not self._model_ready:
print("Model not loaded yet - loading now...")
self.load_model()
return self.model
def load_model(self):
"""Load model from local files"""
try:
print(f"Loading SavedModel from: {self.local_model_path}")
# Check if model exists
if not os.path.exists(self.local_model_path):
raise FileNotFoundError(f"Model not found at: {self.local_model_path}")
# Clear any existing model first (important for resets)
self.model = None
self._model_ready = False
# Load the model
loaded_model = tf.saved_model.load(self.local_model_path)
# Add predict wrapper
def predict_wrapper(input_data):
return loaded_model(input_data)
loaded_model.predict = predict_wrapper
self.model = loaded_model
self._model_ready = True
print("✅ Model loaded successfully from local files")
except Exception as e:
print(f"❌ Error loading model: {e}")
self.model = None
self._model_ready = False
raise
def reset(self):
"""Reset the model loader state - fixes corruption issues"""
with self._lock:
print("🔄 Resetting ModelLoader state...")
self.model = None
self._model_ready = False
# Force garbage collection to free memory
try:
import gc
gc.collect()
except:
pass
print("✅ ModelLoader state reset complete")
def is_model_ready(self):
"""Check if model is loaded and ready"""
with self._lock:
return self.model is not None and self._model_ready |