Update app.py

app.py

@@ -4,8 +4,6 @@ import matplotlib
 matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 from sklearn.preprocessing import MinMaxScaler
-import tensorflow as tf
-from tensorflow.keras import layers, models, regularizers
 import os
 import spaces
 
@@ -27,9 +25,15 @@ CLASS_RISK = {
 }
 COLORS = ['#2ecc71','#3498db','#e74c3c','#f39c12','#9b59b6']
 
-# ── Build model on CPU at startup ─────────────────────────────────────────────
-def build_model():
-    with tf.device('/CPU:0'):
+# ── Load model lazily inside GPU function ─────────────────────────────────────
+model = None
+
+@spaces.GPU
+def predict_ecg(text_input):
+    global model
+    if model is None:
+        import tensorflow as tf
+        from tensorflow.keras import layers, models, regularizers
         m = models.Sequential([
             layers.Input(shape=(187, 1)),
             layers.Conv1D(64, 7, padding='same', activation='relu',
@@ -52,12 +56,8 @@ def build_model():
         ], name="CNN_ECG")
         m.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
         m.load_weights("cnn_weights.weights.h5")
-    return m
-
-model = build_model()
+        model = m
 
-# ── Predict ───────────────────────────────────────────────────────────────────
-def predict_ecg(text_input):
     try:
         cleaned = text_input.replace(",", " ").replace("\n", " ").replace("\t", " ")
         parsed  = [float(x) for x in cleaned.split() if x.strip()]
@@ -71,12 +71,9 @@ def predict_ecg(text_input):
 
         scaler        = MinMaxScaler()
         signal_scaled = scaler.fit_transform(features.reshape(-1, 1)).reshape(1, 187, 1)
-
-        with tf.device('/CPU:0'):
-            probs = model.predict(signal_scaled, verbose=0)[0]
-
-        pred_class = int(np.argmax(probs))
-        confidence = float(np.max(probs)) * 100
+        probs         = model.predict(signal_scaled, verbose=0)[0]
+        pred_class    = int(np.argmax(probs))
+        confidence    = float(np.max(probs)) * 100
 
         fig, axes = plt.subplots(1, 2, figsize=(14, 4))
         fig.patch.set_facecolor('#0e1117')