Enhance model loading: auto-detect head_hidden_dim from checkpoint and streamline checkpoint loading process

app.py

@@ -84,30 +84,40 @@ class MultiModelProfiler:
             # Load model - use MultiTaskSpeakerModel
             from src.models import MultiTaskSpeakerModel
             
+            # Load checkpoint first to detect head_hidden_dim
+            checkpoint_path = model_path / "model.safetensors"
+            pt_path = model_path / "best_model.pt"
+            state_dict = None
+            
+            if checkpoint_path.exists():
+                state_dict = load_safetensors(str(checkpoint_path))
+            elif pt_path.exists():
+                checkpoint = torch.load(pt_path, map_location=self.device, weights_only=False)
+                if "model_state_dict" in checkpoint:
+                    state_dict = checkpoint["model_state_dict"]
+                else:
+                    state_dict = checkpoint
+            
+            # Auto-detect head_hidden_dim from checkpoint
+            head_hidden_dim = 256  # default
+            if state_dict is not None and "gender_head.0.weight" in state_dict:
+                # gender_head.0.weight has shape [head_hidden_dim, hidden_size]
+                head_hidden_dim = state_dict["gender_head.0.weight"].shape[0]
+                print(f"Detected head_hidden_dim: {head_hidden_dim}")
+            
             model = MultiTaskSpeakerModel(
                 model_name=encoder_name,
                 num_genders=2,
                 num_dialects=3,
                 dropout=0.1,
+                head_hidden_dim=head_hidden_dim,
                 freeze_encoder=True
             )
             
-            # Load checkpoint from safetensors
-            checkpoint_path = model_path / "model.safetensors"
-            if checkpoint_path.exists():
-                state_dict = load_safetensors(str(checkpoint_path))
+            # Load checkpoint weights
+            if state_dict is not None:
                 model.load_state_dict(state_dict)
-                print(f"Loaded checkpoint: {checkpoint_path}")
-            else:
-                # Try loading from .pt file
-                pt_path = model_path / "best_model.pt"
-                if pt_path.exists():
-                    checkpoint = torch.load(pt_path, map_location=self.device, weights_only=False)
-                    if "model_state_dict" in checkpoint:
-                        model.load_state_dict(checkpoint["model_state_dict"])
-                    else:
-                        model.load_state_dict(checkpoint)
-                    print(f"Loaded checkpoint: {pt_path}")
+                print(f"Loaded checkpoint: {checkpoint_path if checkpoint_path.exists() else pt_path}")
             
             model.to(self.device)
             model.eval()