remove chinese characters in dualcodec

Author: jiaqili3

models/codec/dualcodec/README.md

@@ -144,6 +144,11 @@ python -m dualcodec.infer.valle.cli_valle_infer --ref_audio <path_to_ref_audio>
 ```
 You can also leave all options empty and it will use the default values.
 
+#### Gradio interface
+```bash
+python -m dualcodec.infer.valle.gradio_valle_demo
+```
+
 ### DualCodec-Voicebox
 #### CLI Inference
 ```bash

models/codec/dualcodec/dualcodec/conf_tts/model/flattened_ar/llama_1x32768.yaml

@@ -0,0 +1,25 @@
+vocab_size: 84644  # ${51866+16384+10}
+speech_vocab_size: 32768
+initial_offset: 10
+
+llama_cfg:
+  _target_: transformers.models.llama.modeling_llama.LlamaConfig
+  vocab_size: ${..vocab_size}
+  hidden_size: 2048
+  intermediate_size: 8192
+  num_hidden_layers: 10
+  num_attention_heads: 16
+  pad_token_id: 0
+  bos_token_id: 1
+  eos_token_id: 2
+
+llm:
+  _target_: transformers.models.llama.modeling_llama.LlamaForCausalLM
+  config: ${..llama_cfg}
+model:
+  _target_: dualcodec.model_tts.flattened_ar.llama_wrapper.LLM
+  llm: ${..llm}
+  config: ${..llama_cfg}
+  speech_vocab_size: ${..speech_vocab_size}
+  initial_offset: ${..initial_offset}
+  sep_token: 3

models/codec/dualcodec/dualcodec/infer/flattened_ar/flatten_patterns.py

@@ -0,0 +1,63 @@
+import torch
+from einops import rearrange
+
+import numpy as np
+
+def offset_codes(semantic_code, offset_sizes):
+    """
+    Applies layer-specific offsets to each codec layer.
+
+    Args:
+        semantic_code (torch.Tensor): Input tensor of shape (batch_size, T, num_codec_layers).
+        offset_sizes (list[int]): List of offsets for each codec layer to distinguish them.
+
+    Returns:
+        torch.Tensor: Offset-applied tensor of shape (batch_size, T, num_codec_layers).
+    """
+    # Calculate cumulative offsets for each layer
+    cumulative_offsets = np.cumsum([0] + offset_sizes[:-1])  # Start with 0 for the first layer
+    # Apply offsets layer by layer
+    offsetted_code = []
+    for i, offset in enumerate(cumulative_offsets):
+        current_layer_code = semantic_code[..., i].clone().detach()  # Extract layer i
+        current_layer_code += offset  # Apply the cumulative offset
+        offsetted_code.append(current_layer_code)
+
+    # Stack all layers along the codec layer dimension
+    offsetted_code = torch.stack(offsetted_code, dim=-1)  # Shape: (batch_size, T, num_codec_layers)
+
+    return offsetted_code
+
+def deoffset_codes(flattened_codes, offset_sizes):
+    """
+    De-offsets a flattened tensor by subtracting the codebook size offsets for each codec layer.
+
+    Args:
+        flattened_codes (torch.Tensor): The offset and flattened tensor of shape (batch_size, T * num_codec_layers).
+        codebook_sizes (list[int]): A list of codebook sizes for each codec layer, used to remove offsets.
+
+    Returns:
+        torch.Tensor: The de-offset tensor of shape (batch_size, T, num_codec_layers).
+    """
+    # Calculate cumulative offsets for each layer
+    cumulative_offsets = np.cumsum([0] + offset_sizes[:-1])  # Start with 0 for the first layer
+
+    # Determine dimensions for reshaping
+    batch_size, flattened_dim = flattened_codes.shape
+    num_codec_layers = len(offset_sizes)
+    T = flattened_dim // num_codec_layers
+
+    # Reshape flattened_codes back to (batch_size, T, num_codec_layers)
+    reshaped_codes = flattened_codes.view(batch_size, T, num_codec_layers)
+
+    # De-offset each layer by subtracting the respective cumulative offset
+    deoffsetted_code = []
+    for i, offset in enumerate(cumulative_offsets):
+        current_layer_code = reshaped_codes[..., i].clone()  # Clone to avoid in-place operation
+        current_layer_code = current_layer_code - offset  # Remove the cumulative offset
+        deoffsetted_code.append(current_layer_code)
+    
+    # Stack all layers along the codec layer dimension
+    deoffsetted_code = torch.stack(deoffsetted_code, dim=-1)  # Shape: (batch_size, T, num_codec_layers)
+
+    return deoffsetted_code

models/codec/dualcodec/dualcodec/infer/flattened_ar/inference_flattened.py

