change probability func

inference.py

@@ -169,16 +169,16 @@ def scaled_sigmoid(x, scale_factor=0.2, linear_property=0.3):
 
 # Apply the scaled sigmoid
 
-
 def get_model(model_type, device):
     """Load the specified model."""
     if model_type == "MERT":
-        #from model import MusicAudioClassifier
-        #model = MusicAudioClassifier(input_dim=768, is_emb=True, mode = 'both', share_parameter = False).to(device)
-        ckpt_file = 'checkpoints/step=003432-val_loss=0.0216-val_acc=0.9963.ckpt'#'mert_finetune_10.pth'
-        model = MERT_AudioCNN.load_from_checkpoint(ckpt_file).to(device)
+        ckpt_file = 'checkpoints/step=003432-val_loss=0.0216-val_acc=0.9963.ckpt'
+        # map_location 추가
+        model = MERT_AudioCNN.load_from_checkpoint(
+            ckpt_file, 
+            map_location=device  # 또는 'cuda:0' 또는 'cpu'
+        ).to(device)
         model.eval()
-        # model.load_state_dict(torch.load(ckpt_file, map_location=device))
         embed_dim = 768
 
     elif model_type == "pure_MERT":
@@ -189,42 +189,33 @@ def get_model(model_type, device):
     else:
         raise ValueError(f"Unknown model type: {model_type}")
    
-    
     model.eval()
     return model, embed_dim
-    
 
 def inference(audio_path):
-    backbone_model, input_dim = get_model('MERT', 'cuda')
-    segments, padding_mask = load_audio(audio_path, sr=24000)
-    segments = segments.to('cuda').to(torch.float32)
-    padding_mask = padding_mask.to('cuda').unsqueeze(0)    
-    logits,embedding = backbone_model(segments.squeeze(1))
-    # test_dataset = FakeMusicCapsDataset([audio_path], [0], target_duration=10.0)
-    # test_data, test_target = test_dataset[0]
-    # test_data = test_data.to('cuda').to(torch.float32)
-    # test_target = test_target.to('cuda')
-    # output, _ = backbone_model(test_data.unsqueeze(0))
+    # device 설정을 명확히 하기
+    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
     
+    backbone_model, input_dim = get_model('MERT', device)
+    segments, padding_mask = load_audio(audio_path, sr=24000)
+    segments = segments.to(device).to(torch.float32)
+    padding_mask = padding_mask.to(device).unsqueeze(0)    
+    logits, embedding = backbone_model(segments.squeeze(1))
 
-
-    # 모델 로드 부분 추가
+    # 모델 로드할 때도 map_location 추가
     model = MusicAudioClassifier.load_from_checkpoint(
-        checkpoint_path = 'checkpoints/EmbeddingModel_MERT_768_2class_weighted-epoch=0014-val_loss=0.0099-val_acc=0.9993-val_f1=0.9978-val_precision=0.9967-val_recall=0.9989.ckpt',
-        input_dim=input_dim, 
-        #emb_model=backbone_model
+        checkpoint_path='checkpoints/EmbeddingModel_MERT_768_2class_weighted-epoch=0014-val_loss=0.0099-val_acc=0.9993-val_f1=0.9978-val_precision=0.9967-val_recall=0.9989.ckpt',
+        input_dim=input_dim,
+        map_location=device  # 이 부분 추가
     )
     
-    
     # Run inference
     print(f"Segments shape: {segments.shape}")
     print("Running inference...")
-    results = run_inference(model, embedding, padding_mask, 'cuda')
+    results = run_inference(model, embedding, padding_mask, device)
     
     # 결과 출력
     print(f"Results: {results}")
-
-
     
     return results