Fix model loading compatibility issues

- Switch from JAX to TensorFlow backend for better model compatibility
- Use traditional HuggingFace Hub download approach instead of keras.saving
- Add huggingface_hub dependency for proper model downloading
- Fix 'weights_store' error by using compatible loading method
- Update documentation to reflect TensorFlow backend usage

Resolves: ValueError: Expected a model.weights.h5 or model.weights.npz file
Resolves: UnboundLocalError: local variable 'weights_store' referenced before assignment

README.md

@@ -17,23 +17,24 @@ AI-powered skin condition analysis using deep learning with EfficientNetV2 archi
 
 ## Features
 
-- **Modern Keras 3.0 with JAX Backend**: Uses the latest Keras with JAX for improved performance
-- **Direct HuggingFace Model Loading**: Loads models directly from HuggingFace Hub
+- **Modern Keras 3.0 with TensorFlow Backend**: Uses Keras 3.0 with TensorFlow for reliable model compatibility
+- **HuggingFace Hub Integration**: Downloads and loads models from HuggingFace Hub
 - **5 Skin Conditions**: Classifies Atopic Dermatitis, Eczema, Psoriasis, Seborrheic Keratoses, and Tinea Ringworm Candidiasis
 - **Medical Disclaimers**: Includes proper medical disclaimers and advice
 
 ## Technical Stack
 
-- **Backend**: JAX (via Keras 3.0)
+- **Backend**: TensorFlow (via Keras 3.0)
 - **Frontend**: Gradio 5.0+
 - **Model Architecture**: EfficientNetV2
-- **Model Source**: Direct loading from `hf://Siraja704/DermaAI`
+- **Model Source**: `Siraja704/DermaAI` via HuggingFace Hub
 
 ## Setup
 
-This application uses Keras 3.0 with JAX backend for optimal performance. The model is loaded directly from HuggingFace Hub using the new `keras.saving.load_model("hf://...")` syntax.
+This application uses Keras 3.0 with TensorFlow backend for reliable model compatibility. The model is downloaded from HuggingFace Hub using the `huggingface_hub` library.
 
 If the model repository is private, you may need to:
+
 1. Go to the Space settings
 2. Add a new secret with key `HF_TOKEN` and your Hugging Face access token as the value
 3. Make sure your token has access to the `Siraja704/DermaAI` model repository

app.py

@@ -1,13 +1,14 @@
-# Available backend options are: "jax", "torch", "tensorflow".
+# Use TensorFlow backend for better compatibility with existing models
 import os
-os.environ["KERAS_BACKEND"] = "jax"
+os.environ["KERAS_BACKEND"] = "tensorflow"
 
 import gradio as gr
-import keras
+import tensorflow as tf
+from tensorflow import keras
 import numpy as np
 from PIL import Image
-import tensorflow as tf  # Keep for preprocessing function
 from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
+from huggingface_hub import hf_hub_download
 
 # Class names for the 5 skin conditions
 CLASS_NAMES = [
@@ -33,20 +34,35 @@ class DermaAIModel:
         self.load_model()
     
     def load_model(self):
-        """Load the DermaAI model from Hugging Face Hub using Keras 3.0"""
+        """Load the DermaAI model from Hugging Face Hub using traditional approach"""
         try:
             print("🔄 Loading DermaAI model from Hugging Face...")
+            
             # Get authentication token from environment variable only (secure)
             hf_token = os.getenv("HF_TOKEN")
             
+            # Download model file using HuggingFace Hub
             if hf_token:
-                # Load model with authentication
-                self.model = keras.saving.load_model("hf://Siraja704/DermaAI", token=hf_token)
-                print("✅ Model loaded successfully with authentication!")
+                print("🔐 Using authentication token...")
+                model_path = hf_hub_download(
+                    repo_id="Siraja704/DermaAI",
+                    filename="DermaAI.keras",
+                    token=hf_token,
+                    cache_dir="./model_cache"
+                )
             else:
-                # Try loading without token (if model becomes public)
-                self.model = keras.saving.load_model("hf://Siraja704/DermaAI")
-                print("✅ Model loaded successfully!")
+                print("📂 Trying without authentication...")
+                model_path = hf_hub_download(
+                    repo_id="Siraja704/DermaAI",
+                    filename="DermaAI.keras",
+                    cache_dir="./model_cache"
+                )
+            
+            # Load the model using TensorFlow/Keras
+            print(f"📁 Loading model from: {model_path}")
+            self.model = keras.models.load_model(model_path)
+            print("✅ Model loaded successfully!")
+            
         except Exception as e:
             error_msg = str(e)
             print(f"❌ Error loading model: {e}")

requirements.txt

@@ -1,6 +1,6 @@
 gradio>=5.0.0
-keras>=3.0.0
-jax[cpu]>=0.4.0
 tensorflow>=2.13.0
+keras>=3.0.0
 pillow>=9.0.0
-numpy>=1.21.0
+numpy>=1.21.0
+huggingface_hub>=0.16.0