@@ -0,0 +1,174 @@
+from cv2 import repeat
+import torch
+from einops import rearrange
+from .flatten_patterns import offset_codes, deoffset_codes
+
+
+class Inference:
+    def __init__(
+        self, model, tokenizer_obj, dualcodec_inference_obj, device="cuda", normalize=False, 
+        half=False, split_paragraph=True, offset_sizes=[16384, 4096, 4096, 4096], **kwargs
+    ) -> None:
+        self.model = model
+        import safetensors.torch
+
+        self.model.to(device)
+        self.model.eval()
+        self.tokenizer = tokenizer_obj
+        self.dualcodec_inference_obj = dualcodec_inference_obj
+        self.device = device
+        self.normalize = normalize
+        self.offset_sizes = offset_sizes
+
+        self.model = self.model.half()
+        
+        self.split_paragraph = split_paragraph
+    
+    @torch.no_grad()
+    def inference(
+        self,
+        speech_24k,
+        prompt_speech,
+        prompt_text,
+        prompt_language,
+        target_text,
+        target_language,
+        use_prompt_text=True,
+        temp=1.0,
+        top_k=1000,
+        top_p=0.85,
+        repeat_penalty=1.1,
+    ):
+        """
+            Generate text given speech and text prompts.
+
+        Args:
+            prompt_speech (str or Tensor): Speech file path or a tensor with shape (n_samples,).
+            prompt_text (str): Text prompt.
+            prompt_language (str): Language of the prompt.
+            target_text (str): Target text to be completed.
+            target_language (str): Language of the target text.
+            use_prompt_text (bool, optional): Whether to use the prompt text as input. Defaults to True.
+            temp (float, optional): Temperature parameter for the distribution. Defaults to 1.0.
+            top_k (int, optional): Number of tokens to keep before applying `top_p`. Defaults to 1000.
+            top_p (float, optional): Probability threshold to use for filtering tokens. Defaults to 0.85.
+
+        Returns:
+            str: Completed text.
+        """
+        self.model.eval()
+        prompt_text = prompt_text.strip()
+        # prompt_text = prompt_text.replace('.',',')
+        # prompt_text = prompt_text.replace('。','，')
+        target_text = target_text.replace("\n", "")
+        target_text = target_text.replace("\t", "")
+        return_values_0 = []
+        return_values_1 = []
+
+        prompt_len_tmp = len(self.tokenizer.encode(prompt_text)) // 2
+
+        if self.split_paragraph:
+            if prompt_language == 'zh':
+                from dualcodec.utils.frontend_utils import split_paragraph
+                texts = split_paragraph(
+                    target_text,
+                    None,
+                    "zh",
+                    token_max_n=60 - prompt_len_tmp,
+                    token_min_n=40 - prompt_len_tmp,
+                    merge_len=20,
+                    comma_split=False,
+                )
+            elif prompt_language == 'ja':
+                from dualcodec.utils.frontend_utils import split_paragraph
+                texts = split_paragraph(
+                    target_text,
+                    None,
+                    "zh",
+                    token_max_n=70,
+                    token_min_n=60,
+                    merge_len=20,
+                    comma_split=False,
+                )
+            elif prompt_language == 'en':
+                from dualcodec.utils.frontend_utils import split_paragraph
+                texts = split_paragraph(
+                    target_text,
+                    self.tokenizer.encode,
+                    "en",
+                    token_max_n=70 - prompt_len_tmp,
+                    token_min_n=60 - prompt_len_tmp,
+                    merge_len=20,
+                    comma_split=True,
+                )
+            else:
+                texts = [target_text]
+        if prompt_language == 'en':
+            texts = [prompt_text + ' ' + t for t in texts]
+        else:
+            texts = [prompt_text + t for t in texts]
+        print(texts)
+
+        all_codes = []
+
+        for text in texts:
+
+            if self.normalize:
+                from dualcodec.dataset.processor import normalize
+                text = list(normalize([{
+                    'language': prompt_language,
+                    'text': text,
+                }], en_punct=True, use_kana=False))[0]['text']
+            print(text)
+
+
+            prompt_text_tokens = torch.tensor(
+                [
+                    [self.tokenizer.to_language_token(prompt_language)]
+                    + self.tokenizer.encode(text)
+                ],
+                dtype=torch.int32,
+                device=self.device,
+            )
+            prompt_text_len = torch.tensor(
+                [prompt_text_tokens.shape[-1]], device=self.device
+            )
+
+            # target_text_tokens = torch.tensor(
+            #     [tokenizer.encode(target_text)], dtype=torch.int32
+            # )
+            # target_text_len = torch.tensor([target_text_tokens.shape[-1]])
+
+            text_token = prompt_text_tokens
+
+            # prompt semantic codes
+            # semantic_code, _ = self._extract_semantic_code(input_features, attention_mask)
+            semantic_codes, acoustic_codes = self.dualcodec_inference_obj.encode(prompt_speech, n_quantizers=4)
+            semantic_codes = rearrange(semantic_codes, 'b t -> b t 1')
+            num_codec_layers = 4
+            semantic_code = torch.cat([semantic_codes, acoustic_codes], dim=-1)[..., :num_codec_layers]
+
+            semantic_code = offset_codes(semantic_code, self.offset_sizes)
+            semantic_code = rearrange(semantic_code, 'b t q -> b (t q)')
+
+            ret_semantic_code = semantic_code.clone().detach()
+
+            out = self.model.inference(
+                text=text_token,
+                text_len=prompt_text_len,
+                prompt_text=None,
+                prompt_text_len=None,
+                prompt_speech_token=semantic_code,
+                prompt_speech_token_len=torch.tensor([semantic_code.shape[-1]]),
+                top_k=top_k,
+                top_p=top_p,
+                repeat_penalty=repeat_penalty,
+                temperature=temp,
+            )
+            out = deoffset_codes(out, self.offset_sizes)
+
+            all_codes.append(out)
+
+        all_codes = torch.cat(all_codes, dim=1) # FIXME not tested
+        out = self.dualcodec_inference_obj.decode(all_codes)
+        return out