Spaces:
Running on Zero
Running on Zero
| # 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) | |
| 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 | |
| def encode_vocal(self, speech): # (B, T) | |
| speech = speech.to(self.device) | |
| return self._extract_coco_codec(speech) | |
| 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] | |