Class RedCodecInfer

fireredtts2/codec/model.py:197–376

class RedCodecInfer(RedCodec):
    def __init__(self, codec: RedCodec):
        super().__init__(
            codec.ssl,
            codec.ssl_adaptor,
            codec.acoustic_encoder,
            codec.downsample,
            codec.rvq,
            codec.upsample,
            codec.semantic_decoder,
            codec.acoustic_decoder,
        )

    @classmethod
    def from_pretrained(cls, conf_path: str, ckpt_path: str) -> "RedCodecInfer":
        with open(conf_path, "r") as f:
            codec = RedCodec.from_config(conf_path)
        ckpt = torch.load(ckpt_path)["generator"]
        codec.load_state_dict(ckpt)
        return cls(codec)

    def _encode_one_batch(self, audio16k: torch.Tensor):
        B, T = audio16k.shape
        audio16k_length = torch.tensor(
            [T] * B, dtype=torch.long, device=audio16k.device
        )
        # Semantic
        ssl, ssl_length = self.ssl.forward(audio16k, audio16k_length)
        ssl = ssl.clone()  # For onnx export
        sem_feats, sem_length = self.ssl_adaptor(ssl, ssl_length)
        # Acoustic
        aco_feats, aco_length = self.acoustic_encoder(audio16k, audio16k_length)
        # VQ
        vq_in_feats = torch.cat([sem_feats, aco_feats], dim=2)
        vq_in_feats, vq_in_length = self.downsample(vq_in_feats, aco_length)
        # RVQ,
        indices = self.rvq.encode_codes(vq_in_feats.transpose(1, 2))  # (nq, B, L)
        indices = indices.permute(1, 0, 2)
        return indices  # (B, nq, L)

    @staticmethod
    def _pad_and_chunk(audio: torch.Tensor, chunk_size: int) -> List[torch.Tensor]:
        pad_len = math.ceil(audio.shape[1] / chunk_size) * chunk_size - audio.shape[1]
        audio = F.pad(audio, (0, pad_len), mode="constant", value=0)
        audio_chunks = audio.split(chunk_size, dim=1)
        return audio_chunks

    @torch.inference_mode()
    def encode(
        self,
        audio16k: torch.Tensor,
        audio16k_length: torch.Tensor = None,
        batch_size: int = 96,
    ):
        """
        Args:
            audio16k: shape (b, t)
            audio16k_length: (b,)