Update src/streamlit_app.py

src/streamlit_app.py

@@ -57,23 +57,49 @@ GEMMA_PATH = "/app/gemma_tomato_lora"
 # ======================
 @st.cache_resource
 def load_cnn():
-    from keras.src.engine.input_layer import InputLayer
+    import tensorflow as tf
+    from tensorflow.keras.layers import InputLayer, Rescaling
+    from tensorflow.keras.utils import custom_object_scope
 
-    original_from_config = InputLayer.from_config
+    # -------------------------
+    # PATCH InputLayer
+    # -------------------------
+    original_input = InputLayer.from_config
 
     @classmethod
-    def patched_from_config(cls, config):
+    def patched_input(cls, config):
         config.pop("batch_shape", None)
         config.pop("optional", None)
-        return original_from_config(config)
+        return original_input(config)
 
-    InputLayer.from_config = patched_from_config
+    InputLayer.from_config = patched_input
 
-    return tf.keras.models.load_model(
-        MODEL_PATH,
-        compile=False
-    )
+    # -------------------------
+    # PATCH Rescaling
+    # -------------------------
+    original_rescaling = Rescaling.from_config
 
+    @classmethod
+    def patched_rescaling(cls, config):
+        if isinstance(config.get("dtype"), dict):
+            config["dtype"] = "float32"
+        return original_rescaling(config)
+
+    Rescaling.from_config = patched_rescaling
+
+    # -------------------------
+    # LOAD MODEL
+    # -------------------------
+    with custom_object_scope({
+        "Rescaling": Rescaling,
+        "InputLayer": InputLayer
+    }):
+        model = tf.keras.models.load_model(
+            MODEL_PATH,
+            compile=False
+        )
+
+    return model
 cnn_model = load_cnn()
 
 # ======================