upd: load model to .pb

api_server.py

@@ -15,8 +15,8 @@ LOCAL = False
 if LOCAL is True:
     model = keras.models.load_model('artifacts/models/mnist_model.h5')
 else:
-    REPO_ID = "1vash/mnist-model"
-    FILENAME = "mnist_model.h5"
+    REPO_ID = "1vash/mnist_demo_model"
+    FILENAME = "saved_model.pb"
     model = keras.models.load_model(hf_hub_download(repo_id=REPO_ID, filename=FILENAME))
 
 # Initialize the Flask application