Explicitly fail on Keras3 (#2107)

* Explicitly fail on Keras3

* code quality

* try to load from tf_keras

* allow saving if tf_keras is installed

* lint

Author: Wauplin

.github/workflows/python-tests.yml

@@ -75,7 +75,7 @@ jobs:
             tensorflow)
               sudo apt update
               sudo apt install -y graphviz
-              uv pip install "huggingface_hub[tensorflow] @ ."
+              uv pip install "huggingface_hub[tensorflow-testing] @ ."
               ;;
 
           esac

setup.py

@@ -45,7 +45,16 @@ def get_version() -> str:
     "fastcore>=1.3.27",
 ]
 
-extras["tensorflow"] = ["tensorflow", "pydot", "graphviz"]
+extras["tensorflow"] = [
+    "tensorflow",
+    "pydot",
+    "graphviz",
+]
+
+extras["tensorflow-testing"] = [
+    "tensorflow",
+    "keras<3.0",
+]
 
 
 extras["testing"] = (

src/huggingface_hub/keras_mixin.py

@@ -2,6 +2,7 @@
 import json
 import os
 import warnings
+from functools import wraps
 from pathlib import Path
 from shutil import copytree
 from typing import Any, Dict, List, Optional, Union
@@ -18,12 +19,37 @@
 from .constants import CONFIG_NAME
 from .hf_api import HfApi
 from .utils import SoftTemporaryDirectory, logging, validate_hf_hub_args
+from .utils._typing import CallableT
 
 
 logger = logging.get_logger(__name__)
 
+keras = None
 if is_tf_available():
-    import tensorflow as tf  # type: ignore
+    # Depending on which version of TensorFlow is installed, we need to import
+    # keras from the correct location.
+    # See https://github.com/tensorflow/tensorflow/releases/tag/v2.16.1.
+    # Note: saving a keras model only works with Keras<3.0.
+    try:
+        import tf_keras as keras  # type: ignore
+    except ImportError:
+        import tensorflow as tf  # type: ignore
+
+        keras = tf.keras
+
+
+def _requires_keras_2_model(fn: CallableT) -> CallableT:
+    # Wrapper to raise if user tries to save a Keras 3.x model
+    @wraps(fn)
+    def _inner(model, *args, **kwargs):
+        if not hasattr(model, "history"):  # hacky way to check if model is Keras 2.x
+            raise NotImplementedError(
+                f"Cannot use '{fn.__name__}': Keras 3.x is not supported."
+                " Please save models manually and upload them using `upload_folder` or `huggingface-cli upload`."
+            )
+        return fn(model, *args, **kwargs)
+
+    return _inner  # type: ignore [return-value]
 
 
 def _flatten_dict(dictionary, parent_key=""):
@@ -62,15 +88,15 @@ def _create_hyperparameter_table(model):
         optimizer_params = model.optimizer.get_config()
         # flatten the configuration
         optimizer_params = _flatten_dict(optimizer_params)
-        optimizer_params["training_precision"] = tf.keras.mixed_precision.global_policy().name
+        optimizer_params["training_precision"] = keras.mixed_precision.global_policy().name
         table = "| Hyperparameters | Value |\n| :-- | :-- |\n"
         for key, value in optimizer_params.items():
             table += f"| {key} | {value} |\n"
     return table
 
 
 def _plot_network(model, save_directory):
-    tf.keras.utils.plot_model(
+    keras.utils.plot_model(
         model,
         to_file=f"{save_directory}/model.png",
         show_shapes=False,
@@ -127,6 +153,7 @@ def _create_model_card(
     readme_path.write_text(model_card)
 
 
+@_requires_keras_2_model
 def save_pretrained_keras(
     model,
     save_directory: Union[str, Path],
@@ -161,9 +188,7 @@ def save_pretrained_keras(
             model_save_kwargs will be passed to
             [`tf.keras.models.save_model()`](https://www.tensorflow.org/api_docs/python/tf/keras/models/save_model).
     """
-    if is_tf_available():
-        import tensorflow as tf
-    else:
+    if keras is None:
         raise ImportError("Called a Tensorflow-specific function but could not import it.")
 
     if not model.built:
@@ -209,7 +234,7 @@ def save_pretrained_keras(
                 json.dump(model.history.history, f, indent=2, sort_keys=True)
 
     _create_model_card(model, save_directory, plot_model, metadata)
-    tf.keras.models.save_model(model, save_directory, include_optimizer=include_optimizer, **model_save_kwargs)
+    keras.models.save_model(model, save_directory, include_optimizer=include_optimizer, **model_save_kwargs)
 
 
 def from_pretrained_keras(*args, **kwargs) -> "KerasModelHubMixin":
@@ -272,6 +297,7 @@ def from_pretrained_keras(*args, **kwargs) -> "KerasModelHubMixin":
 
 
 @validate_hf_hub_args
+@_requires_keras_2_model
 def push_to_hub_keras(
     model,
     repo_id: str,
@@ -452,9 +478,7 @@ def _from_pretrained(
                 TODO - Some args above aren't used since we are calling
                 snapshot_download instead of hf_hub_download.
         """
-        if is_tf_available():
-            import tensorflow as tf
-        else:
+        if keras is None:
             raise ImportError("Called a TensorFlow-specific function but could not import it.")
 
         # Root is either a local filepath matching model_id or a cached snapshot
@@ -470,7 +494,7 @@ def _from_pretrained(
             storage_folder = model_id
 
         # TODO: change this in a future PR. We are not returning a KerasModelHubMixin instance here...
-        model = tf.keras.models.load_model(storage_folder)
+        model = keras.models.load_model(storage_folder)
 
         # For now, we add a new attribute, config, to store the config loaded from the hub/a local dir.
         model.config = config

src/huggingface_hub/utils/_runtime.py

@@ -35,6 +35,7 @@
     "graphviz": {"graphviz"},
     "hf_transfer": {"hf_transfer"},
     "jinja": {"Jinja2"},
+    "keras": {"keras"},
     "minijinja": {"minijinja"},
     "numpy": {"numpy"},
     "pillow": {"Pillow"},
@@ -140,6 +141,15 @@ def get_hf_transfer_version() -> str:
     return _get_version("hf_transfer")
 
 
+# keras
+def is_keras_available() -> bool:
+    return is_package_available("keras")
+
+
+def get_keras_version() -> str:
+    return _get_version("keras")
+
+
 # Minijinja
 def is_minijinja_available() -> bool:
     return is_package_available("minijinja")
@@ -332,6 +342,7 @@ def dump_environment_info() -> Dict[str, Any]:
     info["Torch"] = get_torch_version()
     info["Jinja2"] = get_jinja_version()
     info["Graphviz"] = get_graphviz_version()
+    info["keras"] = get_keras_version()
     info["Pydot"] = get_pydot_version()
     info["Pillow"] = get_pillow_version()
     info["hf_transfer"] = get_hf_transfer_version()