Upload 10 files

soprano/__init__.py

@@ -0,0 +1 @@
+from .tts import SopranoTTS

soprano/backends/base.py

@@ -0,0 +1,20 @@
+class BaseModel:
+    def infer(self,
+            prompts,
+            top_p=0.95,
+            temperature=0.3,
+            repetition_penalty=1.2):
+        '''
+        Takes a list of prompts and returns the output hidden states
+        '''
+        pass
+
+    def stream_infer(self,
+            prompt,
+            top_p=0.95,
+            temperature=0.3,
+            repetition_penalty=1.2):
+        '''
+        Takes a prompt and returns an iterator of the output hidden states
+        '''
+        pass

soprano/backends/lmdeploy.py

@@ -0,0 +1,54 @@
+import torch
+from lmdeploy import pipeline, TurbomindEngineConfig, GenerationConfig
+from .base import BaseModel
+
+
+class LMDeployModel(BaseModel):
+    def __init__(self,
+            device='cuda',
+            cache_size_mb=100,
+            **kwargs):
+        assert device == 'cuda', "lmdeploy only supports cuda devices, consider changing device or using a different backend instead."
+        cache_size_ratio = cache_size_mb * 1024**2 / torch.cuda.get_device_properties('cuda').total_memory
+        backend_config = TurbomindEngineConfig(cache_max_entry_count=cache_size_ratio)
+        self.pipeline = pipeline('ekwek/Soprano-80M',
+            log_level='ERROR',
+            backend_config=backend_config)
+
+    def infer(self,
+            prompts,
+            top_p=0.95,
+            temperature=0.3,
+            repetition_penalty=1.2):
+        gen_config=GenerationConfig(output_last_hidden_state='generation',
+            do_sample=True,
+            top_p=top_p,
+            temperature=temperature,
+            repetition_penalty=repetition_penalty,
+            max_new_tokens=512)
+        responses = self.pipeline(prompts, gen_config=gen_config)
+        res = []
+        for response in responses:
+            res.append({
+                'finish_reason': response.finish_reason,
+                'hidden_state': response.last_hidden_state
+            })
+        return res
+
+    def stream_infer(self,
+            prompt,
+            top_p=0.95,
+            temperature=0.3,
+            repetition_penalty=1.2):
+        gen_config=GenerationConfig(output_last_hidden_state='generation',
+            do_sample=True,
+            top_p=top_p,
+            temperature=temperature,
+            repetition_penalty=repetition_penalty,
+            max_new_tokens=512)
+        responses = self.pipeline.stream_infer([prompt], gen_config=gen_config)
+        for response in responses:
+            yield {
+                'finish_reason': response.finish_reason,
+                'hidden_state': response.last_hidden_state
+            }

soprano/backends/transformers.py

@@ -0,0 +1,68 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from .base import BaseModel
+
+
+class TransformersModel(BaseModel):
+    def __init__(self,
+            device='cuda',
+            **kwargs):
+        self.device = device
+        
+        self.model = AutoModelForCausalLM.from_pretrained(
+            'ekwek/Soprano-80M',
+            torch_dtype=torch.bfloat16 if device == 'cuda' else torch.float32,
+            device_map=device
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained('ekwek/Soprano-80M')
+        self.model.eval()
+
+    def infer(self,
+            prompts,
+            top_p=0.95,
+            temperature=0.3,
+            repetition_penalty=1.2):
+        inputs = self.tokenizer(
+            prompts,
+            return_tensors='pt',
+            padding=True,
+            truncation=True,
+            max_length=512,
+        ).to(self.device)
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                input_ids=inputs['input_ids'],
+                attention_mask=inputs['attention_mask'],
+                max_new_tokens=512,
+                do_sample=True,
+                top_p=top_p,
+                temperature=temperature,
+                repetition_penalty=repetition_penalty,
+                pad_token_id=self.tokenizer.pad_token_id,
+                return_dict_in_generate=True,
+                output_hidden_states=True,
+            )
+        res = []
+        eos_token_id = self.model.config.eos_token_id
+        for i in range(len(prompts)):
+            seq = outputs.sequences[i]
+            hidden_states = []
+            num_output_tokens = len(outputs.hidden_states)
+            for j in range(num_output_tokens):
+                token = seq[j + seq.size(0) - num_output_tokens]
+                if token != eos_token_id: hidden_states.append(outputs.hidden_states[j][-1][i, -1, :])
+            last_hidden_state = torch.stack(hidden_states).squeeze()
+            finish_reason = 'stop' if seq[-1].item() == eos_token_id else 'length'
+            res.append({
+                'finish_reason': finish_reason,
+                'hidden_state': last_hidden_state
+            })
+        return res
+
+    def stream_infer(self,
+            prompt,
+            top_p=0.95,
+            temperature=0.3,
+            repetition_penalty=1.2):
+        raise NotImplementedError("transformers backend does not currently support streaming, please consider using lmdeploy backend instead.")

