vincenthugging/MOSS-TTSD-Enhanced

1import os2import re3 4import torch5import torchaudio6import numpy as np7 8from transformers import AutoTokenizer9from modeling_asteroid import AsteroidTTSInstruct10from XY_Tokenizer.xy_tokenizer.model import XY_Tokenizer11 12MAX_CHANNELS = 813SILENCE_DURATION = 0.0  # Fixed silence duration: 0 seconds14 15def load_model(model_path, spt_config_path, spt_checkpoint_path, torch_dtype=torch.bfloat16, attn_implementation="sdpa"):16    tokenizer = AutoTokenizer.from_pretrained(model_path)17    18    # 尝试使用 FlashAttention2，失败则回退到标准实现19    try:20        model = AsteroidTTSInstruct.from_pretrained(model_path, torch_dtype=torch_dtype, attn_implementation="flash_attention_2")21        print("✅ 使用 FlashAttention2")22    except ImportError:23        print("⚠️ FlashAttention2 不可用，使用标准注意力机制")24        model = AsteroidTTSInstruct.from_pretrained(model_path, torch_dtype=torch_dtype, attn_implementation=attn_implementation)25 26    spt = XY_Tokenizer.load_from_checkpoint(config_path=spt_config_path, ckpt_path=spt_checkpoint_path)27    28    model.eval()29    spt.eval()30    return tokenizer, model, spt31 32 33def process_jsonl_item(item):34    """Process JSONL data items and extract audio and text information according to the new format"""35    base_path = item.get("base_path", "")36    text = item.get("text", "")37    38    prompt_audio = None39    prompt_text = ""40 41    # Process prompt audio and text42    if "prompt_audio" in item and "prompt_text" in item:43        print("Using prompt_audio and prompt_text directly from item.")44        # If prompt_audio and prompt_text exist, use them directly45        prompt_audio_val = item["prompt_audio"]46        if prompt_audio_val: # Only assign if not empty47            prompt_audio = prompt_audio_val48            prompt_text = item["prompt_text"]49        50            # Only perform path joining when prompt_audio is a string path51            if isinstance(prompt_audio, str) and base_path and prompt_audio:52                prompt_audio = os.path.join(base_path, prompt_audio)53    else:54        # Otherwise, merge speaker1 and speaker2 information55        prompt_audio_speaker1 = item.get("prompt_audio_speaker1", "")56        prompt_text_speaker1 = item.get("prompt_text_speaker1", "")57        prompt_audio_speaker2 = item.get("prompt_audio_speaker2", "")58        prompt_text_speaker2 = item.get("prompt_text_speaker2", "")59        60        has_speaker1_audio = (isinstance(prompt_audio_speaker1, str) and prompt_audio_speaker1) or isinstance(prompt_audio_speaker1, tuple)61        has_speaker2_audio = (isinstance(prompt_audio_speaker2, str) and prompt_audio_speaker2) or isinstance(prompt_audio_speaker2, tuple)62 63        if has_speaker1_audio or has_speaker2_audio:64            print("Using speaker1 and speaker2 information for prompt audio and text.")65            # Process audio: if it's a string path, perform path joining; if it's a tuple, use directly66            if isinstance(prompt_audio_speaker1, str):67                speaker1_audio = os.path.join(base_path, prompt_audio_speaker1) if base_path and prompt_audio_speaker1 else prompt_audio_speaker168            else:69                speaker1_audio = prompt_audio_speaker1  # Use tuple directly70                71            if isinstance(prompt_audio_speaker2, str):72                speaker2_audio = os.path.join(base_path, prompt_audio_speaker2) if base_path and prompt_audio_speaker2 else prompt_audio_speaker273            else:74                speaker2_audio = prompt_audio_speaker2  # Use tuple directly75            76            prompt_audio = {77                "speaker1": speaker1_audio,78                "speaker2": speaker2_audio79            }80        81        # Merge text82        temp_prompt_text = ""83        if prompt_text_speaker1:84            temp_prompt_text += f"[S1]{prompt_text_speaker1}"85        if prompt_text_speaker2:86            temp_prompt_text += f"[S2]{prompt_text_speaker2}"87        prompt_text = temp_prompt_text.strip()88    89    return {90        "text": text,91        "prompt_text": prompt_text,92        "prompt_audio": prompt_audio93    }94 95 96def load_audio_data(prompt_audio, target_sample_rate=16000):97    """Load audio data and return processed audio tensor98    99    Args:100        prompt_audio: Can be in the following formats:101            - String: audio file path102            - Tuple: (wav, sr) result from torchaudio.load103            - Dict: {"speaker1": path_or_tuple, "speaker2": path_or_tuple}104    """105    if prompt_audio is None:106        return None107    108    try:109        # Check if prompt_audio is a dictionary (containing speaker1 and speaker2)110        if isinstance(prompt_audio, dict) and "speaker1" in prompt_audio and "speaker2" in prompt_audio:111            # Process audio from both speakers separately112            wav1, sr1 = _load_single_audio(prompt_audio["speaker1"])113            wav2, sr2 = _load_single_audio(prompt_audio["speaker2"])114            # Merge audio from both speakers115            wav = merge_speaker_audios(wav1, sr1, wav2, sr2, target_sample_rate)116            if wav is None:117                return None118        else:119            # Single audio120            wav, sr = _load_single_audio(prompt_audio)121            # Resample to 16k122            if sr != target_sample_rate: 123                wav = torchaudio.functional.resample(wav, sr, target_sample_rate)124            # Ensure mono channel125            if wav.shape[0] > 1:126                wav = wav.mean(dim=0, keepdim=True)  # Convert multi-channel to mono127            if len(wav.shape) == 1: 128                wav = wav.unsqueeze(0)129        130        return wav131    except Exception as e:132        print(f"Error loading audio data: {e}")133        raise134 135 136def _load_single_audio(audio_input):137    """Load single audio, supports file path or (wav, sr) tuple138    139    Args:140        audio_input: String (file path) or tuple (wav, sr)141        142    Returns:143        tuple: (wav, sr)144    """145    if isinstance(audio_input, tuple) and len(audio_input) == 2:146        # Already a (wav, sr) tuple147        wav, sr = audio_input148        return wav, sr149    elif isinstance(audio_input, str):150        # Is a file path, needs to be loaded151        wav, sr = torchaudio.load(audio_input)152        return wav, sr153    else:154        raise ValueError(f"Unsupported audio input format: {type(audio_input)}")155 156 157def merge_speaker_audios(wav1, sr1, wav2, sr2, target_sample_rate=16000):158    """Merge audio data from two speakers"""159    try:160        # Process first audio161        if sr1 != target_sample_rate:162            wav1 = torchaudio.functional.resample(wav1, sr1, target_sample_rate)163        # Ensure mono channel164        if wav1.shape[0] > 1:165            wav1 = wav1.mean(dim=0, keepdim=True)  # Convert multi-channel to mono166        if len(wav1.shape) == 1:167            wav1 = wav1.unsqueeze(0)168        169        # Process second audio  170        if sr2 != target_sample_rate:171            wav2 = torchaudio.functional.resample(wav2, sr2, target_sample_rate)172        # Ensure mono channel173        if wav2.shape[0] > 1:174            wav2 = wav2.mean(dim=0, keepdim=True)  # Convert multi-channel to mono175        if len(wav2.shape) == 1:176            wav2 = wav2.unsqueeze(0)177        178        # Concatenate audio179        merged_wav = torch.cat([wav1, wav2], dim=1)180        return merged_wav181    except Exception as e:182        print(f"Error merging audio: {e}")183        raise184 185 186def process_inputs(tokenizer, spt, prompt, text, device, audio_data=None, max_channels=8, pad_token=1024):187    seq = f"<|begin_of_style|>{prompt}<|end_of_style|>\n<|begin_of_text|>{text}<|end_of_text|>\n<|begin_of_speech|>"188    inputs1 = np.array(tokenizer.encode(seq))189    input_ids = np.full((inputs1.shape[0], max_channels), pad_token)190    input_ids[:, 0] = inputs1191    192    if audio_data is not None:193        try:194            # audio_data should now be a processed audio tensor195            wav = audio_data196            197            # Add fixed 5-second silence at the end of audio (using 16k sample rate)198            silence_samples = int(SILENCE_DURATION * 16000)199            silence = torch.zeros(wav.shape[0], silence_samples)200            wav = torch.cat([wav, silence], dim=1)201            202            with torch.no_grad():203                # Use SPT encoding204                encode_result = spt.encode([wav.squeeze().to(device)])205                audio_token = encode_result["codes_list"][0].permute(1, 0).cpu().numpy()  # Adjust dimension order206                207            # similar to DAC encoding adjustment208            audio_token[:, 0] = audio_token[:, 0] + 151665  # Keep this line if offset is needed, otherwise delete209            input_ids = np.concatenate([input_ids, audio_token])210        except Exception as e:211            print(f"Error processing audio data: {e}")212            raise213    214    return input_ids215 216 217def shifting_inputs(input_ids, tokenizer, pad_token=1024, max_channels=8):218    seq_len = input_ids.shape[0]219    new_seq_len = seq_len + max_channels - 1220    shifted_input_ids = np.full((new_seq_len, max_channels), pad_token, dtype=np.int64)221    shifted_input_ids[:, 0] = np.full(new_seq_len, tokenizer.pad_token_id, dtype=np.int64)222    for i in range(max_channels):223        shifted_input_ids[i : (seq_len + i), i] = input_ids[:, i]224    return shifted_input_ids225 226 227def rpadding(input_ids, channels, tokenizer):228    attention_masks = [np.ones(inputs.shape[0]) for inputs in input_ids]229    max_length = max(ids.shape[0] for ids in input_ids)230    padded_input_ids, padded_attns = [], []231        232    for ids, attn in zip(input_ids, attention_masks):233        pad_len = max_length - ids.shape[0]234        input_pad = np.full((pad_len, channels), 1024)235        input_pad[:, 0] = tokenizer.pad_token_id236        padded_input_ids.append(np.concatenate([input_pad, ids]))237        attn_pad = np.zeros(pad_len)238        padded_attns.append(np.concatenate([attn_pad, attn]))239 240    input_ids = torch.tensor(np.stack(padded_input_ids))241    attention_mask = torch.tensor(np.stack(padded_attns))242 243    return input_ids, attention_mask244 245 246def find_max_valid_positions(C: torch.Tensor, invalid_value=1024) -> torch.Tensor:247    values = C[:, :, 1]248    mask = (values != invalid_value)249    reversed_mask = mask.flip(dims=[1])250    reversed_indices = torch.argmax(reversed_mask.int(), dim=1)251    seq_len = C.size(1)252    original_indices = seq_len - 1 - reversed_indices253    has_valid = mask.any(dim=1)254    original_indices = torch.where(has_valid, original_indices, -1)255    return original_indices256 257 258def normalize_text(text: str) -> str:259    """260    Normalize multi-speaker script.261 262    1. Don't preserve line breaks.263    2. Remove brackets for non-speaker tags (if [] doesn't contain S1/S2...Sx format, remove the brackets themselves).264    3. Remove decorative symbols: 【】《》（）『』「」"-“” .265    4. Internal punctuation ！；：、 → ，；only allow ？ and ，。266    5. Multiple 。 keep only the last one, others → ，。267    6. Replace consecutive "哈" (>=2) with "(笑)".268    7. Auto-recognize [S1] / [S2] … tags; if missing, treat as whole segment.269    8. Merge adjacent identical speaker tags.270    """271    # Replace [1], [2] etc. format with [S1], [S2] etc. format272    text = re.sub(r'\[(\d+)\]', r'[S\1]', text)273 274    # Remove decorative characters275    remove_chars = "【】《》（）『』「」""\"-“”～~"276 277 278    # Remove brackets for non-speaker tags (keep content, only remove brackets themselves)279    text = re.sub(r'\[(?!S\d+\])([^\]]*)\]', r'\1', text)280 281    # Use positive lookahead to split text by speaker tags (tags themselves are still preserved)282    segments = re.split(r'(?=\[S\d+\])', text.replace("\n", " "))283    processed_parts = []284 285    for seg in segments:286        seg = seg.strip()287        if not seg:288            continue289 290        # Extract tags291        m = re.match(r'^(\[S\d+\])\s*(.*)', seg)292        tag, content = m.groups() if m else ('', seg)293 294        # Remove irrelevant symbols295        content = re.sub(f"[{re.escape(remove_chars)}]", "", content)296 297        # Handle consecutive "哈" characters: replace 2 or more with "(笑)"298        content = re.sub(r'哈{2,}', '(笑)', content)299 300        # Handle English laughter (e.g., "haha", "ha ha")301        content = re.sub(r'\b(ha(\s*ha)+)\b', '(laughs)', content, flags=re.IGNORECASE)302 303        # First handle multi-character punctuation marks304        content = content.replace('——', '，')305        content = content.replace('……', '，')306 307        # Handle single-character internal punctuation marks308        internal_punct_map = str.maketrans({309            '！': '，', '!': ',',310            '；': '，', ';': ',',311            '：': '，', ':': ',',312            '、': '，', 313            '？': '，', '?': ','314        })315        content = content.translate(internal_punct_map)316        content = content.strip()317 318        # Keep only the final period319        if len(content) > 1:320            last_ch = "。" if content[-1] == "，" else ("." if content[-1] == "," else content[-1])321            body = content[:-1].replace('。', '，')322            content = body + last_ch323 324        processed_parts.append({'tag': tag, 'content': content})325 326    if not processed_parts:327        return ""328 329    # Merge consecutive same speakers330    merged_lines = []331    current_tag = processed_parts[0]['tag']332    current_content = [processed_parts[0]['content']]333 334    for part in processed_parts[1:]:335        if part['tag'] == current_tag and current_tag:336            current_content.append(part['content'])337        else:338            merged_lines.append(f"{current_tag}{''.join(current_content)}".strip())339            current_tag = part['tag']340            current_content = [part['content']]341 342    merged_lines.append(f"{current_tag}{''.join(current_content)}".strip())343    344    return "".join(merged_lines).replace('‘', "'").replace('’', "'")345 346 347def process_batch(batch_items, tokenizer, model, spt, device, system_prompt, start_idx, use_normalize=False):348    """Process a batch of data items and generate audio, return audio data and metadata"""349    try:350        # Prepare batch data351        batch_size = len(batch_items)352        texts = []353        prompts = [system_prompt] * batch_size354        prompt_audios = []355        actual_texts_data = []  # Store actual text data used356        357        print(f"Processing {batch_size} samples starting from index {start_idx}...")358        359        # Extract text and audio from each sample360        for i, item in enumerate(batch_items):361            # Use new processing function362            processed_item = process_jsonl_item(item)363            364            text = processed_item["text"]365            prompt_text = processed_item["prompt_text"]366            367            # Merge text, if prompt_text is empty, full_text is just text368            full_text = prompt_text + text if prompt_text else text369            original_full_text = full_text  # Save original text370            371            # Apply text normalization based on parameter372            if use_normalize:373                full_text = normalize_text(full_text)374            375            # Replace speaker tags376            final_text = full_text.replace("[S1]", "<speaker1>").replace("[S2]", "<speaker2>")377            texts.append(final_text)378            379            # Save actual text information used380            actual_texts_data.append({381                "index": start_idx + i,382                "original_text": original_full_text,383                "normalized_text": normalize_text(original_full_text) if use_normalize else None,384                "final_text": final_text,385                "use_normalize": use_normalize386            })387            388            # Get reference audio389            prompt_audios.append(processed_item["prompt_audio"])390        391        # Process inputs392        input_ids_list = []393        for i, (text, prompt, audio_path) in enumerate(zip(texts, prompts, prompt_audios)):394            # Load audio data here395            audio_data = load_audio_data(audio_path) if audio_path else None396            inputs = process_inputs(tokenizer, spt, prompt, text, device, audio_data)397            inputs = shifting_inputs(inputs, tokenizer)398            input_ids_list.append(inputs)399        400        # Pad batch inputs401        input_ids, attention_mask = rpadding(input_ids_list, MAX_CHANNELS, tokenizer)402        403        # Batch generation404        print(f"Starting batch audio generation...")405        start = input_ids.shape[1] - MAX_CHANNELS + 1406        407        # Move inputs to GPU408        input_ids = input_ids.to(device)409        attention_mask = attention_mask.to(device)410        411        # Generate model outputs412        outputs = model.generate(413            input_ids=input_ids, 414            attention_mask=attention_mask,415        )416        print(f"Original outputs shape: {outputs.shape}")417        print(f"Start value: {start}")418        print(f"Shape after slicing: {outputs[:, start:].shape}")419        print(f"MAX_CHANNELS: {MAX_CHANNELS}")420        print(f"Calculated seq_len: {outputs.shape[1] - MAX_CHANNELS + 1}")421        # Process outputs422        outputs = outputs[:, start:]423        seq_len = outputs.shape[1] - MAX_CHANNELS + 1424        speech_ids = torch.full((outputs.shape[0], seq_len, MAX_CHANNELS), 0).to(device)425        426        427        # Adjust output format428        for j in range(MAX_CHANNELS):429            speech_ids[..., j] = outputs[:, j : seq_len + j, j]430            if j == 0: 431                speech_ids[..., j] = speech_ids[..., j] - 151665432        433        # Find valid positions for each sample434        li = find_max_valid_positions(speech_ids)435        436        # Store audio result data437        audio_results = []438        439        # Process batch sample results individually440        for i in range(batch_size):441            try:442                # Extract valid speech tokens443                end_idx = li[i] + 1444                if end_idx <= 0:445                    print(f"Sample {start_idx + i} has no valid speech tokens")446                    audio_results.append(None)447                    continue448                    449                this_speech_id = speech_ids[i, :end_idx]450                print(f"Speech token shape for sample {start_idx + i}: {this_speech_id.shape}")451                452                # Decode generated audio453                with torch.no_grad():454                    codes_list = [this_speech_id.permute(1, 0)]  # Convert to SPT expected format455                    decode_result = spt.decode(codes_list, overlap_seconds=10)456                    audio_result = decode_result["syn_wav_list"][0].cpu().detach()457                    458                    if audio_result.ndim == 1:  # If 1D [samples]459                        audio_result = audio_result.unsqueeze(0)  # Convert to 2D [1, samples]460                461                # Save audio data instead of file path462                audio_results.append({463                    "audio_data": audio_result,464                    "sample_rate": spt.output_sample_rate,465                    "index": start_idx + i466                })467                print(f"Audio generation completed: sample {start_idx + i}")468                469            except Exception as e:470                print(f"Error processing sample {start_idx + i}: {str(e)}, skipping...")471                import traceback472                traceback.print_exc()473                audio_results.append(None)474        475        # Clean up GPU memory476        torch.cuda.empty_cache()477        478        # Return text data and audio data479        return actual_texts_data, audio_results480        481    except Exception as e:482        print(f"Error during batch processing: {str(e)}")483        raise484