Upload folder using huggingface_hub

asr_diarization/pipeline.py

@@ -15,13 +15,14 @@ class ASR_Diarization:
     def __init__(self, HF_TOKEN,
                  diar_model="pyannote/speaker-diarization-3.1",
                  asr_model="Capstone04/TrainedWhisper",
-                 model_path = "None"):
+                 model_path=None):  # NEW: model_path parameter
         self.HF_TOKEN = HF_TOKEN
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
         # Load diarization model
         self.diar_pipeline = Pipeline.from_pretrained(diar_model, use_auth_token=HF_TOKEN)
 
+        # MODIFIED: Use custom model_path if provided, otherwise use asr_model
         if model_path and os.path.exists(model_path):
             print(f"🔄 Loading custom ASR model from: {model_path}")
             actual_asr_model = model_path
@@ -29,8 +30,8 @@ class ASR_Diarization:
             print(f"🔄 Loading default ASR model: {asr_model}")
             actual_asr_model = asr_model
 
-        processor = WhisperProcessor.from_pretrained(asr_model, token=HF_TOKEN)
-        model = WhisperForConditionalGeneration.from_pretrained(asr_model, token=HF_TOKEN).to(self.device)
+        processor = WhisperProcessor.from_pretrained(actual_asr_model, token=HF_TOKEN)
+        model = WhisperForConditionalGeneration.from_pretrained(actual_asr_model, token=HF_TOKEN).to(self.device)
 
         self.asr_pipeline = hf_pipeline(
             "automatic-speech-recognition",
@@ -41,14 +42,6 @@ class ASR_Diarization:
             return_timestamps=True
         )
 
-        # model_id = "Capstone04/TrainedWhisper"
-        # self.asr_pipeline = hf_pipeline(
-        #     "automatic-speech-recognition",
-        #     model=model_id,
-        #     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        #     device=0 if torch.cuda.is_available() else -1,
-        # )
-
     def run_diarization(self, audio_path):
         diarization = self.diar_pipeline(audio_path)
         return [