soprano/tts.py

@@ -0,0 +1,172 @@
+from .vocos.decoder import SopranoDecoder
+import torch
+import re
+from unidecode import unidecode
+from scipy.io import wavfile
+from huggingface_hub import hf_hub_download
+import os
+import time
+
+
+class SopranoTTS:
+    def __init__(self,
+            backend='auto',
+            device='cuda',
+            cache_size_mb=10,
+            decoder_batch_size=1):
+        RECOGNIZED_DEVICES = ['cuda']
+        RECOGNIZED_BACKENDS = ['auto', 'lmdeploy', 'transformers']
+        assert device in RECOGNIZED_DEVICES, f"unrecognized device {device}, device must be in {RECOGNIZED_DEVICES}"
+        if backend == 'auto':
+            if device == 'cpu':
+                backend = 'transformers'
+            else:
+                try:
+                    import lmdeploy
+                    backend = 'lmdeploy'
+                except ImportError:
+                    backend='transformers'
+            print(f"Using backend {backend}.")
+        assert backend in RECOGNIZED_BACKENDS, f"unrecognized backend {backend}, backend must be in {RECOGNIZED_BACKENDS}"
+
+        if backend == 'lmdeploy':
+            from .backends.lmdeploy import LMDeployModel
+            self.pipeline = LMDeployModel(device=device, cache_size_mb=cache_size_mb)
+        elif backend == 'transformers':
+            from .backends.transformers import TransformersModel
+            self.pipeline = TransformersModel(device=device)
+
+        self.decoder = SopranoDecoder().cuda()
+        decoder_path = hf_hub_download(repo_id='ekwek/Soprano-80M', filename='decoder.pth')
+        self.decoder.load_state_dict(torch.load(decoder_path))
+        self.decoder_batch_size=decoder_batch_size
+        self.RECEPTIVE_FIELD = 4 # Decoder receptive field
+        self.TOKEN_SIZE = 2048 # Number of samples per audio token
+
+        self.infer("Hello world!") # warmup
+
+    def _preprocess_text(self, texts):
+        '''
+        adds prompt format and sentence/part index
+        '''
+        res = []
+        for text_idx, text in enumerate(texts):
+            text = text.strip()
+            sentences = re.split(r"(?<=[.!?])\s+", text)
+            processed_sentences = []
+            for sentence_idx, sentence in enumerate(sentences):
+                old_len = len(sentence)
+                new_sentence = re.sub(r"[^A-Za-z !\$%&'*+,-./0123456789<>?_]", "", sentence)
+                new_sentence = re.sub(r"[<>/_+]", "", new_sentence)
+                new_sentence = re.sub(r"\.\.[^\.]", ".", new_sentence)
+                new_len = len(new_sentence)
+                if old_len != new_len:
+                    print(f"Warning: unsupported characters found in sentence: {sentence}\n\tThese characters have been removed.")
+                new_sentence = unidecode(new_sentence.strip())
+                processed_sentences.append((f'[STOP][TEXT]{new_sentence}[START]', text_idx, sentence_idx))
+            res.extend(processed_sentences)
+        return res
+
+    def infer(self,
+            text,
+            out_path=None,
+            top_p=0.95,
+            temperature=0.3,
+            repetition_penalty=1.2):
+        results = self.infer_batch([text],
+            top_p=top_p,
+            temperature=temperature,
+            repetition_penalty=repetition_penalty,
+            out_dir=None)[0]
+        if out_path:
+            wavfile.write(out_path, 32000, results.cpu().numpy())
+        return results
+
+    def infer_batch(self,
+            texts,
+            out_dir=None,
+            top_p=0.95,
+            temperature=0.3,
+            repetition_penalty=1.2):
+        sentence_data = self._preprocess_text(texts)
+        prompts = list(map(lambda x: x[0], sentence_data))
+        responses = self.pipeline.infer(prompts,
+            top_p=top_p,
+            temperature=temperature,
+            repetition_penalty=repetition_penalty)
+        hidden_states = []
+        for i, response in enumerate(responses):
+            if response['finish_reason'] != 'stop':
+                print(f"Warning: some sentences did not complete generation, likely due to hallucination.")
+            hidden_state = response['hidden_state']
+            hidden_states.append(hidden_state)
+        combined = list(zip(hidden_states, sentence_data))
+        combined.sort(key=lambda x: -x[0].size(0))
+        hidden_states, sentence_data = zip(*combined)
+
+        num_texts = len(texts)
+        audio_concat = [[] for _ in range(num_texts)]
+        for sentence in sentence_data:
+            audio_concat[sentence[1]].append(None)
+        for idx in range(0, len(hidden_states), self.decoder_batch_size):
+            batch_hidden_states = []
+            lengths = list(map(lambda x: x.size(0), hidden_states[idx:idx+self.decoder_batch_size]))
+            N = len(lengths)
+            for i in range(N):
+                batch_hidden_states.append(torch.cat([
+                    torch.zeros((1, 512, lengths[0]-lengths[i]), device='cuda'),
+                    hidden_states[idx+i].unsqueeze(0).transpose(1,2).cuda().to(torch.float32),
+                ], dim=2))
+            batch_hidden_states = torch.cat(batch_hidden_states)
+            with torch.no_grad():
+                audio = self.decoder(batch_hidden_states)
+            
+            for i in range(N):
+                text_id = sentence_data[idx+i][1]
+                sentence_id = sentence_data[idx+i][2]
+                audio_concat[text_id][sentence_id] = audio[i].squeeze()[-(lengths[i]*self.TOKEN_SIZE-self.TOKEN_SIZE):]
+        audio_concat = [torch.cat(x).cpu() for x in audio_concat]
+        
+        if out_dir:
+            os.makedirs(out_dir, exist_ok=True)
+            for i in range(len(audio_concat)):
+                wavfile.write(f"{out_dir}/{i}.wav", 32000, audio_concat[i].cpu().numpy())
+        return audio_concat
+
+    def infer_stream(self,
+            text,
+            chunk_size=1,
+            top_p=0.95,
+            temperature=0.3,
+            repetition_penalty=1.2):
+        start_time = time.time()
+        sentence_data = self._preprocess_text([text])
+
+        first_chunk = True
+        for sentence, _, _ in sentence_data:
+            responses = self.pipeline.stream_infer(sentence,
+                top_p=top_p,
+                temperature=temperature,
+                repetition_penalty=repetition_penalty)
+            hidden_states_buffer = []
+            chunk_counter = chunk_size
+            for token in responses:
+                finished = token['finish_reason'] is not None
+                if not finished: hidden_states_buffer.append(token['hidden_state'][-1])
+                hidden_states_buffer = hidden_states_buffer[-(2*self.RECEPTIVE_FIELD+chunk_size):]
+                if finished or len(hidden_states_buffer) >= self.RECEPTIVE_FIELD + chunk_size:
+                    if finished or chunk_counter == chunk_size:
+                        batch_hidden_states = torch.stack(hidden_states_buffer)
+                        inp = batch_hidden_states.unsqueeze(0).transpose(1, 2).cuda().to(torch.float32)
+                        with torch.no_grad():
+                            audio = self.decoder(inp)[0]
+                        if finished:
+                            audio_chunk = audio[-((self.RECEPTIVE_FIELD+chunk_counter-1)*self.TOKEN_SIZE-self.TOKEN_SIZE):]
+                        else:
+                            audio_chunk = audio[-((self.RECEPTIVE_FIELD+chunk_size)*self.TOKEN_SIZE-self.TOKEN_SIZE):-(self.RECEPTIVE_FIELD*self.TOKEN_SIZE-self.TOKEN_SIZE)]
+                        chunk_counter = 0
+                        if first_chunk:
+                            print(f"Streaming latency: {1000*(time.time()-start_time):.2f} ms")
+                            first_chunk = False
+                        yield audio_chunk.cpu()
+                    chunk_counter += 1

