Coverage for wrapper/vibevoice/audio_streamer.py: 24%

1# mypy: ignore-errors 

2# Copy from: https://github.com/microsoft/VibeVoice/blob/main/vibevoice/modular/streamer.py 

3 

4from __future__ import annotations 

5 

6import time 

7import torch 

8import asyncio 

9 

10from queue import Queue 

11 

12from typing import Optional 

13from typing import Any 

14 

15from transformers.generation import BaseStreamer 

16 

17 

18class AudioStreamer(BaseStreamer): 

19 """ 

20 Audio streamer that stores audio chunks in queues for each sample in the batch. 

21 This allows streaming audio generation for multiple samples simultaneously. 

22 

23 Parameters: 

24 batch_size (`int`): 

25 The batch size for generation 

26 stop_signal (`any`, *optional*): 

27 The signal to put in the queue when generation ends. Defaults to None. 

28 timeout (`float`, *optional*): 

29 The timeout for the audio queue. If `None`, the queue will block indefinitely. 

30 """ 

31 

32 def __init__( 

33 self, 

34 batch_size: int, 

35 stop_signal: Optional[Any] = None, 

36 timeout: Optional[float] = None, 

37 ): 

38 self.batch_size = batch_size 

39 self.stop_signal = stop_signal 

40 self.timeout = timeout 

41 

42 # Create a queue for each sample in the batch 

43 self.audio_queues: list[Queue] = [Queue() for _ in range(batch_size)] 

44 self.finished_flags: list[bool] = [False for _ in range(batch_size)] 

45 self.sample_indices_map: dict[int, int] = {} # Maps from sample index to queue index 

46 

47 def put(self, audio_chunks: torch.Tensor, sample_indices: torch.Tensor) -> None: # type: ignore[override] 

48 """ 

49 Receives audio chunks and puts them in the appropriate queues. 

50 

51 Args: 

52 audio_chunks: Tensor of shape (num_samples, ...) containing audio chunks 

53 sample_indices: Tensor indicating which samples these chunks belong to 

54 """ 

55 for i, sample_idx in enumerate(sample_indices): 

56 idx = sample_idx.item() 

57 if idx < self.batch_size and not self.finished_flags[idx]: 

58 # Convert to numpy or keep as tensor based on preference 

59 audio_chunk = audio_chunks[i].detach().cpu() 

60 self.audio_queues[idx].put(audio_chunk, timeout=self.timeout) 

61 

62 def end(self, sample_indices: Optional[torch.Tensor] = None) -> None: 

63 """ 

64 Signals the end of generation for specified samples or all samples. 

65 

66 Args: 

67 sample_indices: Optional tensor of sample indices to end. If None, ends all. 

68 """ 

69 if sample_indices is None: 

70 # End all samples 

71 for idx in range(self.batch_size): 

72 if not self.finished_flags[idx]: 

73 self.audio_queues[idx].put(self.stop_signal, timeout=self.timeout) 

74 self.finished_flags[idx] = True 

75 else: 

76 # End specific samples 

77 for sample_idx in sample_indices: 

78 idx = sample_idx.item() if torch.is_tensor(sample_idx) else sample_idx # type: ignore[assignment] 

79 if idx < self.batch_size and not self.finished_flags[idx]: 

80 self.audio_queues[idx].put(self.stop_signal, timeout=self.timeout) 

81 self.finished_flags[idx] = True 

82 

83 def __iter__(self) -> AudioBatchIterator: 

84 """Returns an iterator over the batch of audio streams.""" 

85 return AudioBatchIterator(self) 

86 

87 def get_stream(self, sample_idx: int) -> AudioSampleIterator: 

88 """Get the audio stream for a specific sample.""" 

89 if sample_idx >= self.batch_size: 

90 raise ValueError(f"Sample index {sample_idx} exceeds batch size {self.batch_size}") 

91 return AudioSampleIterator(self, sample_idx) 

92 

93 

94class AudioSampleIterator: 

95 """Iterator for a single audio stream from the batch.""" 

96 

97 def __init__(self, streamer: AudioStreamer, sample_idx: int) -> None: 

98 self.streamer = streamer 

99 self.sample_idx = sample_idx 

100 

101 def __iter__(self) -> AudioSampleIterator: 

102 return self 

103 

104 def __next__(self) -> Any: 

105 value = self.streamer.audio_queues[self.sample_idx].get(timeout=self.streamer.timeout) 

106 if value == self.streamer.stop_signal: 

107 raise StopIteration() 

108 return value 

109 

110 

111class AudioBatchIterator: 

112 """Iterator that yields audio chunks for all samples in the batch.""" 

113 

114 def __init__(self, streamer: AudioStreamer) -> None: 

115 self.streamer = streamer 

116 self.active_samples = set(range(streamer.batch_size)) 

117 

118 def __iter__(self) -> AudioBatchIterator: 

119 return self 

120 

121 def __next__(self) -> dict[int, Any]: 

122 if not self.active_samples: 

123 raise StopIteration() 

124 

125 batch_chunks = {} 

126 samples_to_remove = set() 

127 

128 # Try to get chunks from all active samples 

129 for idx in self.active_samples: 

