Retrieval-based-Voice-Conversion-WebUI/infer/fcpe.py

import torch

from tools.cuda_graph import cuda_graph_enabled, run_cuda_graph


def _is_directml_device(device):
    """Return whether *device* is the PrivateUse1 device registered by DirectML."""
    return getattr(device, "type", None) == "privateuseone" or "privateuseone" in str(
        device
    ).lower()


class FCPEInfer:
    """Project-local FCPE inference adapter with a DirectML execution path.

    DirectML does not support the complex tensor produced by ``torch.stft`` in
    torchfcpe's wav2mel stage.  Keep preprocessing and the small indexed
    decoder on CPU while the FCPE neural network runs on DirectML.  Other
    devices retain torchfcpe's original end-to-end inference path.
    """

    def __init__(self, device):
        from torchfcpe import spawn_bundled_infer_model

        self.device = device
        self.is_directml = _is_directml_device(device)
        if self.is_directml:
            # Loading a checkpoint directly with map_location=privateuseone is
            # not supported consistently.  Load on CPU, leave wav2mel there,
            # and move only the real-valued FCPE network to DirectML.
            self.infer_model = spawn_bundled_infer_model("cpu")
            self.infer_model.wav2mel.eval()
            self.cent_table_cpu = (
                self.infer_model.model.cent_table.detach().float().cpu().clone()
            )
            self.out_dims = int(self.infer_model.model.out_dims)
            self.infer_model.model.to(device).eval()
        else:
            self.infer_model = spawn_bundled_infer_model(device)
            if getattr(device, "type", None) == "cuda" or str(device).startswith(
                "cuda"
            ):
                # torchfcpe creates this tensor on CPU and copies it to CUDA in
                # every local-argmax decode.  Host-to-device copies are forbidden
                # during CUDA Graph capture, so keep the immutable offsets on the
                # target device and use the graph-safe equivalent decoder below.
                self.local_offsets = torch.arange(
                    9, device=device, dtype=torch.long
                ).view(1, 1, 9)

    def _graphable_model_infer(self, mel, decoder_mode, threshold):
        """Run FCPE network and an exactly equivalent capture-safe decoder."""
        model = self.infer_model.model
        latent = model(mel)
        batch, frames, _ = latent.shape
        cents = model.cent_table[None, None, :].expand(batch, frames, -1)

        if decoder_mode == "argmax":
            confidence = torch.max(latent, dim=-1, keepdim=True).values
            decoded = torch.sum(cents * latent, dim=-1, keepdim=True) / torch.sum(
                latent, dim=-1, keepdim=True
            )
        elif decoder_mode == "local_argmax":
            confidence, max_index = torch.max(latent, dim=-1, keepdim=True)
            local_index = self.local_offsets + (max_index - 4)
            local_index = local_index.clamp(0, model.out_dims - 1)
            local_cents = torch.gather(cents, -1, local_index)
            local_latent = torch.gather(latent, -1, local_index)
            decoded = torch.sum(
                local_cents * local_latent, dim=-1, keepdim=True
            ) / torch.sum(local_latent, dim=-1, keepdim=True)
        else:
            raise ValueError("Unknown FCPE decoder mode: %s" % decoder_mode)

        # Match torchfcpe's masking and cent-to-Hz formulas operation-for-operation.
        confidence_mask = torch.ones_like(confidence)
        confidence_mask.masked_fill_(confidence <= threshold, float("-inf"))
        decoded = decoded * confidence_mask
        return 10.0 * torch.pow(2.0, decoded / 1200.0)

    def _decode_on_cpu(self, latent, decoder_mode, threshold):
        """Decode DML network logits on CPU with torchfcpe's exact formulas.

        The current DirectML backend's ``aten::gather`` returns incorrect FCPE
        bin values even though its indices and the neural-network logits match
        CPU.  Decoding is tiny compared with the model, so keep this
        compatibility boundary on CPU as well as the complex STFT.
        """
        latent = latent.detach().float().cpu()
        batch, frames, _ = latent.shape
        cents = self.cent_table_cpu[None, None, :].expand(batch, frames, -1)

        if decoder_mode == "argmax":
            confidence = torch.max(latent, dim=-1, keepdim=True).values
            decoded = torch.sum(cents * latent, dim=-1, keepdim=True) / torch.sum(
                latent, dim=-1, keepdim=True
            )
        elif decoder_mode == "local_argmax":
            confidence, max_index = torch.max(latent, dim=-1, keepdim=True)
            local_index = torch.arange(9, dtype=torch.long) + (max_index - 4)
            local_index.clamp_(0, self.out_dims - 1)
            local_cents = torch.gather(cents, -1, local_index)
            local_latent = torch.gather(latent, -1, local_index)
            decoded = torch.sum(
                local_cents * local_latent, dim=-1, keepdim=True
            ) / torch.sum(local_latent, dim=-1, keepdim=True)
        else:
            raise ValueError(f"Unknown FCPE decoder mode: {decoder_mode}")

        decoded = decoded.masked_fill(confidence <= threshold, float("-inf"))
        return 10.0 * torch.pow(2.0, decoded / 1200.0)

    # torch.no_grad is used instead of inference_mode because DirectML's
    # PrivateUse1 backend still updates version counters in a few operators.
    @torch.no_grad()
    def infer(
        self,
        wav,
        sr,
        decoder_mode="local_argmax",
        threshold=0.006,
    ):
        if not self.is_directml:
            wav = wav.to(self.device)
            if cuda_graph_enabled(wav.device):
                # Wav2MelModule contains tensor-dependent Python conditionals and
                # cannot be captured.  Running it eagerly also creates/caches its
                # STFT window before the graph.  The much larger FCPE neural net
                # and decoder form the stable-shape CUDA Graph boundary.
                mel = self.infer_model.wav2mel(wav, sr)
                return run_cuda_graph(
                    self.infer_model.model,
                    "fcpe-core-%s-%s" % (decoder_mode, threshold),
                    lambda input_mel: self._graphable_model_infer(
                        input_mel,
                        decoder_mode,
                        threshold,
                    ),
                    mel,
                )
            return run_cuda_graph(
                self.infer_model,
                "fcpe-%s-%s-%s" % (sr, decoder_mode, threshold),
                lambda input_wav: self.infer_model.infer(
                    input_wav,
                    sr=sr,
                    decoder_mode=decoder_mode,
                    threshold=threshold,
                ),
                wav,
            )

        wav_cpu = wav.detach().to(device="cpu", dtype=torch.float32)
        mel_cpu = self.infer_model.wav2mel(wav_cpu, sr)
        mel_dml = mel_cpu.to(device=self.device, dtype=torch.float32)
        latent_dml = self.infer_model.model(mel_dml)
        return self._decode_on_cpu(
            latent_dml,
            decoder_mode=decoder_mode,
            threshold=threshold,
        )