Upload folder using huggingface_hub

app.py

@@ -67,6 +67,14 @@ class WatermelonMoEModel(torch.nn.Module):
         print(f"\033[92mINFO\033[0m: Loaded {len(self.models)} models for MoE ensemble")
         print(f"\033[92mINFO\033[0m: Model weights: {self.weights}")
 
+    def to(self, device):
+        """
+        Override to() method to ensure all sub-models are moved to the same device
+        """
+        for model in self.models:
+            model.to(device)
+        return super(WatermelonMoEModel, self).to(device)
+
     def forward(self, mfcc, image):
         """
         Forward pass through the MoE model.
@@ -152,7 +160,8 @@ def predict_sugar_content(audio, image, model_dir="models", weights=None):
         
         # Load MoE model
         moe_model = WatermelonMoEModel(TOP_MODELS, model_dir, weights)
-        moe_model.to(device)
+        # Explicitly move the entire model to device
+        moe_model = moe_model.to(device)
         moe_model.eval()
         print(f"\033[92mINFO\033[0m: Loaded MoE model with {len(moe_model.models)} backbone models")
         
@@ -239,12 +248,19 @@ def predict_sugar_content(audio, image, model_dir="models", weights=None):
         
         # Add batch dimension for inference and move to device
         if mfcc is not None:
+            # Ensure mfcc is on the same device as the model
             mfcc = mfcc.unsqueeze(0).to(device)
-            print(f"\033[92mDEBUG\033[0m: Final MFCC shape with batch dimension: {mfcc.shape}")
+            print(f"\033[92mDEBUG\033[0m: Final MFCC shape with batch dimension: {mfcc.shape}, device: {mfcc.device}")
         
         if processed_image is not None:
+            # Ensure processed_image is on the same device as the model
             processed_image = processed_image.unsqueeze(0).to(device)
-            print(f"\033[92mDEBUG\033[0m: Final image shape with batch dimension: {processed_image.shape}")
+            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}")
         
         # Run inference with MoE model
         print(f"\033[92mDEBUG\033[0m: Running inference with MoE model on device: {device}")

app_moe.py

@@ -67,6 +67,14 @@ class WatermelonMoEModel(torch.nn.Module):
         print(f"\033[92mINFO\033[0m: Loaded {len(self.models)} models for MoE ensemble")
         print(f"\033[92mINFO\033[0m: Model weights: {self.weights}")
 
+    def to(self, device):
+        """
+        Override to() method to ensure all sub-models are moved to the same device
+        """
+        for model in self.models:
+            model.to(device)
+        return super(WatermelonMoEModel, self).to(device)
+
     def forward(self, mfcc, image):
         """
         Forward pass through the MoE model.
@@ -152,7 +160,8 @@ def predict_sugar_content(audio, image, model_dir="models", weights=None):
         
         # Load MoE model
         moe_model = WatermelonMoEModel(TOP_MODELS, model_dir, weights)
-        moe_model.to(device)
+        # Explicitly move the entire model to device
+        moe_model = moe_model.to(device)
         moe_model.eval()
         print(f"\033[92mINFO\033[0m: Loaded MoE model with {len(moe_model.models)} backbone models")
         
@@ -239,12 +248,19 @@ def predict_sugar_content(audio, image, model_dir="models", weights=None):
         
         # Add batch dimension for inference and move to device
         if mfcc is not None:
+            # Ensure mfcc is on the same device as the model
             mfcc = mfcc.unsqueeze(0).to(device)
-            print(f"\033[92mDEBUG\033[0m: Final MFCC shape with batch dimension: {mfcc.shape}")
+            print(f"\033[92mDEBUG\033[0m: Final MFCC shape with batch dimension: {mfcc.shape}, device: {mfcc.device}")
         
         if processed_image is not None:
+            # Ensure processed_image is on the same device as the model
             processed_image = processed_image.unsqueeze(0).to(device)
-            print(f"\033[92mDEBUG\033[0m: Final image shape with batch dimension: {processed_image.shape}")
+            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}")
         
         # Run inference with MoE model
         print(f"\033[92mDEBUG\033[0m: Running inference with MoE model on device: {device}")