Spaces:
Running on Zero
Running on Zero
File size: 4,004 Bytes
5a30378 469e695 5a30378 469e695 5a30378 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | # removed torchvision, torchaudio, safetensors, yaml, math, json, sys, io, warnings,
# random, T5Tokenizer, T5EncoderModel, tqdm, MelSpectrogram. not needed for inference.
import librosa
import torch
import accelerate
import numpy as np
from librosa.feature import chroma_stft
from anyaccomp.fmt_model import FlowMatchingTransformerConcat
from models.codec.amphion_codec.vocos import Vocos
from models.codec.coco.rep_coco_model import CocoStyle
from utils.util import load_config
class Sing2SongInferencePipeline:
"""
Wraps the three model components needed for inference:
1. CocoStyle – encodes the vocal chromagram into discrete style tokens
2. FlowMatchingTransformerConcat – diffuses those tokens into a mel spectrogram
3. Vocos – decodes the mel spectrogram into a waveform
"""
def __init__(
self,
checkpoint_path,
cfg_path,
vocoder_checkpoint_path,
vocoder_cfg_path,
device="cuda",
):
self.cfg = load_config(cfg_path)
self.device = device
self.checkpoint_path = checkpoint_path
self._load_model(checkpoint_path) # flow matching transformer
self._build_input_model() # chromagram encoder (CocoStyle)
self.vocoder_checkpoint_path = vocoder_checkpoint_path
self.vocoder_cfg = load_config(vocoder_cfg_path)
self._build_output_model() # vocoder (Vocos)
print("Output model built")
def _load_model(self, checkpoint_path):
self.model = FlowMatchingTransformerConcat(
cfg=self.cfg.model.flow_matching_transformer
)
accelerate.load_checkpoint_and_dispatch(self.model, checkpoint_path)
self.model.eval().to(self.device)
print(
f"model Params: {round(sum(p.numel() for p in self.model.parameters() if p.requires_grad)/1e6, 2)}M"
)
print(f"Loaded model from {checkpoint_path}")
def _build_input_model(self):
# construct_only_for_quantizer=True skips building the decoder half of CocoStyle —
# we only need the encoder + quantizer for inference.
self.coco_model = CocoStyle(
cfg=self.cfg.model.coco, construct_only_for_quantizer=True
)
self.coco_model.eval()
self.coco_model.to(self.device)
accelerate.load_checkpoint_and_dispatch(
self.coco_model, self.cfg.model.coco.pretrained_path
)
def _build_output_model(self):
self.vocoder = Vocos(cfg=self.vocoder_cfg.model.vocos)
accelerate.load_checkpoint_and_dispatch(
self.vocoder, self.vocoder_checkpoint_path
)
self.vocoder = self.vocoder.eval().to(self.device)
@torch.no_grad()
@torch.amp.autocast("cuda", dtype=torch.bfloat16)
def _extract_coco_codec(self, speech):
"""Compute a chromagram from the waveform, then quantize it into discrete tokens."""
target_chroma_dim = self.cfg.model.coco.chromagram_dim
speech = speech.cpu().numpy().squeeze()
# librosa returns [n_chroma, T]; transpose to [T, n_chroma] for the model.
chromagram = chroma_stft(
y=speech,
sr=self.cfg.preprocess.chromagram.sample_rate,
n_fft=self.cfg.preprocess.chromagram.n_fft,
hop_length=self.cfg.preprocess.chromagram.hop_size,
win_length=self.cfg.preprocess.chromagram.win_size,
n_chroma=target_chroma_dim,
).T
chromagram_feats = torch.tensor(chromagram).unsqueeze(0).to(self.device)
codecs, _ = self.coco_model.quantize(chromagram_feats)
return codecs
@torch.no_grad()
def encode_vocal(self, speech): # (B, T)
speech = speech.to(self.device)
return self._extract_coco_codec(speech)
@torch.no_grad()
def _generate_audio(self, mel):
# mel is [B, T, C]; vocoder expects [B, C, T].
return (self.vocoder(mel.transpose(1, 2)).detach().cpu())[0]
|