Upload folder using huggingface_hub

app.py

@@ -34,10 +34,11 @@ class WatermelonMoEModel(torch.nn.Module):
             weights: Optional list of weights for each model (None for equal weighting)
         """
         super(WatermelonMoEModel, self).__init__()
-        self.models = []
+        self.models = torch.nn.ModuleList()  # Use ModuleList instead of regular list
         self.model_configs = model_configs
         
         # Load each model
+        loaded_count = 0
         for config in model_configs:
             img_backbone = config["image_backbone"]
             audio_backbone = config["audio_backbone"]
@@ -49,22 +50,31 @@ class WatermelonMoEModel(torch.nn.Module):
             model_path = os.path.join(model_dir, f"{img_backbone}_{audio_backbone}_model.pt")
             if os.path.exists(model_path):
                 print(f"\033[92mINFO\033[0m: Loading model {img_backbone}_{audio_backbone} from {model_path}")
-                model.load_state_dict(torch.load(model_path, map_location='cpu'))
+                try:
+                    model.load_state_dict(torch.load(model_path, map_location='cpu'))
+                    model.eval()  # Set to evaluation mode
+                    self.models.append(model)
+                    loaded_count += 1
+                except Exception as e:
+                    print(f"\033[91mERR!\033[0m: Failed to load model from {model_path}: {e}")
+                    continue
             else:
                 print(f"\033[91mERR!\033[0m: Model checkpoint not found at {model_path}")
                 continue
-                
-            model.eval()  # Set to evaluation mode
-            self.models.append(model)
         
+        # Add a dummy parameter if no models were loaded to prevent StopIteration
+        if loaded_count == 0:
+            print(f"\033[91mERR!\033[0m: No models were successfully loaded!")
+            self.dummy_param = torch.nn.Parameter(torch.zeros(1))
+            
         # Set model weights (uniform by default)
-        if weights:
+        if weights and loaded_count > 0:
             assert len(weights) == len(self.models), "Number of weights must match number of models"
             self.weights = weights
         else:
-            self.weights = [1.0 / len(self.models)] * len(self.models)
+            self.weights = [1.0 / max(loaded_count, 1)] * max(loaded_count, 1)
             
-        print(f"\033[92mINFO\033[0m: Loaded {len(self.models)} models for MoE ensemble")
+        print(f"\033[92mINFO\033[0m: Loaded {loaded_count} models for MoE ensemble")
         print(f"\033[92mINFO\033[0m: Model weights: {self.weights}")
 
     def to(self, device):
@@ -80,6 +90,11 @@ class WatermelonMoEModel(torch.nn.Module):
         Forward pass through the MoE model.
         Returns the weighted average of all model outputs.
         """
+        # Check if we have models loaded
+        if not self.models:
+            print(f"\033[91mERR!\033[0m: No models available for inference!")
+            return torch.tensor([0.0], device=mfcc.device)  # Return a default value
+            
         outputs = []
         
         # Get outputs from each model
@@ -258,9 +273,19 @@ def predict_sugar_content(audio, image, model_dir="models", weights=None):
             print(f"\033[92mDEBUG\033[0m: Final image shape with batch dimension: {processed_image.shape}, device: {processed_image.device}")
         
         # Double-check model is on the correct device
-        print(f"\033[92mDEBUG\033[0m: MoE model device: {next(moe_model.parameters()).device}")
-        for i, model in enumerate(moe_model.models):
-            print(f"\033[92mDEBUG\033[0m: Model {i} device: {next(model.parameters()).device}")
+        try:
+            param = next(moe_model.parameters())
+            print(f"\033[92mDEBUG\033[0m: MoE model device: {param.device}")
+            
+            # Check individual models
+            for i, model in enumerate(moe_model.models):
+                try:
+                    model_param = next(model.parameters())
+                    print(f"\033[92mDEBUG\033[0m: Model {i} device: {model_param.device}")
+                except StopIteration:
+                    print(f"\033[91mERR!\033[0m: Model {i} has no parameters!")
+        except StopIteration:
+            print(f"\033[91mERR!\033[0m: MoE model has no parameters!")
         
         # Run inference with MoE model
         print(f"\033[92mDEBUG\033[0m: Running inference with MoE model on device: {device}")

app_moe.py

@@ -34,10 +34,11 @@ class WatermelonMoEModel(torch.nn.Module):
             weights: Optional list of weights for each model (None for equal weighting)
         """
         super(WatermelonMoEModel, self).__init__()
-        self.models = []
+        self.models = torch.nn.ModuleList()  # Use ModuleList instead of regular list
         self.model_configs = model_configs
         
         # Load each model
+        loaded_count = 0
         for config in model_configs:
             img_backbone = config["image_backbone"]
             audio_backbone = config["audio_backbone"]
@@ -49,22 +50,31 @@ class WatermelonMoEModel(torch.nn.Module):
             model_path = os.path.join(model_dir, f"{img_backbone}_{audio_backbone}_model.pt")
             if os.path.exists(model_path):
                 print(f"\033[92mINFO\033[0m: Loading model {img_backbone}_{audio_backbone} from {model_path}")
-                model.load_state_dict(torch.load(model_path, map_location='cpu'))
+                try:
+                    model.load_state_dict(torch.load(model_path, map_location='cpu'))
+                    model.eval()  # Set to evaluation mode
+                    self.models.append(model)
+                    loaded_count += 1
+                except Exception as e:
+                    print(f"\033[91mERR!\033[0m: Failed to load model from {model_path}: {e}")
+                    continue
             else:
                 print(f"\033[91mERR!\033[0m: Model checkpoint not found at {model_path}")
                 continue
-                
-            model.eval()  # Set to evaluation mode
-            self.models.append(model)
         
+        # Add a dummy parameter if no models were loaded to prevent StopIteration
+        if loaded_count == 0:
+            print(f"\033[91mERR!\033[0m: No models were successfully loaded!")
+            self.dummy_param = torch.nn.Parameter(torch.zeros(1))
+            
         # Set model weights (uniform by default)
-        if weights:
+        if weights and loaded_count > 0:
             assert len(weights) == len(self.models), "Number of weights must match number of models"
             self.weights = weights
         else:
-            self.weights = [1.0 / len(self.models)] * len(self.models)
+            self.weights = [1.0 / max(loaded_count, 1)] * max(loaded_count, 1)
             
-        print(f"\033[92mINFO\033[0m: Loaded {len(self.models)} models for MoE ensemble")
+        print(f"\033[92mINFO\033[0m: Loaded {loaded_count} models for MoE ensemble")
         print(f"\033[92mINFO\033[0m: Model weights: {self.weights}")
 
     def to(self, device):
@@ -80,6 +90,11 @@ class WatermelonMoEModel(torch.nn.Module):
         Forward pass through the MoE model.
         Returns the weighted average of all model outputs.
         """
+        # Check if we have models loaded
+        if not self.models:
+            print(f"\033[91mERR!\033[0m: No models available for inference!")
+            return torch.tensor([0.0], device=mfcc.device)  # Return a default value
+            
         outputs = []
         
         # Get outputs from each model