| 53 | REPLAY_SECONDS = 3 # What the user hears as context. |
| 54 | |
| 55 | class AudioProcessor: |
| 56 | def __init__(self, model, prompt_path): |
| 57 | self.model = model |
| 58 | self.prompt_path = prompt_path |
| 59 | self.initialize_state(prompt_path) |
| 60 | |
| 61 | def initialize_state(self, prompt_path): |
| 62 | loaded_audio, sr = torchaudio.load(prompt_path) |
| 63 | self.replay_seconds = REPLAY_SECONDS |
| 64 | |
| 65 | if sr != SAMPLE_RATE: |
| 66 | resampler = torchaudio.transforms.Resample(sr, SAMPLE_RATE) |
| 67 | loaded_audio = resampler(loaded_audio) |
| 68 | |
| 69 | if loaded_audio.shape[0] == 1: |
| 70 | loaded_audio = loaded_audio.repeat(2, 1) |
| 71 | |
| 72 | audio_length = loaded_audio.shape[-1] |
| 73 | num_chunks = audio_length // 2000 |
| 74 | loaded_audio = loaded_audio[..., :num_chunks * 2000] |
| 75 | |
| 76 | self.loaded_audio = loaded_audio.to(device) |
| 77 | |
| 78 | with T.autocast(device_type=device, dtype=T.bfloat16), T.inference_mode(): |
| 79 | self.model.init_cache(bsize=1, device=device, dtype=T.bfloat16, length=1024) |
| 80 | self.next_model_audio = self.model.next_audio_from_audio(self.loaded_audio.unsqueeze(0), temps=TEMPS) |
| 81 | self.prompt_buffer = None |
| 82 | self.prompt_position = 0 |
| 83 | self.chunks_until_live = int(self.replay_seconds * 8) |
| 84 | self.initialize_prompt_buffer() |
| 85 | print_colored("AudioProcessor state initialized", "green") |
| 86 | |
| 87 | def initialize_prompt_buffer(self): |
| 88 | self.recorded_audio = self.loaded_audio |
| 89 | prompt_audio = self.loaded_audio.reshape(1, 2, -1) |
| 90 | prompt_audio = prompt_audio[:, :, -(16000*self.replay_seconds):].cpu().numpy() |
| 91 | prompt_audio_mono = prompt_audio.mean(axis=1) |
| 92 | self.prompt_buffer = np.array_split(prompt_audio_mono[0], int(self.replay_seconds * 8)) |
| 93 | print_colored(f"Initialized prompt buffer with {len(self.prompt_buffer)} chunks", "grey") |
| 94 | |
| 95 | async def process_audio(self, audio_data): |
| 96 | if self.chunks_until_live > 0: |
| 97 | print_colored(f"Serving from prompt buffer, {self.chunks_until_live} chunks left", "grey") |
| 98 | chunk = self.prompt_buffer[int(self.replay_seconds * 8) - self.chunks_until_live] |
| 99 | self.chunks_until_live -= 1 |
| 100 | |
| 101 | if self.chunks_until_live == 0: |
| 102 | print_colored("Switching to live processing mode", "green") |
| 103 | |
| 104 | time.sleep(0.05) |
| 105 | return chunk |
| 106 | |
| 107 | audio_tensor = T.from_numpy(audio_data).to(device) |
| 108 | audio_tensor = audio_tensor.reshape(1, 1, -1) |
| 109 | audio_tensor = T.cat([audio_tensor, self.next_model_audio], dim=1) |
| 110 | |
| 111 | with T.autocast(device_type=device, dtype=T.bfloat16), T.inference_mode(): |
| 112 | curr_model_audio = self.model.next_audio_from_audio( |