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
@@ -93,13 +93,14 @@ class SAMAudio(BaseModel):
         self.timestep_emb = SinusoidalEmbedding(cfg.transformer.dim)
         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):
@@ -256,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