xscdvfaaqqq/DiffRhythm

1import torch2import librosa3import torchaudio4import random5import json6from muq import MuQMuLan, MuQ7from mutagen.mp3 import MP38import os9import numpy as np10from huggingface_hub import hf_hub_download11from hydra.utils import instantiate12from omegaconf import OmegaConf13from safetensors.torch import load_file14 15from diffrhythm.model import DiT, CFM16 17def vae_sample(mean, scale):18    stdev = torch.nn.functional.softplus(scale) + 1e-419    var = stdev * stdev20    logvar = torch.log(var)21    latents = torch.randn_like(mean) * stdev + mean22 23    kl = (mean * mean + var - logvar - 1).sum(1).mean()24 25    return latents, kl26 27def normalize_audio(y, target_dbfs=0):28    max_amplitude = torch.max(torch.abs(y))29 30    target_amplitude = 10.0**(target_dbfs / 20.0)31    scale_factor = target_amplitude / max_amplitude32 33    normalized_audio = y * scale_factor34 35    return normalized_audio36 37def set_audio_channels(audio, target_channels):38    if target_channels == 1:39        # Convert to mono40        audio = audio.mean(1, keepdim=True)41    elif target_channels == 2:42        # Convert to stereo43        if audio.shape[1] == 1:44            audio = audio.repeat(1, 2, 1)45        elif audio.shape[1] > 2:46            audio = audio[:, :2, :]47    return audio48 49class PadCrop(torch.nn.Module):50    def __init__(self, n_samples, randomize=True):51        super().__init__()52        self.n_samples = n_samples53        self.randomize = randomize54 55    def __call__(self, signal):56        n, s = signal.shape57        start = 0 if (not self.randomize) else torch.randint(0, max(0, s - self.n_samples) + 1, []).item()58        end = start + self.n_samples59        output = signal.new_zeros([n, self.n_samples])60        output[:, :min(s, self.n_samples)] = signal[:, start:end]61        return output62 63def prepare_audio(audio, in_sr, target_sr, target_length, target_channels, device):64    65    audio = audio.to(device)66 67    if in_sr != target_sr:68        resample_tf = torchaudio.transforms.Resample(in_sr, target_sr).to(device)69        audio = resample_tf(audio)70    if target_length is None:71        target_length = audio.shape[-1]72    audio = PadCrop(target_length, randomize=False)(audio)73 74    # Add batch dimension75    if audio.dim() == 1:76        audio = audio.unsqueeze(0).unsqueeze(0)77    elif audio.dim() == 2:78        audio = audio.unsqueeze(0)79 80    audio = set_audio_channels(audio, target_channels)81 82    return audio83 84def decode_audio(latents, vae_model, chunked=False, overlap=32, chunk_size=128):85    downsampling_ratio = 204886    io_channels = 287    if not chunked:88        return vae_model.decode_export(latents)89    else:90        # chunked decoding91        hop_size = chunk_size - overlap92        total_size = latents.shape[2]93        batch_size = latents.shape[0]94        chunks = []95        i = 096        for i in range(0, total_size - chunk_size + 1, hop_size):97            chunk = latents[:, :, i : i + chunk_size]98            chunks.append(chunk)99        if i + chunk_size != total_size:100            # Final chunk101            chunk = latents[:, :, -chunk_size:]102            chunks.append(chunk)103        chunks = torch.stack(chunks)104        num_chunks = chunks.shape[0]105        # samples_per_latent is just the downsampling ratio106        samples_per_latent = downsampling_ratio107        # Create an empty waveform, we will populate it with chunks as decode them108        y_size = total_size * samples_per_latent109        y_final = torch.zeros((batch_size, io_channels, y_size)).to(latents.device)110        for i in range(num_chunks):111            x_chunk = chunks[i, :]112            # decode the chunk113            y_chunk = vae_model.decode_export(x_chunk)114            # figure out where to put the audio along the time domain115            if i == num_chunks - 1:116                # final chunk always goes at the end117                t_end = y_size118                t_start = t_end - y_chunk.shape[2]119            else:120                t_start = i * hop_size * samples_per_latent121                t_end = t_start + chunk_size * samples_per_latent122            #  remove the edges of the overlaps123            ol = (overlap // 2) * samples_per_latent124            chunk_start = 0125            chunk_end = y_chunk.shape[2]126            if i > 0:127                # no overlap for the start of the first chunk128                t_start += ol129                chunk_start += ol130            if i < num_chunks - 1:131                # no overlap for the end of the last chunk132                t_end -= ol133                chunk_end -= ol134            # paste the chunked audio into our y_final output audio135            y_final[:, :, t_start:t_end] = y_chunk[:, :, chunk_start:chunk_end]136        return y_final137 138def encode_audio(audio, vae_model, chunked=False, overlap=32, chunk_size=128):139    downsampling_ratio = 2048140    latent_dim = 128141    if not chunked:142        # default behavior. Encode the entire audio in parallel143        return vae_model.encode_export(audio)144    else:145        # CHUNKED ENCODING146        # samples_per_latent is just the downsampling ratio (which is also the upsampling ratio)147        samples_per_latent = downsampling_ratio148        total_size = audio.shape[2] # in samples149        batch_size = audio.shape[0]150        chunk_size *= samples_per_latent # converting metric in latents to samples151        overlap *= samples_per_latent # converting metric in latents to samples152        hop_size = chunk_size - overlap153        chunks = []154        for i in range(0, total_size - chunk_size + 1, hop_size):155            chunk = audio[:,:,i:i+chunk_size]156            chunks.append(chunk)157        if i+chunk_size != total_size:158            # Final chunk159            chunk = audio[:,:,-chunk_size:]160            chunks.append(chunk)161        chunks = torch.stack(chunks)162        num_chunks = chunks.shape[0]163        # Note: y_size might be a different value from the latent length used in diffusion training164        # because we can encode audio of varying lengths165        # However, the audio should've been padded to a multiple of samples_per_latent by now.166        y_size = total_size // samples_per_latent167        # Create an empty latent, we will populate it with chunks as we encode them168        y_final = torch.zeros((batch_size,latent_dim,y_size)).to(audio.device)169        for i in range(num_chunks):170            x_chunk = chunks[i,:]171            # encode the chunk172            y_chunk = vae_model.encode_export(x_chunk)173            # figure out where to put the audio along the time domain174            if i == num_chunks-1:175                # final chunk always goes at the end176                t_end = y_size177                t_start = t_end - y_chunk.shape[2]178            else:179                t_start = i * hop_size // samples_per_latent180                t_end = t_start + chunk_size // samples_per_latent181            #  remove the edges of the overlaps182            ol = overlap//samples_per_latent//2183            chunk_start = 0184            chunk_end = y_chunk.shape[2]185            if i > 0:186                # no overlap for the start of the first chunk187                t_start += ol188                chunk_start += ol189            if i < num_chunks-1:190                # no overlap for the end of the last chunk191                t_end -= ol192                chunk_end -= ol193            # paste the chunked audio into our y_final output audio194            y_final[:,:,t_start:t_end] = y_chunk[:,:,chunk_start:chunk_end]195        return y_final196 197def prepare_model(max_frames, device):198    # prepare cfm model199    if max_frames == 2048:200        dit_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-1_2", filename="cfm_model.pt")201    else:202        dit_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-1_2-full", filename="cfm_model.pt")203        204    dit_config_path = "./diffrhythm/config/config.json"205    with open(dit_config_path) as f:206        model_config = json.load(f)207    dit_model_cls = DiT208    cfm = CFM(209                transformer=dit_model_cls(**model_config["model"], max_frames=max_frames),210                num_channels=model_config["model"]['mel_dim'],211                max_frames=max_frames212             )213    cfm = cfm.to(device)214    cfm = load_checkpoint(cfm, dit_ckpt_path, device=device, use_ema=False)215 216    # prepare tokenizer217    tokenizer = CNENTokenizer()218 219    # prepare muq220    muq = MuQMuLan.from_pretrained("OpenMuQ/MuQ-MuLan-large", cache_dir="./pretrained")221    muq = muq.to(device).eval()222 223    # prepare vae224    vae_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-vae", filename="vae_model.pt")225    vae = torch.jit.load(vae_ckpt_path, map_location="cpu").to(device)226    227    228    # prepare eval model229    train_config = OmegaConf.load("./pretrained/eval.yaml")230    checkpoint_path = "./pretrained/eval.safetensors"231 232    eval_model = instantiate(train_config.generator).to(device).eval()233    state_dict = load_file(checkpoint_path, device="cpu")234    eval_model.load_state_dict(state_dict)235 236    eval_muq = MuQ.from_pretrained("OpenMuQ/MuQ-large-msd-iter")237    eval_muq = eval_muq.to(device).eval()238 239    return cfm, tokenizer, muq, vae, eval_model, eval_muq240 241 242# for song edit, will be added in the future243def get_reference_latent(device, max_frames, edit, pred_segments, ref_song, vae_model):244    sampling_rate = 44100245    downsample_rate = 2048246    io_channels = 2247    if edit:248        input_audio, in_sr = torchaudio.load(ref_song)249        input_audio = prepare_audio(input_audio, in_sr=in_sr, target_sr=sampling_rate, target_length=None, target_channels=io_channels, device=device)250        input_audio = normalize_audio(input_audio, -6)251        252        with torch.no_grad():253            latent = encode_audio(input_audio, vae_model, chunked=True) # [b d t]254            mean, scale = latent.chunk(2, dim=1)255            prompt, _ = vae_sample(mean, scale)256            prompt = prompt.transpose(1, 2) # [b t d]257            prompt = prompt[:,:max_frames,:] if prompt.shape[1] >= max_frames else torch.nn.functional.pad(prompt, (0, 0, 0, max_frames - prompt.shape[1]), mode="constant", value=0)258        259        pred_segments = json.loads(pred_segments)260        # import pdb; pdb.set_trace()261        pred_frames = []262        for st, et in pred_segments:263            sf = 0 if st == -1 else int(st * sampling_rate / downsample_rate)264            # if st == -1:265            #     sf = 0266            # else:267            #     sf = int(st * sampling_rate / downsample_rate )268            269            ef = max_frames if et == -1 else int(et * sampling_rate / downsample_rate)270            # if et == -1:271            #     ef = max_frames272            # else:273            #     ef = int(et * sampling_rate / downsample_rate )274            pred_frames.append((sf, ef))275        # import pdb; pdb.set_trace()276        return prompt, pred_frames277    else:278        prompt = torch.zeros(1, max_frames, 64).to(device)279        pred_frames = [(0, max_frames)]280        return prompt, pred_frames281 282 283def get_negative_style_prompt(device):284    file_path = "./src/negative_prompt.npy"285    vocal_stlye = np.load(file_path)286 287    vocal_stlye = torch.from_numpy(vocal_stlye).to(device)  # [1, 512]288    vocal_stlye = vocal_stlye.half()289 290    return vocal_stlye291 292@torch.no_grad()293def eval_song(eval_model, eval_muq, songs):294    295    resampled_songs = [torchaudio.functional.resample(song.mean(dim=0, keepdim=True), 44100, 24000) for song in songs]296    ssl_list = []297    for i in range(len(resampled_songs)):298        output = eval_muq(resampled_songs[i], output_hidden_states=True)299        muq_ssl = output["hidden_states"][6]300        ssl_list.append(muq_ssl.squeeze(0))301 302    ssl = torch.stack(ssl_list)303    scores_g = eval_model(ssl)304    score = torch.mean(scores_g, dim=1)305    idx = score.argmax(dim=0)306    307    return songs[idx]308    309 310@torch.no_grad()311def get_audio_style_prompt(model, wav_path):312    vocal_flag = False313    mulan = model314    audio, _ = librosa.load(wav_path, sr=24000)315    audio_len = librosa.get_duration(y=audio, sr=24000)316    317    if audio_len <= 1:318        vocal_flag = True319    320    if audio_len > 10:321        start_time = int(audio_len // 2 - 5)322        wav = audio[start_time*24000:(start_time+10)*24000]323    324    else:325        wav = audio326    wav = torch.tensor(wav).unsqueeze(0).to(model.device)327    328    with torch.no_grad():329        audio_emb = mulan(wavs = wav) # [1, 512]330        331    audio_emb = audio_emb.half()332 333    return audio_emb, vocal_flag334 335 336@torch.no_grad()337def get_text_style_prompt(model, text_prompt):338    mulan = model339    340    with torch.no_grad():341        text_emb = mulan(texts = text_prompt) # [1, 512]342    text_emb = text_emb.half()343 344    return text_emb345 346 347@torch.no_grad()348def get_style_prompt(model, wav_path=None, prompt=None):349    mulan = model350 351    if prompt is not None:352        return mulan(texts=prompt).half()353 354    ext = os.path.splitext(wav_path)[-1].lower()355    if ext == ".mp3":356        meta = MP3(wav_path)357        audio_len = meta.info.length358    elif ext in [".wav", ".flac"]:359        audio_len = librosa.get_duration(path=wav_path)360    else:361        raise ValueError("Unsupported file format: {}".format(ext))362 363    if audio_len < 10:364        print(365            f"Warning: The audio file {wav_path} is too short ({audio_len:.2f} seconds). Expected at least 10 seconds."366        )367 368    assert audio_len >= 10369 370    mid_time = audio_len // 2371    start_time = mid_time - 5372    wav, _ = librosa.load(wav_path, sr=24000, offset=start_time, duration=10)373 374    wav = torch.tensor(wav).unsqueeze(0).to(model.device)375 376    with torch.no_grad():377        audio_emb = mulan(wavs=wav)  # [1, 512]378 379    audio_emb = audio_emb380    audio_emb = audio_emb.half()381 382    return audio_emb383 384def parse_lyrics(lyrics: str):385    lyrics_with_time = []386    lyrics = lyrics.strip()387    for line in lyrics.split("\n"):388        try:389            time, lyric = line[1:9], line[10:]390            lyric = lyric.strip()391            mins, secs = time.split(":")392            secs = int(mins) * 60 + float(secs)393            lyrics_with_time.append((secs, lyric))394        except:395            continue396    return lyrics_with_time397 398 399class CNENTokenizer:400    def __init__(self):401        with open("./diffrhythm/g2p/g2p/vocab.json", "r", encoding='utf-8') as file:402            self.phone2id: dict = json.load(file)["vocab"]403        self.id2phone = {v: k for (k, v) in self.phone2id.items()}404        from diffrhythm.g2p.g2p_generation import chn_eng_g2p405 406        self.tokenizer = chn_eng_g2p407 408    def encode(self, text):409        phone, token = self.tokenizer(text)410        token = [x + 1 for x in token]411        return token412 413    def decode(self, token):414        return "|".join([self.id2phone[x - 1] for x in token])415 416 417def get_lrc_token(max_frames, text, tokenizer, max_secs, device):418 419    lyrics_shift = 0420    sampling_rate = 44100421    downsample_rate = 2048422 423    comma_token_id = 1424    period_token_id = 2425 426    lrc_with_time = parse_lyrics(text)427 428    modified_lrc_with_time = []429    for i in range(len(lrc_with_time)):430        time, line = lrc_with_time[i]431        line_token = tokenizer.encode(line)432        modified_lrc_with_time.append((time, line_token))433    lrc_with_time = modified_lrc_with_time434 435    lrc_with_time = [436        (time_start, line)437        for (time_start, line) in lrc_with_time438        if time_start < max_secs439    ]440    if max_frames == 2048:441        lrc_with_time = lrc_with_time[:-1] if len(lrc_with_time) >= 1 else lrc_with_time442        443    end_frame = max_frames if max_frames == 2048 else int(max_secs * (sampling_rate / downsample_rate))444    end_frame = min(end_frame, max_frames) 445 446    normalized_start_time = 0.0447    448    normalized_duration = end_frame / max_frames449 450    lrc = torch.zeros((end_frame,), dtype=torch.long)451 452    tokens_count = 0453    last_end_pos = 0454    for time_start, line in lrc_with_time:455        tokens = [456            token if token != period_token_id else comma_token_id for token in line457        ] + [period_token_id]458        tokens = torch.tensor(tokens, dtype=torch.long)459        num_tokens = tokens.shape[0]460 461        gt_frame_start = int(time_start * sampling_rate / downsample_rate)462 463        frame_shift = random.randint(int(-lyrics_shift), int(lyrics_shift))464 465        frame_start = max(gt_frame_start - frame_shift, last_end_pos)466        frame_len = min(num_tokens, end_frame - frame_start)467 468        lrc[frame_start : frame_start + frame_len] = tokens[:frame_len]469 470        tokens_count += num_tokens471        last_end_pos = frame_start + frame_len472 473    lrc_emb = lrc.unsqueeze(0).to(device)474 475    normalized_start_time = torch.tensor(normalized_start_time).unsqueeze(0).to(device)476    normalized_start_time = normalized_start_time.half()477    478    normalized_duration = torch.tensor(normalized_duration).unsqueeze(0).to(device)479    normalized_duration = normalized_duration.half()480 481    return lrc_emb, normalized_start_time, end_frame, normalized_duration482 483 484def load_checkpoint(model, ckpt_path, device, use_ema=True):485    model = model.half()486 487    ckpt_type = ckpt_path.split(".")[-1]488    if ckpt_type == "safetensors":489        from safetensors.torch import load_file490 491        checkpoint = load_file(ckpt_path)492    else:493        checkpoint = torch.load(ckpt_path, weights_only=True)494 495    if use_ema:496        if ckpt_type == "safetensors":497            checkpoint = {"ema_model_state_dict": checkpoint}498        checkpoint["model_state_dict"] = {499            k.replace("ema_model.", ""): v500            for k, v in checkpoint["ema_model_state_dict"].items()501            if k not in ["initted", "step"]502        }503        model.load_state_dict(checkpoint["model_state_dict"], strict=False)504    else:505        if ckpt_type == "safetensors":506            checkpoint = {"model_state_dict": checkpoint}507        model.load_state_dict(checkpoint["model_state_dict"], strict=False)508 509    return model.to(device)510