xscdvfaaqqq/DiffRhythm

1import torch2import torchaudio3from einops import rearrange4import argparse5import os6import time7import random8 9import torch10import torchaudio11import numpy as np12from einops import rearrange13import io14import pydub15 16from diffrhythm.infer.infer_utils import (17    decode_audio,18    get_lrc_token,19    get_negative_style_prompt,20    get_reference_latent,21    get_style_prompt,22    prepare_model,23    eval_song,24)25 26 27def inference(28    cfm_model,29    vae_model,30    eval_model,31    eval_muq,32    cond,33    text,34    duration,35    style_prompt,36    negative_style_prompt,37    steps,38    cfg_strength,39    sway_sampling_coef,40    start_time,41    file_type,42    vocal_flag,43    odeint_method,44    pred_frames,45    batch_infer_num,46    song_duration,47    chunked=True,48):49    with torch.inference_mode():50        latents, _ = cfm_model.sample(51            cond=cond,52            text=text,53            duration=duration,54            max_duration=duration,55            style_prompt=style_prompt,56            negative_style_prompt=negative_style_prompt,57            steps=steps,58            cfg_strength=cfg_strength,59            sway_sampling_coef=sway_sampling_coef,60            start_time=start_time,61            vocal_flag=vocal_flag,62            odeint_method=odeint_method,63            latent_pred_segments=pred_frames,64            batch_infer_num=batch_infer_num,65            song_duration=song_duration66        )67 68        outputs = []69        for latent in latents:70            latent = latent.to(torch.float32)71            latent = latent.transpose(1, 2)  # [b d t]72 73            output = decode_audio(latent, vae_model, chunked=chunked)74 75            # Rearrange audio batch to a single sequence76            output = rearrange(output, "b d n -> d (b n)")77            78            outputs.append(output)79        if batch_infer_num > 1:80            generated_song = eval_song(eval_model, eval_muq, outputs)81        else:82            generated_song = outputs[0]83        output_tensor = generated_song.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).cpu()84        output_np = output_tensor.numpy().T.astype(np.float32)85        if file_type == 'wav':86            return (44100, output_np)87        else:88            buffer = io.BytesIO()89            output_np = np.int16(output_np * 2**15)90            song = pydub.AudioSegment(output_np.tobytes(), frame_rate=44100, sample_width=2, channels=2)91            if file_type == 'mp3':92                song.export(buffer, format="mp3", bitrate="320k")93            else:94                song.export(buffer, format="ogg", bitrate="320k")95            return buffer.getvalue()96        97 98 99if __name__ == "__main__":100    parser = argparse.ArgumentParser()101    parser.add_argument(102        "--lrc-path",103        type=str,104        help="lyrics of target song",105    )  # lyrics of target song106    parser.add_argument(107        "--ref-prompt",108        type=str,109        help="reference prompt as style prompt for target song",110        required=False,111    )  # reference prompt as style prompt for target song112    parser.add_argument(113        "--ref-audio-path",114        type=str,115        help="reference audio as style prompt for target song",116        required=False,117    )  # reference audio as style prompt for target song118    parser.add_argument(119        "--chunked",120        action="store_true",121        help="whether to use chunked decoding",122    )  # whether to use chunked decoding123    parser.add_argument(124        "--audio-length",125        type=int,126        default=95,127        choices=[95, 285],128        help="length of generated song",129    )  # length of target song130    parser.add_argument(131        "--repo-id", type=str, default="ASLP-lab/DiffRhythm-base", help="target model"132    )133    parser.add_argument(134        "--output-dir",135        type=str,136        default="infer/example/output",137        help="output directory fo generated song",138    )  # output directory of target song139    parser.add_argument(140        "--edit",141        action="store_true",142        help="whether to open edit mode",143    )  # edit flag144    parser.add_argument(145        "--ref-song",146        type=str,147        required=False,148        help="reference prompt as latent prompt for editing",149    )  # reference prompt as latent prompt for editing150    parser.add_argument(151        "--edit-segments",152        type=str,153        required=False,154        help="edit segments o target song",155    )  # edit segments o target song156    args = parser.parse_args()157 158    assert (159        args.ref_prompt or args.ref_audio_path160    ), "either ref_prompt or ref_audio_path should be provided"161    assert not (162        args.ref_prompt and args.ref_audio_path163    ), "only one of them should be provided"164    if args.edit:165        assert (166            args.ref_song and args.edit_segments167        ), "reference song and edit segments should be provided for editing"168 169    device = "cpu"170    if torch.cuda.is_available():171        device = "cuda"172    elif torch.mps.is_available():173        device = "mps"174 175    audio_length = args.audio_length176    if audio_length == 95:177        max_frames = 2048178    elif audio_length == 285:179        max_frames = 6144180 181    cfm, tokenizer, muq, vae, eval_model, eval_muq = prepare_model(max_frames, device, repo_id=args.repo_id)182 183    if args.lrc_path:184        with open(args.lrc_path, "r", encoding='utf-8') as f:185            lrc = f.read()186    else:187        lrc = ""188    lrc_prompt, start_time = get_lrc_token(max_frames, lrc, tokenizer, device)189 190    if args.ref_audio_path:191        style_prompt = get_style_prompt(muq, args.ref_audio_path)192    else:193        style_prompt = get_style_prompt(muq, prompt=args.ref_prompt)194 195    negative_style_prompt = get_negative_style_prompt(device)196 197    latent_prompt, pred_frames = get_reference_latent(device, max_frames, args.edit, args.edit_segments, args.ref_song, vae)198 199    s_t = time.time()200    generated_songs = inference(201        cfm_model=cfm,202        vae_model=vae,203        cond=latent_prompt,204        text=lrc_prompt,205        duration=max_frames,206        style_prompt=style_prompt,207        negative_style_prompt=negative_style_prompt,208        start_time=start_time,209        pred_frames=pred_frames,210        chunked=args.chunked,211    )212    213    214    215    generated_song = eval_song(eval_model, eval_muq, generated_songs)216    217    # Peak normalize, clip, convert to int16, and save to file218    generated_song = (219        generated_song.to(torch.float32)220        .div(torch.max(torch.abs(generated_song)))221        .clamp(-1, 1)222        .mul(32767)223        .to(torch.int16)224        .cpu()225    )226    e_t = time.time() - s_t227    print(f"inference cost {e_t:.2f} seconds")228    output_dir = args.output_dir229    os.makedirs(output_dir, exist_ok=True)230 231    output_path = os.path.join(output_dir, "output.wav")232    torchaudio.save(output_path, generated_song, sample_rate=44100)233