soprano/vocos/decoder.py

@@ -0,0 +1,45 @@
+import torch
+from torch import nn
+
+from .models import VocosBackbone
+from .heads import ISTFTHead
+
+
+class SopranoDecoder(nn.Module):
+    def __init__(self,
+                 num_input_channels=512,
+                 decoder_num_layers=8,
+                 decoder_dim=512,
+                 decoder_intermediate_dim=None,
+                 hop_length=512,
+                 n_fft=2048,
+                 upscale=4,
+                 dw_kernel=3,
+                ):
+        super().__init__()
+        self.decoder_initial_channels = num_input_channels
+        self.num_layers = decoder_num_layers
+        self.dim = decoder_dim
+        self.intermediate_dim = decoder_intermediate_dim if decoder_intermediate_dim else decoder_dim*3
+        self.hop_length = hop_length
+        self.n_fft = n_fft
+        self.upscale = upscale
+        self.dw_kernel = dw_kernel
+        
+        self.decoder = VocosBackbone(input_channels=self.decoder_initial_channels,
+                                      dim=self.dim,
+                                      intermediate_dim=self.intermediate_dim,
+                                      num_layers=self.num_layers,
+                                      input_kernel_size=dw_kernel,
+                                      dw_kernel_size=dw_kernel,
+                                      )
+        self.head = ISTFTHead(dim=self.dim,
+                            n_fft=self.n_fft,
+                            hop_length=self.hop_length)
+
+    def forward(self, x):
+        T = x.size(2)
+        x = torch.nn.functional.interpolate(x, size=self.upscale*(T-1)+1, mode='linear', align_corners=True)
+        x = self.decoder(x)
+        reconstructed = self.head(x)
+        return reconstructed

