| 31 | |
| 32 | |
| 33 | def process_inputs( |
| 34 | tokenizer, |
| 35 | spt, |
| 36 | prompt, |
| 37 | text, |
| 38 | device, |
| 39 | audio_data=None, |
| 40 | reference_audio=None, |
| 41 | main_audio=None, |
| 42 | max_channels=8, |
| 43 | pad_token=1024, |
| 44 | ): |
| 45 | # Decompose template into multiple parts |
| 46 | # 1. Style prompt part |
| 47 | seg1 = f"<|begin_of_style|>{prompt}<|end_of_style|>\n<|begin_of_text|>" |
| 48 | inputs1 = np.array(tokenizer.encode(seg1)) |
| 49 | inputs_expanded1 = np.full((len(inputs1), max_channels), pad_token) |
| 50 | inputs_expanded1[:, 0] = inputs1 |
| 51 | labels1 = np.full( |
| 52 | inputs_expanded1.shape, -100 |
| 53 | ) # Style prompt does not compute loss |
| 54 | |
| 55 | # 2. Text part |
| 56 | text_tokens = tokenizer.encode(text, add_special_tokens=False) |
| 57 | inputs2 = np.array(text_tokens) |
| 58 | inputs_expanded2 = np.full((len(inputs2), max_channels), pad_token) |
| 59 | inputs_expanded2[:, 0] = inputs2 |
| 60 | labels2 = np.full(inputs_expanded2.shape, -100) # Text does not compute loss |
| 61 | |
| 62 | # 3. Text end/speech begin part |
| 63 | seg3 = f"<|end_of_text|>\n<|begin_of_speech|>" |
| 64 | inputs3 = np.array(tokenizer.encode(seg3)) |
| 65 | inputs_expanded3 = np.full((len(inputs3), max_channels), pad_token) |
| 66 | inputs_expanded3[:, 0] = inputs3 |
| 67 | labels3 = np.full( |
| 68 | inputs_expanded3.shape, -100 |
| 69 | ) # Start marker does not compute loss |
| 70 | |
| 71 | # 4. Audio processing part |
| 72 | audio_token = None |
| 73 | if reference_audio is not None and main_audio is not None: |
| 74 | # New format: process two audio files separately and then concatenate tokens |
| 75 | try: |
| 76 | # Add silence to the end of each audio |
| 77 | silence_samples = int(SILENCE_DURATION * 16000) |
| 78 | silence = torch.zeros(1, silence_samples) |
| 79 | |
| 80 | # Ensure audio has correct shape [1, samples] |
| 81 | if len(reference_audio.shape) == 1: |
| 82 | reference_audio = reference_audio.unsqueeze(0) |
| 83 | if len(main_audio.shape) == 1: |
| 84 | main_audio = main_audio.unsqueeze(0) |
| 85 | |
| 86 | # Add silence to each audio |
| 87 | ref_audio_with_silence = torch.cat([reference_audio, silence], dim=1) |
| 88 | main_audio_with_silence = torch.cat([main_audio, silence], dim=1) |
| 89 | |
| 90 | with torch.no_grad(): |