from huggingface_hub import hf_hub_download
import tensorflow as tf
from tensorflow import keras

# Replace 'your-username/your-model-name' with your actual Hugging Face model repository ID.
repo_id = "your-username/your-model-name"
# Replace 'my_keras_model.keras' with the name of the file you uploaded.
filename = "my_keras_model.keras"

# Download the model file from the Hugging Face Hub.
model_path = hf_hub_download(repo_id=repo_id, filename=filename)

# Load the model using Keras's built-in function.
# The 'safe_mode=False' argument is often necessary when loading models saved from older TensorFlow versions
# or if the model contains custom layers.
model = keras.models.load_model(model_path, safe_mode=False)

# Now you can use the loaded model for inference.
# Example: Load a single MNIST test image and make a prediction.
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_test = x_test.astype("float32") / 255.0
x_test = tf.expand_dims(x_test, -1)
image_to_predict = x_test[0:1]

# Get the model's prediction.
predictions = model.predict(image_to_predict)

# Print the predicted class (the one with the highest probability).
predicted_class = tf.argmax(predictions[0]).numpy()
print(f"Predicted class: {predicted_class}")

# Display the model summary to confirm it's loaded correctly.
model.summary()