soprano/vocos/heads.py

@@ -0,0 +1,50 @@
+import torch
+from torch import nn
+from .spectral_ops import ISTFT
+
+
+class ISTFTHead(nn.Module):
+    """
+    ISTFT Head module for predicting STFT complex coefficients.
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames, which should align with
+                          the resolution of the input features.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "center"):
+        super().__init__()
+        out_dim = n_fft + 2
+        self.out = torch.nn.Linear(dim, out_dim)
+        self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
+
+    @torch.compiler.disable
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the ISTFTHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x = self.out(x.transpose(1,2)).transpose(1, 2)
+        mag, p = x.chunk(2, dim=1)
+        mag = torch.exp(mag)
+        mag = torch.clip(mag, max=1e2)  # safeguard to prevent excessively large magnitudes
+        # wrapping happens here. These two lines produce real and imaginary value
+        x = torch.cos(p)
+        y = torch.sin(p)
+        # recalculating phase here does not produce anything new
+        # only costs time
+        # phase = torch.atan2(y, x)
+        # S = mag * torch.exp(phase * 1j)
+        # better directly produce the complex value 
+        S = mag * (x + 1j * y)
+        audio = self.istft(S)
+        return audio

soprano/vocos/models.py

@@ -0,0 +1,61 @@
+from typing import Optional
+
+import torch
+from torch import nn
+
+from .modules import ConvNeXtBlock
+
+class VocosBackbone(nn.Module):
+    """
+    Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization
+
+    Args:
+        input_channels (int): Number of input features channels.
+        dim (int): Hidden dimension of the model.
+        intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
+        num_layers (int): Number of ConvNeXtBlock layers.
+        layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`.
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        dim: int,
+        intermediate_dim: int,
+        num_layers: int,
+        input_kernel_size: int = 9,
+        dw_kernel_size: int = 9,
+        layer_scale_init_value: Optional[float] = None,
+        pad: str = 'zeros',
+    ):
+        super().__init__()
+        self.embed = nn.Conv1d(input_channels, dim, kernel_size=input_kernel_size, padding=input_kernel_size//2, padding_mode=pad)
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.convnext = nn.ModuleList(
+            [
+                ConvNeXtBlock(
+                    dim=dim,
+                    intermediate_dim=intermediate_dim,
+                    dw_kernel_size=dw_kernel_size,
+                    layer_scale_init_value=layer_scale_init_value or 1 / num_layers**0.5,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv1d, nn.Linear)):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None: nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.embed(x) # (B, C, L)
+        x = self.norm(x.transpose(1, 2))
+        x = x.transpose(1, 2)
+        for conv_block in self.convnext:
+            x = conv_block(x)
+        x = self.final_layer_norm(x.transpose(1, 2))
+        x = x.transpose(1, 2)
+        return x

soprano/vocos/modules.py

@@ -0,0 +1,47 @@
+import torch
+from torch import nn
+
+
+class ConvNeXtBlock(nn.Module):
+    """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
+
+    Args:
+        dim (int): Number of input channels.
+        intermediate_dim (int): Dimensionality of the intermediate layer.
+        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+            Defaults to None.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        layer_scale_init_value: float,
+        dw_kernel_size: int = 9,
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv1d(dim, dim, kernel_size=dw_kernel_size, padding=dw_kernel_size//2, groups=dim)  # depthwise conv
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)
+
+        x = residual + x
+        return x

soprano/vocos/spectral_ops.py

@@ -0,0 +1,74 @@
+import torch
+from torch import nn
+
+class ISTFT(nn.Module):
+    """
+    Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
+    windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
+    See issue: https://github.com/pytorch/pytorch/issues/62323
+    Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
+    The NOLA constraint is met as we trim padded samples anyway.
+
+    Args:
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames.
+        win_length (int): The size of window frame and STFT filter.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        window = torch.hann_window(win_length).to('cuda')
+        self.register_buffer("window", window)
+
+    def forward(self, spec: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
+
+        Args:
+            spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
+                            N is the number of frequency bins, and T is the number of time frames.
+
+        Returns:
+            Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
+        """
+        if self.padding == "center":
+            spec[:,0] = 0 # fixes some strange bug where first/last freqs don't matter when bs<16 which causes exploding gradients
+            spec[:,-1] = 0
+            # Fallback to pytorch native implementation
+            return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True)
+        elif self.padding == "same":
+            pad = (self.win_length - self.hop_length) // 2
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        assert spec.dim() == 3, "Expected a 3D tensor as input"
+        B, N, T = spec.shape
+
+        # Inverse FFT
+        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
+        ifft = ifft * self.window[None, :, None]
+
+        # Overlap and Add
+        output_size = (T - 1) * self.hop_length + self.win_length
+        y = torch.nn.functional.fold(
+            ifft, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length),
+        )[:, 0, 0, pad:-pad]
+
+        # Window envelope
+        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
+        window_envelope = torch.nn.functional.fold(
+            window_sq, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length),
+        ).squeeze()[pad:-pad]
+
+        # Normalize
+        assert (window_envelope > 1e-11).all()
+        y = y / window_envelope
+
+        return y