cocktailpeanut/DiffRhythm

1import torch2import torchaudio3from einops import rearrange4import argparse5import json6import os7from tqdm import tqdm8import random9import numpy as np10import time11import io12import gc13import pydub14 15from diffrhythm.infer.infer_utils import (16    get_reference_latent,17    get_lrc_token,18    get_style_prompt,19    get_audio_style_prompt,20    prepare_model,21    get_negative_style_prompt22)23 24def decode_audio(latents, vae_model, chunked=False, overlap=32, chunk_size=128):25    downsampling_ratio = 204826    io_channels = 227    if not chunked:28        # default behavior. Decode the entire latent in parallel29        return vae_model.decode_export(latents)30    else:31        # chunked decoding32        hop_size = chunk_size - overlap33        total_size = latents.shape[2]34        batch_size = latents.shape[0]35        chunks = []36        i = 037        for i in range(0, total_size - chunk_size + 1, hop_size):38            chunk = latents[:,:,i:i+chunk_size]39            chunks.append(chunk)40        if i+chunk_size != total_size:41            # Final chunk42            chunk = latents[:,:,-chunk_size:]43            chunks.append(chunk)44        chunks = torch.stack(chunks)45        num_chunks = chunks.shape[0]46        # samples_per_latent is just the downsampling ratio47        samples_per_latent = downsampling_ratio48        # Create an empty waveform, we will populate it with chunks as decode them49        y_size = total_size * samples_per_latent50        y_final = torch.zeros((batch_size,io_channels,y_size)).to(latents.device)51        for i in range(num_chunks):52            x_chunk = chunks[i,:]53            # decode the chunk54            y_chunk = vae_model.decode_export(x_chunk)55            # figure out where to put the audio along the time domain56            if i == num_chunks-1:57                # final chunk always goes at the end58                t_end = y_size59                t_start = t_end - y_chunk.shape[2]60            else:61                t_start = i * hop_size * samples_per_latent62                t_end = t_start + chunk_size * samples_per_latent63            #  remove the edges of the overlaps64            ol = (overlap//2) * samples_per_latent65            chunk_start = 066            chunk_end = y_chunk.shape[2]67            if i > 0:68                # no overlap for the start of the first chunk69                t_start += ol70                chunk_start += ol71            if i < num_chunks-1:72                # no overlap for the end of the last chunk73                t_end -= ol74                chunk_end -= ol75            # paste the chunked audio into our y_final output audio76            y_final[:,:,t_start:t_end] = y_chunk[:,:,chunk_start:chunk_end]77        return y_final78 79def inference(cfm_model, vae_model, cond, text, duration, style_prompt, negative_style_prompt, steps, cfg_strength, sway_sampling_coef, start_time, file_type, vocal_flag, odeint_method):80 81    with torch.inference_mode():82        print(">1")83        generated, _ = cfm_model.sample(84            cond=cond,85            text=text,86            duration=duration,87            style_prompt=style_prompt,88            negative_style_prompt=negative_style_prompt,89            steps=steps,90            cfg_strength=cfg_strength,91            sway_sampling_coef=sway_sampling_coef,92            start_time=start_time,93            vocal_flag=vocal_flag,94            odeint_method=odeint_method,95        )96        if torch.cuda.is_available():97            torch.cuda.empty_cache()98        elif torch.mps.is_available():99            torch.mps.empty_cache()100        gc.collect()101 102        103        print(">2")104        generated = generated.to(torch.float32)105        print(">3")106        latent = generated.transpose(1, 2) # [b d t]107        print(">4")108        output = decode_audio(latent, vae_model, chunked=True)109        print(">5")110 111        del latent, generated112        if torch.cuda.is_available():113            torch.cuda.empty_cache()114        elif torch.mps.is_available():115            torch.mps.empty_cache()116        gc.collect()117 118        print(">6")119 120        # Rearrange audio batch to a single sequence121        output = rearrange(output, "b d n -> d (b n)")122        output_tensor = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).cpu()123        output_np = output_tensor.numpy().T.astype(np.float32)124        125        if file_type == 'wav':126            return (44100, output_np)127        else:128            buffer = io.BytesIO()129            output_np = np.int16(output_np * 2**15)130            song = pydub.AudioSegment(output_np.tobytes(), frame_rate=44100, sample_width=2, channels=2)131            if file_type == 'mp3':132                song.export(buffer, format="mp3", bitrate="320k")133            else:134                song.export(buffer, format="ogg", bitrate="320k")135            return buffer.getvalue()136    137        138if __name__ == "__main__":139    parser = argparse.ArgumentParser()140    parser.add_argument('--lrc-path', type=str, default="example/eg.lrc") # lyrics of target song141    parser.add_argument('--ref-audio-path', type=str, default="example/eg.mp3") # reference audio as style prompt for target song142    parser.add_argument('--audio-length', type=int, default=95) # length of target song143    parser.add_argument('--output-dir', type=str, default="example/output")144    args = parser.parse_args()145    146    device = "cpu"147    if torch.cuda.is_available():148        device = "cuda"149    elif torch.mps.is_available():150        device = "mps"151    152    audio_length = args.audio_length153    if audio_length == 95:154        max_frames = 2048155    elif audio_length == 285:156        max_frames = 6144157    158    cfm, tokenizer, muq, vae = prepare_model(device)159    160    with open(args.lrc_path, 'r') as f:161        lrc = f.read()162    lrc_prompt, start_time = get_lrc_token(lrc, tokenizer, device)163    164    style_prompt = get_style_prompt(muq, args.ref_audio_path)165    166    negative_style_prompt = get_negative_style_prompt(device)167    168    latent_prompt = get_reference_latent(device, max_frames)169    170    s_t = time.time()171    generated_song = inference(cfm_model=cfm, 172                               vae_model=vae, 173                               cond=latent_prompt, 174                               text=lrc_prompt, 175                               duration=max_frames, 176                               style_prompt=style_prompt,177                               negative_style_prompt=negative_style_prompt,178                               start_time=start_time179                               )180    e_t = time.time() - s_t181    print(f"inference cost {e_t} seconds")182    183    output_dir = args.output_dir184    os.makedirs(output_dir, exist_ok=True)185    186    output_path = os.path.join(output_dir, "output.wav")187    torchaudio.save(output_path, generated_song, sample_rate=44100)188    189