Update app.py

app.py

@@ -44,9 +44,13 @@ class MultiAttributeClassifier:
                     except Exception as e1:
                         print(f"   ⚠️ Normal loading failed: {e1}")
                         try:
-                            # Try with compile=False
-                            self.models[category] = tf.keras.models.load_model(model_path, compile=False)
-                            print(f"   ✅ Loaded {category} with compile=False")
+                            # Try with compile=False and custom objects
+                            self.models[category] = tf.keras.models.load_model(
+                                model_path, 
+                                compile=False,
+                                custom_objects=CUSTOM_OBJECTS
+                            )
+                            print(f"   ✅ Loaded {category} with custom_objects")
                         except Exception as e2:
                             print(f"   ❌ Failed to load {category}: {e2}")
                             continue
@@ -157,15 +161,17 @@ class MultiAttributeClassifier:
                             except Exception as e2:
                                 print(f"   ⚠️ compile=False failed: {e2}")
                                 try:
-                                    # Third try: Load with custom objects (for any custom layers)
+                                    # Third try: Load with custom objects (CRITICAL for your models)
                                     self.gan_models[model_name] = tf.keras.models.load_model(
                                         model_path, 
                                         compile=False,
-                                        custom_objects={'tf': tf}
+                                        custom_objects=CUSTOM_OBJECTS
                                     )
-                                    print(f"   ✅ Loaded with custom_objects")
+                                    print(f"   ✅ Loaded with custom_objects (InstanceNormalization)")
                                 except Exception as e3:
                                     print(f"   ❌ All loading methods failed: {e3}")
+                                    # Print the actual error for debugging
+                                    print(f"   🔍 Error details: {str(e3)}")
                                     raise e3
                         
                         print(f"✅ Loaded GAN: {model_name} from {model_path}")