Custom Streamer for streaming text_ids and audio_ids The model will call put() twice: 1. streamer.put(next_tokens) - Text tokens, shape (batch_size,) 2. streamer.put(next_speech_tokens) - Audio tokens, shape (batch_size, group_size)
| 22 | register_funaudiochat() |
| 23 | |
| 24 | class FunaudioChatStreamer(BaseStreamer): |
| 25 | """ |
| 26 | Custom Streamer for streaming text_ids and audio_ids |
| 27 | |
| 28 | The model will call put() twice: |
| 29 | 1. streamer.put(next_tokens) - Text tokens, shape (batch_size,) |
| 30 | 2. streamer.put(next_speech_tokens) - Audio tokens, shape (batch_size, group_size) |
| 31 | """ |
| 32 | def __init__(self, processor, skip_prompt=True, group_size=5, **decode_kwargs): |
| 33 | self.processor = processor |
| 34 | self.skip_prompt = skip_prompt |
| 35 | self.decode_kwargs = decode_kwargs |
| 36 | self.group_size = group_size |
| 37 | |
| 38 | self.step_results = [] |
| 39 | self.text_token_cache = [] |
| 40 | self.audio_token_cache = [] |
| 41 | self.pending_text_ids = None |
| 42 | self.pending_text_str = None |
| 43 | |
| 44 | self.done = False |
| 45 | self.prompt_length = 0 |
| 46 | |
| 47 | def put(self, value): |
| 48 | """ |
| 49 | Receiving generated tokens |
| 50 | """ |
| 51 | is_text_token = False |
| 52 | is_audio_token = False |
| 53 | |
| 54 | if value.dim() == 1: |
| 55 | is_text_token = True |
| 56 | text_ids = value.unsqueeze(-1) # (batch_size,) -> (batch_size, 1) |
| 57 | elif value.dim() == 2: |
| 58 | if value.shape[1] == self.group_size or value.shape[1] > 1: |
| 59 | is_audio_token = True |
| 60 | audio_ids = value |
| 61 | else: |
| 62 | is_text_token = True |
| 63 | text_ids = value |
| 64 | |
| 65 | # handle text tokens |
| 66 | if is_text_token: |
| 67 | # mark the prompt length for the first call |
| 68 | if len(self.text_token_cache) == 0 and self.skip_prompt and text_ids.shape[1] > 1: |
| 69 | self.prompt_length = text_ids.shape[1] - 1 |
| 70 | new_text_ids = text_ids[:, -1:] |
| 71 | else: |
| 72 | new_text_ids = text_ids |
| 73 | |
| 74 | self.text_token_cache.append(new_text_ids.clone()) |
| 75 | |
| 76 | new_text_str = "" |
| 77 | try: |
| 78 | new_text_str = self.processor.decode(new_text_ids[0], **self.decode_kwargs) |
| 79 | except Exception as e: |
| 80 | new_text_str = f"[Error: {e}]" |
| 81 |