130 try: 

131 value = self.streamer.audio_queues[idx].get(block=False) 

132 if value == self.streamer.stop_signal: 

133 samples_to_remove.add(idx) 

134 else: 

135 batch_chunks[idx] = value 

136 except Exception: 

137 # Queue is empty for this sample, skip it this iteration 

138 pass 

139 

140 # Remove finished samples 

141 self.active_samples -= samples_to_remove 

142 

143 if batch_chunks: 

144 return batch_chunks 

145 elif self.active_samples: 

146 # If no chunks were ready but we still have active samples, 

147 # wait a bit and try again 

148 time.sleep(0.01) 

149 return self.__next__() 

150 else: 

151 raise StopIteration() 

152 

153 

154class AsyncAudioStreamer(AudioStreamer): 

155 """ 

156 Async version of AudioStreamer for use in async contexts. 

157 """ 

158 

159 def __init__( 

160 self, 

161 batch_size: int, 

162 stop_signal: Optional[Any] = None, 

163 timeout: Optional[float] = None, 

164 ) -> None: 

165 super().__init__(batch_size, stop_signal, timeout) 

166 # Replace regular queues with async queues 

167 self.audio_queues: list[asyncio.Queue] = [ # type: ignore[assignment] 

168 asyncio.Queue() for _ in range(batch_size) 

169 ] 

170 self.loop = asyncio.get_running_loop() 

171 

172 def put(self, audio_chunks: torch.Tensor, sample_indices: torch.Tensor) -> None: # type: ignore[override] 

173 """Put audio chunks in the appropriate async queues.""" 

174 for i, sample_idx in enumerate(sample_indices): 

175 idx = sample_idx.item() 

176 if idx < self.batch_size and not self.finished_flags[idx]: 

177 audio_chunk = audio_chunks[i].detach().cpu() 

178 self.loop.call_soon_threadsafe( 

179 self.audio_queues[idx].put_nowait, audio_chunk 

180 ) 

181 

182 def end(self, sample_indices: Optional[torch.Tensor] = None) -> None: 

183 """Signal the end of generation for specified samples.""" 

184 if sample_indices is None: 

185 indices_to_end = range(self.batch_size) 

186 else: 

187 indices_to_end = [s.item() if torch.is_tensor(s) else s for s in sample_indices] # type: ignore[assignment] 

188 

189 for idx in indices_to_end: 

190 if idx < self.batch_size and not self.finished_flags[idx]: 

191 self.loop.call_soon_threadsafe( 

192 self.audio_queues[idx].put_nowait, self.stop_signal 

193 ) 

194 self.finished_flags[idx] = True 

195 

196 async def get_stream(self, sample_idx: int) -> Any: 

197 """Get async iterator for a specific sample's audio stream.""" 

198 if sample_idx >= self.batch_size: 

199 raise ValueError(f"Sample index {sample_idx} exceeds batch size {self.batch_size}") 

200 

201 while True: 

202 value = await self.audio_queues[sample_idx].get() 

203 if value == self.stop_signal: 

204 break 

205 yield value 

206 

207 def __aiter__(self) -> AsyncAudioBatchIterator: 

208 """Returns an async iterator over all audio streams.""" 

209 return AsyncAudioBatchIterator(self) 

210 

211 

212class AsyncAudioBatchIterator: 

213 """Async iterator for batch audio streaming.""" 

214 

215 def __init__(self, streamer: AsyncAudioStreamer): 

216 self.streamer = streamer 

217 self.active_samples = set(range(streamer.batch_size)) 

218 

219 def __aiter__(self) -> AsyncAudioBatchIterator: 

220 return self 

221 

222 async def __anext__(self) -> dict[int, Any]: 

223 if not self.active_samples: 

224 raise StopAsyncIteration() 

225 

226 batch_chunks = {} 

227 samples_to_remove = set() 

228 

229 # Create tasks for all active samples 

230 tasks = { 

231 idx: asyncio.create_task(self._get_chunk(idx)) 

232 for idx in self.active_samples 

233 } 

234 

235 # Wait for at least one chunk to be ready 

236 done, pending = await asyncio.wait( 

237 tasks.values(), 

238 return_when=asyncio.FIRST_COMPLETED, 

239 timeout=self.streamer.timeout 

240 ) 

241 

242 # Cancel pending tasks 

243 for task in pending: 

244 task.cancel() 

245 

246 # Process completed tasks 

247 for idx, task in tasks.items(): 

248 if task in done: 

249 try: 

250 value = await task 

251 if value == self.streamer.stop_signal: 

252 samples_to_remove.add(idx) 

253 else: 

254 batch_chunks[idx] = value 

255 except asyncio.CancelledError: 

256 pass 

257 

258 self.active_samples -= samples_to_remove 

259 

260 if batch_chunks: 

261 return batch_chunks 

262 elif self.active_samples: 

263 # Try again if we still have active samples 

264 return await self.__anext__() 

265 else: 

266 raise StopAsyncIteration() 

267 

268 async def _get_chunk(self, idx: int) -> Any: 

269 """Helper to get a chunk from a specific queue.""" 

270 return await self.streamer.audio_queues[idx].get()