MarkDaniel212 commited on
Commit
e13a882
·
verified ·
1 Parent(s): 257b104

Perf: load audio tokenizer directly in bf16 (skip fp32 intermediate)

Browse files
Files changed (1) hide show
  1. src/mimo_audio/mimo_audio.py +5 -2
src/mimo_audio/mimo_audio.py CHANGED
@@ -115,9 +115,12 @@ class MimoAudio:
115
  }
116
 
117
  start_loading_mimo_audio_tokenizer_time = time.monotonic()
118
- self.mimo_audio_tokenizer = MiMoAudioTokenizer.from_pretrained(self.mimo_audio_tokenizer_path)
 
 
 
119
 
120
- self.mimo_audio_tokenizer.eval().bfloat16().to(self.device)
121
  print(
122
  f"MiMo-Audio Tokenizer loaded in {time.monotonic() - start_loading_mimo_audio_tokenizer_time:.2f} seconds, device: {self.device}"
123
  )
 
115
  }
116
 
117
  start_loading_mimo_audio_tokenizer_time = time.monotonic()
118
+ self.mimo_audio_tokenizer = MiMoAudioTokenizer.from_pretrained(
119
+ self.mimo_audio_tokenizer_path,
120
+ torch_dtype=torch.bfloat16,
121
+ )
122
 
123
+ self.mimo_audio_tokenizer.eval().to(self.device)
124
  print(
125
  f"MiMo-Audio Tokenizer loaded in {time.monotonic() - start_loading_mimo_audio_tokenizer_time:.2f} seconds, device: {self.device}"
126
  )