Update sam_audio/model/model.py

sam_audio/model/model.py

@@ -6,7 +6,7 @@ from dataclasses import dataclass
 from typing import Any, Dict, Optional
 
 import torch
-#from core.audio_visual_encoder import PEAudioFrame, PEAudioFrameTransform
+from core.audio_visual_encoder import PEAudioFrame, PEAudioFrameTransform
 from torchdiffeq import odeint
 
 from sam_audio.model.align import AlignModalities
@@ -94,13 +94,13 @@ class SAMAudio(BaseModel):
         self.visual_ranker = create_ranker(cfg.visual_ranker)
         self.text_ranker = create_ranker(cfg.text_ranker)
         
-        #if cfg.span_predictor is not None:
-        #    self.span_predictor = PEAudioFrame.from_config(
-        #        cfg.span_predictor, pretrained=True
-        #    )
-        #    self.span_predictor_transform = PEAudioFrameTransform.from_config(
-        #        cfg.span_predictor
-        #    )
+        if cfg.span_predictor is not None:
+            self.span_predictor = PEAudioFrame.from_config(
+                cfg.span_predictor, pretrained=True
+            )
+            self.span_predictor_transform = PEAudioFrameTransform.from_config(
+                cfg.span_predictor
+            )
 
     @property
     def sample_rate(self):
@@ -257,16 +257,16 @@ class SAMAudio(BaseModel):
         # Encode audio
         forward_args = self._get_forward_args(batch, candidates=reranking_candidates)
 
-        #if predict_spans and hasattr(self, "span_predictor") and batch.anchors is None:
-        #    batch = self.predict_spans(
-        #        batch=batch,
-        #        audio_features=self._unrepeat_from_reranking(
-        #            forward_args["audio_features"], reranking_candidates
-        #        ),
-        #        audio_pad_mask=self._unrepeat_from_reranking(
-        #           forward_args["audio_pad_mask"], reranking_candidates
-        #        ),
-        #    )
+        if predict_spans and hasattr(self, "span_predictor") and batch.anchors is None:
+            batch = self.predict_spans(
+                batch=batch,
+                audio_features=self._unrepeat_from_reranking(
+                    forward_args["audio_features"], reranking_candidates
+                ),
+                audio_pad_mask=self._unrepeat_from_reranking(
+                   forward_args["audio_pad_mask"], reranking_candidates
+                ),
+            )
 
         audio_features = forward_args["audio_features"]
         B, T, C = audio_features.shape