cocktailpeanut/DiffRhythm

1import torch2import librosa3import random4import json5from muq import MuQMuLan6from mutagen.mp3 import MP37import os8import numpy as np9from huggingface_hub import hf_hub_download10from diffrhythm.model import DiT, CFM11 12 13def prepare_model(device):14    # prepare cfm model15    dit_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-base", filename="cfm_model.pt")16    #dit_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-full", filename="cfm_model.pt")17    dit_config_path = "./diffrhythm/config/diffrhythm-1b.json"18    with open(dit_config_path, encoding="utf-8") as f:19        model_config = json.load(f)20    dit_model_cls = DiT21    cfm = CFM(22                #transformer=dit_model_cls(**model_config["model"], use_style_prompt=True, max_pos=6144),23                transformer=dit_model_cls(**model_config["model"], use_style_prompt=True, max_pos=2048),24                num_channels=model_config["model"]['mel_dim'],25                use_style_prompt=True26             )27    cfm = cfm.to(device)28    cfm = load_checkpoint(cfm, dit_ckpt_path, device=device, use_ema=False)29    30    # prepare tokenizer31    tokenizer = CNENTokenizer()32    33    # prepare muq34    muq = MuQMuLan.from_pretrained("OpenMuQ/MuQ-MuLan-large")35    muq = muq.to(device).eval()36    37    # prepare vae38    vae_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-vae", filename="vae_model.pt")39    vae = torch.jit.load(vae_ckpt_path, map_location='cpu').to(device)40    return cfm, tokenizer, muq, vae41    42 43# for song edit, will be added in the future44def get_reference_latent(device, max_frames):45    return torch.zeros(1, max_frames, 64).to(device)46 47def get_negative_style_prompt(device):48    file_path = "./src/negative_prompt.npy"49    vocal_stlye = np.load(file_path)50    51    vocal_stlye = torch.from_numpy(vocal_stlye).to(device) # [1, 512]52    vocal_stlye = vocal_stlye.half()53    54    return vocal_stlye55 56 57def get_audio_style_prompt(model, wav_path):58    vocal_flag = False59    mulan = model60    audio, _ = librosa.load(wav_path, sr=24000)61    audio_len = librosa.get_duration(y=audio, sr=24000)62 63    if audio_len <= 1:64        vocal_flag = True65 66    if audio_len > 10:67        start_time = int(audio_len // 2 - 5)68        wav = audio[start_time*24000:(start_time+10)*24000]69 70    else:71        wav = audio72    wav = torch.tensor(wav).unsqueeze(0).to(model.device)73 74    with torch.no_grad():75        audio_emb = mulan(wavs = wav) # [1, 512]76 77    audio_emb = audio_emb.half()78 79    return audio_emb, vocal_flag80 81def get_text_style_prompt(model, text_prompt):82    mulan = model83 84    with torch.no_grad():85        text_emb = mulan(texts = text_prompt) # [1, 512]86    text_emb = text_emb.half()87 88    return text_emb89 90 91@torch.no_grad()92def get_style_prompt(model, wav_path, prompt):93    mulan = model94 95    if prompt is not None:96        return mulan(texts=prompt).half()97    98    ext = os.path.splitext(wav_path)[-1].lower()99    if ext == '.mp3':100        meta = MP3(wav_path)101        audio_len = meta.info.length102        src_sr = meta.info.sample_rate103    elif ext == '.wav':104        audio, sr = librosa.load(wav_path, sr=None)105        audio_len = librosa.get_duration(y=audio, sr=sr)106        src_sr = sr107    else:108        raise ValueError("Unsupported file format: {}".format(ext))109    110    assert(audio_len >= 10)111    112    mid_time = audio_len // 2113    start_time = mid_time - 5114    wav, sr = librosa.load(wav_path, sr=None, offset=start_time, duration=10)115    116    resampled_wav = librosa.resample(wav, orig_sr=src_sr, target_sr=24000)117    resampled_wav = torch.tensor(resampled_wav).unsqueeze(0).to(model.device)118    119    with torch.no_grad():120        audio_emb = mulan(wavs = resampled_wav) # [1, 512]121        122    audio_emb = audio_emb123    audio_emb = audio_emb.half()124 125    return audio_emb126 127def parse_lyrics(lyrics: str):128    lyrics_with_time = []129    lyrics = lyrics.strip()130    for line in lyrics.split('\n'):131        try:132            time, lyric = line[1:9], line[10:]133            lyric = lyric.strip()134            mins, secs = time.split(':')135            secs = int(mins) * 60 + float(secs)136            lyrics_with_time.append((secs, lyric))137        except:138            continue139    return lyrics_with_time140 141class CNENTokenizer():142    def __init__(self):143        with open('./diffrhythm/g2p/g2p/vocab.json', 'r', encoding="utf-8") as file:144            self.phone2id:dict = json.load(file)['vocab']145        self.id2phone = {v:k for (k, v) in self.phone2id.items()}146        # from f5_tts.g2p.g2p_generation import chn_eng_g2p147        from diffrhythm.g2p.g2p_generation import chn_eng_g2p148        self.tokenizer = chn_eng_g2p149    def encode(self, text):150        phone, token = self.tokenizer(text)151        token = [x+1 for x in token]152        return token153    def decode(self, token):154        return "|".join([self.id2phone[x-1] for x in token])155    156def get_lrc_token(max_frames, text, tokenizer, device):157 158#    max_frames = 2048159    lyrics_shift = 0160    sampling_rate = 44100161    downsample_rate = 2048162    max_secs = max_frames / (sampling_rate / downsample_rate)163   164    pad_token_id = 0165    comma_token_id = 1166    period_token_id = 2    167 168    if text == "":169        return torch.zeros((max_frames,), dtype=torch.long).unsqueeze(0).to(device), torch.tensor(0.).unsqueeze(0).to(device).half()170 171    lrc_with_time = parse_lyrics(text)172    173    modified_lrc_with_time = []174    for i in range(len(lrc_with_time)):175        time, line = lrc_with_time[i]176        line_token = tokenizer.encode(line)177        modified_lrc_with_time.append((time, line_token))178    lrc_with_time = modified_lrc_with_time179 180    lrc_with_time = [(time_start, line) for (time_start, line) in lrc_with_time if time_start < max_secs]181#    lrc_with_time = lrc_with_time[:-1] if len(lrc_with_time) >= 1 else lrc_with_time182    183    normalized_start_time = 0.184 185    lrc = torch.zeros((max_frames,), dtype=torch.long)186 187    tokens_count = 0188    last_end_pos = 0189    for time_start, line in lrc_with_time:190        tokens = [token if token != period_token_id else comma_token_id for token in line] + [period_token_id]191        tokens = torch.tensor(tokens, dtype=torch.long)192        num_tokens = tokens.shape[0]193 194        gt_frame_start = int(time_start * sampling_rate / downsample_rate)195        196        frame_shift = random.randint(int(lyrics_shift), int(lyrics_shift))197        198        frame_start = max(gt_frame_start - frame_shift, last_end_pos)199        frame_len = min(num_tokens, max_frames - frame_start)200 201        #print(gt_frame_start, frame_shift, frame_start, frame_len, tokens_count, last_end_pos, full_pos_emb.shape)202 203        lrc[frame_start:frame_start + frame_len] = tokens[:frame_len]204 205        tokens_count += num_tokens206        last_end_pos = frame_start + frame_len   207        208    lrc_emb = lrc.unsqueeze(0).to(device)209    210    normalized_start_time = torch.tensor(normalized_start_time).unsqueeze(0).to(device)211    normalized_start_time = normalized_start_time.half()212    213    return lrc_emb, normalized_start_time214 215def load_checkpoint(model, ckpt_path, device, use_ema=True):216    model = model.half()217 218    ckpt_type = ckpt_path.split(".")[-1]219    if ckpt_type == "safetensors":220        from safetensors.torch import load_file221 222        checkpoint = load_file(ckpt_path)223    else:224        checkpoint = torch.load(ckpt_path, weights_only=True)225 226    if use_ema:227        if ckpt_type == "safetensors":228            checkpoint = {"ema_model_state_dict": checkpoint}229        checkpoint["model_state_dict"] = {230            k.replace("ema_model.", ""): v231            for k, v in checkpoint["ema_model_state_dict"].items()232            if k not in ["initted", "step"]233        }234        model.load_state_dict(checkpoint["model_state_dict"], strict=False)235    else:236        if ckpt_type == "safetensors":237            checkpoint = {"model_state_dict": checkpoint}238        model.load_state_dict(checkpoint["model_state_dict"], strict=False)239 240    return model.to(device)241